langgraph/libs/langgraph/tests/test_pregel_async.py

import asyncio
import enum
import functools
import gc
import logging
import operator
import random
import sys
import uuid
from collections import Counter, deque
from contextlib import asynccontextmanager, contextmanager
from dataclasses import replace
from time import perf_counter
from typing import (
    Annotated,
    Any,
    AsyncGenerator,
    AsyncIterator,
    Dict,
    Generator,
    List,
    Literal,
    Optional,
    Tuple,
    Union,
)
from uuid import UUID

import httpx
import pytest
from langchain_core.language_models import GenericFakeChatModel
from langchain_core.runnables import RunnableConfig, RunnableLambda, RunnablePassthrough
from langchain_core.utils.aiter import aclosing
from pytest_mock import MockerFixture
from syrupy import SnapshotAssertion
from typing_extensions import TypedDict

from langgraph.channels.base import BaseChannel
from langgraph.channels.binop import BinaryOperatorAggregate
from langgraph.channels.context import Context
from langgraph.channels.last_value import LastValue
from langgraph.channels.topic import Topic
from langgraph.checkpoint.base import (
    ChannelVersions,
    Checkpoint,
    CheckpointMetadata,
    CheckpointTuple,
)
from langgraph.checkpoint.memory import InMemorySaver, MemorySaver
from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, PUSH, START
from langgraph.errors import InvalidUpdateError, NodeInterrupt
from langgraph.func import entrypoint, task
from langgraph.graph import END, Graph, StateGraph
from langgraph.graph.message import MessagesState, add_messages
from langgraph.prebuilt.tool_node import ToolNode
from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot
from langgraph.pregel.loop import AsyncPregelLoop
from langgraph.pregel.retry import RetryPolicy
from langgraph.pregel.runner import PregelRunner
from langgraph.store.base import BaseStore
from langgraph.types import (
    Command,
    Interrupt,
    PregelTask,
    Send,
    StateUpdate,
    StreamWriter,
    interrupt,
)
from tests.any_str import AnyStr, AnyVersion, FloatBetween, UnsortedSequence
from tests.conftest import (
    ALL_CHECKPOINTERS_ASYNC,
    ALL_CHECKPOINTERS_ASYNC_PLUS_NONE,
    ALL_STORES_ASYNC,
    REGULAR_CHECKPOINTERS_ASYNC,
    SHOULD_CHECK_SNAPSHOTS,
    awith_checkpointer,
    awith_store,
)
from tests.fake_tracer import FakeTracer
from tests.memory_assert import MemorySaverNoPending
from tests.messages import (
    _AnyIdAIMessage,
    _AnyIdAIMessageChunk,
    _AnyIdHumanMessage,
    _AnyIdToolMessage,
)

logger = logging.getLogger(__name__)

pytestmark = pytest.mark.anyio

NEEDS_CONTEXTVARS = pytest.mark.skipif(
    sys.version_info < (3, 11),
    reason="Python 3.11+ is required for async contextvars support",
)


async def test_checkpoint_errors() -> None:
    class FaultyGetCheckpointer(InMemorySaver):
        async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
            raise ValueError("Faulty get_tuple")

    class FaultyPutCheckpointer(InMemorySaver):
        async def aput(
            self,
            config: RunnableConfig,
            checkpoint: Checkpoint,
            metadata: CheckpointMetadata,
            new_versions: ChannelVersions,
        ) -> RunnableConfig:
            raise ValueError("Faulty put")

    class FaultyPutWritesCheckpointer(InMemorySaver):
        async def aput_writes(
            self, config: RunnableConfig, writes: List[Tuple[str, Any]], task_id: str
        ) -> RunnableConfig:
            raise ValueError("Faulty put_writes")

    class FaultyVersionCheckpointer(InMemorySaver):
        def get_next_version(self, current: Optional[int], channel: BaseChannel) -> int:
            raise ValueError("Faulty get_next_version")

    def logic(inp: str) -> str:
        return ""

    builder = StateGraph(Annotated[str, operator.add])
    builder.add_node("agent", logic)
    builder.add_edge(START, "agent")

    graph = builder.compile(checkpointer=FaultyGetCheckpointer())
    with pytest.raises(ValueError, match="Faulty get_tuple"):
        await graph.ainvoke("", {"configurable": {"thread_id": "thread-1"}})
    with pytest.raises(ValueError, match="Faulty get_tuple"):
        async for _ in graph.astream("", {"configurable": {"thread_id": "thread-2"}}):
            pass
    with pytest.raises(ValueError, match="Faulty get_tuple"):
        async for _ in graph.astream_events(
            "", {"configurable": {"thread_id": "thread-3"}}, version="v2"
        ):
            pass

    graph = builder.compile(checkpointer=FaultyPutCheckpointer())
    with pytest.raises(ValueError, match="Faulty put"):
        await graph.ainvoke("", {"configurable": {"thread_id": "thread-1"}})
    with pytest.raises(ValueError, match="Faulty put"):
        async for _ in graph.astream("", {"configurable": {"thread_id": "thread-2"}}):
            pass
    with pytest.raises(ValueError, match="Faulty put"):
        async for _ in graph.astream_events(
            "", {"configurable": {"thread_id": "thread-3"}}, version="v2"
        ):
            pass

    graph = builder.compile(checkpointer=FaultyVersionCheckpointer())
    with pytest.raises(ValueError, match="Faulty get_next_version"):
        await graph.ainvoke("", {"configurable": {"thread_id": "thread-1"}})
    with pytest.raises(ValueError, match="Faulty get_next_version"):
        async for _ in graph.astream("", {"configurable": {"thread_id": "thread-2"}}):
            pass
    with pytest.raises(ValueError, match="Faulty get_next_version"):
        async for _ in graph.astream_events(
            "", {"configurable": {"thread_id": "thread-3"}}, version="v2"
        ):
            pass

    # add a parallel node
    builder.add_node("parallel", logic)
    builder.add_edge(START, "parallel")
    graph = builder.compile(checkpointer=FaultyPutWritesCheckpointer())
    with pytest.raises(ValueError, match="Faulty put_writes"):
        await graph.ainvoke("", {"configurable": {"thread_id": "thread-1"}})
    with pytest.raises(ValueError, match="Faulty put_writes"):
        async for _ in graph.astream("", {"configurable": {"thread_id": "thread-2"}}):
            pass
    with pytest.raises(ValueError, match="Faulty put_writes"):
        async for _ in graph.astream_events(
            "", {"configurable": {"thread_id": "thread-3"}}, version="v2"
        ):
            pass


async def test_py_async_with_cancel_behavior() -> None:
    """This test confirms that in all versions of Python we support, __aexit__
    is not cancelled when the coroutine containing the async with block is cancelled."""

    logs: list[str] = []

    class MyContextManager:
        async def __aenter__(self):
            logs.append("Entering")
            return self

        async def __aexit__(self, exc_type, exc_val, exc_tb):
            logs.append("Starting exit")
            try:
                # Simulate some cleanup work
                await asyncio.sleep(2)
                logs.append("Cleanup completed")
            except asyncio.CancelledError:
                logs.append("Cleanup was cancelled!")
                raise
            logs.append("Exit finished")

    async def main():
        try:
            async with MyContextManager():
                logs.append("In context")
                await asyncio.sleep(1)
                logs.append("This won't print if cancelled")
        except asyncio.CancelledError:
            logs.append("Context was cancelled")
            raise

    # create task
    t = asyncio.create_task(main())
    # cancel after 0.2 seconds
    await asyncio.sleep(0.2)
    t.cancel()
    # check logs before cancellation is handled
    assert logs == [
        "Entering",
        "In context",
    ], "Cancelled before cleanup started"
    # wait for task to finish
    try:
        await t
    except asyncio.CancelledError:
        # check logs after cancellation is handled
        assert logs == [
            "Entering",
            "In context",
            "Starting exit",
            "Cleanup completed",
            "Exit finished",
            "Context was cancelled",
        ], "Cleanup started and finished after cancellation"
    else:
        assert False, "Task should be cancelled"


async def test_checkpoint_put_after_cancellation() -> None:
    logs: list[str] = []

    class LongPutCheckpointer(MemorySaver):
        async def aput(
            self,
            config: RunnableConfig,
            checkpoint: Checkpoint,
            metadata: CheckpointMetadata,
            new_versions: ChannelVersions,
        ) -> RunnableConfig:
            logs.append("checkpoint.aput.start")
            try:
                await asyncio.sleep(1)
                return await super().aput(config, checkpoint, metadata, new_versions)
            finally:
                logs.append("checkpoint.aput.end")

    inner_task_cancelled = False

    async def awhile(input: Any) -> None:
        logs.append("awhile.start")
        try:
            await asyncio.sleep(1)
        except asyncio.CancelledError:
            nonlocal inner_task_cancelled
            inner_task_cancelled = True
            raise
        finally:
            logs.append("awhile.end")

    builder = Graph()
    builder.add_node("agent", awhile)
    builder.set_entry_point("agent")
    builder.set_finish_point("agent")

    graph = builder.compile(checkpointer=LongPutCheckpointer())
    thread1 = {"configurable": {"thread_id": "1"}}

    # start the task
    t = asyncio.create_task(graph.ainvoke(1, thread1))
    # cancel after 0.2 seconds
    await asyncio.sleep(0.2)
    t.cancel()
    # check logs before cancellation is handled
    assert sorted(logs) == [
        "awhile.start",
        "checkpoint.aput.start",
    ], "Cancelled before checkpoint put started"
    # wait for task to finish
    try:
        await t
    except asyncio.CancelledError:
        # check logs after cancellation is handled
        assert sorted(logs) == [
            "awhile.end",
            "awhile.start",
            "checkpoint.aput.end",
            "checkpoint.aput.start",
        ], "Checkpoint put is not cancelled"
    else:
        assert False, "Task should be cancelled"


async def test_checkpoint_put_after_cancellation_stream_anext() -> None:
    logs: list[str] = []

    class LongPutCheckpointer(MemorySaver):
        async def aput(
            self,
            config: RunnableConfig,
            checkpoint: Checkpoint,
            metadata: CheckpointMetadata,
            new_versions: ChannelVersions,
        ) -> RunnableConfig:
            logs.append("checkpoint.aput.start")
            try:
                await asyncio.sleep(1)
                return await super().aput(config, checkpoint, metadata, new_versions)
            finally:
                logs.append("checkpoint.aput.end")

    inner_task_cancelled = False

    async def awhile(input: Any) -> None:
        logs.append("awhile.start")
        try:
            await asyncio.sleep(1)
        except asyncio.CancelledError:
            nonlocal inner_task_cancelled
            inner_task_cancelled = True
            raise
        finally:
            logs.append("awhile.end")

    builder = Graph()
    builder.add_node("agent", awhile)
    builder.set_entry_point("agent")
    builder.set_finish_point("agent")

    graph = builder.compile(checkpointer=LongPutCheckpointer())
    thread1 = {"configurable": {"thread_id": "1"}}

    # start the task
    s = graph.astream(1, thread1)
    t = asyncio.create_task(s.__anext__())
    # cancel after 0.2 seconds
    await asyncio.sleep(0.2)
    t.cancel()
    # check logs before cancellation is handled
    assert sorted(logs) == [
        "awhile.start",
        "checkpoint.aput.start",
    ], "Cancelled before checkpoint put started"
    # wait for task to finish
    try:
        await t
    except asyncio.CancelledError:
        # check logs after cancellation is handled
        assert sorted(logs) == [
            "awhile.end",
            "awhile.start",
            "checkpoint.aput.end",
            "checkpoint.aput.start",
        ], "Checkpoint put is not cancelled"
    else:
        assert False, "Task should be cancelled"


async def test_checkpoint_put_after_cancellation_stream_events_anext() -> None:
    logs: list[str] = []

    class LongPutCheckpointer(MemorySaver):
        async def aput(
            self,
            config: RunnableConfig,
            checkpoint: Checkpoint,
            metadata: CheckpointMetadata,
            new_versions: ChannelVersions,
        ) -> RunnableConfig:
            logs.append("checkpoint.aput.start")
            try:
                await asyncio.sleep(1)
                return await super().aput(config, checkpoint, metadata, new_versions)
            finally:
                logs.append("checkpoint.aput.end")

    inner_task_cancelled = False

    async def awhile(input: Any) -> None:
        logs.append("awhile.start")
        try:
            await asyncio.sleep(1)
        except asyncio.CancelledError:
            nonlocal inner_task_cancelled
            inner_task_cancelled = True
            raise
        finally:
            logs.append("awhile.end")

    builder = Graph()
    builder.add_node("agent", awhile)
    builder.set_entry_point("agent")
    builder.set_finish_point("agent")

    graph = builder.compile(checkpointer=LongPutCheckpointer())
    thread1 = {"configurable": {"thread_id": "1"}}

    # start the task
    s = graph.astream_events(1, thread1, version="v2", include_names=["LangGraph"])
    # skip first event (happens right away)
    await s.__anext__()
    # start the task for 2nd event
    t = asyncio.create_task(s.__anext__())
    # cancel after 0.2 seconds
    await asyncio.sleep(0.2)
    t.cancel()
    # check logs before cancellation is handled
    assert logs == [
        "checkpoint.aput.start",
        "awhile.start",
    ], "Cancelled before checkpoint put started"
    # wait for task to finish
    try:
        await t
    except asyncio.CancelledError:
        # check logs after cancellation is handled
        assert logs == [
            "checkpoint.aput.start",
            "awhile.start",
            "awhile.end",
            "checkpoint.aput.end",
        ], "Checkpoint put is not cancelled"
    else:
        assert False, "Task should be cancelled"


async def test_node_cancellation_on_external_cancel() -> None:
    inner_task_cancelled = False

    async def awhile(input: Any) -> None:
        try:
            await asyncio.sleep(1)
        except asyncio.CancelledError:
            nonlocal inner_task_cancelled
            inner_task_cancelled = True
            raise

    builder = Graph()
    builder.add_node("agent", awhile)
    builder.set_entry_point("agent")
    builder.set_finish_point("agent")

    graph = builder.compile()

    with pytest.raises(asyncio.TimeoutError):
        await asyncio.wait_for(graph.ainvoke(1), 0.5)

    assert inner_task_cancelled


async def test_node_cancellation_on_other_node_exception() -> None:
    inner_task_cancelled = False

    async def awhile(input: Any) -> None:
        try:
            await asyncio.sleep(1)
        except asyncio.CancelledError:
            nonlocal inner_task_cancelled
            inner_task_cancelled = True
            raise

    async def iambad(input: Any) -> None:
        raise ValueError("I am bad")

    builder = Graph()
    builder.add_node("agent", awhile)
    builder.add_node("bad", iambad)
    builder.set_conditional_entry_point(lambda _: ["agent", "bad"], then=END)

    graph = builder.compile()

    with pytest.raises(ValueError, match="I am bad"):
        # This will raise ValueError, not TimeoutError
        await asyncio.wait_for(graph.ainvoke(1), 0.5)

    assert inner_task_cancelled


async def test_node_cancellation_on_other_node_exception_two() -> None:
    async def awhile(input: Any) -> None:
        await asyncio.sleep(1)

    async def iambad(input: Any) -> None:
        raise ValueError("I am bad")

    builder = Graph()
    builder.add_node("agent", awhile)
    builder.add_node("bad", iambad)
    builder.set_conditional_entry_point(lambda _: ["agent", "bad"], then=END)

    graph = builder.compile()

    with pytest.raises(ValueError, match="I am bad"):
        # This will raise ValueError, not CancelledError
        await graph.ainvoke(1)


@NEEDS_CONTEXTVARS
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_dynamic_interrupt(checkpointer_name: str) -> None:
    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":
            answer = interrupt("Just because...")
        else:
            answer = " all good"
        return {"my_key": answer}

    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",
    }

    async with awith_checkpointer(checkpointer_name) as checkpointer:
        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"})

        # flow: interrupt -> resume with answer
        thread2 = {"configurable": {"thread_id": "2"}}
        # stop when about to enter node
        assert [
            c
            async for c in tool_two.astream(
                {"my_key": "value ⛰️", "market": "DE"}, thread2
            )
        ] == [
            {
                "__interrupt__": (
                    Interrupt(
                        value="Just because...",
                        resumable=True,
                        ns=[AnyStr("tool_two:")],
                    ),
                )
            },
        ]
        # resume with answer
        assert [
            c async for c in tool_two.astream(Command(resume=" my answer"), thread2)
        ] == [
            {"tool_two": {"my_key": " my answer"}},
        ]

        # flow: interrupt -> clear
        thread1 = {"configurable": {"thread_id": "1"}}
        # stop when about to enter node
        assert [
            c
            async for c in tool_two.astream(
                {"my_key": "value ⛰️", "market": "DE"}, thread1
            )
        ] == [
            {
                "__interrupt__": (
                    Interrupt(
                        value="Just because...",
                        resumable=True,
                        ns=[AnyStr("tool_two:")],
                    ),
                )
            },
        ]
        if "shallow" not in checkpointer_name:
            assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [
                {
                    "parents": {},
                    "source": "loop",
                    "step": 0,
                    "writes": None,
                    "thread_id": "1",
                },
                {
                    "parents": {},
                    "source": "input",
                    "step": -1,
                    "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}},
                    "thread_id": "1",
                },
            ]
        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",
                    (PULL, "tool_two"),
                    interrupts=(
                        Interrupt(
                            value="Just because...",
                            resumable=True,
                            ns=[AnyStr("tool_two:")],
                        ),
                    ),
                ),
            ),
            config=tup.config,
            created_at=tup.checkpoint["ts"],
            metadata={
                "parents": {},
                "source": "loop",
                "step": 0,
                "writes": None,
                "thread_id": "1",
            },
            parent_config=(
                None
                if "shallow" in checkpointer_name
                else [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][
                    -1
                ].config
            ),
        )

        # clear the interrupt and next tasks
        await tool_two.aupdate_state(thread1, None, as_node=END)
        # interrupt is cleared, as well as the next tasks
        tup = await tool_two.checkpointer.aget_tuple(thread1)
        assert await tool_two.aget_state(thread1) == StateSnapshot(
            values={"my_key": "value ⛰️", "market": "DE"},
            next=(),
            tasks=(),
            config=tup.config,
            created_at=tup.checkpoint["ts"],
            metadata={
                "parents": {},
                "source": "update",
                "step": 1,
                "writes": {},
                "thread_id": "1",
            },
            parent_config=(
                None
                if "shallow" in checkpointer_name
                else [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][
                    -1
                ].config
            ),
        )


@NEEDS_CONTEXTVARS
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_dynamic_interrupt_subgraph(checkpointer_name: str) -> None:
    class SubgraphState(TypedDict):
        my_key: str
        market: str

    tool_two_node_count = 0

    def tool_two_node(s: SubgraphState) -> SubgraphState:
        nonlocal tool_two_node_count
        tool_two_node_count += 1
        if s["market"] == "DE":
            answer = interrupt("Just because...")
        else:
            answer = " all good"
        return {"my_key": answer}

    subgraph = StateGraph(SubgraphState)
    subgraph.add_node("do", tool_two_node, retry=RetryPolicy())
    subgraph.add_edge(START, "do")

    class State(TypedDict):
        my_key: Annotated[str, operator.add]
        market: str

    tool_two_graph = StateGraph(State)
    tool_two_graph.add_node("tool_two", subgraph.compile())
    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",
    }

    async with awith_checkpointer(checkpointer_name) as checkpointer:
        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"})

        # flow: interrupt -> resume with answer
        thread2 = {"configurable": {"thread_id": "2"}}
        # stop when about to enter node
        assert [
            c
            async for c in tool_two.astream(
                {"my_key": "value ⛰️", "market": "DE"}, thread2
            )
        ] == [
            {
                "__interrupt__": (
                    Interrupt(
                        value="Just because...",
                        resumable=True,
                        ns=[AnyStr("tool_two:"), AnyStr("do:")],
                    ),
                )
            },
        ]
        # resume with answer
        assert [
            c async for c in tool_two.astream(Command(resume=" my answer"), thread2)
        ] == [
            {"tool_two": {"my_key": " my answer", "market": "DE"}},
        ]

        # flow: interrupt -> clear
        thread1 = {"configurable": {"thread_id": "1"}}
        thread1root = {"configurable": {"thread_id": "1", "checkpoint_ns": ""}}
        # stop when about to enter node
        assert [
            c
            async for c in tool_two.astream(
                {"my_key": "value ⛰️", "market": "DE"}, thread1
            )
        ] == [
            {
                "__interrupt__": (
                    Interrupt(
                        value="Just because...",
                        resumable=True,
                        ns=[AnyStr("tool_two:"), AnyStr("do:")],
                    ),
                )
            },
        ]
        if "shallow" not in checkpointer_name:
            assert [
                c.metadata async for c in tool_two.checkpointer.alist(thread1root)
            ] == [
                {
                    "parents": {},
                    "source": "loop",
                    "step": 0,
                    "writes": None,
                    "thread_id": "1",
                },
                {
                    "parents": {},
                    "source": "input",
                    "step": -1,
                    "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}},
                    "thread_id": "1",
                },
            ]
        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",
                    (PULL, "tool_two"),
                    interrupts=(
                        Interrupt(
                            value="Just because...",
                            resumable=True,
                            ns=[AnyStr("tool_two:"), AnyStr("do:")],
                        ),
                    ),
                    state={
                        "configurable": {
                            "thread_id": "1",
                            "checkpoint_ns": AnyStr("tool_two:"),
                        }
                    },
                ),
            ),
            config=tup.config,
            created_at=tup.checkpoint["ts"],
            metadata={
                "parents": {},
                "source": "loop",
                "step": 0,
                "writes": None,
                "thread_id": "1",
            },
            parent_config=(
                None
                if "shallow" in checkpointer_name
                else [
                    c async for c in tool_two.checkpointer.alist(thread1root, limit=2)
                ][-1].config
            ),
        )

        # clear the interrupt and next tasks
        await tool_two.aupdate_state(thread1, None, as_node=END)
        # interrupt is cleared, as well as the next tasks
        tup = await tool_two.checkpointer.aget_tuple(thread1)
        assert await tool_two.aget_state(thread1) == StateSnapshot(
            values={"my_key": "value ⛰️", "market": "DE"},
            next=(),
            tasks=(),
            config=tup.config,
            created_at=tup.checkpoint["ts"],
            metadata={
                "parents": {},
                "source": "update",
                "step": 1,
                "writes": {},
                "thread_id": "1",
            },
            parent_config=(
                None
                if "shallow" in checkpointer_name
                else [
                    c async for c in tool_two.checkpointer.alist(thread1root, limit=2)
                ][-1].config
            ),
        )


@NEEDS_CONTEXTVARS
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_copy_checkpoint(checkpointer_name: str) -> None:
    class State(TypedDict):
        my_key: Annotated[str, operator.add]
        market: str

    def tool_one(s: State) -> State:
        return {"my_key": " one"}

    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":
            answer = interrupt("Just because...")
        else:
            answer = " all good"
        return {"my_key": answer}

    def start(state: State) -> list[Union[Send, str]]:
        return ["tool_two", Send("tool_one", state)]

    tool_two_graph = StateGraph(State)
    tool_two_graph.add_node("tool_two", tool_two_node, retry=RetryPolicy())
    tool_two_graph.add_node("tool_one", tool_one)
    tool_two_graph.set_conditional_entry_point(start)
    tool_two = tool_two_graph.compile()

    tracer = FakeTracer()
    assert await tool_two.ainvoke(
        {"my_key": "value", "market": "DE"}, {"callbacks": [tracer]}, debug=True
    ) == {
        "my_key": "value one",
        "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 one"}

    assert await tool_two.ainvoke({"my_key": "value", "market": "US"}) == {
        "my_key": "value all good one",
        "market": "US",
    }

    async with awith_checkpointer(checkpointer_name) as checkpointer:
        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"})

        # flow: interrupt -> resume with answer
        thread2 = {"configurable": {"thread_id": "2"}}
        # stop when about to enter node
        assert [
            c
            async for c in tool_two.astream(
                {"my_key": "value ⛰️", "market": "DE"}, thread2
            )
        ] == UnsortedSequence(
            {
                "__interrupt__": (
                    Interrupt(
                        value="Just because...",
                        resumable=True,
                        ns=[AnyStr("tool_two:")],
                    ),
                )
            },
            {
                "tool_one": {"my_key": " one"},
            },
        )
        # resume with answer
        assert [
            c async for c in tool_two.astream(Command(resume=" my answer"), thread2)
        ] == [
            {
                "__metadata__": {"cached": True},
                "tool_one": {"my_key": " one"},
            },
            {"tool_two": {"my_key": " my answer"}},
        ]

        # flow: interrupt -> clear tasks
        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 ⛰️ one",
            "market": "DE",
        }

        if "shallow" not in checkpointer_name:
            assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [
                {
                    "parents": {},
                    "source": "loop",
                    "step": 0,
                    "writes": None,
                    "thread_id": "1",
                },
                {
                    "parents": {},
                    "source": "input",
                    "step": -1,
                    "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}},
                    "thread_id": "1",
                },
            ]

        tup = await tool_two.checkpointer.aget_tuple(thread1)
        assert await tool_two.aget_state(thread1) == StateSnapshot(
            values={"my_key": "value ⛰️ one", "market": "DE"},
            next=("tool_two",),
            tasks=(
                PregelTask(
                    AnyStr(),
                    name="tool_one",
                    path=("__pregel_push", 0),
                    error=None,
                    interrupts=(),
                    state=None,
                    result={"my_key": " one"},
                ),
                PregelTask(
                    AnyStr(),
                    "tool_two",
                    (PULL, "tool_two"),
                    interrupts=(
                        Interrupt(
                            value="Just because...",
                            resumable=True,
                            ns=[AnyStr("tool_two:")],
                        ),
                    ),
                ),
            ),
            config=tup.config,
            created_at=tup.checkpoint["ts"],
            metadata={
                "parents": {},
                "source": "loop",
                "step": 0,
                "writes": None,
                "thread_id": "1",
            },
            parent_config=(
                None
                if "shallow" in checkpointer_name
                else [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][
                    -1
                ].config
            ),
        )

        if "shallow" in checkpointer_name:
            # shallow checkpointer doesn't support copy
            return

        # clear the interrupt and next tasks
        await tool_two.aupdate_state(thread1, None, as_node="__copy__")
        # interrupt is cleared, next task is kept
        tup = await tool_two.checkpointer.aget_tuple(thread1)
        assert await tool_two.aget_state(thread1) == StateSnapshot(
            values={"my_key": "value ⛰️", "market": "DE"},
            next=("tool_one", "tool_two"),
            tasks=(
                PregelTask(
                    AnyStr(),
                    "tool_one",
                    (PUSH, 0),
                    result=None,
                ),
                PregelTask(
                    AnyStr(),
                    "tool_two",
                    (PULL, "tool_two"),
                    interrupts=(),
                ),
            ),
            config=tup.config,
            created_at=tup.checkpoint["ts"],
            metadata={
                "parents": {},
                "source": "fork",
                "step": 1,
                "writes": None,
                "thread_id": "1",
            },
            parent_config=(
                [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][
                    -1
                ].parent_config
            ),
        )


@NEEDS_CONTEXTVARS
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_node_not_cancelled_on_other_node_interrupted(
    checkpointer_name: str,
) -> None:
    class State(TypedDict):
        hello: Annotated[str, operator.add]

    awhiles = 0
    inner_task_cancelled = False

    async def awhile(input: State) -> None:
        nonlocal awhiles

        awhiles += 1
        try:
            await asyncio.sleep(1)
            return {"hello": " again"}
        except asyncio.CancelledError:
            nonlocal inner_task_cancelled
            inner_task_cancelled = True
            raise

    async def iambad(input: State) -> None:
        return {"hello": interrupt("I am bad")}

    builder = StateGraph(State)
    builder.add_node("agent", awhile)
    builder.add_node("bad", iambad)
    builder.set_conditional_entry_point(lambda _: ["agent", "bad"], then=END)

    async with awith_checkpointer(checkpointer_name) as checkpointer:
        graph = builder.compile(checkpointer=checkpointer)
        thread = {"configurable": {"thread_id": "1"}}

        # writes from "awhile" are applied to last chunk
        assert await graph.ainvoke({"hello": "world"}, thread) == {
            "hello": "world again"
        }

        assert not inner_task_cancelled
        assert awhiles == 1

        assert await graph.ainvoke(None, thread, debug=True) == {"hello": "world again"}

        assert not inner_task_cancelled
        assert awhiles == 1

        # resume with answer
        assert await graph.ainvoke(Command(resume=" okay"), thread) == {
            "hello": "world again okay"
        }

        assert not inner_task_cancelled
        assert awhiles == 1


@pytest.mark.parametrize("stream_hang_s", [0.3, 0.6])
async def test_step_timeout_on_stream_hang(stream_hang_s: float) -> None:
    inner_task_cancelled = False

    async def awhile(input: Any) -> None:
        try:
            await asyncio.sleep(1.5)
        except asyncio.CancelledError:
            nonlocal inner_task_cancelled
            inner_task_cancelled = True
            raise

    async def alittlewhile(input: Any) -> None:
        await asyncio.sleep(0.6)
        return "1"

    builder = Graph()
    builder.add_node(awhile)
    builder.add_node(alittlewhile)
    builder.set_conditional_entry_point(lambda _: ["awhile", "alittlewhile"], then=END)
    graph = builder.compile()
    graph.step_timeout = 1

    with pytest.raises(asyncio.TimeoutError):
        async for chunk in graph.astream(1, stream_mode="updates"):
            assert chunk == {"alittlewhile": {"alittlewhile": "1"}}
            await asyncio.sleep(stream_hang_s)

    assert inner_task_cancelled


@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC_PLUS_NONE)
async def test_cancel_graph_astream(checkpointer_name: str) -> None:
    class State(TypedDict):
        value: Annotated[int, operator.add]

    class AwhileMaker:
        def __init__(self) -> None:
            self.reset()

        async def __call__(self, input: State) -> Any:
            self.started = True
            try:
                await asyncio.sleep(1.5)
            except asyncio.CancelledError:
                self.cancelled = True
                raise

        def reset(self):
            self.started = False
            self.cancelled = False

    async def alittlewhile(input: State) -> None:
        await asyncio.sleep(0.6)
        return {"value": 2}

    awhile = AwhileMaker()
    aparallelwhile = AwhileMaker()
    builder = StateGraph(State)
    builder.add_node("awhile", awhile)
    builder.add_node("aparallelwhile", aparallelwhile)
    builder.add_node(alittlewhile)
    builder.add_edge(START, "alittlewhile")
    builder.add_edge(START, "aparallelwhile")
    builder.add_edge("alittlewhile", "awhile")

    async with awith_checkpointer(checkpointer_name) as checkpointer:
        graph = builder.compile(checkpointer=checkpointer)

        # test interrupting astream
        got_event = False
        thread1: RunnableConfig = {"configurable": {"thread_id": "1"}}
        async with aclosing(graph.astream({"value": 1}, thread1)) as stream:
            async for chunk in stream:
                assert chunk == {"alittlewhile": {"value": 2}}
                got_event = True
                break

        assert got_event

        # node aparallelwhile should start, but be cancelled
        assert aparallelwhile.started is True
        assert aparallelwhile.cancelled is True

        # node "awhile" should never start
        assert awhile.started is False

        # checkpoint with output of "alittlewhile" should not be saved
        # but we should have applied pending writes
        if checkpointer is not None:
            state = await graph.aget_state(thread1)
            assert state is not None
            assert state.values == {"value": 3}  # 1 + 2
            assert state.next == ("aparallelwhile",)
            assert state.metadata == {
                "parents": {},
                "source": "loop",
                "step": 0,
                "writes": None,
                "thread_id": "1",
            }


@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC_PLUS_NONE)
async def test_cancel_graph_astream_events_v2(checkpointer_name: Optional[str]) -> None:
    class State(TypedDict):
        value: int

    class AwhileMaker:
        def __init__(self) -> None:
            self.reset()

        async def __call__(self, input: State) -> Any:
            self.started = True
            try:
                await asyncio.sleep(1.5)
            except asyncio.CancelledError:
                self.cancelled = True
                raise

        def reset(self):
            self.started = False
            self.cancelled = False

    async def alittlewhile(input: State) -> None:
        await asyncio.sleep(0.6)
        return {"value": 2}

    awhile = AwhileMaker()
    anotherwhile = AwhileMaker()
    builder = StateGraph(State)
    builder.add_node(alittlewhile)
    builder.add_node("awhile", awhile)
    builder.add_node("anotherwhile", anotherwhile)
    builder.add_edge(START, "alittlewhile")
    builder.add_edge("alittlewhile", "awhile")
    builder.add_edge("awhile", "anotherwhile")

    async with awith_checkpointer(checkpointer_name) as checkpointer:
        graph = builder.compile(checkpointer=checkpointer)

        # test interrupting astream_events v2
        got_event = False
        thread2: RunnableConfig = {"configurable": {"thread_id": "2"}}
        async with aclosing(
            graph.astream_events({"value": 1}, thread2, version="v2")
        ) as stream:
            async for chunk in stream:
                if chunk["event"] == "on_chain_stream" and not chunk["parent_ids"]:
                    got_event = True
                    assert chunk["data"]["chunk"] == {"alittlewhile": {"value": 2}}
                    await asyncio.sleep(0.1)
                    break

        # did break
        assert got_event

        # node "awhile" maybe starts (impl detail of astream_events)
        # if it does start, it must be cancelled
        if awhile.started:
            assert awhile.cancelled is True

        # node "anotherwhile" should never start
        assert anotherwhile.started is False

        # checkpoint with output of "alittlewhile" should not be saved
        if checkpointer is not None:
            state = await graph.aget_state(thread2)
            assert state is not None
            assert state.values == {"value": 2}
            assert state.next == ("awhile",)
            assert state.metadata == {
                "parents": {},
                "source": "loop",
                "step": 1,
                "writes": {"alittlewhile": {"value": 2}},
                "thread_id": "2",
            }


async def test_node_schemas_custom_output() -> None:
    class State(TypedDict):
        hello: str
        bye: str
        messages: Annotated[list[str], add_messages]

    class Output(TypedDict):
        messages: list[str]

    class StateForA(TypedDict):
        hello: str
        messages: Annotated[list[str], add_messages]

    async def node_a(state: StateForA):
        assert state == {
            "hello": "there",
            "messages": [_AnyIdHumanMessage(content="hello")],
        }

    class StateForB(TypedDict):
        bye: str
        now: int

    async def node_b(state: StateForB):
        assert state == {
            "bye": "world",
        }
        return {
            "now": 123,
            "hello": "again",
        }

    class StateForC(TypedDict):
        hello: str
        now: int

    async def node_c(state: StateForC):
        assert state == {
            "hello": "again",
            "now": 123,
        }

    builder = StateGraph(State, output=Output)
    builder.add_node("a", node_a)
    builder.add_node("b", node_b)
    builder.add_node("c", node_c)
    builder.add_edge(START, "a")
    builder.add_edge("a", "b")
    builder.add_edge("b", "c")
    graph = builder.compile()

    assert await graph.ainvoke(
        {"hello": "there", "bye": "world", "messages": "hello"}
    ) == {
        "messages": [_AnyIdHumanMessage(content="hello")],
    }

    builder = StateGraph(State, output=Output)
    builder.add_node("a", node_a)
    builder.add_node("b", node_b)
    builder.add_node("c", node_c)
    builder.add_edge(START, "a")
    builder.add_edge("a", "b")
    builder.add_edge("b", "c")
    graph = builder.compile()

    assert await graph.ainvoke(
        {
            "hello": "there",
            "bye": "world",
            "messages": "hello",
            "now": 345,  # ignored because not in input schema
        }
    ) == {
        "messages": [_AnyIdHumanMessage(content="hello")],
    }

    assert [
        c
        async for c in graph.astream(
            {
                "hello": "there",
                "bye": "world",
                "messages": "hello",
                "now": 345,  # ignored because not in input schema
            }
        )
    ] == [
        {"a": None},
        {"b": {"hello": "again", "now": 123}},
        {"c": None},
    ]


async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
    add_one = mocker.Mock(side_effect=lambda x: x + 1)
    chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output")

    app = Pregel(
        nodes={
            "one": chain,
        },
        channels={
            "input": LastValue(int),
            "output": LastValue(int),
        },
        input_channels="input",
        output_channels="output",
    )
    graph = Graph()
    graph.add_node("add_one", add_one)
    graph.set_entry_point("add_one")
    graph.set_finish_point("add_one")
    gapp = graph.compile()

    if SHOULD_CHECK_SNAPSHOTS:
        assert app.input_schema.model_json_schema() == {
            "title": "LangGraphInput",
            "type": "integer",
        }
        assert app.output_schema.model_json_schema() == {
            "title": "LangGraphOutput",
            "type": "integer",
        }
    assert await app.ainvoke(2) == 3
    assert await app.ainvoke(2, output_keys=["output"]) == {"output": 3}

    assert await gapp.ainvoke(2) == 3


@pytest.mark.parametrize(
    "falsy_value",
    [None, False, 0, "", [], {}, set(), frozenset(), 0.0, 0j],
)
async def test_invoke_single_process_in_out_falsy_values(falsy_value: Any) -> None:
    graph = Graph()
    graph.add_node("return_falsy_const", lambda *args, **kwargs: falsy_value)
    graph.set_entry_point("return_falsy_const")
    graph.set_finish_point("return_falsy_const")
    gapp = graph.compile()
    assert falsy_value == await gapp.ainvoke(1)


async def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> None:
    add_one = mocker.Mock(side_effect=lambda x: x + 1)
    chain = (
        Channel.subscribe_to("input")
        | add_one
        | Channel.write_to("output", fixed=5, output_plus_one=lambda x: x + 1)
    )

    app = Pregel(
        nodes={"one": chain},
        channels={
            "input": LastValue(int),
            "output": LastValue(int),
            "fixed": LastValue(int),
            "output_plus_one": LastValue(int),
        },
        output_channels=["output", "fixed", "output_plus_one"],
        input_channels="input",
    )

    if SHOULD_CHECK_SNAPSHOTS:
        assert app.input_schema.model_json_schema() == {
            "title": "LangGraphInput",
            "type": "integer",
        }
        assert app.output_schema.model_json_schema() == {
            "title": "LangGraphOutput",
            "type": "object",
            "properties": {
                "output": {"title": "Output", "type": "integer", "default": None},
                "fixed": {"title": "Fixed", "type": "integer", "default": None},
                "output_plus_one": {
                    "title": "Output Plus One",
                    "type": "integer",
                    "default": None,
                },
            },
        }
    assert await app.ainvoke(2) == {"output": 3, "fixed": 5, "output_plus_one": 4}


async def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None:
    add_one = mocker.Mock(side_effect=lambda x: x + 1)
    chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output")

    app = Pregel(
        nodes={"one": chain},
        channels={"input": LastValue(int), "output": LastValue(int)},
        input_channels="input",
        output_channels=["output"],
    )

    if SHOULD_CHECK_SNAPSHOTS:
        assert app.input_schema.model_json_schema() == {
            "title": "LangGraphInput",
            "type": "integer",
        }
        assert app.output_schema.model_json_schema() == {
            "title": "LangGraphOutput",
            "type": "object",
            "properties": {
                "output": {"title": "Output", "type": "integer", "default": None}
            },
        }
    assert await app.ainvoke(2) == {"output": 3}


async def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> None:
    add_one = mocker.Mock(side_effect=lambda x: x + 1)
    chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output")

    app = Pregel(
        nodes={"one": chain},
        channels={"input": LastValue(int), "output": LastValue(int)},
        input_channels=["input"],
        output_channels=["output"],
    )

    if SHOULD_CHECK_SNAPSHOTS:
        assert app.input_schema.model_json_schema() == {
            "title": "LangGraphInput",
            "type": "object",
            "properties": {
                "input": {"title": "Input", "type": "integer", "default": None}
            },
        }
        assert app.output_schema.model_json_schema() == {
            "title": "LangGraphOutput",
            "type": "object",
            "properties": {
                "output": {"title": "Output", "type": "integer", "default": None}
            },
        }
    assert await app.ainvoke({"input": 2}) == {"output": 3}


async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
    add_one = mocker.Mock(side_effect=lambda x: x + 1)
    one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
    two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output")

    app = Pregel(
        nodes={"one": one, "two": two},
        channels={
            "inbox": LastValue(int),
            "output": LastValue(int),
            "input": LastValue(int),
        },
        input_channels="input",
        output_channels="output",
        stream_channels=["inbox", "output"],
    )

    assert await app.ainvoke(2) == 4

    with pytest.raises(GraphRecursionError):
        await app.ainvoke(2, {"recursion_limit": 1})

    step = 0
    async for values in app.astream(2):
        step += 1
        if step == 1:
            assert values == {
                "inbox": 3,
            }
        elif step == 2:
            assert values == {
                "inbox": 3,
                "output": 4,
            }
    assert step == 2

    graph = Graph()
    graph.add_node("add_one", add_one)
    graph.add_node("add_one_more", add_one)
    graph.set_entry_point("add_one")
    graph.set_finish_point("add_one_more")
    graph.add_edge("add_one", "add_one_more")
    gapp = graph.compile()

    assert await gapp.ainvoke(2) == 4

    step = 0
    async for values in gapp.astream(2):
        step += 1
        if step == 1:
            assert values == {
                "add_one": 3,
            }
        elif step == 2:
            assert values == {
                "add_one_more": 4,
            }
    assert step == 2


async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
    add_one = mocker.Mock(side_effect=lambda x: x + 1)
    one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
    two = (
        Channel.subscribe_to("inbox")
        | RunnableLambda(add_one).abatch
        | Channel.write_to("output").abatch
    )

    app = Pregel(
        nodes={"one": one, "two": two},
        channels={
            "inbox": Topic(int),
            "output": LastValue(int),
            "input": LastValue(int),
        },
        input_channels=["input", "inbox"],
        stream_channels=["output", "inbox"],
        output_channels=["output"],
    )

    # [12 + 1, 2 + 1 + 1]
    assert [
        c
        async for c in app.astream(
            {"input": 2, "inbox": 12}, output_keys="output", stream_mode="updates"
        )
    ] == [
        {"one": None},
        {"two": 13},
        {"two": 4},
    ]
    assert [
        c async for c in app.astream({"input": 2, "inbox": 12}, output_keys="output")
    ] == [13, 4]

    assert [
        c async for c in app.astream({"input": 2, "inbox": 12}, stream_mode="updates")
    ] == [
        {"one": {"inbox": 3}},
        {"two": {"output": 13}},
        {"two": {"output": 4}},
    ]
    assert [c async for c in app.astream({"input": 2, "inbox": 12})] == [
        {"inbox": [3], "output": 13},
        {"output": 4},
    ]
    assert [
        c async for c in app.astream({"input": 2, "inbox": 12}, stream_mode="debug")
    ] == [
        {
            "type": "task",
            "timestamp": AnyStr(),
            "step": 0,
            "payload": {
                "id": AnyStr(),
                "name": "one",
                "input": 2,
                "triggers": ("input",),
            },
        },
        {
            "type": "task",
            "timestamp": AnyStr(),
            "step": 0,
            "payload": {
                "id": AnyStr(),
                "name": "two",
                "input": [12],
                "triggers": ("inbox",),
            },
        },
        {
            "type": "task_result",
            "timestamp": AnyStr(),
            "step": 0,
            "payload": {
                "id": AnyStr(),
                "name": "one",
                "result": [("inbox", 3)],
                "error": None,
                "interrupts": [],
            },
        },
        {
            "type": "task_result",
            "timestamp": AnyStr(),
            "step": 0,
            "payload": {
                "id": AnyStr(),
                "name": "two",
                "result": [("output", 13)],
                "error": None,
                "interrupts": [],
            },
        },
        {
            "type": "task",
            "timestamp": AnyStr(),
            "step": 1,
            "payload": {
                "id": AnyStr(),
                "name": "two",
                "input": [3],
                "triggers": ("inbox",),
            },
        },
        {
            "type": "task_result",
            "timestamp": AnyStr(),
            "step": 1,
            "payload": {
                "id": AnyStr(),
                "name": "two",
                "result": [("output", 4)],
                "error": None,
                "interrupts": [],
            },
        },
    ]


async def test_batch_two_processes_in_out() -> None:
    async def add_one_with_delay(inp: int) -> int:
        await asyncio.sleep(inp / 10)
        return inp + 1

    one = Channel.subscribe_to("input") | add_one_with_delay | Channel.write_to("one")
    two = Channel.subscribe_to("one") | add_one_with_delay | Channel.write_to("output")

    app = Pregel(
        nodes={"one": one, "two": two},
        channels={
            "one": LastValue(int),
            "output": LastValue(int),
            "input": LastValue(int),
        },
        input_channels="input",
        output_channels="output",
    )

    assert await app.abatch([3, 2, 1, 3, 5]) == [5, 4, 3, 5, 7]
    assert await app.abatch([3, 2, 1, 3, 5], output_keys=["output"]) == [
        {"output": 5},
        {"output": 4},
        {"output": 3},
        {"output": 5},
        {"output": 7},
    ]

    graph = Graph()
    graph.add_node("add_one", add_one_with_delay)
    graph.add_node("add_one_more", add_one_with_delay)
    graph.set_entry_point("add_one")
    graph.set_finish_point("add_one_more")
    graph.add_edge("add_one", "add_one_more")
    gapp = graph.compile()

    assert await gapp.abatch([3, 2, 1, 3, 5]) == [5, 4, 3, 5, 7]


async def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None:
    test_size = 100
    add_one = mocker.Mock(side_effect=lambda x: x + 1)

    nodes = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")}
    for i in range(test_size - 2):
        nodes[str(i)] = (
            Channel.subscribe_to(str(i - 1)) | add_one | Channel.write_to(str(i))
        )
    nodes["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output")

    app = Pregel(
        nodes=nodes,
        channels={str(i): LastValue(int) for i in range(-1, test_size - 2)}
        | {"input": LastValue(int), "output": LastValue(int)},
        input_channels="input",
        output_channels="output",
    )

    # No state is left over from previous invocations
    for _ in range(10):
        assert await app.ainvoke(2, {"recursion_limit": test_size}) == 2 + test_size

    # Concurrent invocations do not interfere with each other
    assert await asyncio.gather(
        *(app.ainvoke(2, {"recursion_limit": test_size}) for _ in range(10))
    ) == [2 + test_size for _ in range(10)]


async def test_batch_many_processes_in_out(mocker: MockerFixture) -> None:
    test_size = 100
    add_one = mocker.Mock(side_effect=lambda x: x + 1)

    nodes = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")}
    for i in range(test_size - 2):
        nodes[str(i)] = (
            Channel.subscribe_to(str(i - 1)) | add_one | Channel.write_to(str(i))
        )
    nodes["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output")

    app = Pregel(
        nodes=nodes,
        channels={str(i): LastValue(int) for i in range(-1, test_size - 2)}
        | {"input": LastValue(int), "output": LastValue(int)},
        input_channels="input",
        output_channels="output",
    )

    # No state is left over from previous invocations
    for _ in range(3):
        # Then invoke pubsub
        assert await app.abatch([2, 1, 3, 4, 5], {"recursion_limit": test_size}) == [
            2 + test_size,
            1 + test_size,
            3 + test_size,
            4 + test_size,
            5 + test_size,
        ]

    # Concurrent invocations do not interfere with each other
    assert await asyncio.gather(
        *(app.abatch([2, 1, 3, 4, 5], {"recursion_limit": test_size}) for _ in range(3))
    ) == [
        [2 + test_size, 1 + test_size, 3 + test_size, 4 + test_size, 5 + test_size]
        for _ in range(3)
    ]


async def test_invoke_two_processes_two_in_two_out_invalid(
    mocker: MockerFixture,
) -> None:
    add_one = mocker.Mock(side_effect=lambda x: x + 1)

    one = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
    two = Channel.subscribe_to("input") | add_one | Channel.write_to("output")

    app = Pregel(
        nodes={"one": one, "two": two},
        channels={"output": LastValue(int), "input": LastValue(int)},
        input_channels="input",
        output_channels="output",
    )

    with pytest.raises(InvalidUpdateError):
        # LastValue channels can only be updated once per iteration
        await app.ainvoke(2)


async def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> None:
    add_one = mocker.Mock(side_effect=lambda x: x + 1)

    one = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
    two = Channel.subscribe_to("input") | add_one | Channel.write_to("output")

    app = Pregel(
        nodes={"one": one, "two": two},
        channels={
            "input": LastValue(int),
            "output": Topic(int),
        },
        input_channels="input",
        output_channels="output",
    )

    # An Topic channel accumulates updates into a sequence
    assert await app.ainvoke(2) == [3, 3]


@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_invoke_checkpoint(mocker: MockerFixture, checkpointer_name: str) -> None:
    add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"])
    errored_once = False

    def raise_if_above_10(input: int) -> int:
        nonlocal errored_once
        if input > 4:
            if errored_once:
                pass
            else:
                errored_once = True
                raise ConnectionError("I will be retried")
        if input > 10:
            raise ValueError("Input is too large")
        return input

    one = (
        Channel.subscribe_to(["input"]).join(["total"])
        | add_one
        | Channel.write_to("output", "total")
        | raise_if_above_10
    )

    async with awith_checkpointer(checkpointer_name) as checkpointer:
        app = Pregel(
            nodes={"one": one},
            channels={
                "total": BinaryOperatorAggregate(int, operator.add),
                "input": LastValue(int),
                "output": LastValue(int),
            },
            input_channels="input",
            output_channels="output",
            checkpointer=checkpointer,
            retry_policy=RetryPolicy(),
        )

        # total starts out as 0, so output is 0+2=2
        assert await app.ainvoke(2, {"configurable": {"thread_id": "1"}}) == 2
        checkpoint = await checkpointer.aget({"configurable": {"thread_id": "1"}})
        assert checkpoint is not None
        assert checkpoint["channel_values"].get("total") == 2
        # total is now 2, so output is 2+3=5
        assert await app.ainvoke(3, {"configurable": {"thread_id": "1"}}) == 5
        assert errored_once, "errored and retried"
        checkpoint = await checkpointer.aget({"configurable": {"thread_id": "1"}})
        assert checkpoint is not None
        assert checkpoint["channel_values"].get("total") == 7
        # total is now 2+5=7, so output would be 7+4=11, but raises ValueError
        with pytest.raises(ValueError):
            await app.ainvoke(4, {"configurable": {"thread_id": "1"}})
        # checkpoint is not updated
        checkpoint = await checkpointer.aget({"configurable": {"thread_id": "1"}})
        assert checkpoint is not None
        assert checkpoint["channel_values"].get("total") == 7
        # on a new thread, total starts out as 0, so output is 0+5=5
        assert await app.ainvoke(5, {"configurable": {"thread_id": "2"}}) == 5
        checkpoint = await checkpointer.aget({"configurable": {"thread_id": "1"}})
        assert checkpoint is not None
        assert checkpoint["channel_values"].get("total") == 7
        checkpoint = await checkpointer.aget({"configurable": {"thread_id": "2"}})
        assert checkpoint is not None
        assert checkpoint["channel_values"].get("total") == 5


@pytest.mark.parametrize("checkpoint_during", [True, False])
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_pending_writes_resume(
    checkpointer_name: str, checkpoint_during: bool
) -> None:
    if not checkpoint_during and "shallow" in checkpointer_name:
        pytest.skip("Checkpointing during execution not supported")

    class State(TypedDict):
        value: Annotated[int, operator.add]

    class AwhileMaker:
        def __init__(self, sleep: float, rtn: Union[Dict, Exception]) -> None:
            self.sleep = sleep
            self.rtn = rtn
            self.reset()

        async def __call__(self, input: State) -> Any:
            self.calls += 1
            await asyncio.sleep(self.sleep)
            if isinstance(self.rtn, Exception):
                raise self.rtn
            else:
                return self.rtn

        def reset(self):
            self.calls = 0

    one = AwhileMaker(0.1, {"value": 2})
    two = AwhileMaker(0.2, ConnectionError("I'm not good"))
    builder = StateGraph(State)
    builder.add_node("one", one)
    builder.add_node(
        "two", two, retry=RetryPolicy(max_attempts=2, initial_interval=0, jitter=False)
    )
    builder.add_edge(START, "one")
    builder.add_edge(START, "two")
    async with awith_checkpointer(checkpointer_name) as checkpointer:
        graph = builder.compile(checkpointer=checkpointer)

        thread1: RunnableConfig = {"configurable": {"thread_id": "1"}}
        with pytest.raises(ConnectionError, match="I'm not good"):
            await graph.ainvoke(
                {"value": 1}, thread1, checkpoint_during=checkpoint_during
            )

        # both nodes should have been called once
        assert one.calls == 1
        assert two.calls == 2

        # latest checkpoint should be before nodes "one", "two"
        # but we should have applied pending writes from "one"
        state = await graph.aget_state(thread1)
        assert state is not None
        assert state.values == {"value": 3}
        assert state.next == ("two",)
        assert state.tasks == (
            PregelTask(AnyStr(), "one", (PULL, "one"), result={"value": 2}),
            PregelTask(
                AnyStr(),
                "two",
                (PULL, "two"),
                'ConnectionError("I\'m not good")',
            ),
        )
        assert state.metadata == {
            "parents": {},
            "source": "loop",
            "step": 0,
            "writes": None,
            "thread_id": "1",
        }
        # get_state with checkpoint_id should not apply any pending writes
        state = await graph.aget_state(state.config)
        assert state is not None
        assert state.values == {"value": 1}
        assert state.next == ("one", "two")
        # should contain pending write of "one"
        checkpoint = await checkpointer.aget_tuple(thread1)
        assert checkpoint is not None
        # should contain error from "two"
        expected_writes = [
            (AnyStr(), "value", 2),
            (AnyStr(), ERROR, 'ConnectionError("I\'m not good")'),
        ]
        assert len(checkpoint.pending_writes) == 2
        assert all(w in expected_writes for w in checkpoint.pending_writes)
        # both non-error pending writes come from same task
        non_error_writes = [w for w in checkpoint.pending_writes if w[1] != ERROR]
        # error write is from the other task
        error_write = next(w for w in checkpoint.pending_writes if w[1] == ERROR)
        assert error_write[0] != non_error_writes[0][0]

        # resume execution
        with pytest.raises(ConnectionError, match="I'm not good"):
            await graph.ainvoke(None, thread1, checkpoint_during=checkpoint_during)

        # node "one" succeeded previously, so shouldn't be called again
        assert one.calls == 1
        # node "two" should have been called once again
        assert two.calls == 4

        # confirm no new checkpoints saved
        state_two = await graph.aget_state(thread1)
        assert state_two.metadata == state.metadata

        # resume execution, without exception
        two.rtn = {"value": 3}
        # both the pending write and the new write were applied, 1 + 2 + 3 = 6
        assert await graph.ainvoke(
            None, thread1, checkpoint_during=checkpoint_during
        ) == {"value": 6}

        if "shallow" in checkpointer_name:
            assert len([c async for c in checkpointer.alist(thread1)]) == 1
            return

        # check all final checkpoints
        checkpoints = [c async for c in checkpointer.alist(thread1)]
        # we should have 3
        assert len(checkpoints) == (3 if checkpoint_during else 2)
        # 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": 3,
                "id": AnyStr(),
                "ts": AnyStr(),
                "pending_sends": [],
                "versions_seen": {
                    "one": {
                        "branch:to:one": AnyVersion(),
                    },
                    "two": {
                        "branch:to:two": AnyVersion(),
                    },
                    "__input__": {},
                    "__start__": {
                        "__start__": AnyVersion(),
                    },
                    "__interrupt__": {
                        "value": AnyVersion(),
                        "__start__": AnyVersion(),
                        "branch:to:one": AnyVersion(),
                        "branch:to:two": AnyVersion(),
                    },
                },
                "channel_versions": {
                    "value": AnyVersion(),
                    "__start__": AnyVersion(),
                    "branch:to:one": AnyVersion(),
                    "branch:to:two": AnyVersion(),
                },
                "channel_values": {"value": 6},
            },
            metadata={
                "parents": {},
                "step": 1,
                "source": "loop",
                "writes": {"one": {"value": 2}, "two": {"value": 3}},
                "thread_id": "1",
            },
            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": 3,
                "id": AnyStr(),
                "ts": AnyStr(),
                "pending_sends": [],
                "versions_seen": {
                    "__input__": {},
                    "__start__": {
                        "__start__": AnyVersion(),
                    },
                },
                "channel_versions": {
                    "value": AnyVersion(),
                    "__start__": AnyVersion(),
                    "branch:to:one": AnyVersion(),
                    "branch:to:two": AnyVersion(),
                },
                "channel_values": {
                    "value": 1,
                    "branch:to:one": None,
                    "branch:to:two": None,
                },
            },
            metadata={
                "parents": {},
                "step": 0,
                "source": "loop",
                "writes": None,
                "thread_id": "1",
            },
            parent_config={
                "configurable": {
                    "thread_id": "1",
                    "checkpoint_ns": "",
                    "checkpoint_id": checkpoints[2].config["configurable"][
                        "checkpoint_id"
                    ]
                    if checkpoint_during
                    else AnyStr(),
                }
            },
            pending_writes=UnsortedSequence(
                (AnyStr(), "value", 2),
                (AnyStr(), "__error__", 'ConnectionError("I\'m not good")'),
                (AnyStr(), "value", 3),
            )
            if checkpoint_during
            else UnsortedSequence(
                (AnyStr(), "value", 2),
                (AnyStr(), "__error__", 'ConnectionError("I\'m not good")'),
                # the write against the previous checkpoint is not saved, as it is
                # produced in a run where only the next checkpoint (the last) is saved
            ),
        )
        if not checkpoint_during:
            return
        assert checkpoints[2] == CheckpointTuple(
            config={
                "configurable": {
                    "thread_id": "1",
                    "checkpoint_ns": "",
                    "checkpoint_id": AnyStr(),
                }
            },
            checkpoint={
                "v": 3,
                "id": AnyStr(),
                "ts": AnyStr(),
                "pending_sends": [],
                "versions_seen": {"__input__": {}},
                "channel_versions": {
                    "__start__": AnyVersion(),
                },
                "channel_values": {"__start__": {"value": 1}},
            },
            metadata={
                "parents": {},
                "step": -1,
                "source": "input",
                "writes": {"__start__": {"value": 1}},
                "thread_id": "1",
            },
            parent_config=None,
            pending_writes=UnsortedSequence(
                (AnyStr(), "value", 1),
                (AnyStr(), "branch:to:one", None),
                (AnyStr(), "branch:to:two", None),
            ),
        )


@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC)
async def test_run_from_checkpoint_id_retains_previous_writes(
    checkpointer_name: str,
) -> None:
    class MyState(TypedDict):
        myval: Annotated[int, operator.add]
        otherval: bool

    class Anode:
        def __init__(self):
            self.switch = False

        async def __call__(self, state: MyState):
            self.switch = not self.switch
            return {"myval": 2 if self.switch else 1, "otherval": self.switch}

    builder = StateGraph(MyState)
    thenode = Anode()  # Fun.
    builder.add_node("node_one", thenode)
    builder.add_node("node_two", thenode)
    builder.add_edge(START, "node_one")

    def _getedge(src: str):
        swap = "node_one" if src == "node_two" else "node_two"

        def _edge(st: MyState) -> Literal["__end__", "node_one", "node_two"]:
            if st["myval"] > 3:
                return END
            if st["otherval"]:
                return swap
            return src

        return _edge

    builder.add_conditional_edges("node_one", _getedge("node_one"))
    builder.add_conditional_edges("node_two", _getedge("node_two"))
    async with awith_checkpointer(checkpointer_name) as checkpointer:
        graph = builder.compile(checkpointer=checkpointer)

        thread_id = uuid.uuid4()
        thread1 = {"configurable": {"thread_id": str(thread_id)}}

        result = await graph.ainvoke({"myval": 1}, thread1)
        assert result["myval"] == 4
        history = [c async for c in graph.aget_state_history(thread1)]

        assert len(history) == 4
        assert history[0].values == {"myval": 4, "otherval": False}
        assert history[-1].values == {"myval": 0}

        second_run_config = {
            **thread1,
            "configurable": {
                **thread1["configurable"],
                "checkpoint_id": history[1].config["configurable"]["checkpoint_id"],
            },
        }
        second_result = await graph.ainvoke(None, second_run_config)
        assert second_result == {"myval": 5, "otherval": True}

        new_history = [
            c
            async for c in graph.aget_state_history(
                {"configurable": {"thread_id": str(thread_id), "checkpoint_ns": ""}}
            )
        ]

        assert len(new_history) == len(history) + 1
        for original, new in zip(history, new_history[1:]):
            assert original.values == new.values
            assert original.next == new.next
            assert original.metadata["step"] == new.metadata["step"]

        def _get_tasks(hist: list, start: int):
            return [h.tasks for h in hist[start:]]

        assert _get_tasks(new_history, 1) == _get_tasks(history, 0)


async def test_cond_edge_after_send() -> None:
    class Node:
        def __init__(self, name: str):
            self.name = name
            setattr(self, "__name__", name)

        async def __call__(self, state):
            return [self.name]

    async def send_for_fun(state):
        return [Send("2", state), Send("2", state)]

    async def route_to_three(state) -> Literal["3"]:
        return "3"

    builder = StateGraph(Annotated[list, operator.add])
    builder.add_node(Node("1"))
    builder.add_node(Node("2"))
    builder.add_node(Node("3"))
    builder.add_edge(START, "1")
    builder.add_conditional_edges("1", send_for_fun)
    builder.add_conditional_edges("2", route_to_three)
    graph = builder.compile()

    assert await graph.ainvoke(["0"]) == ["0", "1", "2", "2", "3"]


async def test_concurrent_emit_sends() -> None:
    class Node:
        def __init__(self, name: str):
            self.name = name
            setattr(self, "__name__", name)

        async def __call__(self, state):
            return (
                [self.name]
                if isinstance(state, list)
                else ["|".join((self.name, str(state)))]
            )

    async def send_for_fun(state):
        return [Send("2", 1), Send("2", 2), "3.1"]

    async def send_for_profit(state):
        return [Send("2", 3), Send("2", 4)]

    async def route_to_three(state) -> Literal["3"]:
        return "3"

    builder = StateGraph(Annotated[list, operator.add])
    builder.add_node(Node("1"))
    builder.add_node(Node("1.1"))
    builder.add_node(Node("2"))
    builder.add_node(Node("3"))
    builder.add_node(Node("3.1"))
    builder.add_edge(START, "1")
    builder.add_edge(START, "1.1")
    builder.add_conditional_edges("1", send_for_fun)
    builder.add_conditional_edges("1.1", send_for_profit)
    builder.add_conditional_edges("2", route_to_three)
    graph = builder.compile()
    assert await graph.ainvoke(["0"]) == (
        [
            "0",
            "1",
            "1.1",
            "3.1",
            "2|1",
            "2|2",
            "2|3",
            "2|4",
            "3",
        ]
    )


@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_send_sequences(checkpointer_name: str) -> None:
    class Node:
        def __init__(self, name: str):
            self.name = name
            setattr(self, "__name__", name)

        async def __call__(self, state):
            update = (
                [self.name]
                if isinstance(state, list)  # or isinstance(state, Control)
                else ["|".join((self.name, str(state)))]
            )
            if isinstance(state, Command):
                return replace(state, update=update)
            else:
                return update

    async def send_for_fun(state):
        return [
            Send("2", Command(goto=Send("2", 3))),
            Send("2", Command(goto=Send("2", 4))),
            "3.1",
        ]

    async def route_to_three(state) -> Literal["3"]:
        return "3"

    builder = StateGraph(Annotated[list, operator.add])
    builder.add_node(Node("1"))
    builder.add_node(Node("2"))
    builder.add_node(Node("3"))
    builder.add_node(Node("3.1"))
    builder.add_edge(START, "1")
    builder.add_conditional_edges("1", send_for_fun)
    builder.add_conditional_edges("2", route_to_three)
    graph = builder.compile()
    assert await graph.ainvoke(["0"]) == [
        "0",
        "1",
        "3.1",
        "2|Command(goto=Send(node='2', arg=3))",
        "2|Command(goto=Send(node='2', arg=4))",
        "3",
        "2|3",
        "2|4",
        "3",
    ]

    async with awith_checkpointer(checkpointer_name) as checkpointer:
        graph = builder.compile(checkpointer=checkpointer, interrupt_before=["3.1"])
        thread1 = {"configurable": {"thread_id": "1"}}
        assert await graph.ainvoke(["0"], thread1) == [
            "0",
            "1",
        ]
        assert await graph.ainvoke(None, thread1) == [
            "0",
            "1",
            "3.1",
            "2|Command(goto=Send(node='2', arg=3))",
            "2|Command(goto=Send(node='2', arg=4))",
            "3",
            "2|3",
            "2|4",
            "3",
        ]


@NEEDS_CONTEXTVARS
@pytest.mark.parametrize("checkpoint_during", [True, False])
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_imp_task(checkpointer_name: str, checkpoint_during: bool) -> None:
    if not checkpoint_during and "shallow" in checkpointer_name:
        pytest.skip("Checkpointing during execution not supported")

    async with awith_checkpointer(checkpointer_name) as checkpointer:
        mapper_calls = 0

        @task()
        async def mapper(input: int) -> str:
            nonlocal mapper_calls
            mapper_calls += 1
            await asyncio.sleep(0.1 * input)
            return str(input) * 2

        @entrypoint(checkpointer=checkpointer)
        async def graph(input: list[int]) -> list[str]:
            futures = [mapper(i) for i in input]
            mapped = await asyncio.gather(*futures)
            answer = interrupt("question")
            return [m + answer for m in mapped]

        tracer = FakeTracer()
        thread1 = {"configurable": {"thread_id": "1"}, "callbacks": [tracer]}
        assert [
            c
            async for c in graph.astream(
                [0, 1], thread1, checkpoint_during=checkpoint_during
            )
        ] == [
            {"mapper": "00"},
            {"mapper": "11"},
            {
                "__interrupt__": (
                    Interrupt(
                        value="question",
                        resumable=True,
                        ns=[AnyStr("graph:")],
                        when="during",
                    ),
                )
            },
        ]
        assert mapper_calls == 2
        assert len(tracer.runs) == 1
        assert len(tracer.runs[0].child_runs) == 1
        entrypoint_run = tracer.runs[0].child_runs[0]
        assert entrypoint_run.name == "graph"
        mapper_runs = [r for r in entrypoint_run.child_runs if r.name == "mapper"]
        assert len(mapper_runs) == 2
        assert any(r.inputs == {"input": 0} for r in mapper_runs)
        assert any(r.inputs == {"input": 1} for r in mapper_runs)

        assert await graph.ainvoke(
            Command(resume="answer"), thread1, checkpoint_during=checkpoint_during
        ) == [
            "00answer",
            "11answer",
        ]
        assert mapper_calls == 2


@NEEDS_CONTEXTVARS
@pytest.mark.parametrize("checkpoint_during", [True, False])
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_imp_nested(checkpointer_name: str, checkpoint_during: bool) -> None:
    if not checkpoint_during and "shallow" in checkpointer_name:
        pytest.skip("Checkpointing during execution not supported")

    async def mynode(input: list[str]) -> list[str]:
        return [it + "a" for it in input]

    builder = StateGraph(list[str])
    builder.add_node(mynode)
    builder.add_edge(START, "mynode")
    add_a = builder.compile()

    @task
    def submapper(input: int) -> str:
        return str(input)

    @task
    async def mapper(input: int) -> str:
        await asyncio.sleep(input / 100)
        return await submapper(input) * 2

    async with awith_checkpointer(checkpointer_name) as checkpointer:

        @entrypoint(checkpointer=checkpointer)
        async def graph(input: list[int]) -> list[str]:
            futures = [mapper(i) for i in input]
            mapped = await asyncio.gather(*futures)
            answer = interrupt("question")
            final = [m + answer for m in mapped]
            return await add_a.ainvoke(final)

        assert graph.get_input_jsonschema() == {
            "type": "array",
            "items": {"type": "integer"},
            "title": "LangGraphInput",
        }
        assert graph.get_output_jsonschema() == {
            "type": "array",
            "items": {"type": "string"},
            "title": "LangGraphOutput",
        }

        thread1 = {"configurable": {"thread_id": "1"}}
        assert [
            c
            async for c in graph.astream(
                [0, 1], thread1, checkpoint_during=checkpoint_during
            )
        ] == [
            {"submapper": "0"},
            {"mapper": "00"},
            {"submapper": "1"},
            {"mapper": "11"},
            {
                "__interrupt__": (
                    Interrupt(
                        value="question",
                        resumable=True,
                        ns=[AnyStr("graph:")],
                        when="during",
                    ),
                )
            },
        ]

        assert await graph.ainvoke(
            Command(resume="answer"), thread1, checkpoint_during=checkpoint_during
        ) == [
            "00answera",
            "11answera",
        ]


@NEEDS_CONTEXTVARS
@pytest.mark.parametrize("checkpoint_during", [True, False])
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_imp_task_cancel(checkpointer_name: str, checkpoint_during: bool) -> None:
    if not checkpoint_during and "shallow" in checkpointer_name:
        pytest.skip("Checkpointing during execution not supported")

    async with awith_checkpointer(checkpointer_name) as checkpointer:
        mapper_calls = 0
        mapper_cancels = 0

        @task()
        async def mapper(input: int) -> str:
            nonlocal mapper_calls, mapper_cancels
            mapper_calls += 1
            try:
                await asyncio.sleep(1)
            except asyncio.CancelledError:
                mapper_cancels += 1
                raise
            return str(input) * 2

        @entrypoint(checkpointer=checkpointer)
        async def graph(input: list[int]) -> list[str]:
            futures = [mapper(i) for i in input]
            await asyncio.sleep(0.1)
            futures.pop().cancel()  # cancel one
            mapped = await asyncio.gather(*futures)
            answer = interrupt("question")
            return [m + answer for m in mapped]

        thread1 = {"configurable": {"thread_id": "1"}}
        assert [
            c
            async for c in graph.astream(
                [0, 1], thread1, checkpoint_during=checkpoint_during
            )
        ] == [
            {"mapper": "00"},
            {
                "__interrupt__": (
                    Interrupt(
                        value="question",
                        resumable=True,
                        ns=[AnyStr("graph:")],
                        when="during",
                    ),
                )
            },
        ]
        assert mapper_calls == 2
        assert mapper_cancels == 1

        assert await graph.ainvoke(
            Command(resume="answer"), thread1, checkpoint_during=checkpoint_during
        ) == [
            "00answer",
        ]
        assert mapper_calls == 3
        assert mapper_cancels == 2


@NEEDS_CONTEXTVARS
@pytest.mark.parametrize("checkpoint_during", [True, False])
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_imp_sync_from_async(
    checkpointer_name: str, checkpoint_during: bool
) -> None:
    if not checkpoint_during and "shallow" in checkpointer_name:
        pytest.skip("Checkpointing during execution not supported")

    async with awith_checkpointer(checkpointer_name) as checkpointer:

        @task()
        def foo(state: dict) -> dict:
            return {"a": state["a"] + "foo", "b": "bar"}

        @task
        def bar(a: str, b: str, c: Optional[str] = None) -> dict:
            return {"a": a + b, "c": (c or "") + "bark"}

        @task()
        def baz(state: dict) -> dict:
            return {"a": state["a"] + "baz", "c": "something else"}

        @entrypoint(checkpointer=checkpointer)
        def graph(state: dict) -> dict:
            foo_result = foo(state).result()
            fut_bar = bar(foo_result["a"], foo_result["b"])
            fut_baz = baz(fut_bar.result())
            return fut_baz.result()

        thread1 = {"configurable": {"thread_id": "1"}}
        assert [
            c
            async for c in graph.astream(
                {"a": "0"}, thread1, checkpoint_during=checkpoint_during
            )
        ] == [
            {"foo": {"a": "0foo", "b": "bar"}},
            {"bar": {"a": "0foobar", "c": "bark"}},
            {"baz": {"a": "0foobarbaz", "c": "something else"}},
            {"graph": {"a": "0foobarbaz", "c": "something else"}},
        ]


@NEEDS_CONTEXTVARS
@pytest.mark.parametrize("checkpoint_during", [True, False])
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_imp_stream_order(
    checkpointer_name: str, checkpoint_during: bool
) -> None:
    if not checkpoint_during and "shallow" in checkpointer_name:
        pytest.skip("Checkpointing during execution not supported")

    async with awith_checkpointer(checkpointer_name) as checkpointer:

        @task()
        async def foo(state: dict) -> dict:
            return {"a": state["a"] + "foo", "b": "bar"}

        @task
        async def bar(a: str, b: str, c: Optional[str] = None) -> dict:
            return {"a": a + b, "c": (c or "") + "bark"}

        @task()
        async def baz(state: dict) -> dict:
            return {"a": state["a"] + "baz", "c": "something else"}

        @entrypoint(checkpointer=checkpointer)
        async def graph(state: dict) -> dict:
            foo_res = await foo(state)

            fut_bar = bar(foo_res["a"], foo_res["b"])
            fut_baz = baz(await fut_bar)
            return await fut_baz

        thread1 = {"configurable": {"thread_id": "1"}}
        assert [
            c
            async for c in graph.astream(
                {"a": "0"}, thread1, checkpoint_during=checkpoint_during
            )
        ] == [
            {"foo": {"a": "0foo", "b": "bar"}},
            {"bar": {"a": "0foobar", "c": "bark"}},
            {"baz": {"a": "0foobarbaz", "c": "something else"}},
            {"graph": {"a": "0foobarbaz", "c": "something else"}},
        ]


@pytest.mark.parametrize("checkpoint_during", [True, False])
@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC)
async def test_send_dedupe_on_resume(
    checkpointer_name: str, checkpoint_during: bool
) -> None:
    class InterruptOnce:
        ticks: int = 0

        def __call__(self, state):
            self.ticks += 1
            if self.ticks == 1:
                raise NodeInterrupt("Bahh")
            return ["|".join(("flaky", str(state)))]

    class Node:
        def __init__(self, name: str):
            self.name = name
            self.ticks = 0
            setattr(self, "__name__", name)

        def __call__(self, state):
            self.ticks += 1
            update = (
                [self.name]
                if isinstance(state, list)
                else ["|".join((self.name, str(state)))]
            )
            if isinstance(state, Command):
                return replace(state, update=update)
            else:
                return update

    def send_for_fun(state):
        return [
            Send("2", Command(goto=Send("2", 3))),
            Send("2", Command(goto=Send("flaky", 4))),
            "3.1",
        ]

    def route_to_three(state) -> Literal["3"]:
        return "3"

    builder = StateGraph(Annotated[list, operator.add])
    builder.add_node(Node("1"))
    builder.add_node(Node("2"))
    builder.add_node(Node("3"))
    builder.add_node(Node("3.1"))
    builder.add_node("flaky", InterruptOnce())
    builder.add_edge(START, "1")
    builder.add_conditional_edges("1", send_for_fun)
    builder.add_conditional_edges("2", route_to_three)

    async with awith_checkpointer(checkpointer_name) as checkpointer:
        graph = builder.compile(checkpointer=checkpointer)
        thread1 = {"configurable": {"thread_id": "1"}}
        assert await graph.ainvoke(
            ["0"], thread1, checkpoint_during=checkpoint_during
        ) == [
            "0",
            "1",
            "3.1",
            "2|Command(goto=Send(node='2', arg=3))",
            "2|Command(goto=Send(node='flaky', arg=4))",
            "3",
            "2|3",
        ]
        assert builder.nodes["2"].runnable.func.ticks == 3
        assert builder.nodes["flaky"].runnable.func.ticks == 1
        # resume execution
        assert await graph.ainvoke(
            None, thread1, checkpoint_during=checkpoint_during
        ) == [
            "0",
            "1",
            "3.1",
            "2|Command(goto=Send(node='2', arg=3))",
            "2|Command(goto=Send(node='flaky', arg=4))",
            "3",
            "2|3",
            "flaky|4",
            "3",
        ]
        # node "2" doesn't get called again, as we recover writes saved before
        assert builder.nodes["2"].runnable.func.ticks == 3
        # node "flaky" gets called again, as it was interrupted
        assert builder.nodes["flaky"].runnable.func.ticks == 2
        # check history
        history = [c async for c in graph.aget_state_history(thread1)]
        assert len(history) == (6 if checkpoint_during else 2)
        expected_history = [
            StateSnapshot(
                values=[
                    "0",
                    "1",
                    "3.1",
                    "2|Command(goto=Send(node='2', arg=3))",
                    "2|Command(goto=Send(node='flaky', arg=4))",
                    "3",
                    "2|3",
                    "flaky|4",
                    "3",
                ],
                next=(),
                config={
                    "configurable": {
                        "thread_id": "1",
                        "checkpoint_ns": "",
                        "checkpoint_id": AnyStr(),
                    }
                },
                metadata={
                    "source": "loop",
                    "writes": {"3": ["3"]},
                    "thread_id": "1",
                    "step": 4,
                    "parents": {},
                },
                created_at=AnyStr(),
                parent_config={
                    "configurable": {
                        "thread_id": "1",
                        "checkpoint_ns": "",
                        "checkpoint_id": AnyStr(),
                    }
                },
                tasks=(),
            ),
            StateSnapshot(
                values=[
                    "0",
                    "1",
                    "3.1",
                    "2|Command(goto=Send(node='2', arg=3))",
                    "2|Command(goto=Send(node='flaky', arg=4))",
                    "3",
                    "2|3",
                    "flaky|4",
                ],
                next=("3",),
                config={
                    "configurable": {
                        "thread_id": "1",
                        "checkpoint_ns": "",
                        "checkpoint_id": AnyStr(),
                    }
                },
                metadata={
                    "source": "loop",
                    "writes": {"2": ["2|3"], "3": ["3"], "flaky": ["flaky|4"]},
                    "thread_id": "1",
                    "step": 3,
                    "parents": {},
                },
                created_at=AnyStr(),
                parent_config={
                    "configurable": {
                        "thread_id": "1",
                        "checkpoint_ns": "",
                        "checkpoint_id": AnyStr(),
                    }
                },
                tasks=(
                    PregelTask(
                        id=AnyStr(),
                        name="3",
                        path=("__pregel_pull", "3"),
                        error=None,
                        interrupts=(),
                        state=None,
                        result=["3"],
                    ),
                ),
            ),
            StateSnapshot(
                values=[
                    "0",
                    "1",
                    "3.1",
                    "2|Command(goto=Send(node='2', arg=3))",
                    "2|Command(goto=Send(node='flaky', arg=4))",
                ],
                next=("2", "flaky", "3"),
                config={
                    "configurable": {
                        "thread_id": "1",
                        "checkpoint_ns": "",
                        "checkpoint_id": AnyStr(),
                    }
                },
                metadata={
                    "source": "loop",
                    "writes": {
                        "2": [
                            ["2|Command(goto=Send(node='2', arg=3))"],
                            ["2|Command(goto=Send(node='flaky', arg=4))"],
                        ],
                        "3.1": ["3.1"],
                    },
                    "thread_id": "1",
                    "step": 2,
                    "parents": {},
                },
                created_at=AnyStr(),
                parent_config={
                    "configurable": {
                        "thread_id": "1",
                        "checkpoint_ns": "",
                        "checkpoint_id": AnyStr(),
                    }
                },
                tasks=(
                    PregelTask(
                        id=AnyStr(),
                        name="2",
                        path=("__pregel_push", 0),
                        error=None,
                        interrupts=(),
                        state=None,
                        result=["2|3"],
                    ),
                    PregelTask(
                        id=AnyStr(),
                        name="flaky",
                        path=("__pregel_push", 1),
                        error=None,
                        interrupts=(Interrupt(value="Bahh", resumable=False, ns=None),),
                        state=None,
                        result=["flaky|4"] if checkpoint_during else None,
                    ),
                    PregelTask(
                        id=AnyStr(),
                        name="3",
                        path=("__pregel_pull", "3"),
                        error=None,
                        interrupts=(),
                        state=None,
                        result=["3"],
                    ),
                ),
            ),
            StateSnapshot(
                values=["0", "1"],
                next=("2", "2", "3.1"),
                config={
                    "configurable": {
                        "thread_id": "1",
                        "checkpoint_ns": "",
                        "checkpoint_id": AnyStr(),
                    }
                },
                metadata={
                    "source": "loop",
                    "writes": {"1": ["1"]},
                    "thread_id": "1",
                    "step": 1,
                    "parents": {},
                },
                created_at=AnyStr(),
                parent_config={
                    "configurable": {
                        "thread_id": "1",
                        "checkpoint_ns": "",
                        "checkpoint_id": AnyStr(),
                    }
                },
                tasks=(
                    PregelTask(
                        id=AnyStr(),
                        name="2",
                        path=("__pregel_push", 0),
                        error=None,
                        interrupts=(),
                        state=None,
                        result=["2|Command(goto=Send(node='2', arg=3))"],
                    ),
                    PregelTask(
                        id=AnyStr(),
                        name="2",
                        path=("__pregel_push", 1),
                        error=None,
                        interrupts=(),
                        state=None,
                        result=["2|Command(goto=Send(node='flaky', arg=4))"],
                    ),
                    PregelTask(
                        id=AnyStr(),
                        name="3.1",
                        path=("__pregel_pull", "3.1"),
                        error=None,
                        interrupts=(),
                        state=None,
                        result=["3.1"],
                    ),
                ),
            ),
            StateSnapshot(
                values=["0"],
                next=("1",),
                config={
                    "configurable": {
                        "thread_id": "1",
                        "checkpoint_ns": "",
                        "checkpoint_id": AnyStr(),
                    }
                },
                metadata={
                    "source": "loop",
                    "writes": None,
                    "thread_id": "1",
                    "step": 0,
                    "parents": {},
                },
                created_at=AnyStr(),
                parent_config={
                    "configurable": {
                        "thread_id": "1",
                        "checkpoint_ns": "",
                        "checkpoint_id": AnyStr(),
                    }
                },
                tasks=(
                    PregelTask(
                        id=AnyStr(),
                        name="1",
                        path=("__pregel_pull", "1"),
                        error=None,
                        interrupts=(),
                        state=None,
                        result=["1"],
                    ),
                ),
            ),
            StateSnapshot(
                values=[],
                next=("__start__",),
                config={
                    "configurable": {
                        "thread_id": "1",
                        "checkpoint_ns": "",
                        "checkpoint_id": AnyStr(),
                    }
                },
                metadata={
                    "source": "input",
                    "writes": {"__start__": ["0"]},
                    "thread_id": "1",
                    "step": -1,
                    "parents": {},
                },
                created_at=AnyStr(),
                parent_config=None,
                tasks=(
                    PregelTask(
                        id=AnyStr(),
                        name="__start__",
                        path=("__pregel_pull", "__start__"),
                        error=None,
                        interrupts=(),
                        state=None,
                        result=["0"],
                    ),
                ),
            ),
        ]
        if checkpoint_during:
            assert history == expected_history
        else:
            assert history[0] == expected_history[0]
            assert history[1] == expected_history[2]


@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_send_react_interrupt(checkpointer_name: str) -> None:
    from langchain_core.messages import AIMessage, HumanMessage, ToolCall, ToolMessage

    ai_message = AIMessage(
        "",
        id="ai1",
        tool_calls=[ToolCall(name="foo", args={"hi": [1, 2, 3]}, id=AnyStr())],
    )

    async def agent(state):
        return {"messages": ai_message}

    def route(state):
        if isinstance(state["messages"][-1], AIMessage):
            return [
                Send(call["name"], call) for call in state["messages"][-1].tool_calls
            ]

    foo_called = 0

    async def foo(call: ToolCall):
        nonlocal foo_called
        foo_called += 1
        return {"messages": ToolMessage(str(call["args"]), tool_call_id=call["id"])}

    builder = StateGraph(MessagesState)
    builder.add_node(agent)
    builder.add_node(foo)
    builder.add_edge(START, "agent")
    builder.add_conditional_edges("agent", route)
    graph = builder.compile()

    assert await graph.ainvoke({"messages": [HumanMessage("hello")]}) == {
        "messages": [
            _AnyIdHumanMessage(content="hello"),
            _AnyIdAIMessage(
                content="",
                tool_calls=[
                    {
                        "name": "foo",
                        "args": {"hi": [1, 2, 3]},
                        "id": "",
                        "type": "tool_call",
                    }
                ],
            ),
            _AnyIdToolMessage(
                content="{'hi': [1, 2, 3]}",
                tool_call_id=AnyStr(),
            ),
        ]
    }
    assert foo_called == 1

    async with awith_checkpointer(checkpointer_name) as checkpointer:
        # simple interrupt-resume flow
        foo_called = 0
        graph = builder.compile(checkpointer=checkpointer, interrupt_before=["foo"])
        thread1 = {"configurable": {"thread_id": "1"}}
        assert await graph.ainvoke({"messages": [HumanMessage("hello")]}, thread1) == {
            "messages": [
                _AnyIdHumanMessage(content="hello"),
                _AnyIdAIMessage(
                    content="",
                    tool_calls=[
                        {
                            "name": "foo",
                            "args": {"hi": [1, 2, 3]},
                            "id": "",
                            "type": "tool_call",
                        }
                    ],
                ),
            ]
        }
        assert foo_called == 0
        assert await graph.ainvoke(None, thread1) == {
            "messages": [
                _AnyIdHumanMessage(content="hello"),
                _AnyIdAIMessage(
                    content="",
                    tool_calls=[
                        {
                            "name": "foo",
                            "args": {"hi": [1, 2, 3]},
                            "id": "",
                            "type": "tool_call",
                        }
                    ],
                ),
                _AnyIdToolMessage(
                    content="{'hi': [1, 2, 3]}",
                    tool_call_id=AnyStr(),
                ),
            ]
        }
        assert foo_called == 1

        # interrupt-update-resume flow
        foo_called = 0
        graph = builder.compile(checkpointer=checkpointer, interrupt_before=["foo"])
        thread1 = {"configurable": {"thread_id": "2"}}
        assert await graph.ainvoke({"messages": [HumanMessage("hello")]}, thread1) == {
            "messages": [
                _AnyIdHumanMessage(content="hello"),
                _AnyIdAIMessage(
                    content="",
                    tool_calls=[
                        {
                            "name": "foo",
                            "args": {"hi": [1, 2, 3]},
                            "id": "",
                            "type": "tool_call",
                        }
                    ],
                ),
            ]
        }
        assert foo_called == 0

        # get state should show the pending task
        state = await graph.aget_state(thread1)
        assert state == StateSnapshot(
            values={
                "messages": [
                    _AnyIdHumanMessage(content="hello"),
                    _AnyIdAIMessage(
                        content="",
                        tool_calls=[
                            {
                                "name": "foo",
                                "args": {"hi": [1, 2, 3]},
                                "id": "",
                                "type": "tool_call",
                            }
                        ],
                    ),
                ]
            },
            next=("foo",),
            config={
                "configurable": {
                    "thread_id": "2",
                    "checkpoint_ns": "",
                    "checkpoint_id": AnyStr(),
                }
            },
            metadata={
                "step": 1,
                "source": "loop",
                "writes": {
                    "agent": {
                        "messages": AIMessage(
                            content="",
                            id="ai1",
                            tool_calls=[
                                {
                                    "name": "foo",
                                    "args": {"hi": [1, 2, 3]},
                                    "id": "",
                                    "type": "tool_call",
                                }
                            ],
                        )
                    }
                },
                "parents": {},
                "thread_id": "2",
            },
            created_at=AnyStr(),
            parent_config=(
                None
                if "shallow" in checkpointer_name
                else {
                    "configurable": {
                        "thread_id": "2",
                        "checkpoint_ns": "",
                        "checkpoint_id": AnyStr(),
                    }
                }
            ),
            tasks=(
                PregelTask(
                    id=AnyStr(),
                    name="foo",
                    path=("__pregel_push", 0),
                    error=None,
                    interrupts=(),
                    state=None,
                    result=None,
                ),
            ),
        )

        # remove the tool call, clearing the pending task
        await graph.aupdate_state(
            thread1, {"messages": AIMessage("Bye now", id=ai_message.id, tool_calls=[])}
        )

        # tool call no longer in pending tasks
        assert await graph.aget_state(thread1) == StateSnapshot(
            values={
                "messages": [
                    _AnyIdHumanMessage(content="hello"),
                    _AnyIdAIMessage(
                        content="Bye now",
                        tool_calls=[],
                    ),
                ]
            },
            next=(),
            config={
                "configurable": {
                    "thread_id": "2",
                    "checkpoint_ns": "",
                    "checkpoint_id": AnyStr(),
                }
            },
            metadata={
                "step": 2,
                "source": "update",
                "writes": {
                    "agent": {
                        "messages": _AnyIdAIMessage(
                            content="Bye now",
                            tool_calls=[],
                        )
                    }
                },
                "parents": {},
                "thread_id": "2",
            },
            created_at=AnyStr(),
            parent_config=(
                None
                if "shallow" in checkpointer_name
                else {
                    "configurable": {
                        "thread_id": "2",
                        "checkpoint_ns": "",
                        "checkpoint_id": AnyStr(),
                    }
                }
            ),
            tasks=(),
        )

        # tool call not executed
        assert await graph.ainvoke(None, thread1) == {
            "messages": [
                _AnyIdHumanMessage(content="hello"),
                _AnyIdAIMessage(content="Bye now"),
            ]
        }
        assert foo_called == 0

        # interrupt-update-resume flow, creating new Send in update call
        foo_called = 0
        graph = builder.compile(checkpointer=checkpointer, interrupt_before=["foo"])
        thread1 = {"configurable": {"thread_id": "3"}}
        assert await graph.ainvoke({"messages": [HumanMessage("hello")]}, thread1) == {
            "messages": [
                _AnyIdHumanMessage(content="hello"),
                _AnyIdAIMessage(
                    content="",
                    tool_calls=[
                        {
                            "name": "foo",
                            "args": {"hi": [1, 2, 3]},
                            "id": "",
                            "type": "tool_call",
                        }
                    ],
                ),
            ]
        }
        assert foo_called == 0

        # get state should show the pending task
        state = await graph.aget_state(thread1)
        assert state == StateSnapshot(
            values={
                "messages": [
                    _AnyIdHumanMessage(content="hello"),
                    _AnyIdAIMessage(
                        content="",
                        tool_calls=[
                            {
                                "name": "foo",
                                "args": {"hi": [1, 2, 3]},
                                "id": "",
                                "type": "tool_call",
                            }
                        ],
                    ),
                ]
            },
            next=("foo",),
            config={
                "configurable": {
                    "thread_id": "3",
                    "checkpoint_ns": "",
                    "checkpoint_id": AnyStr(),
                }
            },
            metadata={
                "step": 1,
                "source": "loop",
                "writes": {
                    "agent": {
                        "messages": AIMessage(
                            content="",
                            id="ai1",
                            tool_calls=[
                                {
                                    "name": "foo",
                                    "args": {"hi": [1, 2, 3]},
                                    "id": "",
                                    "type": "tool_call",
                                }
                            ],
                        )
                    }
                },
                "parents": {},
                "thread_id": "3",
            },
            created_at=AnyStr(),
            parent_config=(
                None
                if "shallow" in checkpointer_name
                else {
                    "configurable": {
                        "thread_id": "3",
                        "checkpoint_ns": "",
                        "checkpoint_id": AnyStr(),
                    }
                }
            ),
            tasks=(
                PregelTask(
                    id=AnyStr(),
                    name="foo",
                    path=("__pregel_push", 0),
                    error=None,
                    interrupts=(),
                    state=None,
                    result=None,
                ),
            ),
        )

        # replace the tool call, should clear previous send, create new one
        await graph.aupdate_state(
            thread1,
            {
                "messages": AIMessage(
                    "",
                    id=ai_message.id,
                    tool_calls=[
                        {
                            "name": "foo",
                            "args": {"hi": [4, 5, 6]},
                            "id": "tool1",
                            "type": "tool_call",
                        }
                    ],
                )
            },
        )

        # prev tool call no longer in pending tasks, new tool call is
        assert await graph.aget_state(thread1) == StateSnapshot(
            values={
                "messages": [
                    _AnyIdHumanMessage(content="hello"),
                    _AnyIdAIMessage(
                        content="",
                        tool_calls=[
                            {
                                "name": "foo",
                                "args": {"hi": [4, 5, 6]},
                                "id": "tool1",
                                "type": "tool_call",
                            }
                        ],
                    ),
                ]
            },
            next=("foo",),
            config={
                "configurable": {
                    "thread_id": "3",
                    "checkpoint_ns": "",
                    "checkpoint_id": AnyStr(),
                }
            },
            metadata={
                "step": 2,
                "source": "update",
                "writes": {
                    "agent": {
                        "messages": _AnyIdAIMessage(
                            content="",
                            tool_calls=[
                                {
                                    "name": "foo",
                                    "args": {"hi": [4, 5, 6]},
                                    "id": "tool1",
                                    "type": "tool_call",
                                }
                            ],
                        )
                    }
                },
                "parents": {},
                "thread_id": "3",
            },
            created_at=AnyStr(),
            parent_config=(
                None
                if "shallow" in checkpointer_name
                else {
                    "configurable": {
                        "thread_id": "3",
                        "checkpoint_ns": "",
                        "checkpoint_id": AnyStr(),
                    }
                }
            ),
            tasks=(
                PregelTask(
                    id=AnyStr(),
                    name="foo",
                    path=("__pregel_push", 0),
                    error=None,
                    interrupts=(),
                    state=None,
                    result=None,
                ),
            ),
        )

        # prev tool call not executed, new tool call is
        assert await graph.ainvoke(None, thread1) == {
            "messages": [
                _AnyIdHumanMessage(content="hello"),
                AIMessage(
                    "",
                    id="ai1",
                    tool_calls=[
                        {
                            "name": "foo",
                            "args": {"hi": [4, 5, 6]},
                            "id": "tool1",
                            "type": "tool_call",
                        }
                    ],
                ),
                _AnyIdToolMessage(content="{'hi': [4, 5, 6]}", tool_call_id="tool1"),
            ]
        }
        assert foo_called == 1


@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_send_react_interrupt_control(
    checkpointer_name: str, snapshot: SnapshotAssertion
) -> None:
    from langchain_core.messages import AIMessage, HumanMessage, ToolCall, ToolMessage

    ai_message = AIMessage(
        "",
        id="ai1",
        tool_calls=[ToolCall(name="foo", args={"hi": [1, 2, 3]}, id=AnyStr())],
    )

    async def agent(state) -> Command[Literal["foo"]]:
        return Command(
            update={"messages": ai_message},
            goto=[Send(call["name"], call) for call in ai_message.tool_calls],
        )

    foo_called = 0

    async def foo(call: ToolCall):
        nonlocal foo_called
        foo_called += 1
        return {"messages": ToolMessage(str(call["args"]), tool_call_id=call["id"])}

    builder = StateGraph(MessagesState)
    builder.add_node(agent)
    builder.add_node(foo)
    builder.add_edge(START, "agent")
    graph = builder.compile()
    if checkpointer_name == "memory":
        assert graph.get_graph().draw_mermaid() == snapshot

    assert await graph.ainvoke({"messages": [HumanMessage("hello")]}) == {
        "messages": [
            _AnyIdHumanMessage(content="hello"),
            _AnyIdAIMessage(
                content="",
                tool_calls=[
                    {
                        "name": "foo",
                        "args": {"hi": [1, 2, 3]},
                        "id": "",
                        "type": "tool_call",
                    }
                ],
            ),
            _AnyIdToolMessage(
                content="{'hi': [1, 2, 3]}",
                tool_call_id=AnyStr(),
            ),
        ]
    }
    assert foo_called == 1

    async with awith_checkpointer(checkpointer_name) as checkpointer:
        # simple interrupt-resume flow
        foo_called = 0
        graph = builder.compile(checkpointer=checkpointer, interrupt_before=["foo"])
        thread1 = {"configurable": {"thread_id": "1"}}
        assert await graph.ainvoke({"messages": [HumanMessage("hello")]}, thread1) == {
            "messages": [
                _AnyIdHumanMessage(content="hello"),
                _AnyIdAIMessage(
                    content="",
                    tool_calls=[
                        {
                            "name": "foo",
                            "args": {"hi": [1, 2, 3]},
                            "id": "",
                            "type": "tool_call",
                        }
                    ],
                ),
            ]
        }
        assert foo_called == 0
        assert await graph.ainvoke(None, thread1) == {
            "messages": [
                _AnyIdHumanMessage(content="hello"),
                _AnyIdAIMessage(
                    content="",
                    tool_calls=[
                        {
                            "name": "foo",
                            "args": {"hi": [1, 2, 3]},
                            "id": "",
                            "type": "tool_call",
                        }
                    ],
                ),
                _AnyIdToolMessage(
                    content="{'hi': [1, 2, 3]}",
                    tool_call_id=AnyStr(),
                ),
            ]
        }
        assert foo_called == 1

        # interrupt-update-resume flow
        foo_called = 0
        graph = builder.compile(checkpointer=checkpointer, interrupt_before=["foo"])
        thread1 = {"configurable": {"thread_id": "2"}}
        assert await graph.ainvoke({"messages": [HumanMessage("hello")]}, thread1) == {
            "messages": [
                _AnyIdHumanMessage(content="hello"),
                _AnyIdAIMessage(
                    content="",
                    tool_calls=[
                        {
                            "name": "foo",
                            "args": {"hi": [1, 2, 3]},
                            "id": "",
                            "type": "tool_call",
                        }
                    ],
                ),
            ]
        }
        assert foo_called == 0

        # get state should show the pending task
        state = await graph.aget_state(thread1)
        assert state == StateSnapshot(
            values={
                "messages": [
                    _AnyIdHumanMessage(content="hello"),
                    _AnyIdAIMessage(
                        content="",
                        tool_calls=[
                            {
                                "name": "foo",
                                "args": {"hi": [1, 2, 3]},
                                "id": "",
                                "type": "tool_call",
                            }
                        ],
                    ),
                ]
            },
            next=("foo",),
            config={
                "configurable": {
                    "thread_id": "2",
                    "checkpoint_ns": "",
                    "checkpoint_id": AnyStr(),
                }
            },
            metadata={
                "step": 1,
                "source": "loop",
                "writes": {
                    "agent": {
                        "messages": AIMessage(
                            content="",
                            id="ai1",
                            tool_calls=[
                                {
                                    "name": "foo",
                                    "args": {"hi": [1, 2, 3]},
                                    "id": "",
                                    "type": "tool_call",
                                }
                            ],
                        )
                    }
                },
                "parents": {},
                "thread_id": "2",
            },
            created_at=AnyStr(),
            parent_config=(
                None
                if "shallow" in checkpointer_name
                else {
                    "configurable": {
                        "thread_id": "2",
                        "checkpoint_ns": "",
                        "checkpoint_id": AnyStr(),
                    }
                }
            ),
            tasks=(
                PregelTask(
                    id=AnyStr(),
                    name="foo",
                    path=("__pregel_push", 0),
                    error=None,
                    interrupts=(),
                    state=None,
                    result=None,
                ),
            ),
        )

        # remove the tool call, clearing the pending task
        await graph.aupdate_state(
            thread1, {"messages": AIMessage("Bye now", id=ai_message.id, tool_calls=[])}
        )

        # tool call no longer in pending tasks
        assert await graph.aget_state(thread1) == StateSnapshot(
            values={
                "messages": [
                    _AnyIdHumanMessage(content="hello"),
                    _AnyIdAIMessage(
                        content="Bye now",
                        tool_calls=[],
                    ),
                ]
            },
            next=(),
            config={
                "configurable": {
                    "thread_id": "2",
                    "checkpoint_ns": "",
                    "checkpoint_id": AnyStr(),
                }
            },
            metadata={
                "step": 2,
                "source": "update",
                "writes": {
                    "agent": {
                        "messages": _AnyIdAIMessage(
                            content="Bye now",
                            tool_calls=[],
                        )
                    }
                },
                "parents": {},
                "thread_id": "2",
            },
            created_at=AnyStr(),
            parent_config=(
                None
                if "shallow" in checkpointer_name
                else {
                    "configurable": {
                        "thread_id": "2",
                        "checkpoint_ns": "",
                        "checkpoint_id": AnyStr(),
                    }
                }
            ),
            tasks=(),
        )

        # tool call not executed
        assert await graph.ainvoke(None, thread1) == {
            "messages": [
                _AnyIdHumanMessage(content="hello"),
                _AnyIdAIMessage(content="Bye now"),
            ]
        }
        assert foo_called == 0


@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_max_concurrency(checkpointer_name: str) -> None:
    class Node:
        def __init__(self, name: str):
            self.name = name
            setattr(self, "__name__", name)
            self.currently = 0
            self.max_currently = 0

        async def __call__(self, state):
            self.currently += 1
            if self.currently > self.max_currently:
                self.max_currently = self.currently
            await asyncio.sleep(random.random() / 10)
            self.currently -= 1
            return [state]

    def one(state):
        return ["1"]

    def three(state):
        return ["3"]

    async def send_to_many(state):
        return [Send("2", idx) for idx in range(100)]

    async def route_to_three(state) -> Literal["3"]:
        return "3"

    node2 = Node("2")
    builder = StateGraph(Annotated[list, operator.add])
    builder.add_node("1", one)
    builder.add_node(node2)
    builder.add_node("3", three)
    builder.add_edge(START, "1")
    builder.add_conditional_edges("1", send_to_many)
    builder.add_conditional_edges("2", route_to_three)
    graph = builder.compile()

    assert await graph.ainvoke(["0"]) == ["0", "1", *range(100), "3"]
    assert node2.max_currently == 100
    assert node2.currently == 0
    node2.max_currently = 0

    assert await graph.ainvoke(["0"], {"max_concurrency": 10}) == [
        "0",
        "1",
        *range(100),
        "3",
    ]
    assert node2.max_currently == 10
    assert node2.currently == 0

    async with awith_checkpointer(checkpointer_name) as checkpointer:
        graph = builder.compile(checkpointer=checkpointer, interrupt_before=["2"])
        thread1 = {"max_concurrency": 10, "configurable": {"thread_id": "1"}}

        assert await graph.ainvoke(["0"], thread1, debug=True) == ["0", "1"]
        state = await graph.aget_state(thread1)
        assert state.values == ["0", "1"]
        assert await graph.ainvoke(None, thread1) == ["0", "1", *range(100), "3"]


@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_max_concurrency_control(checkpointer_name: str) -> None:
    async def node1(state) -> Command[Literal["2"]]:
        return Command(update=["1"], goto=[Send("2", idx) for idx in range(100)])

    node2_currently = 0
    node2_max_currently = 0

    async def node2(state) -> Command[Literal["3"]]:
        nonlocal node2_currently, node2_max_currently
        node2_currently += 1
        if node2_currently > node2_max_currently:
            node2_max_currently = node2_currently
        await asyncio.sleep(0.1)
        node2_currently -= 1

        return Command(update=[state], goto="3")

    async def node3(state) -> Literal["3"]:
        return ["3"]

    builder = StateGraph(Annotated[list, operator.add])
    builder.add_node("1", node1)
    builder.add_node("2", node2)
    builder.add_node("3", node3)
    builder.add_edge(START, "1")
    graph = builder.compile()

    if checkpointer_name == "memory":
        assert (
            graph.get_graph().draw_mermaid()
            == """---
config:
  flowchart:
    curve: linear
---
graph TD;
	__start__([<p>__start__</p>]):::first
	1(1)
	2(2)
	3(3)
	__end__([<p>__end__</p>]):::last
	1 -.-> 2;
	2 -.-> 3;
	__start__ --> 1;
	3 --> __end__;
	classDef default fill:#f2f0ff,line-height:1.2
	classDef first fill-opacity:0
	classDef last fill:#bfb6fc
"""
        )

    assert await graph.ainvoke(["0"], debug=True) == ["0", "1", *range(100), "3"]
    assert node2_max_currently == 100
    assert node2_currently == 0
    node2_max_currently = 0

    assert await graph.ainvoke(["0"], {"max_concurrency": 10}) == [
        "0",
        "1",
        *range(100),
        "3",
    ]
    assert node2_max_currently == 10
    assert node2_currently == 0

    async with awith_checkpointer(checkpointer_name) as checkpointer:
        graph = builder.compile(checkpointer=checkpointer, interrupt_before=["2"])
        thread1 = {"max_concurrency": 10, "configurable": {"thread_id": "1"}}

        assert await graph.ainvoke(["0"], thread1) == ["0", "1"]
        assert await graph.ainvoke(None, thread1) == ["0", "1", *range(100), "3"]


@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_invoke_checkpoint_three(
    mocker: MockerFixture, checkpointer_name: str
) -> None:
    add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"])

    def raise_if_above_10(input: int) -> int:
        if input > 10:
            raise ValueError("Input is too large")
        return input

    one = (
        Channel.subscribe_to(["input"]).join(["total"])
        | add_one
        | Channel.write_to("output", "total")
        | raise_if_above_10
    )

    async with awith_checkpointer(checkpointer_name) as checkpointer:
        app = Pregel(
            nodes={"one": one},
            channels={
                "total": BinaryOperatorAggregate(int, operator.add),
                "input": LastValue(int),
                "output": LastValue(int),
            },
            input_channels="input",
            output_channels="output",
            checkpointer=checkpointer,
            debug=True,
        )

        thread_1 = {"configurable": {"thread_id": "1"}}
        # total starts out as 0, so output is 0+2=2
        assert await app.ainvoke(2, thread_1) == 2
        state = await app.aget_state(thread_1)
        assert state is not None
        assert state.values.get("total") == 2
        assert (
            state.config["configurable"]["checkpoint_id"]
            == (await checkpointer.aget(thread_1))["id"]
        )
        # total is now 2, so output is 2+3=5
        assert await app.ainvoke(3, thread_1) == 5
        state = await app.aget_state(thread_1)
        assert state is not None
        assert state.values.get("total") == 7
        assert (
            state.config["configurable"]["checkpoint_id"]
            == (await checkpointer.aget(thread_1))["id"]
        )
        # total is now 2+5=7, so output would be 7+4=11, but raises ValueError
        with pytest.raises(ValueError):
            await app.ainvoke(4, thread_1)
        # checkpoint is not updated
        state = await app.aget_state(thread_1)
        assert state is not None
        assert state.values.get("total") == 7
        assert state.next == ("one",)
        """we checkpoint inputs and it failed on "one", so the next node is one"""
        # we can recover from error by sending new inputs
        assert await app.ainvoke(2, thread_1) == 9
        state = await app.aget_state(thread_1)
        assert state is not None
        assert state.values.get("total") == 16, "total is now 7+9=16"
        assert state.next == ()

        thread_2 = {"configurable": {"thread_id": "2"}}
        # on a new thread, total starts out as 0, so output is 0+5=5
        assert await app.ainvoke(5, thread_2) == 5
        state = await app.aget_state({"configurable": {"thread_id": "1"}})
        assert state is not None
        assert state.values.get("total") == 16
        assert state.next == ()
        state = await app.aget_state(thread_2)
        assert state is not None
        assert state.values.get("total") == 5
        assert state.next == ()

        if "shallow" in checkpointer_name:
            return

        assert len([c async for c in app.aget_state_history(thread_1, limit=1)]) == 1
        # list all checkpoints for thread 1
        thread_1_history = [c async for c in app.aget_state_history(thread_1)]
        # there are 7 checkpoints
        assert len(thread_1_history) == 7
        assert Counter(c.metadata["source"] for c in thread_1_history) == {
            "input": 4,
            "loop": 3,
        }
        # sorted descending
        assert (
            thread_1_history[0].config["configurable"]["checkpoint_id"]
            > thread_1_history[1].config["configurable"]["checkpoint_id"]
        )
        # cursor pagination
        cursored = [
            c
            async for c in app.aget_state_history(
                thread_1, limit=1, before=thread_1_history[0].config
            )
        ]
        assert len(cursored) == 1
        assert cursored[0].config == thread_1_history[1].config
        # the last checkpoint
        assert thread_1_history[0].values["total"] == 16
        # the first "loop" checkpoint
        assert thread_1_history[-2].values["total"] == 2
        # can get each checkpoint using aget with config
        assert (await checkpointer.aget(thread_1_history[0].config))[
            "id"
        ] == thread_1_history[0].config["configurable"]["checkpoint_id"]
        assert (await checkpointer.aget(thread_1_history[1].config))[
            "id"
        ] == thread_1_history[1].config["configurable"]["checkpoint_id"]

        thread_1_next_config = await app.aupdate_state(thread_1_history[1].config, 10)
        # update creates a new checkpoint
        assert (
            thread_1_next_config["configurable"]["checkpoint_id"]
            > thread_1_history[0].config["configurable"]["checkpoint_id"]
        )
        # 1 more checkpoint in history
        assert len([c async for c in app.aget_state_history(thread_1)]) == 8
        assert Counter(
            [c.metadata["source"] async for c in app.aget_state_history(thread_1)]
        ) == {
            "update": 1,
            "input": 4,
            "loop": 3,
        }
        # the latest checkpoint is the updated one
        assert await app.aget_state(thread_1) == await app.aget_state(
            thread_1_next_config
        )


async def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None:
    add_one = mocker.Mock(side_effect=lambda x: x + 1)
    add_10_each = mocker.Mock(side_effect=lambda x: sorted(y + 10 for y in x))

    one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
    chain_three = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
    chain_four = (
        Channel.subscribe_to("inbox") | add_10_each | Channel.write_to("output")
    )

    app = Pregel(
        nodes={
            "one": one,
            "chain_three": chain_three,
            "chain_four": chain_four,
        },
        channels={
            "inbox": Topic(int),
            "output": LastValue(int),
            "input": LastValue(int),
        },
        input_channels="input",
        output_channels="output",
    )

    # Then invoke app
    # We get a single array result as chain_four waits for all publishers to finish
    # before operating on all elements published to topic_two as an array
    for _ in range(100):
        assert await app.ainvoke(2) == [13, 13]

    assert await asyncio.gather(*(app.ainvoke(2) for _ in range(100))) == [
        [13, 13] for _ in range(100)
    ]


@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_invoke_join_then_call_other_pregel(
    mocker: MockerFixture, checkpointer_name: str
) -> None:
    add_one = mocker.Mock(side_effect=lambda x: x + 1)
    add_10_each = mocker.Mock(side_effect=lambda x: [y + 10 for y in x])

    inner_app = Pregel(
        nodes={
            "one": Channel.subscribe_to("input") | add_one | Channel.write_to("output")
        },
        channels={
            "output": LastValue(int),
            "input": LastValue(int),
        },
        input_channels="input",
        output_channels="output",
    )

    one = (
        Channel.subscribe_to("input")
        | add_10_each
        | Channel.write_to("inbox_one").map()
    )
    two = (
        Channel.subscribe_to("inbox_one")
        | inner_app.map()
        | sorted
        | Channel.write_to("outbox_one")
    )
    chain_three = Channel.subscribe_to("outbox_one") | sum | Channel.write_to("output")

    app = Pregel(
        nodes={
            "one": one,
            "two": two,
            "chain_three": chain_three,
        },
        channels={
            "inbox_one": Topic(int),
            "outbox_one": LastValue(int),
            "output": LastValue(int),
            "input": LastValue(int),
        },
        input_channels="input",
        output_channels="output",
    )

    # Then invoke pubsub
    for _ in range(10):
        assert await app.ainvoke([2, 3]) == 27

    assert await asyncio.gather(*(app.ainvoke([2, 3]) for _ in range(10))) == [
        27 for _ in range(10)
    ]

    async with awith_checkpointer(checkpointer_name) as checkpointer:
        # add checkpointer
        app.checkpointer = checkpointer
        # subgraph is called twice, and that works
        assert await app.ainvoke([2, 3], {"configurable": {"thread_id": "1"}}) == 27

        # set inner graph checkpointer NeverCheckpoint
        inner_app.checkpointer = False
        # subgraph still called twice, but checkpointing for inner graph is disabled
        assert await app.ainvoke([2, 3], {"configurable": {"thread_id": "1"}}) == 27


async def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None:
    add_one = mocker.Mock(side_effect=lambda x: x + 1)

    one = (
        Channel.subscribe_to("input") | add_one | Channel.write_to("output", "between")
    )
    two = Channel.subscribe_to("between") | add_one | Channel.write_to("output")

    app = Pregel(
        nodes={"one": one, "two": two},
        channels={
            "input": LastValue(int),
            "between": LastValue(int),
            "output": LastValue(int),
        },
        stream_channels=["output", "between"],
        input_channels="input",
        output_channels="output",
    )

    # Then invoke pubsub
    assert [c async for c in app.astream(2)] == [
        {"between": 3, "output": 3},
        {"between": 3, "output": 4},
    ]


async def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None:
    add_one = mocker.Mock(side_effect=lambda x: x + 1)
    one = Channel.subscribe_to("input") | add_one | Channel.write_to("between")
    two = Channel.subscribe_to("between") | add_one

    app = Pregel(
        nodes={"one": one, "two": two},
        channels={
            "input": LastValue(int),
            "between": LastValue(int),
            "output": LastValue(int),
        },
        input_channels="input",
        output_channels="output",
    )

    # It finishes executing (once no more messages being published)
    # but returns nothing, as nothing was published to "output" topic
    assert await app.ainvoke(2) is None


async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None:
    setup_sync = mocker.Mock()
    cleanup_sync = mocker.Mock()
    setup_async = mocker.Mock()
    cleanup_async = mocker.Mock()

    @contextmanager
    def an_int() -> Generator[int, None, None]:
        setup_sync()
        try:
            yield 5
        finally:
            cleanup_sync()

    @asynccontextmanager
    async def an_int_async() -> AsyncGenerator[int, None]:
        setup_async()
        try:
            yield 5
        finally:
            cleanup_async()

    add_one = mocker.Mock(side_effect=lambda x: x + 1)
    one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
    two = (
        Channel.subscribe_to("inbox")
        | RunnableLambda(add_one).abatch
        | Channel.write_to("output").abatch
    )

    app = Pregel(
        nodes={"one": one, "two": two},
        channels={
            "input": LastValue(int),
            "output": LastValue(int),
            "inbox": Topic(int),
            "ctx": Context(an_int, an_int_async),
        },
        input_channels="input",
        output_channels=["inbox", "output"],
        stream_channels=["inbox", "output"],
    )

    async def aenumerate(aiter: AsyncIterator[Any]) -> AsyncIterator[tuple[int, Any]]:
        i = 0
        async for chunk in aiter:
            yield i, chunk
            i += 1

    assert setup_sync.call_count == 0
    assert cleanup_sync.call_count == 0
    assert setup_async.call_count == 0
    assert cleanup_async.call_count == 0
    async for i, chunk in aenumerate(app.astream(2)):
        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"
        if i == 0:
            assert chunk == {"inbox": [3]}
        elif i == 1:
            assert chunk == {"output": 4}
        else:
            assert False, "Expected only two chunks"
    assert setup_sync.call_count == 0
    assert cleanup_sync.call_count == 0
    assert setup_async.call_count == 1, "Expected setup to be called once"
    assert cleanup_async.call_count == 1, "Expected cleanup to be called once"


async def test_conditional_entrypoint_graph() -> None:
    async def left(data: str) -> str:
        return data + "->left"

    async def right(data: str) -> str:
        return data + "->right"

    def should_start(data: str) -> str:
        # Logic to decide where to start
        if len(data) > 10:
            return "go-right"
        else:
            return "go-left"

    # Define a new graph
    workflow = Graph()

    workflow.add_node("left", left)
    workflow.add_node("right", right)

    workflow.set_conditional_entry_point(
        should_start, {"go-left": "left", "go-right": "right"}
    )

    workflow.add_conditional_edges("left", lambda data: END)
    workflow.add_edge("right", END)

    app = workflow.compile()

    assert await app.ainvoke("what is weather in sf") == "what is weather in sf->right"

    assert [c async for c in app.astream("what is weather in sf")] == [
        {"right": "what is weather in sf->right"},
    ]


async def test_conditional_entrypoint_graph_state() -> None:
    class AgentState(TypedDict, total=False):
        input: str
        output: str
        steps: Annotated[list[str], operator.add]

    async def left(data: AgentState) -> AgentState:
        return {"output": data["input"] + "->left"}

    async def right(data: AgentState) -> AgentState:
        return {"output": data["input"] + "->right"}

    def should_start(data: AgentState) -> str:
        assert data["steps"] == [], "Expected input to be read from the state"
        # Logic to decide where to start
        if len(data["input"]) > 10:
            return "go-right"
        else:
            return "go-left"

    # Define a new graph
    workflow = StateGraph(AgentState)

    workflow.add_node("left", left)
    workflow.add_node("right", right)

    workflow.set_conditional_entry_point(
        should_start, {"go-left": "left", "go-right": "right"}
    )

    workflow.add_conditional_edges("left", lambda data: END)
    workflow.add_edge("right", END)

    app = workflow.compile()

    assert await app.ainvoke({"input": "what is weather in sf"}) == {
        "input": "what is weather in sf",
        "output": "what is weather in sf->right",
        "steps": [],
    }

    assert [c async for c in app.astream({"input": "what is weather in sf"})] == [
        {"right": {"output": "what is weather in sf->right"}},
    ]


@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_in_one_fan_out_state_graph_waiting_edge(checkpointer_name: str) -> None:
    def sorted_add(
        x: list[str], y: Union[list[str], list[tuple[str, str]]]
    ) -> list[str]:
        if isinstance(y[0], tuple):
            for rem, _ in y:
                x.remove(rem)
            y = [t[1] for t in y]
        return sorted(operator.add(x, y))

    class State(TypedDict, total=False):
        query: str
        answer: str
        docs: Annotated[list[str], sorted_add]

    async def rewrite_query(data: State) -> State:
        return {"query": f"query: {data['query']}"}

    async def analyzer_one(data: State) -> State:
        return {"query": f"analyzed: {data['query']}"}

    async def retriever_one(data: State) -> State:
        return {"docs": ["doc1", "doc2"]}

    async def retriever_two(data: State) -> State:
        await asyncio.sleep(0.1)
        return {"docs": ["doc3", "doc4"]}

    async def qa(data: State) -> State:
        return {"answer": ",".join(data["docs"])}

    workflow = StateGraph(State)

    workflow.add_node("rewrite_query", rewrite_query)
    workflow.add_node("analyzer_one", analyzer_one)
    workflow.add_node("retriever_one", retriever_one)
    workflow.add_node("retriever_two", retriever_two)
    workflow.add_node("qa", qa)

    workflow.set_entry_point("rewrite_query")
    workflow.add_edge("rewrite_query", "analyzer_one")
    workflow.add_edge("analyzer_one", "retriever_one")
    workflow.add_edge("rewrite_query", "retriever_two")
    workflow.add_edge(["retriever_one", "retriever_two"], "qa")
    workflow.set_finish_point("qa")

    app = workflow.compile()

    assert await app.ainvoke({"query": "what is weather in sf"}) == {
        "query": "analyzed: query: what is weather in sf",
        "docs": ["doc1", "doc2", "doc3", "doc4"],
        "answer": "doc1,doc2,doc3,doc4",
    }

    assert [c async for c in app.astream({"query": "what is weather in sf"})] == [
        {"rewrite_query": {"query": "query: what is weather in sf"}},
        {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
        {"retriever_two": {"docs": ["doc3", "doc4"]}},
        {"retriever_one": {"docs": ["doc1", "doc2"]}},
        {"qa": {"answer": "doc1,doc2,doc3,doc4"}},
    ]

    async with awith_checkpointer(checkpointer_name) as checkpointer:
        app_w_interrupt = workflow.compile(
            checkpointer=checkpointer,
            interrupt_after=["retriever_one"],
        )
        config = {"configurable": {"thread_id": "1"}}

        assert [
            c
            async for c in app_w_interrupt.astream(
                {"query": "what is weather in sf"}, config
            )
        ] == [
            {"rewrite_query": {"query": "query: what is weather in sf"}},
            {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
            {"retriever_two": {"docs": ["doc3", "doc4"]}},
            {"retriever_one": {"docs": ["doc1", "doc2"]}},
            {"__interrupt__": ()},
        ]

        assert [c async for c in app_w_interrupt.astream(None, config)] == [
            {"qa": {"answer": "doc1,doc2,doc3,doc4"}},
        ]


@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_in_one_fan_out_state_graph_waiting_edge_via_branch(
    snapshot: SnapshotAssertion, checkpointer_name: str
) -> None:
    def sorted_add(
        x: list[str], y: Union[list[str], list[tuple[str, str]]]
    ) -> list[str]:
        if isinstance(y[0], tuple):
            for rem, _ in y:
                x.remove(rem)
            y = [t[1] for t in y]
        return sorted(operator.add(x, y))

    class State(TypedDict, total=False):
        query: str
        answer: str
        docs: Annotated[list[str], sorted_add]

    async def rewrite_query(data: State) -> State:
        return {"query": f"query: {data['query']}"}

    async def analyzer_one(data: State) -> State:
        return {"query": f"analyzed: {data['query']}"}

    async def retriever_one(data: State) -> State:
        return {"docs": ["doc1", "doc2"]}

    async def retriever_two(data: State) -> State:
        await asyncio.sleep(0.1)
        return {"docs": ["doc3", "doc4"]}

    async def qa(data: State) -> State:
        return {"answer": ",".join(data["docs"])}

    workflow = StateGraph(State)

    workflow.add_node("rewrite_query", rewrite_query)
    workflow.add_node("analyzer_one", analyzer_one)
    workflow.add_node("retriever_one", retriever_one)
    workflow.add_node("retriever_two", retriever_two)
    workflow.add_node("qa", qa)

    workflow.set_entry_point("rewrite_query")
    workflow.add_edge("rewrite_query", "analyzer_one")
    workflow.add_edge("analyzer_one", "retriever_one")
    workflow.add_conditional_edges(
        "rewrite_query", lambda _: "retriever_two", {"retriever_two": "retriever_two"}
    )
    workflow.add_edge(["retriever_one", "retriever_two"], "qa")
    workflow.set_finish_point("qa")

    app = workflow.compile()

    assert await app.ainvoke({"query": "what is weather in sf"}, debug=True) == {
        "query": "analyzed: query: what is weather in sf",
        "docs": ["doc1", "doc2", "doc3", "doc4"],
        "answer": "doc1,doc2,doc3,doc4",
    }

    assert [c async for c in app.astream({"query": "what is weather in sf"})] == [
        {"rewrite_query": {"query": "query: what is weather in sf"}},
        {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
        {"retriever_two": {"docs": ["doc3", "doc4"]}},
        {"retriever_one": {"docs": ["doc1", "doc2"]}},
        {"qa": {"answer": "doc1,doc2,doc3,doc4"}},
    ]

    async with awith_checkpointer(checkpointer_name) as checkpointer:
        app_w_interrupt = workflow.compile(
            checkpointer=checkpointer,
            interrupt_after=["retriever_one"],
        )
        config = {"configurable": {"thread_id": "1"}}

        assert [
            c
            async for c in app_w_interrupt.astream(
                {"query": "what is weather in sf"}, config
            )
        ] == [
            {"rewrite_query": {"query": "query: what is weather in sf"}},
            {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
            {"retriever_two": {"docs": ["doc3", "doc4"]}},
            {"retriever_one": {"docs": ["doc1", "doc2"]}},
            {"__interrupt__": ()},
        ]

        assert [c async for c in app_w_interrupt.astream(None, config)] == [
            {"qa": {"answer": "doc1,doc2,doc3,doc4"}},
        ]


@pytest.mark.parametrize("version", ["v1", "v2"])
async def test_nested_pydantic_models(version: str) -> None:
    """Test that nested Pydantic models are properly constructed from leaf nodes up."""

    # Define nested Pydantic models
    if version == "v1":
        from pydantic.v1 import BaseModel, Field
    else:
        from pydantic import BaseModel, Field

    class NestedModel(BaseModel):
        value: int
        name: str
        something: Optional[str] = None

    # Forward reference model
    class RecursiveModel(BaseModel):
        value: str
        child: Optional["RecursiveModel"] = None

    # Discriminated union models
    class Cat(BaseModel):
        pet_type: Literal["cat"]
        meow: str

    class Dog(BaseModel):
        pet_type: Literal["dog"]
        bark: str

    # Cyclic reference model
    class Person(BaseModel):
        id: str
        name: str
        friends: list[str] = Field(default_factory=list)  # IDs of friends

    class MyEnum(enum.Enum):
        A = 1
        B = 2

    class MyTypedDict(TypedDict):
        x: int
        my_enum: MyEnum

    class State(BaseModel):
        # Basic nested model tests
        top_level: str
        nested: NestedModel
        optional_nested: Optional[NestedModel] = None
        dict_nested: dict[str, NestedModel]
        my_set: set[int]
        another_set: set
        my_enum: MyEnum
        list_nested: Annotated[
            Union[dict, list[dict[str, NestedModel]]], lambda x, y: (x or []) + [y]
        ]
        list_nested_reversed: Annotated[
            Union[list[dict[str, NestedModel]], NestedModel, dict, list],
            lambda x, y: (x or []) + [y],
        ]
        tuple_nested: tuple[str, NestedModel]
        tuple_list_nested: list[tuple[int, NestedModel]]
        complex_tuple: tuple[str, dict[str, tuple[int, NestedModel]]]
        my_typed_dict: MyTypedDict

        # Forward reference test
        recursive: RecursiveModel

        # Discriminated union test
        pet: Union[Cat, Dog]

        # Cyclic reference test
        people: dict[str, Person]  # Map of ID -> Person

    inputs = {
        # Basic nested models
        "top_level": "initial",
        "nested": {"value": 42, "name": "test"},
        "optional_nested": {"value": 10, "name": "optional"},
        "my_set": [1, 2, 7],
        "another_set": ["foo", 3],
        "my_enum": MyEnum.B,
        "my_typed_dict": {"x": 1, "my_enum": MyEnum.A},
        "dict_nested": {"a": {"value": 5, "name": "a"}},
        "list_nested": [{"a": {"value": 6, "name": "b"}}],
        "list_nested_reversed": ["foo", "bar"],
        "tuple_nested": ["tuple-key", {"value": 7, "name": "tuple-value"}],
        "tuple_list_nested": [[1, {"value": 8, "name": "tuple-in-list"}]],
        "complex_tuple": [
            "complex",
            {"nested": [9, {"value": 10, "name": "deep"}]},
        ],
        # Forward reference
        "recursive": {"value": "parent", "child": {"value": "child", "child": None}},
        # Discriminated union (using a cat in this case)
        "pet": {"pet_type": "cat", "meow": "meow!"},
        # Cyclic references
        "people": {
            "1": {
                "id": "1",
                "name": "Alice",
                "friends": ["2", "3"],  # Alice is friends with Bob and Charlie
            },
            "2": {
                "id": "2",
                "name": "Bob",
                "friends": ["1"],  # Bob is friends with Alice
            },
            "3": {
                "id": "3",
                "name": "Charlie",
                "friends": ["1", "2"],  # Charlie is friends with Alice and Bob
            },
        },
    }

    update = {"top_level": "updated", "nested": {"value": 100, "name": "updated"}}

    async def node_fn(state: State) -> dict:
        assert state == State(**inputs)
        return update

    builder = StateGraph(State)
    builder.add_node("process", node_fn)
    builder.set_entry_point("process")
    builder.set_finish_point("process")
    graph = builder.compile()

    result = await graph.ainvoke(inputs.copy())

    assert result == {**inputs, **update}


@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class(
    snapshot: SnapshotAssertion, mocker: MockerFixture, checkpointer_name: str
) -> None:
    from pydantic.v1 import BaseModel, ValidationError

    setup = mocker.Mock()
    teardown = mocker.Mock()

    @asynccontextmanager
    async def assert_ctx_once() -> AsyncIterator[None]:
        assert setup.call_count == 0
        assert teardown.call_count == 0
        try:
            yield
        finally:
            assert setup.call_count == 1
            assert teardown.call_count == 1
            setup.reset_mock()
            teardown.reset_mock()

    @asynccontextmanager
    async def make_httpx_client() -> AsyncIterator[httpx.AsyncClient]:
        setup()
        async with httpx.AsyncClient() as client:
            try:
                yield client
            finally:
                teardown()

    def sorted_add(
        x: list[str], y: Union[list[str], list[tuple[str, str]]]
    ) -> list[str]:
        if isinstance(y[0], tuple):
            for rem, _ in y:
                x.remove(rem)
            y = [t[1] for t in y]
        return sorted(operator.add(x, y))

    class State(BaseModel):
        class Config:
            arbitrary_types_allowed = True

        query: str
        answer: Optional[str] = None
        docs: Annotated[list[str], sorted_add]
        client: Annotated[httpx.AsyncClient, Context(make_httpx_client)]

    class Input(BaseModel):
        query: str

    class Output(BaseModel):
        answer: str
        docs: list[str]

    class StateUpdate(BaseModel):
        query: Optional[str] = None
        answer: Optional[str] = None
        docs: Optional[list[str]] = None

    async def rewrite_query(data: State) -> State:
        return {"query": f"query: {data.query}"}

    async def analyzer_one(data: State) -> State:
        return StateUpdate(query=f"analyzed: {data.query}")

    async def retriever_one(data: State) -> State:
        return {"docs": ["doc1", "doc2"]}

    async def retriever_two(data: State) -> State:
        await asyncio.sleep(0.1)
        return {"docs": ["doc3", "doc4"]}

    async def qa(data: State) -> State:
        return {"answer": ",".join(data.docs)}

    async def decider(data: State) -> str:
        assert isinstance(data, State)
        return "retriever_two"

    workflow = StateGraph(State, input=Input, output=Output)

    workflow.add_node("rewrite_query", rewrite_query)
    workflow.add_node("analyzer_one", analyzer_one)
    workflow.add_node("retriever_one", retriever_one)
    workflow.add_node("retriever_two", retriever_two)
    workflow.add_node("qa", qa)

    workflow.set_entry_point("rewrite_query")
    workflow.add_edge("rewrite_query", "analyzer_one")
    workflow.add_edge("analyzer_one", "retriever_one")
    workflow.add_conditional_edges(
        "rewrite_query", decider, {"retriever_two": "retriever_two"}
    )
    workflow.add_edge(["retriever_one", "retriever_two"], "qa")
    workflow.set_finish_point("qa")

    app = workflow.compile()

    async with assert_ctx_once():
        with pytest.raises(ValidationError):
            await app.ainvoke({"query": {}})

    async with assert_ctx_once():
        assert await app.ainvoke({"query": "what is weather in sf"}) == {
            "docs": ["doc1", "doc2", "doc3", "doc4"],
            "answer": "doc1,doc2,doc3,doc4",
        }

    async with assert_ctx_once():
        assert [c async for c in app.astream({"query": "what is weather in sf"})] == [
            {"rewrite_query": {"query": "query: what is weather in sf"}},
            {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
            {"retriever_two": {"docs": ["doc3", "doc4"]}},
            {"retriever_one": {"docs": ["doc1", "doc2"]}},
            {"qa": {"answer": "doc1,doc2,doc3,doc4"}},
        ]

    async with awith_checkpointer(checkpointer_name) as checkpointer:
        app_w_interrupt = workflow.compile(
            checkpointer=checkpointer,
            interrupt_after=["retriever_one"],
        )
        config = {"configurable": {"thread_id": "1"}}

        async with assert_ctx_once():
            assert [
                c
                async for c in app_w_interrupt.astream(
                    {"query": "what is weather in sf"}, config
                )
            ] == [
                {"rewrite_query": {"query": "query: what is weather in sf"}},
                {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
                {"retriever_two": {"docs": ["doc3", "doc4"]}},
                {"retriever_one": {"docs": ["doc1", "doc2"]}},
                {"__interrupt__": ()},
            ]

        async with assert_ctx_once():
            assert [c async for c in app_w_interrupt.astream(None, config)] == [
                {"qa": {"answer": "doc1,doc2,doc3,doc4"}},
            ]

        assert await app_w_interrupt.aget_state(config) == StateSnapshot(
            values={
                "query": "analyzed: query: what is weather in sf",
                "answer": "doc1,doc2,doc3,doc4",
                "docs": ["doc1", "doc2", "doc3", "doc4"],
            },
            tasks=(),
            next=(),
            config={
                "configurable": {
                    "thread_id": "1",
                    "checkpoint_ns": "",
                    "checkpoint_id": AnyStr(),
                }
            },
            metadata={
                "parents": {},
                "source": "loop",
                "writes": {"qa": {"answer": "doc1,doc2,doc3,doc4"}},
                "step": 4,
                "thread_id": "1",
            },
            created_at=AnyStr(),
            parent_config=(
                None
                if "shallow" in checkpointer_name
                else {
                    "configurable": {
                        "thread_id": "1",
                        "checkpoint_ns": "",
                        "checkpoint_id": AnyStr(),
                    }
                }
            ),
        )

        async with assert_ctx_once():
            assert await app_w_interrupt.aupdate_state(
                config, {"docs": ["doc5"]}, as_node="rewrite_query"
            ) == {
                "configurable": {
                    "thread_id": "1",
                    "checkpoint_id": AnyStr(),
                    "checkpoint_ns": "",
                }
            }


@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2(
    snapshot: SnapshotAssertion, checkpointer_name: str
) -> None:
    from pydantic import BaseModel, ValidationError

    def sorted_add(
        x: list[str], y: Union[list[str], list[tuple[str, str]]]
    ) -> list[str]:
        if isinstance(y[0], tuple):
            for rem, _ in y:
                x.remove(rem)
            y = [t[1] for t in y]
        return sorted(operator.add(x, y))

    class InnerObject(BaseModel):
        yo: int

    class State(BaseModel):
        query: str
        inner: InnerObject
        answer: Optional[str] = None
        docs: Annotated[list[str], sorted_add]

    class StateUpdate(BaseModel):
        query: Optional[str] = None
        answer: Optional[str] = None
        docs: Optional[list[str]] = None

    async def rewrite_query(data: State) -> State:
        return {"query": f"query: {data.query}"}

    async def analyzer_one(data: State) -> State:
        return StateUpdate(query=f"analyzed: {data.query}")

    async def retriever_one(data: State) -> State:
        return {"docs": ["doc1", "doc2"]}

    async def retriever_two(data: State) -> State:
        await asyncio.sleep(0.1)
        return {"docs": ["doc3", "doc4"]}

    async def qa(data: State) -> State:
        return {"answer": ",".join(data.docs)}

    async def decider(data: State) -> str:
        assert isinstance(data, State)
        return "retriever_two"

    workflow = StateGraph(State)

    workflow.add_node("rewrite_query", rewrite_query)
    workflow.add_node("analyzer_one", analyzer_one)
    workflow.add_node("retriever_one", retriever_one)
    workflow.add_node("retriever_two", retriever_two)
    workflow.add_node("qa", qa)

    workflow.set_entry_point("rewrite_query")
    workflow.add_edge("rewrite_query", "analyzer_one")
    workflow.add_edge("analyzer_one", "retriever_one")
    workflow.add_conditional_edges(
        "rewrite_query", decider, {"retriever_two": "retriever_two"}
    )
    workflow.add_edge(["retriever_one", "retriever_two"], "qa")
    workflow.set_finish_point("qa")

    app = workflow.compile()

    if SHOULD_CHECK_SNAPSHOTS and checkpointer_name == "memory":
        assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
        assert app.get_input_schema().model_json_schema() == snapshot
        assert app.get_output_schema().model_json_schema() == snapshot

    with pytest.raises(ValidationError):
        await app.ainvoke({"query": {}})

    assert await app.ainvoke(
        {"query": "what is weather in sf", "inner": {"yo": 1}}
    ) == {
        "query": "analyzed: query: what is weather in sf",
        "docs": ["doc1", "doc2", "doc3", "doc4"],
        "answer": "doc1,doc2,doc3,doc4",
        "inner": {"yo": 1},
    }

    assert [
        c
        async for c in app.astream(
            {"query": "what is weather in sf", "inner": {"yo": 1}}
        )
    ] == [
        {"rewrite_query": {"query": "query: what is weather in sf"}},
        {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
        {"retriever_two": {"docs": ["doc3", "doc4"]}},
        {"retriever_one": {"docs": ["doc1", "doc2"]}},
        {"qa": {"answer": "doc1,doc2,doc3,doc4"}},
    ]

    async with awith_checkpointer(checkpointer_name) as checkpointer:
        app_w_interrupt = workflow.compile(
            checkpointer=checkpointer,
            interrupt_after=["retriever_one"],
        )
        config = {"configurable": {"thread_id": "1"}}

        assert [
            c
            async for c in app_w_interrupt.astream(
                {"query": "what is weather in sf", "inner": {"yo": 1}}, config
            )
        ] == [
            {"rewrite_query": {"query": "query: what is weather in sf"}},
            {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
            {"retriever_two": {"docs": ["doc3", "doc4"]}},
            {"retriever_one": {"docs": ["doc1", "doc2"]}},
            {"__interrupt__": ()},
        ]

        assert [c async for c in app_w_interrupt.astream(None, config)] == [
            {"qa": {"answer": "doc1,doc2,doc3,doc4"}},
        ]

        assert await app_w_interrupt.aupdate_state(
            config, {"docs": ["doc5"]}, as_node="rewrite_query"
        ) == {
            "configurable": {
                "thread_id": "1",
                "checkpoint_id": AnyStr(),
                "checkpoint_ns": "",
            }
        }


@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_in_one_fan_out_state_graph_waiting_edge_plus_regular(
    checkpointer_name: str,
) -> None:
    def sorted_add(
        x: list[str], y: Union[list[str], list[tuple[str, str]]]
    ) -> list[str]:
        if isinstance(y[0], tuple):
            for rem, _ in y:
                x.remove(rem)
            y = [t[1] for t in y]
        return sorted(operator.add(x, y))

    class State(TypedDict, total=False):
        query: str
        answer: str
        docs: Annotated[list[str], sorted_add]

    async def rewrite_query(data: State) -> State:
        return {"query": f"query: {data['query']}"}

    async def analyzer_one(data: State) -> State:
        await asyncio.sleep(0.1)
        return {"query": f"analyzed: {data['query']}"}

    async def retriever_one(data: State) -> State:
        return {"docs": ["doc1", "doc2"]}

    async def retriever_two(data: State) -> State:
        await asyncio.sleep(0.2)
        return {"docs": ["doc3", "doc4"]}

    async def qa(data: State) -> State:
        return {"answer": ",".join(data["docs"])}

    workflow = StateGraph(State)

    workflow.add_node("rewrite_query", rewrite_query)
    workflow.add_node("analyzer_one", analyzer_one)
    workflow.add_node("retriever_one", retriever_one)
    workflow.add_node("retriever_two", retriever_two)
    workflow.add_node("qa", qa)

    workflow.set_entry_point("rewrite_query")
    workflow.add_edge("rewrite_query", "analyzer_one")
    workflow.add_edge("analyzer_one", "retriever_one")
    workflow.add_edge("rewrite_query", "retriever_two")
    workflow.add_edge(["retriever_one", "retriever_two"], "qa")
    workflow.set_finish_point("qa")

    # silly edge, to make sure having been triggered before doesn't break
    # semantics of named barrier (== waiting edges)
    workflow.add_edge("rewrite_query", "qa")

    app = workflow.compile()

    assert await app.ainvoke({"query": "what is weather in sf"}) == {
        "query": "analyzed: query: what is weather in sf",
        "docs": ["doc1", "doc2", "doc3", "doc4"],
        "answer": "doc1,doc2,doc3,doc4",
    }

    assert [c async for c in app.astream({"query": "what is weather in sf"})] == [
        {"rewrite_query": {"query": "query: what is weather in sf"}},
        {"qa": {"answer": ""}},
        {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
        {"retriever_two": {"docs": ["doc3", "doc4"]}},
        {"retriever_one": {"docs": ["doc1", "doc2"]}},
        {"qa": {"answer": "doc1,doc2,doc3,doc4"}},
    ]

    async with awith_checkpointer(checkpointer_name) as checkpointer:
        app_w_interrupt = workflow.compile(
            checkpointer=checkpointer,
            interrupt_after=["retriever_one"],
        )
        config = {"configurable": {"thread_id": "1"}}

        assert [
            c
            async for c in app_w_interrupt.astream(
                {"query": "what is weather in sf"}, config
            )
        ] == [
            {"rewrite_query": {"query": "query: what is weather in sf"}},
            {"qa": {"answer": ""}},
            {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
            {"retriever_two": {"docs": ["doc3", "doc4"]}},
            {"retriever_one": {"docs": ["doc1", "doc2"]}},
            {"__interrupt__": ()},
        ]

        assert [c async for c in app_w_interrupt.astream(None, config)] == [
            {"qa": {"answer": "doc1,doc2,doc3,doc4"}},
        ]


async def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None:
    def sorted_add(
        x: list[str], y: Union[list[str], list[tuple[str, str]]]
    ) -> list[str]:
        if isinstance(y[0], tuple):
            for rem, _ in y:
                x.remove(rem)
            y = [t[1] for t in y]
        return sorted(operator.add(x, y))

    class State(TypedDict, total=False):
        query: str
        answer: str
        docs: Annotated[list[str], sorted_add]

    async def rewrite_query(data: State) -> State:
        return {"query": f"query: {data['query']}"}

    async def analyzer_one(data: State) -> State:
        return {"query": f"analyzed: {data['query']}"}

    async def retriever_one(data: State) -> State:
        return {"docs": ["doc1", "doc2"]}

    async def retriever_two(data: State) -> State:
        await asyncio.sleep(0.1)
        return {"docs": ["doc3", "doc4"]}

    async def qa(data: State) -> State:
        return {"answer": ",".join(data["docs"])}

    async def decider(data: State) -> None:
        return None

    def decider_cond(data: State) -> str:
        if data["query"].count("analyzed") > 1:
            return "qa"
        else:
            return "rewrite_query"

    workflow = StateGraph(State)

    workflow.add_node("rewrite_query", rewrite_query)
    workflow.add_node("analyzer_one", analyzer_one)
    workflow.add_node("retriever_one", retriever_one)
    workflow.add_node("retriever_two", retriever_two)
    workflow.add_node("decider", decider)
    workflow.add_node("qa", qa)

    workflow.set_entry_point("rewrite_query")
    workflow.add_edge("rewrite_query", "analyzer_one")
    workflow.add_edge("analyzer_one", "retriever_one")
    workflow.add_edge("rewrite_query", "retriever_two")
    workflow.add_edge(["retriever_one", "retriever_two"], "decider")
    workflow.add_conditional_edges("decider", decider_cond)
    workflow.set_finish_point("qa")

    app = workflow.compile()

    assert await app.ainvoke({"query": "what is weather in sf"}) == {
        "query": "analyzed: query: analyzed: query: what is weather in sf",
        "answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4",
        "docs": ["doc1", "doc1", "doc2", "doc2", "doc3", "doc3", "doc4", "doc4"],
    }

    assert [c async for c in app.astream({"query": "what is weather in sf"})] == [
        {"rewrite_query": {"query": "query: what is weather in sf"}},
        {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
        {"retriever_two": {"docs": ["doc3", "doc4"]}},
        {"retriever_one": {"docs": ["doc1", "doc2"]}},
        {"decider": None},
        {"rewrite_query": {"query": "query: analyzed: query: what is weather in sf"}},
        {
            "analyzer_one": {
                "query": "analyzed: query: analyzed: query: what is weather in sf"
            }
        },
        {"retriever_two": {"docs": ["doc3", "doc4"]}},
        {"retriever_one": {"docs": ["doc1", "doc2"]}},
        {"decider": None},
        {"qa": {"answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4"}},
    ]


async def test_in_one_fan_out_state_graph_waiting_edge_multiple_cond_edge() -> None:
    def sorted_add(
        x: list[str], y: Union[list[str], list[tuple[str, str]]]
    ) -> list[str]:
        if isinstance(y[0], tuple):
            for rem, _ in y:
                x.remove(rem)
            y = [t[1] for t in y]
        return sorted(operator.add(x, y))

    class State(TypedDict, total=False):
        query: str
        answer: str
        docs: Annotated[list[str], sorted_add]

    async def rewrite_query(data: State) -> State:
        return {"query": f"query: {data['query']}"}

    async def retriever_picker(data: State) -> list[str]:
        return ["analyzer_one", "retriever_two"]

    async def analyzer_one(data: State) -> State:
        return {"query": f"analyzed: {data['query']}"}

    async def retriever_one(data: State) -> State:
        return {"docs": ["doc1", "doc2"]}

    async def retriever_two(data: State) -> State:
        await asyncio.sleep(0.1)
        return {"docs": ["doc3", "doc4"]}

    async def qa(data: State) -> State:
        return {"answer": ",".join(data["docs"])}

    async def decider(data: State) -> None:
        return None

    def decider_cond(data: State) -> str:
        if data["query"].count("analyzed") > 1:
            return "qa"
        else:
            return "rewrite_query"

    workflow = StateGraph(State)

    workflow.add_node("rewrite_query", rewrite_query)
    workflow.add_node("analyzer_one", analyzer_one)
    workflow.add_node("retriever_one", retriever_one)
    workflow.add_node("retriever_two", retriever_two)
    workflow.add_node("decider", decider)
    workflow.add_node("qa", qa)

    workflow.set_entry_point("rewrite_query")
    workflow.add_conditional_edges("rewrite_query", retriever_picker)
    workflow.add_edge("analyzer_one", "retriever_one")
    workflow.add_edge(["retriever_one", "retriever_two"], "decider")
    workflow.add_conditional_edges("decider", decider_cond)
    workflow.set_finish_point("qa")

    app = workflow.compile()

    assert await app.ainvoke({"query": "what is weather in sf"}) == {
        "query": "analyzed: query: analyzed: query: what is weather in sf",
        "answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4",
        "docs": ["doc1", "doc1", "doc2", "doc2", "doc3", "doc3", "doc4", "doc4"],
    }

    assert [c async for c in app.astream({"query": "what is weather in sf"})] == [
        {"rewrite_query": {"query": "query: what is weather in sf"}},
        {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
        {"retriever_two": {"docs": ["doc3", "doc4"]}},
        {"retriever_one": {"docs": ["doc1", "doc2"]}},
        {"decider": None},
        {"rewrite_query": {"query": "query: analyzed: query: what is weather in sf"}},
        {
            "analyzer_one": {
                "query": "analyzed: query: analyzed: query: what is weather in sf"
            }
        },
        {"retriever_two": {"docs": ["doc3", "doc4"]}},
        {"retriever_one": {"docs": ["doc1", "doc2"]}},
        {"decider": None},
        {"qa": {"answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4"}},
    ]


async def test_nested_graph(snapshot: SnapshotAssertion) -> None:
    def never_called_fn(state: Any):
        assert 0, "This function should never be called"

    never_called = RunnableLambda(never_called_fn)

    class InnerState(TypedDict):
        my_key: str
        my_other_key: str

    def up(state: InnerState):
        return {"my_key": state["my_key"] + " there", "my_other_key": state["my_key"]}

    inner = StateGraph(InnerState)
    inner.add_node("up", up)
    inner.set_entry_point("up")
    inner.set_finish_point("up")

    class State(TypedDict):
        my_key: str
        never_called: Any

    async def side(state: State):
        return {"my_key": state["my_key"] + " and back again"}

    graph = StateGraph(State)
    graph.add_node("inner", inner.compile())
    graph.add_node("side", side)
    graph.set_entry_point("inner")
    graph.add_edge("inner", "side")
    graph.set_finish_point("side")

    app = graph.compile()

    assert await app.ainvoke({"my_key": "my value", "never_called": never_called}) == {
        "my_key": "my value there and back again",
        "never_called": never_called,
    }
    assert [
        chunk
        async for chunk in app.astream(
            {"my_key": "my value", "never_called": never_called}
        )
    ] == [
        {"inner": {"my_key": "my value there"}},
        {"side": {"my_key": "my value there and back again"}},
    ]
    assert [
        chunk
        async for chunk in app.astream(
            {"my_key": "my value", "never_called": never_called}, stream_mode="values"
        )
    ] == [
        {"my_key": "my value", "never_called": never_called},
        {"my_key": "my value there", "never_called": never_called},
        {"my_key": "my value there and back again", "never_called": never_called},
    ]
    times_called = 0
    async for event in app.astream_events(
        {"my_key": "my value", "never_called": never_called},
        version="v2",
        config={"run_id": UUID(int=0)},
        stream_mode="values",
    ):
        if event["event"] == "on_chain_end" and event["run_id"] == str(UUID(int=0)):
            times_called += 1
            assert event["data"] == {
                "output": {
                    "my_key": "my value there and back again",
                    "never_called": never_called,
                }
            }
    assert times_called == 1
    times_called = 0
    async for event in app.astream_events(
        {"my_key": "my value", "never_called": never_called},
        version="v2",
        config={"run_id": UUID(int=0)},
    ):
        if event["event"] == "on_chain_end" and event["run_id"] == str(UUID(int=0)):
            times_called += 1
            assert event["data"] == {
                "output": {
                    "my_key": "my value there and back again",
                    "never_called": never_called,
                }
            }
    assert times_called == 1

    chain = app | RunnablePassthrough()

    assert await chain.ainvoke(
        {"my_key": "my value", "never_called": never_called}
    ) == {
        "my_key": "my value there and back again",
        "never_called": never_called,
    }
    assert [
        chunk
        async for chunk in chain.astream(
            {"my_key": "my value", "never_called": never_called}
        )
    ] == [
        {"inner": {"my_key": "my value there"}},
        {"side": {"my_key": "my value there and back again"}},
    ]
    times_called = 0
    async for event in chain.astream_events(
        {"my_key": "my value", "never_called": never_called},
        version="v2",
        config={"run_id": UUID(int=0)},
    ):
        if event["event"] == "on_chain_end" and event["run_id"] == str(UUID(int=0)):
            times_called += 1
            assert event["data"] == {
                "output": [
                    {"inner": {"my_key": "my value there"}},
                    {"side": {"my_key": "my value there and back again"}},
                ]
            }
    assert times_called == 1


@pytest.mark.parametrize("checkpoint_during", [True, False])
@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC)
async def test_subgraph_checkpoint_true(
    checkpointer_name: str, checkpoint_during: bool
) -> None:
    class InnerState(TypedDict):
        my_key: Annotated[str, operator.add]
        my_other_key: str

    def inner_1(state: InnerState):
        return {"my_key": " got here", "my_other_key": state["my_key"]}

    def inner_2(state: InnerState):
        return {"my_key": " and there"}

    inner = StateGraph(InnerState)
    inner.add_node("inner_1", inner_1)
    inner.add_node("inner_2", inner_2)
    inner.add_edge("inner_1", "inner_2")
    inner.set_entry_point("inner_1")
    inner.set_finish_point("inner_2")

    class State(TypedDict):
        my_key: str

    graph = StateGraph(State)
    graph.add_node("inner", inner.compile(checkpointer=True))
    graph.add_edge(START, "inner")
    graph.add_conditional_edges(
        "inner", lambda s: "inner" if s["my_key"].count("there") < 2 else END
    )

    async with awith_checkpointer(checkpointer_name) as checkpointer:
        app = graph.compile(checkpointer=checkpointer)

        config = {"configurable": {"thread_id": "2"}}
        assert [
            c
            async for c in app.astream(
                {"my_key": ""},
                config,
                subgraphs=True,
                checkpoint_during=checkpoint_during,
            )
        ] == [
            (("inner",), {"inner_1": {"my_key": " got here", "my_other_key": ""}}),
            (("inner",), {"inner_2": {"my_key": " and there"}}),
            ((), {"inner": {"my_key": " got here and there"}}),
            (
                ("inner",),
                {
                    "inner_1": {
                        "my_key": " got here",
                        "my_other_key": " got here and there got here and there",
                    }
                },
            ),
            (("inner",), {"inner_2": {"my_key": " and there"}}),
            (
                (),
                {
                    "inner": {
                        "my_key": " got here and there got here and there got here and there"
                    }
                },
            ),
        ]


@NEEDS_CONTEXTVARS
@pytest.mark.parametrize("checkpoint_during", [True, False])
@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC)
async def test_subgraph_checkpoint_true_interrupt(
    checkpointer_name: str, checkpoint_during: bool
) -> None:
    # Define subgraph
    class SubgraphState(TypedDict):
        # note that none of these keys are shared with the parent graph state
        bar: str
        baz: str

    def subgraph_node_1(state: SubgraphState):
        baz_value = interrupt("Provide baz value")
        return {"baz": baz_value}

    def subgraph_node_2(state: SubgraphState):
        return {"bar": state["bar"] + state["baz"]}

    subgraph_builder = StateGraph(SubgraphState)
    subgraph_builder.add_node(subgraph_node_1)
    subgraph_builder.add_node(subgraph_node_2)
    subgraph_builder.add_edge(START, "subgraph_node_1")
    subgraph_builder.add_edge("subgraph_node_1", "subgraph_node_2")
    subgraph = subgraph_builder.compile(checkpointer=True)

    class ParentState(TypedDict):
        foo: str

    def node_1(state: ParentState):
        return {"foo": "hi! " + state["foo"]}

    async def node_2(state: ParentState, config: RunnableConfig):
        response = await subgraph.ainvoke({"bar": state["foo"]})
        return {"foo": response["bar"]}

    builder = StateGraph(ParentState)
    builder.add_node("node_1", node_1)
    builder.add_node("node_2", node_2)
    builder.add_edge(START, "node_1")
    builder.add_edge("node_1", "node_2")

    async with awith_checkpointer(checkpointer_name) as checkpointer:
        graph = builder.compile(checkpointer=checkpointer)
        config = {"configurable": {"thread_id": "1"}}

        assert await graph.ainvoke(
            {"foo": "foo"}, config, checkpoint_during=checkpoint_during
        ) == {"foo": "hi! foo"}
        assert (await graph.aget_state(config, subgraphs=True)).tasks[
            0
        ].state.values == {"bar": "hi! foo"}
        assert await graph.ainvoke(
            Command(resume="baz"), config, checkpoint_during=checkpoint_during
        ) == {"foo": "hi! foobaz"}


@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_stream_subgraphs_during_execution(checkpointer_name: str) -> None:
    class InnerState(TypedDict):
        my_key: Annotated[str, operator.add]
        my_other_key: str

    async def inner_1(state: InnerState):
        return {"my_key": "got here", "my_other_key": state["my_key"]}

    async def inner_2(state: InnerState):
        await asyncio.sleep(0.5)
        return {
            "my_key": " and there",
            "my_other_key": state["my_key"],
        }

    inner = StateGraph(InnerState)
    inner.add_node("inner_1", inner_1)
    inner.add_node("inner_2", inner_2)
    inner.add_edge("inner_1", "inner_2")
    inner.set_entry_point("inner_1")
    inner.set_finish_point("inner_2")

    class State(TypedDict):
        my_key: Annotated[str, operator.add]

    async def outer_1(state: State):
        await asyncio.sleep(0.2)
        return {"my_key": " and parallel"}

    async def outer_2(state: State):
        return {"my_key": " and back again"}

    graph = StateGraph(State)
    graph.add_node("inner", inner.compile())
    graph.add_node("outer_1", outer_1)
    graph.add_node("outer_2", outer_2)

    graph.add_edge(START, "inner")
    graph.add_edge(START, "outer_1")
    graph.add_edge(["inner", "outer_1"], "outer_2")
    graph.add_edge("outer_2", END)

    async with awith_checkpointer(checkpointer_name) as checkpointer:
        app = graph.compile(checkpointer=checkpointer)

        start = perf_counter()
        chunks: list[tuple[float, Any]] = []
        config = {"configurable": {"thread_id": "2"}}
        async for c in app.astream({"my_key": ""}, config, subgraphs=True):
            chunks.append((round(perf_counter() - start, 1), c))
        for idx in range(len(chunks)):
            elapsed, c = chunks[idx]
            chunks[idx] = (round(elapsed - chunks[0][0], 1), c)

        assert chunks == [
            # arrives before "inner" finishes
            (
                FloatBetween(0.0, 0.1),
                (
                    (AnyStr("inner:"),),
                    {"inner_1": {"my_key": "got here", "my_other_key": ""}},
                ),
            ),
            (FloatBetween(0.2, 0.4), ((), {"outer_1": {"my_key": " and parallel"}})),
            (
                FloatBetween(0.5, 0.8),
                (
                    (AnyStr("inner:"),),
                    {"inner_2": {"my_key": " and there", "my_other_key": "got here"}},
                ),
            ),
            (FloatBetween(0.5, 0.8), ((), {"inner": {"my_key": "got here and there"}})),
            (FloatBetween(0.5, 0.8), ((), {"outer_2": {"my_key": " and back again"}})),
        ]


@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_stream_buffering_single_node(checkpointer_name: str) -> None:
    class State(TypedDict):
        my_key: Annotated[str, operator.add]

    async def node(state: State, writer: StreamWriter):
        writer("Before sleep")
        await asyncio.sleep(0.2)
        writer("After sleep")
        return {"my_key": "got here"}

    builder = StateGraph(State)
    builder.add_node("node", node)
    builder.add_edge(START, "node")
    builder.add_edge("node", END)

    async with awith_checkpointer(checkpointer_name) as checkpointer:
        graph = builder.compile(checkpointer=checkpointer)

        start = perf_counter()
        chunks: list[tuple[float, Any]] = []
        config = {"configurable": {"thread_id": "2"}}
        async for c in graph.astream({"my_key": ""}, config, stream_mode="custom"):
            chunks.append((round(perf_counter() - start, 1), c))

        assert chunks == [
            (FloatBetween(0.0, 0.1), "Before sleep"),
            (FloatBetween(0.2, 0.3), "After sleep"),
        ]


@pytest.mark.parametrize("checkpoint_during", [True, False])
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_nested_graph_interrupts_parallel(
    checkpointer_name: str, checkpoint_during: bool
) -> None:
    class InnerState(TypedDict):
        my_key: Annotated[str, operator.add]
        my_other_key: str

    async def inner_1(state: InnerState):
        await asyncio.sleep(0.1)
        return {"my_key": "got here", "my_other_key": state["my_key"]}

    async def inner_2(state: InnerState):
        return {
            "my_key": " and there",
            "my_other_key": state["my_key"],
        }

    inner = StateGraph(InnerState)
    inner.add_node("inner_1", inner_1)
    inner.add_node("inner_2", inner_2)
    inner.add_edge("inner_1", "inner_2")
    inner.set_entry_point("inner_1")
    inner.set_finish_point("inner_2")

    class State(TypedDict):
        my_key: Annotated[str, operator.add]

    async def outer_1(state: State):
        return {"my_key": " and parallel"}

    async def outer_2(state: State):
        return {"my_key": " and back again"}

    graph = StateGraph(State)
    graph.add_node(
        "inner",
        inner.compile(interrupt_before=["inner_2"]),
    )
    graph.add_node("outer_1", outer_1)
    graph.add_node("outer_2", outer_2)

    graph.add_edge(START, "inner")
    graph.add_edge(START, "outer_1")
    graph.add_edge(["inner", "outer_1"], "outer_2")
    graph.set_finish_point("outer_2")

    async with awith_checkpointer(checkpointer_name) as checkpointer:
        app = graph.compile(checkpointer=checkpointer)

        # test invoke w/ nested interrupt
        config = {"configurable": {"thread_id": "1"}}
        assert await app.ainvoke(
            {"my_key": ""}, config, checkpoint_during=checkpoint_during
        ) == {
            "my_key": " and parallel",
        }

        assert await app.ainvoke(None, config, checkpoint_during=checkpoint_during) == {
            "my_key": "got here and there and parallel and back again",
        }

        # below combo of assertions is asserting two things
        # - outer_1 finishes before inner interrupts (because we see its output in stream, which only happens after node finishes)
        # - the writes of outer are persisted in 1st call and used in 2nd call, ie outer isn't called again (because we dont see outer_1 output again in 2nd stream)
        # test stream updates w/ nested interrupt
        config = {"configurable": {"thread_id": "2"}}
        assert [
            c
            async for c in app.astream(
                {"my_key": ""},
                config,
                subgraphs=True,
                checkpoint_during=checkpoint_during,
            )
        ] == [
            # we got to parallel node first
            ((), {"outer_1": {"my_key": " and parallel"}}),
            (
                (AnyStr("inner:"),),
                {"inner_1": {"my_key": "got here", "my_other_key": ""}},
            ),
            ((), {"__interrupt__": ()}),
        ]
        assert [
            c
            async for c in app.astream(
                None, config, checkpoint_during=checkpoint_during
            )
        ] == [
            {"outer_1": {"my_key": " and parallel"}, "__metadata__": {"cached": True}},
            {"inner": {"my_key": "got here and there"}},
            {"outer_2": {"my_key": " and back again"}},
        ]

        # test stream values w/ nested interrupt
        config = {"configurable": {"thread_id": "3"}}
        assert [
            c
            async for c in app.astream(
                {"my_key": ""},
                config,
                stream_mode="values",
                checkpoint_during=checkpoint_during,
            )
        ] == [
            {"my_key": ""},
            {"my_key": " and parallel"},
        ]
        assert [
            c
            async for c in app.astream(
                None, config, stream_mode="values", checkpoint_during=checkpoint_during
            )
        ] == [
            {"my_key": ""},
            {"my_key": "got here and there and parallel"},
            {"my_key": "got here and there and parallel and back again"},
        ]

        # # test interrupts BEFORE the parallel node
        app = graph.compile(checkpointer=checkpointer, interrupt_before=["outer_1"])
        config = {"configurable": {"thread_id": "4"}}
        assert [
            c
            async for c in app.astream(
                {"my_key": ""},
                config,
                stream_mode="values",
                checkpoint_during=checkpoint_during,
            )
        ] == [
            {"my_key": ""},
        ]
        # while we're waiting for the node w/ interrupt inside to finish
        assert [
            c
            async for c in app.astream(
                None, config, stream_mode="values", checkpoint_during=checkpoint_during
            )
        ] == [
            {"my_key": ""},
            {"my_key": " and parallel"},
        ]
        assert [
            c
            async for c in app.astream(
                None, config, stream_mode="values", checkpoint_during=checkpoint_during
            )
        ] == [
            {"my_key": ""},
            {"my_key": "got here and there and parallel"},
            {"my_key": "got here and there and parallel and back again"},
        ]

        # test interrupts AFTER the parallel node
        app = graph.compile(checkpointer=checkpointer, interrupt_after=["outer_1"])
        config = {"configurable": {"thread_id": "5"}}
        assert [
            c
            async for c in app.astream(
                {"my_key": ""},
                config,
                stream_mode="values",
                checkpoint_during=checkpoint_during,
            )
        ] == [
            {"my_key": ""},
            {"my_key": " and parallel"},
        ]
        assert [
            c
            async for c in app.astream(
                None, config, stream_mode="values", checkpoint_during=checkpoint_during
            )
        ] == [
            {"my_key": ""},
            {"my_key": "got here and there and parallel"},
        ]
        assert [
            c
            async for c in app.astream(
                None, config, stream_mode="values", checkpoint_during=checkpoint_during
            )
        ] == [
            {"my_key": "got here and there and parallel"},
            {"my_key": "got here and there and parallel and back again"},
        ]


@pytest.mark.parametrize("checkpoint_during", [True, False])
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_doubly_nested_graph_interrupts(
    checkpointer_name: str, checkpoint_during: bool
) -> None:
    class State(TypedDict):
        my_key: str

    class ChildState(TypedDict):
        my_key: str

    class GrandChildState(TypedDict):
        my_key: str

    async def grandchild_1(state: ChildState):
        return {"my_key": state["my_key"] + " here"}

    async def grandchild_2(state: ChildState):
        return {
            "my_key": state["my_key"] + " and there",
        }

    grandchild = StateGraph(GrandChildState)
    grandchild.add_node("grandchild_1", grandchild_1)
    grandchild.add_node("grandchild_2", grandchild_2)
    grandchild.add_edge("grandchild_1", "grandchild_2")
    grandchild.set_entry_point("grandchild_1")
    grandchild.set_finish_point("grandchild_2")

    child = StateGraph(ChildState)
    child.add_node(
        "child_1",
        grandchild.compile(interrupt_before=["grandchild_2"]),
    )
    child.set_entry_point("child_1")
    child.set_finish_point("child_1")

    async def parent_1(state: State):
        return {"my_key": "hi " + state["my_key"]}

    async def parent_2(state: State):
        return {"my_key": state["my_key"] + " and back again"}

    graph = StateGraph(State)
    graph.add_node("parent_1", parent_1)
    graph.add_node("child", child.compile())
    graph.add_node("parent_2", parent_2)
    graph.set_entry_point("parent_1")
    graph.add_edge("parent_1", "child")
    graph.add_edge("child", "parent_2")
    graph.set_finish_point("parent_2")

    async with awith_checkpointer(checkpointer_name) as checkpointer:
        app = graph.compile(checkpointer=checkpointer)

        # test invoke w/ nested interrupt
        config = {"configurable": {"thread_id": "1"}}
        assert await app.ainvoke(
            {"my_key": "my value"}, config, checkpoint_during=checkpoint_during
        ) == {
            "my_key": "hi my value",
        }

        assert await app.ainvoke(None, config, checkpoint_during=checkpoint_during) == {
            "my_key": "hi my value here and there and back again",
        }

        # test stream updates w/ nested interrupt
        nodes: list[str] = []
        config = {
            "configurable": {"thread_id": "2", CONFIG_KEY_NODE_FINISHED: nodes.append}
        }
        assert [
            c
            async for c in app.astream(
                {"my_key": "my value"}, config, checkpoint_during=checkpoint_during
            )
        ] == [
            {"parent_1": {"my_key": "hi my value"}},
            {"__interrupt__": ()},
        ]
        assert nodes == ["parent_1", "grandchild_1"]
        assert [
            c
            async for c in app.astream(
                None, config, checkpoint_during=checkpoint_during
            )
        ] == [
            {"child": {"my_key": "hi my value here and there"}},
            {"parent_2": {"my_key": "hi my value here and there and back again"}},
        ]
        assert nodes == [
            "parent_1",
            "grandchild_1",
            "grandchild_2",
            "child_1",
            "child",
            "parent_2",
        ]

        # test stream values w/ nested interrupt
        config = {"configurable": {"thread_id": "3"}}
        assert [
            c
            async for c in app.astream(
                {"my_key": "my value"},
                config,
                stream_mode="values",
                checkpoint_during=checkpoint_during,
            )
        ] == [
            {"my_key": "my value"},
            {"my_key": "hi my value"},
        ]
        assert [
            c
            async for c in app.astream(
                None, config, stream_mode="values", checkpoint_during=checkpoint_during
            )
        ] == [
            {"my_key": "hi my value"},
            {"my_key": "hi my value here and there"},
            {"my_key": "hi my value here and there and back again"},
        ]


async def test_checkpoint_metadata() -> None:
    """This test verifies that a run's configurable fields are merged with the
    previous checkpoint config for each step in the run.
    """
    # set up test
    from langchain_core.language_models.fake_chat_models import (
        FakeMessagesListChatModel,
    )
    from langchain_core.messages import AIMessage, AnyMessage
    from langchain_core.prompts import ChatPromptTemplate
    from langchain_core.tools import tool

    # graph state
    class BaseState(TypedDict):
        messages: Annotated[list[AnyMessage], add_messages]

    # initialize graph nodes
    @tool()
    def search_api(query: str) -> str:
        """Searches the API for the query."""
        return f"result for {query}"

    tools = [search_api]

    prompt = ChatPromptTemplate.from_messages(
        [
            ("system", "You are a nice assistant."),
            ("placeholder", "{messages}"),
        ]
    )

    model = FakeMessagesListChatModel(
        responses=[
            AIMessage(
                content="",
                tool_calls=[
                    {
                        "id": "tool_call123",
                        "name": "search_api",
                        "args": {"query": "query"},
                    },
                ],
            ),
            AIMessage(content="answer"),
        ]
    )

    def agent(state: BaseState, config: RunnableConfig) -> BaseState:
        formatted = prompt.invoke(state)
        response = model.invoke(formatted)
        return {"messages": response}

    def should_continue(data: BaseState) -> str:
        # Logic to decide whether to continue in the loop or exit
        if not data["messages"][-1].tool_calls:
            return "exit"
        else:
            return "continue"

    # define graphs w/ and w/o interrupt
    workflow = StateGraph(BaseState)
    workflow.add_node("agent", agent)
    workflow.add_node("tools", ToolNode(tools))
    workflow.set_entry_point("agent")
    workflow.add_conditional_edges(
        "agent", should_continue, {"continue": "tools", "exit": END}
    )
    workflow.add_edge("tools", "agent")

    # graph w/o interrupt
    checkpointer_1 = InMemorySaver()
    app = workflow.compile(checkpointer=checkpointer_1)

    # graph w/ interrupt
    checkpointer_2 = InMemorySaver()
    app_w_interrupt = workflow.compile(
        checkpointer=checkpointer_2, interrupt_before=["tools"]
    )

    # assertions

    # invoke graph w/o interrupt
    await app.ainvoke(
        {"messages": ["what is weather in sf"]},
        {
            "configurable": {
                "thread_id": "1",
                "test_config_1": "foo",
                "test_config_2": "bar",
            },
        },
    )

    config = {"configurable": {"thread_id": "1"}}

    # assert that checkpoint metadata contains the run's configurable fields
    chkpnt_metadata_1 = (await checkpointer_1.aget_tuple(config)).metadata
    assert chkpnt_metadata_1["thread_id"] == "1"
    assert chkpnt_metadata_1["test_config_1"] == "foo"
    assert chkpnt_metadata_1["test_config_2"] == "bar"

    # Verify that all checkpoint metadata have the expected keys. This check
    # is needed because a run may have an arbitrary number of steps depending
    # on how the graph is constructed.
    chkpnt_tuples_1 = checkpointer_1.alist(config)
    async for chkpnt_tuple in chkpnt_tuples_1:
        assert chkpnt_tuple.metadata["thread_id"] == "1"
        assert chkpnt_tuple.metadata["test_config_1"] == "foo"
        assert chkpnt_tuple.metadata["test_config_2"] == "bar"

    # invoke graph, but interrupt before tool call
    await app_w_interrupt.ainvoke(
        {"messages": ["what is weather in sf"]},
        {
            "configurable": {
                "thread_id": "2",
                "test_config_3": "foo",
                "test_config_4": "bar",
            },
        },
    )

    config = {"configurable": {"thread_id": "2"}}

    # assert that checkpoint metadata contains the run's configurable fields
    chkpnt_metadata_2 = (await checkpointer_2.aget_tuple(config)).metadata
    assert chkpnt_metadata_2["thread_id"] == "2"
    assert chkpnt_metadata_2["test_config_3"] == "foo"
    assert chkpnt_metadata_2["test_config_4"] == "bar"

    # resume graph execution
    await app_w_interrupt.ainvoke(
        input=None,
        config={
            "configurable": {
                "thread_id": "2",
                "test_config_3": "foo",
                "test_config_4": "bar",
            }
        },
    )

    # assert that checkpoint metadata contains the run's configurable fields
    chkpnt_metadata_3 = (await checkpointer_2.aget_tuple(config)).metadata
    assert chkpnt_metadata_3["thread_id"] == "2"
    assert chkpnt_metadata_3["test_config_3"] == "foo"
    assert chkpnt_metadata_3["test_config_4"] == "bar"

    # Verify that all checkpoint metadata have the expected keys. This check
    # is needed because a run may have an arbitrary number of steps depending
    # on how the graph is constructed.
    chkpnt_tuples_2 = checkpointer_2.alist(config)
    async for chkpnt_tuple in chkpnt_tuples_2:
        assert chkpnt_tuple.metadata["thread_id"] == "2"
        assert chkpnt_tuple.metadata["test_config_3"] == "foo"
        assert chkpnt_tuple.metadata["test_config_4"] == "bar"


async def test_checkpointer_null_pending_writes() -> None:
    class Node:
        def __init__(self, name: str):
            self.name = name
            setattr(self, "__name__", name)

        def __call__(self, state):
            return [self.name]

    builder = StateGraph(Annotated[list, operator.add])
    builder.add_node(Node("1"))
    builder.add_edge(START, "1")
    graph = builder.compile(checkpointer=MemorySaverNoPending())
    assert graph.invoke([], {"configurable": {"thread_id": "foo"}}) == ["1"]
    assert graph.invoke([], {"configurable": {"thread_id": "foo"}}) == ["1"] * 2
    assert (await graph.ainvoke([], {"configurable": {"thread_id": "foo"}})) == [
        "1"
    ] * 3
    assert (await graph.ainvoke([], {"configurable": {"thread_id": "foo"}})) == [
        "1"
    ] * 4


@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
@pytest.mark.parametrize("store_name", ALL_STORES_ASYNC)
async def test_store_injected_async(checkpointer_name: str, store_name: str) -> None:
    class State(TypedDict):
        count: Annotated[int, operator.add]

    doc_id = str(uuid.uuid4())
    doc = {"some-key": "this-is-a-val"}
    uid = uuid.uuid4().hex
    namespace = (f"foo-{uid}", "bar")
    thread_1 = str(uuid.uuid4())
    thread_2 = str(uuid.uuid4())

    class Node:
        def __init__(self, i: Optional[int] = None):
            self.i = i

        async def __call__(
            self, inputs: State, config: RunnableConfig, store: BaseStore
        ):
            assert isinstance(store, BaseStore)
            await store.aput(
                (
                    namespace
                    if self.i is not None
                    and config["configurable"]["thread_id"] in (thread_1, thread_2)
                    else (f"foo_{self.i}", "bar")
                ),
                doc_id,
                {
                    **doc,
                    "from_thread": config["configurable"]["thread_id"],
                    "some_val": inputs["count"],
                },
            )
            return {"count": 1}

    def other_node(inputs: State, config: RunnableConfig, store: BaseStore):
        assert isinstance(store, BaseStore)
        store.put(("not", "interesting"), "key", {"val": "val"})
        item = store.get(("not", "interesting"), "key")
        assert item is not None
        assert item.value == {"val": "val"}
        return {"count": 0}

    builder = StateGraph(State)
    builder.add_node("node", Node())
    builder.add_node("other_node", other_node)
    builder.add_edge("__start__", "node")
    builder.add_edge("node", "other_node")

    N = 500
    M = 1

    for i in range(N):
        builder.add_node(f"node_{i}", Node(i))
        builder.add_edge("__start__", f"node_{i}")

    async with (
        awith_checkpointer(checkpointer_name) as checkpointer,
        awith_store(store_name) as the_store,
    ):
        graph = builder.compile(store=the_store, checkpointer=checkpointer)

        # Test batch operations with multiple threads
        results = await graph.abatch(
            [{"count": 0}] * M,
            ([{"configurable": {"thread_id": str(uuid.uuid4())}}] * (M - 1))
            + [{"configurable": {"thread_id": thread_1}}],
        )
        result = results[-1]
        assert result == {"count": N + 1}
        returned_doc = (await the_store.aget(namespace, doc_id)).value
        assert returned_doc == {**doc, "from_thread": thread_1, "some_val": 0}
        assert len((await the_store.asearch(namespace))) == 1

        # Check results after another turn of the same thread
        result = await graph.ainvoke(
            {"count": 0}, {"configurable": {"thread_id": thread_1}}
        )
        assert result == {"count": (N + 1) * 2}
        returned_doc = (await the_store.aget(namespace, doc_id)).value
        assert returned_doc == {**doc, "from_thread": thread_1, "some_val": N + 1}
        assert len((await the_store.asearch(namespace))) == 1

        # Test with a different thread
        result = await graph.ainvoke(
            {"count": 0}, {"configurable": {"thread_id": thread_2}}
        )
        assert result == {"count": N + 1}
        returned_doc = (await the_store.aget(namespace, doc_id)).value
        assert returned_doc == {
            **doc,
            "from_thread": thread_2,
            "some_val": 0,
        }  # Overwrites the whole doc
        assert (
            len((await the_store.asearch(namespace))) == 1
        )  # still overwriting the same one


async def test_debug_retry():
    class State(TypedDict):
        messages: Annotated[list[str], operator.add]

    def node(name):
        async def _node(state: State):
            return {"messages": [f"entered {name} node"]}

        return _node

    builder = StateGraph(State)
    builder.add_node("one", node("one"))
    builder.add_node("two", node("two"))
    builder.add_edge(START, "one")
    builder.add_edge("one", "two")
    builder.add_edge("two", END)

    saver = InMemorySaver()

    graph = builder.compile(checkpointer=saver)

    config = {"configurable": {"thread_id": "1"}}
    await graph.ainvoke({"messages": []}, config=config)

    # re-run step: 1
    async for c in saver.alist(config):
        if c.metadata["step"] == 1:
            target_config = c.parent_config
            break
    assert target_config is not None

    update_config = await graph.aupdate_state(target_config, values=None)

    events = [
        c async for c in graph.astream(None, config=update_config, stream_mode="debug")
    ]

    checkpoint_events = list(
        reversed([e["payload"] for e in events if e["type"] == "checkpoint"])
    )

    checkpoint_history = {
        c.config["configurable"]["checkpoint_id"]: c
        async for c in graph.aget_state_history(config)
    }

    def lax_normalize_config(config: Optional[dict]) -> Optional[dict]:
        if config is None:
            return None
        return config["configurable"]

    for stream in checkpoint_events:
        stream_conf = lax_normalize_config(stream["config"])
        stream_parent_conf = lax_normalize_config(stream["parent_config"])
        assert stream_conf != stream_parent_conf

        # ensure the streamed checkpoint == checkpoint from checkpointer.list()
        history = checkpoint_history[stream["config"]["configurable"]["checkpoint_id"]]
        history_conf = lax_normalize_config(history.config)
        assert stream_conf == history_conf

        history_parent_conf = lax_normalize_config(history.parent_config)
        assert stream_parent_conf == history_parent_conf


async def test_debug_subgraphs():
    class State(TypedDict):
        messages: Annotated[list[str], operator.add]

    def node(name):
        async def _node(state: State):
            return {"messages": [f"entered {name} node"]}

        return _node

    parent = StateGraph(State)
    child = StateGraph(State)

    child.add_node("c_one", node("c_one"))
    child.add_node("c_two", node("c_two"))
    child.add_edge(START, "c_one")
    child.add_edge("c_one", "c_two")
    child.add_edge("c_two", END)

    parent.add_node("p_one", node("p_one"))
    parent.add_node("p_two", child.compile())
    parent.add_edge(START, "p_one")
    parent.add_edge("p_one", "p_two")
    parent.add_edge("p_two", END)

    graph = parent.compile(checkpointer=InMemorySaver())

    config = {"configurable": {"thread_id": "1"}}
    events = [
        c
        async for c in graph.astream(
            {"messages": []},
            config=config,
            stream_mode="debug",
        )
    ]

    checkpoint_events = list(
        reversed([e["payload"] for e in events if e["type"] == "checkpoint"])
    )
    checkpoint_history = [c async for c in graph.aget_state_history(config)]

    assert len(checkpoint_events) == len(checkpoint_history)

    def normalize_config(config: Optional[dict]) -> Optional[dict]:
        if config is None:
            return None
        return config["configurable"]

    for stream, history in zip(checkpoint_events, checkpoint_history):
        assert stream["values"] == history.values
        assert stream["next"] == list(history.next)
        assert normalize_config(stream["config"]) == normalize_config(history.config)
        assert normalize_config(stream["parent_config"]) == normalize_config(
            history.parent_config
        )

        assert len(stream["tasks"]) == len(history.tasks)
        for stream_task, history_task in zip(stream["tasks"], history.tasks):
            assert stream_task["id"] == history_task.id
            assert stream_task["name"] == history_task.name
            assert stream_task["interrupts"] == history_task.interrupts
            assert stream_task.get("error") == history_task.error
            assert stream_task.get("state") == history_task.state


async def test_debug_nested_subgraphs():
    from collections import defaultdict

    class State(TypedDict):
        messages: Annotated[list[str], operator.add]

    def node(name):
        async def _node(state: State):
            return {"messages": [f"entered {name} node"]}

        return _node

    grand_parent = StateGraph(State)
    parent = StateGraph(State)
    child = StateGraph(State)

    child.add_node("c_one", node("c_one"))
    child.add_node("c_two", node("c_two"))
    child.add_edge(START, "c_one")
    child.add_edge("c_one", "c_two")
    child.add_edge("c_two", END)

    parent.add_node("p_one", node("p_one"))
    parent.add_node("p_two", child.compile())
    parent.add_edge(START, "p_one")
    parent.add_edge("p_one", "p_two")
    parent.add_edge("p_two", END)

    grand_parent.add_node("gp_one", node("gp_one"))
    grand_parent.add_node("gp_two", parent.compile())
    grand_parent.add_edge(START, "gp_one")
    grand_parent.add_edge("gp_one", "gp_two")
    grand_parent.add_edge("gp_two", END)

    graph = grand_parent.compile(checkpointer=InMemorySaver())

    config = {"configurable": {"thread_id": "1"}}
    events = [
        c
        async for c in graph.astream(
            {"messages": []},
            config=config,
            stream_mode="debug",
            subgraphs=True,
        )
    ]

    stream_ns: dict[tuple, dict] = defaultdict(list)
    for ns, e in events:
        if e["type"] == "checkpoint":
            stream_ns[ns].append(e["payload"])

    assert list(stream_ns.keys()) == [
        (),
        (AnyStr("gp_two:"),),
        (AnyStr("gp_two:"), AnyStr("p_two:")),
    ]

    history_ns = {}
    for ns in stream_ns.keys():

        async def get_history():
            history = [
                c
                async for c in graph.aget_state_history(
                    {"configurable": {"thread_id": "1", "checkpoint_ns": "|".join(ns)}}
                )
            ]
            return history[::-1]

        history_ns[ns] = await get_history()

    def normalize_config(config: Optional[dict]) -> Optional[dict]:
        if config is None:
            return None

        clean_config = {}
        clean_config["thread_id"] = config["configurable"]["thread_id"]
        clean_config["checkpoint_id"] = config["configurable"]["checkpoint_id"]
        clean_config["checkpoint_ns"] = config["configurable"]["checkpoint_ns"]
        if "checkpoint_map" in config["configurable"]:
            clean_config["checkpoint_map"] = config["configurable"]["checkpoint_map"]

        return clean_config

    for checkpoint_events, checkpoint_history in zip(
        stream_ns.values(), history_ns.values()
    ):
        for stream, history in zip(checkpoint_events, checkpoint_history):
            assert stream["values"] == history.values
            assert stream["next"] == list(history.next)
            assert normalize_config(stream["config"]) == normalize_config(
                history.config
            )
            assert normalize_config(stream["parent_config"]) == normalize_config(
                history.parent_config
            )

            assert len(stream["tasks"]) == len(history.tasks)
            for stream_task, history_task in zip(stream["tasks"], history.tasks):
                assert stream_task["id"] == history_task.id
                assert stream_task["name"] == history_task.name
                assert stream_task["interrupts"] == history_task.interrupts
                assert stream_task.get("error") == history_task.error
                assert stream_task.get("state") == history_task.state


@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_parent_command(checkpointer_name: str) -> None:
    from langchain_core.messages import BaseMessage
    from langchain_core.tools import tool

    @tool(return_direct=True)
    def get_user_name() -> Command:
        """Retrieve user name"""
        return Command(update={"user_name": "Meow"}, graph=Command.PARENT)

    subgraph_builder = StateGraph(MessagesState)
    subgraph_builder.add_node("tool", get_user_name)
    subgraph_builder.add_edge(START, "tool")
    subgraph = subgraph_builder.compile()

    class CustomParentState(TypedDict):
        messages: Annotated[list[BaseMessage], add_messages]
        # this key is not available to the child graph
        user_name: str

    builder = StateGraph(CustomParentState)
    builder.add_node("alice", subgraph)
    builder.add_edge(START, "alice")
    async with awith_checkpointer(checkpointer_name) as checkpointer:
        graph = builder.compile(checkpointer=checkpointer)

        config = {"configurable": {"thread_id": "1"}}

        assert await graph.ainvoke(
            {"messages": [("user", "get user name")]}, config
        ) == {
            "messages": [
                _AnyIdHumanMessage(
                    content="get user name", additional_kwargs={}, response_metadata={}
                ),
            ],
            "user_name": "Meow",
        }
        assert await graph.aget_state(config) == StateSnapshot(
            values={
                "messages": [
                    _AnyIdHumanMessage(
                        content="get user name",
                        additional_kwargs={},
                        response_metadata={},
                    ),
                ],
                "user_name": "Meow",
            },
            next=(),
            config={
                "configurable": {
                    "thread_id": "1",
                    "checkpoint_ns": "",
                    "checkpoint_id": AnyStr(),
                }
            },
            metadata={
                "source": "loop",
                "writes": {
                    "alice": {
                        "user_name": "Meow",
                    }
                },
                "thread_id": "1",
                "step": 1,
                "parents": {},
            },
            created_at=AnyStr(),
            parent_config=(
                None
                if "shallow" in checkpointer_name
                else {
                    "configurable": {
                        "thread_id": "1",
                        "checkpoint_ns": "",
                        "checkpoint_id": AnyStr(),
                    }
                }
            ),
            tasks=(),
        )


@NEEDS_CONTEXTVARS
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_interrupt_subgraph(checkpointer_name: str):
    class State(TypedDict):
        baz: str

    def foo(state):
        return {"baz": "foo"}

    def bar(state):
        value = interrupt("Please provide baz value:")
        return {"baz": value}

    child_builder = StateGraph(State)
    child_builder.add_node(bar)
    child_builder.add_edge(START, "bar")

    builder = StateGraph(State)
    builder.add_node(foo)
    builder.add_node("bar", child_builder.compile())
    builder.add_edge(START, "foo")
    builder.add_edge("foo", "bar")

    async with awith_checkpointer(checkpointer_name) as checkpointer:
        graph = builder.compile(checkpointer=checkpointer)

        thread1 = {"configurable": {"thread_id": "1"}}
        # First run, interrupted at bar
        assert await graph.ainvoke({"baz": ""}, thread1)
        # Resume with answer
        assert await graph.ainvoke(Command(resume="bar"), thread1)


@NEEDS_CONTEXTVARS
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_interrupt_multiple(checkpointer_name: str):
    class State(TypedDict):
        my_key: Annotated[str, operator.add]

    async def node(s: State) -> State:
        answer = interrupt({"value": 1})
        answer2 = interrupt({"value": 2})
        return {"my_key": answer + " " + answer2}

    builder = StateGraph(State)
    builder.add_node("node", node)
    builder.add_edge(START, "node")

    async with awith_checkpointer(checkpointer_name) as checkpointer:
        graph = builder.compile(checkpointer=checkpointer)
        thread1 = {"configurable": {"thread_id": "1"}}

        assert [
            e async for e in graph.astream({"my_key": "DE", "market": "DE"}, thread1)
        ] == [
            {
                "__interrupt__": (
                    Interrupt(
                        value={"value": 1},
                        resumable=True,
                        ns=[AnyStr("node:")],
                        when="during",
                    ),
                )
            }
        ]

        assert [
            event
            async for event in graph.astream(
                Command(resume="answer 1", update={"my_key": "foofoo"}),
                thread1,
                stream_mode="updates",
            )
        ] == [
            {
                "__interrupt__": (
                    Interrupt(
                        value={"value": 2},
                        resumable=True,
                        ns=[AnyStr("node:")],
                        when="during",
                    ),
                )
            }
        ]

        assert [
            event
            async for event in graph.astream(
                Command(resume="answer 2"), thread1, stream_mode="updates"
            )
        ] == [
            {"node": {"my_key": "answer 1 answer 2"}},
        ]


@NEEDS_CONTEXTVARS
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_interrupt_loop(checkpointer_name: str):
    class State(TypedDict):
        age: int
        other: str

    async def ask_age(s: State):
        """Ask an expert for help."""
        question = "How old are you?"
        value = None
        for _ in range(10):
            value: str = interrupt(question)
            if not value.isdigit() or int(value) < 18:
                question = "invalid response"
                value = None
            else:
                break

        return {"age": int(value)}

    builder = StateGraph(State)
    builder.add_node("node", ask_age)
    builder.add_edge(START, "node")

    async with awith_checkpointer(checkpointer_name) as checkpointer:
        graph = builder.compile(checkpointer=checkpointer)
        thread1 = {"configurable": {"thread_id": "1"}}

        assert [e async for e in graph.astream({"other": ""}, thread1)] == [
            {
                "__interrupt__": (
                    Interrupt(
                        value="How old are you?",
                        resumable=True,
                        ns=[AnyStr("node:")],
                        when="during",
                    ),
                )
            }
        ]

        assert [
            event
            async for event in graph.astream(
                Command(resume="13"),
                thread1,
            )
        ] == [
            {
                "__interrupt__": (
                    Interrupt(
                        value="invalid response",
                        resumable=True,
                        ns=[AnyStr("node:")],
                        when="during",
                    ),
                )
            }
        ]

        assert [
            event
            async for event in graph.astream(
                Command(resume="15"),
                thread1,
            )
        ] == [
            {
                "__interrupt__": (
                    Interrupt(
                        value="invalid response",
                        resumable=True,
                        ns=[AnyStr("node:")],
                        when="during",
                    ),
                )
            }
        ]

        assert [
            event async for event in graph.astream(Command(resume="19"), thread1)
        ] == [
            {"node": {"age": 19}},
        ]


@NEEDS_CONTEXTVARS
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_interrupt_functional(checkpointer_name: str) -> None:
    @task
    async def foo(state: dict) -> dict:
        return {"a": state["a"] + "foo"}

    @task
    async def bar(state: dict) -> dict:
        return {"a": state["a"] + "bar", "b": state["b"]}

    async with awith_checkpointer(checkpointer_name) as checkpointer:

        @entrypoint(checkpointer=checkpointer)
        async def graph(inputs: dict) -> dict:
            foo_result = await foo(inputs)
            value = interrupt("Provide value for bar:")
            bar_input = {**foo_result, "b": value}
            bar_result = await bar(bar_input)
            return bar_result

        config = {"configurable": {"thread_id": "1"}}
        # First run, interrupted at bar
        await graph.ainvoke({"a": ""}, config)
        # Resume with an answer
        res = await graph.ainvoke(Command(resume="bar"), config)
        assert res == {"a": "foobar", "b": "bar"}


@NEEDS_CONTEXTVARS
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_interrupt_task_functional(checkpointer_name: str) -> None:
    @task
    async def foo(state: dict) -> dict:
        return {"a": state["a"] + "foo"}

    @task
    async def bar(state: dict) -> dict:
        value = interrupt("Provide value for bar:")
        return {"a": state["a"] + value}

    async with awith_checkpointer(checkpointer_name) as checkpointer:

        @entrypoint(checkpointer=checkpointer)
        async def graph(inputs: dict) -> dict:
            foo_result = await foo(inputs)
            bar_result = await bar(foo_result)
            return bar_result

        config = {"configurable": {"thread_id": "1"}}
        # First run, interrupted at bar
        await graph.ainvoke({"a": ""}, config)
        # Resume with an answer
        res = await graph.ainvoke(Command(resume="bar"), config)
        assert res == {"a": "foobar"}


@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_command_with_static_breakpoints(checkpointer_name: str) -> None:
    """Test that we can use Command to resume and update with static breakpoints."""

    class State(TypedDict):
        """The graph state."""

        foo: str

    def node1(state: State):
        return {
            "foo": state["foo"] + "|node-1",
        }

    def node2(state: State):
        return {
            "foo": state["foo"] + "|node-2",
        }

    builder = StateGraph(State)
    builder.add_node("node1", node1)
    builder.add_node("node2", node2)
    builder.add_edge(START, "node1")
    builder.add_edge("node1", "node2")

    async with awith_checkpointer(checkpointer_name) as checkpointer:
        graph = builder.compile(checkpointer=checkpointer, interrupt_before=["node1"])
        config = {"configurable": {"thread_id": str(uuid.uuid4())}}

        # Start the graph and interrupt at the first node
        await graph.ainvoke({"foo": "abc"}, config)
        result = await graph.ainvoke(Command(update={"foo": "def"}), config)
        assert result == {"foo": "def|node-1|node-2"}


@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_multistep_plan(checkpointer_name: str):
    from langchain_core.messages import AnyMessage

    class State(TypedDict, total=False):
        plan: list[Union[str, list[str]]]
        messages: Annotated[list[AnyMessage], add_messages]

    def planner(state: State):
        if state.get("plan") is None:
            # create plan somehow
            plan = ["step1", ["step2", "step3"], "step4"]
            # pick the first step to execute next
            first_step, *plan = plan
            # put the rest of plan in state
            return Command(goto=first_step, update={"plan": plan})
        elif state["plan"]:
            # go to the next step of the plan
            next_step, *next_plan = state["plan"]
            return Command(goto=next_step, update={"plan": next_plan})
        else:
            # the end of the plan
            pass

    def step1(state: State):
        return Command(goto="planner", update={"messages": [("human", "step1")]})

    def step2(state: State):
        return Command(goto="planner", update={"messages": [("human", "step2")]})

    def step3(state: State):
        return Command(goto="planner", update={"messages": [("human", "step3")]})

    def step4(state: State):
        return Command(goto="planner", update={"messages": [("human", "step4")]})

    builder = StateGraph(State)
    builder.add_node(planner)
    builder.add_node(step1)
    builder.add_node(step2)
    builder.add_node(step3)
    builder.add_node(step4)
    builder.add_edge(START, "planner")

    async with awith_checkpointer(checkpointer_name) as checkpointer:
        graph = builder.compile(checkpointer=checkpointer)

        config = {"configurable": {"thread_id": "1"}}

        assert await graph.ainvoke({"messages": [("human", "start")]}, config) == {
            "messages": [
                _AnyIdHumanMessage(content="start"),
                _AnyIdHumanMessage(content="step1"),
                _AnyIdHumanMessage(content="step2"),
                _AnyIdHumanMessage(content="step3"),
                _AnyIdHumanMessage(content="step4"),
            ],
            "plan": [],
        }


@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_command_goto_with_static_breakpoints(checkpointer_name: str) -> None:
    """Use Command goto with static breakpoints."""

    class State(TypedDict):
        """The graph state."""

        foo: Annotated[str, operator.add]

    def node1(state: State):
        return {
            "foo": "|node-1",
        }

    def node2(state: State):
        return {
            "foo": "|node-2",
        }

    builder = StateGraph(State)
    builder.add_node("node1", node1)
    builder.add_node("node2", node2)
    builder.add_edge(START, "node1")
    builder.add_edge("node1", "node2")

    async with awith_checkpointer(checkpointer_name) as checkpointer:
        graph = builder.compile(checkpointer=checkpointer, interrupt_before=["node1"])

        config = {"configurable": {"thread_id": str(uuid.uuid4())}}

        # Start the graph and interrupt at the first node
        await graph.ainvoke({"foo": "abc"}, config)
        result = await graph.ainvoke(Command(goto=["node2"]), config)
        assert result == {"foo": "abc|node-1|node-2|node-2"}


async def test_parallel_node_execution():
    """Test that parallel nodes execute concurrently."""

    class State(TypedDict):
        results: Annotated[list[str], operator.add]

    async def slow_node(state: State):
        await asyncio.sleep(1)
        return {"results": ["slow"]}

    async def fast_node(state: State):
        await asyncio.sleep(2)
        return {"results": ["fast"]}

    builder = StateGraph(State)
    builder.add_node("slow", slow_node)
    builder.add_node("fast", fast_node)
    builder.add_edge(START, "slow")
    builder.add_edge(START, "fast")

    graph = builder.compile()

    start = perf_counter()
    result = await graph.ainvoke({"results": []})
    duration = perf_counter() - start

    # Fast node result should be available first
    assert "fast" in result["results"][0]

    # Total duration should be less than sum of both nodes
    assert duration < 3.0


@NEEDS_CONTEXTVARS
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_multiple_interrupt_state_persistence(checkpointer_name: str) -> None:
    """Test that state is preserved correctly across multiple interrupts."""

    class State(TypedDict):
        steps: Annotated[list[str], operator.add]

    def interruptible_node(state: State):
        first = interrupt("First interrupt")
        second = interrupt("Second interrupt")
        return {"steps": [first, second]}

    builder = StateGraph(State)
    builder.add_node("node", interruptible_node)
    builder.add_edge(START, "node")

    async with awith_checkpointer(checkpointer_name) as checkpointer:
        app = builder.compile(checkpointer=checkpointer)
        config = {"configurable": {"thread_id": "1"}}

        # First execution - should hit first interrupt
        await app.ainvoke({"steps": []}, config)

        # State should still be empty since node hasn't returned
        state = await app.aget_state(config)
        assert state.values == {"steps": []}

        # Resume after first interrupt - should hit second interrupt
        await app.ainvoke(Command(resume="step1"), config)

        # State should still be empty since node hasn't returned
        state = await app.aget_state(config)
        assert state.values == {"steps": []}

        # Resume after second interrupt - node should complete
        result = await app.ainvoke(Command(resume="step2"), config)

        # Now state should contain both steps since node returned
        assert result["steps"] == ["step1", "step2"]
        state = await app.aget_state(config)
        assert state.values["steps"] == ["step1", "step2"]


async def test_concurrent_execution():
    """Test concurrent execution with async nodes."""

    class State(TypedDict):
        counter: Annotated[int, operator.add]

    results = deque()

    async def slow_node(state: State):
        await asyncio.sleep(0.1)
        return {"counter": 1}

    builder = StateGraph(State)
    builder.add_node("node", slow_node)
    builder.add_edge(START, "node")
    graph = builder.compile()

    async def run_graph():
        result = await graph.ainvoke({"counter": 0})
        results.append(result)

    # Create and gather tasks
    tasks = [run_graph() for _ in range(10)]
    await asyncio.gather(*tasks)

    # Verify results are independent
    assert len(results) == 10
    for result in results:
        assert result["counter"] == 1


@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_checkpoint_recovery_async(checkpointer_name: str):
    """Test recovery from checkpoints after failures with async nodes."""

    class State(TypedDict):
        steps: Annotated[list[str], operator.add]
        attempt: int  # Track number of attempts

    async def failing_node(state: State):
        # Fail on first attempt, succeed on retry
        if state["attempt"] == 1:
            raise RuntimeError("Simulated failure")
        await asyncio.sleep(0.1)  # Simulate async work
        return {"steps": ["node1"]}

    async def second_node(state: State):
        await asyncio.sleep(0.1)  # Simulate async work
        return {"steps": ["node2"]}

    builder = StateGraph(State)
    builder.add_node("node1", failing_node)
    builder.add_node("node2", second_node)
    builder.add_edge(START, "node1")
    builder.add_edge("node1", "node2")

    async with awith_checkpointer(checkpointer_name) as checkpointer:
        graph = builder.compile(checkpointer=checkpointer)
        config = {"configurable": {"thread_id": "1"}}

        # First attempt should fail
        with pytest.raises(RuntimeError):
            await graph.ainvoke({"steps": ["start"], "attempt": 1}, config)

        # Verify checkpoint state
        state = await graph.aget_state(config)
        assert state is not None
        assert state.values == {"steps": ["start"], "attempt": 1}  # input state saved
        assert state.next == ("node1",)  # Should retry failed node

        # Retry with updated attempt count
        result = await graph.ainvoke({"steps": [], "attempt": 2}, config)
        assert result == {"steps": ["start", "node1", "node2"], "attempt": 2}

        if "shallow" in checkpointer_name:
            return

        # Verify checkpoint history shows both attempts
        history = [c async for c in graph.aget_state_history(config)]
        assert len(history) == 6  # Initial + failed attempt + successful attempt

        # Verify the error was recorded in checkpoint
        failed_checkpoint = next(c for c in history if c.tasks and c.tasks[0].error)
        assert "RuntimeError('Simulated failure')" in failed_checkpoint.tasks[0].error


async def test_multiple_updates_root() -> None:
    def node_a(state):
        return [Command(update="a1"), Command(update="a2")]

    def node_b(state):
        return "b"

    graph = (
        StateGraph(Annotated[str, operator.add])
        .add_sequence([node_a, node_b])
        .add_edge(START, "node_a")
        .compile()
    )

    assert await graph.ainvoke("") == "a1a2b"

    # only streams the last update from node_a
    assert [c async for c in graph.astream("", stream_mode="updates")] == [
        {"node_a": ["a1", "a2"]},
        {"node_b": "b"},
    ]


async def test_multiple_updates() -> None:
    class State(TypedDict):
        foo: Annotated[str, operator.add]

    def node_a(state):
        return [Command(update={"foo": "a1"}), Command(update={"foo": "a2"})]

    def node_b(state):
        return {"foo": "b"}

    graph = (
        StateGraph(State)
        .add_sequence([node_a, node_b])
        .add_edge(START, "node_a")
        .compile()
    )

    assert await graph.ainvoke({"foo": ""}) == {
        "foo": "a1a2b",
    }

    # only streams the last update from node_a
    assert [c async for c in graph.astream({"foo": ""}, stream_mode="updates")] == [
        {"node_a": [{"foo": "a1"}, {"foo": "a2"}]},
        {"node_b": {"foo": "b"}},
    ]


@NEEDS_CONTEXTVARS
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_falsy_return_from_task(checkpointer_name: str) -> None:
    """Test with a falsy return from a task."""

    @task
    async def falsy_task() -> bool:
        return False

    async with awith_checkpointer(checkpointer_name) as checkpointer:

        @entrypoint(checkpointer=checkpointer)
        async def graph(state: dict) -> dict:
            """React tool."""
            await falsy_task()
            interrupt("test")

        configurable = {"configurable": {"thread_id": str(uuid.uuid4())}}
        await graph.ainvoke({"a": 5}, configurable)
        await graph.ainvoke(Command(resume="123"), configurable)


@NEEDS_CONTEXTVARS
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_multiple_interrupts_functional(checkpointer_name: str) -> None:
    """Test multiple interrupts with functional API."""
    from langgraph.func import entrypoint, task

    counter = 0

    @task
    async def double(x: int) -> int:
        """Increment the counter."""
        nonlocal counter
        counter += 1
        return 2 * x

    async with awith_checkpointer(checkpointer_name) as checkpointer:

        @entrypoint(checkpointer=checkpointer)
        async def graph(state: dict) -> dict:
            """React tool."""

            values = []

            for idx in [1, 2, 3]:
                values.extend([await double(idx), interrupt({"a": "boo"})])

            return {"values": values}

        configurable = {"configurable": {"thread_id": str(uuid.uuid4())}}
        await graph.ainvoke({}, configurable)
        await graph.ainvoke(Command(resume="a"), configurable)
        await graph.ainvoke(Command(resume="b"), configurable)
        result = await graph.ainvoke(Command(resume="c"), configurable)
        # `double` value should be cached appropriately when used w/ `interrupt`
        assert result == {
            "values": [2, "a", 4, "b", 6, "c"],
        }
        assert counter == 3


@NEEDS_CONTEXTVARS
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_double_interrupt_subgraph(checkpointer_name: str) -> None:
    class AgentState(TypedDict):
        input: str

    def node_1(state: AgentState):
        result = interrupt("interrupt node 1")
        return {"input": result}

    def node_2(state: AgentState):
        result = interrupt("interrupt node 2")
        return {"input": result}

    subgraph_builder = (
        StateGraph(AgentState)
        .add_node("node_1", node_1)
        .add_node("node_2", node_2)
        .add_edge(START, "node_1")
        .add_edge("node_1", "node_2")
        .add_edge("node_2", END)
    )

    async with awith_checkpointer(checkpointer_name) as checkpointer:
        # invoke the sub graph
        subgraph = subgraph_builder.compile(checkpointer=checkpointer)
        thread = {"configurable": {"thread_id": str(uuid.uuid4())}}
        assert [c async for c in subgraph.astream({"input": "test"}, thread)] == [
            {
                "__interrupt__": (
                    Interrupt(
                        value="interrupt node 1",
                        resumable=True,
                        ns=[AnyStr("node_1:")],
                        when="during",
                    ),
                )
            },
        ]
        # resume from the first interrupt
        assert [c async for c in subgraph.astream(Command(resume="123"), thread)] == [
            {
                "node_1": {"input": "123"},
            },
            {
                "__interrupt__": (
                    Interrupt(
                        value="interrupt node 2",
                        resumable=True,
                        ns=[AnyStr("node_2:")],
                        when="during",
                    ),
                )
            },
        ]
        # resume from the second interrupt
        assert [c async for c in subgraph.astream(Command(resume="123"), thread)] == [
            {
                "node_2": {"input": "123"},
            },
        ]

        subgraph = subgraph_builder.compile()

        def invoke_sub_agent(state: AgentState):
            return subgraph.invoke(state)

        thread = {"configurable": {"thread_id": str(uuid.uuid4())}}

        parent_agent = (
            StateGraph(AgentState)
            .add_node("invoke_sub_agent", invoke_sub_agent)
            .add_edge(START, "invoke_sub_agent")
            .add_edge("invoke_sub_agent", END)
            .compile(checkpointer=checkpointer)
        )

        assert [c async for c in parent_agent.astream({"input": "test"}, thread)] == [
            {
                "__interrupt__": (
                    Interrupt(
                        value="interrupt node 1",
                        resumable=True,
                        ns=[AnyStr("invoke_sub_agent:"), AnyStr("node_1:")],
                        when="during",
                    ),
                )
            },
        ]

        # resume from the first interrupt
        assert [
            c async for c in parent_agent.astream(Command(resume=True), thread)
        ] == [
            {
                "__interrupt__": (
                    Interrupt(
                        value="interrupt node 2",
                        resumable=True,
                        ns=[AnyStr("invoke_sub_agent:"), AnyStr("node_2:")],
                        when="during",
                    ),
                )
            }
        ]

        # resume from 2nd interrupt
        assert [
            c async for c in parent_agent.astream(Command(resume=True), thread)
        ] == [
            {
                "invoke_sub_agent": {"input": True},
            },
        ]


@NEEDS_CONTEXTVARS
async def test_async_streaming_with_functional_api() -> None:
    """Test streaming with functional API.

    This test verifies that we're able to stream results as they're being generated
    rather than have all the results arrive at once after the graph has completed.

    The time of arrival between the two updates corresponding to the two `slow` tasks
    should be greater than the time delay between the two tasks.
    """

    time_delay = 0.01

    @task()
    async def slow() -> dict:
        await asyncio.sleep(time_delay)  # Simulate a delay of 10 ms
        return {"tic": asyncio.get_running_loop().time()}

    @entrypoint()
    async def graph(inputs: dict) -> list:
        first = await slow()
        second = await slow()
        return [first, second]

    arrival_times = []

    async for chunk in graph.astream({}):
        if "slow" not in chunk:  # We'll just look at the updates from `slow`
            continue
        arrival_times.append(asyncio.get_running_loop().time())

    assert len(arrival_times) == 2
    delta = arrival_times[1] - arrival_times[0]
    # Delta cannot be less than 10 ms if it is streaming as results are generated.
    assert delta > time_delay


@NEEDS_CONTEXTVARS
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_multiple_subgraphs(checkpointer_name: str) -> None:
    class State(TypedDict):
        a: int
        b: int

    class Output(TypedDict):
        result: int

    async with awith_checkpointer(checkpointer_name) as checkpointer:
        # Define the subgraphs
        async def add(state):
            return {"result": state["a"] + state["b"]}

        add_subgraph = (
            StateGraph(State, output=Output)
            .add_node(add)
            .add_edge(START, "add")
            .compile()
        )

        async def multiply(state):
            return {"result": state["a"] * state["b"]}

        multiply_subgraph = (
            StateGraph(State, output=Output)
            .add_node(multiply)
            .add_edge(START, "multiply")
            .compile()
        )

        # Test calling the same subgraph multiple times
        async def call_same_subgraph(state):
            result = await add_subgraph.ainvoke(state)
            another_result = await add_subgraph.ainvoke(
                {"a": result["result"], "b": 10}
            )
            return another_result

        parent_call_same_subgraph = (
            StateGraph(State, output=Output)
            .add_node(call_same_subgraph)
            .add_edge(START, "call_same_subgraph")
            .compile(checkpointer=checkpointer)
        )
        config = {"configurable": {"thread_id": "1"}}
        assert await parent_call_same_subgraph.ainvoke({"a": 2, "b": 3}, config) == {
            "result": 15
        }

        # Test calling multiple subgraphs
        class Output(TypedDict):
            add_result: int
            multiply_result: int

        async def call_multiple_subgraphs(state):
            add_result = await add_subgraph.ainvoke(state)
            multiply_result = await multiply_subgraph.ainvoke(state)
            return {
                "add_result": add_result["result"],
                "multiply_result": multiply_result["result"],
            }

        parent_call_multiple_subgraphs = (
            StateGraph(State, output=Output)
            .add_node(call_multiple_subgraphs)
            .add_edge(START, "call_multiple_subgraphs")
            .compile(checkpointer=checkpointer)
        )
        config = {"configurable": {"thread_id": "2"}}
        assert await parent_call_multiple_subgraphs.ainvoke(
            {"a": 2, "b": 3}, config
        ) == {
            "add_result": 5,
            "multiply_result": 6,
        }


@NEEDS_CONTEXTVARS
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_multiple_subgraphs_functional(checkpointer_name: str) -> None:
    async with awith_checkpointer(checkpointer_name) as checkpointer:
        # Define addition subgraph
        @entrypoint()
        async def add(inputs):
            a, b = inputs
            return a + b

        # Define multiplication subgraph using tasks
        @task
        async def multiply_task(a, b):
            return a * b

        @entrypoint()
        async def multiply(inputs):
            return await multiply_task(*inputs)

        # Test calling the same subgraph multiple times
        @task
        async def call_same_subgraph(a, b):
            result = await add.ainvoke([a, b])
            another_result = await add.ainvoke([result, 10])
            return another_result

        @entrypoint(checkpointer=checkpointer)
        async def parent_call_same_subgraph(inputs):
            return await call_same_subgraph(*inputs)

        config = {"configurable": {"thread_id": "1"}}
        assert await parent_call_same_subgraph.ainvoke([2, 3], config) == 15

        # Test calling multiple subgraphs
        @task
        async def call_multiple_subgraphs(a, b):
            add_result = await add.ainvoke([a, b])
            multiply_result = await multiply.ainvoke([a, b])
            return [add_result, multiply_result]

        @entrypoint(checkpointer=checkpointer)
        async def parent_call_multiple_subgraphs(inputs):
            return await call_multiple_subgraphs(*inputs)

        config = {"configurable": {"thread_id": "2"}}
        assert await parent_call_multiple_subgraphs.ainvoke([2, 3], config) == [5, 6]


@NEEDS_CONTEXTVARS
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_multiple_subgraphs_mixed_entrypoint(checkpointer_name: str) -> None:
    """Test calling multiple StateGraph subgraphs from an entrypoint."""

    class State(TypedDict):
        a: int
        b: int

    class Output(TypedDict):
        result: int

    async with awith_checkpointer(checkpointer_name) as checkpointer:
        # Define the subgraphs
        async def add(state):
            return {"result": state["a"] + state["b"]}

        add_subgraph = (
            StateGraph(State, output=Output)
            .add_node(add)
            .add_edge(START, "add")
            .compile()
        )

        async def multiply(state):
            return {"result": state["a"] * state["b"]}

        multiply_subgraph = (
            StateGraph(State, output=Output)
            .add_node(multiply)
            .add_edge(START, "multiply")
            .compile()
        )

        # Test calling the same subgraph multiple times
        @task
        async def call_same_subgraph(a, b):
            result = (await add_subgraph.ainvoke({"a": a, "b": b}))["result"]
            another_result = (await add_subgraph.ainvoke({"a": result, "b": 10}))[
                "result"
            ]
            return another_result

        @entrypoint(checkpointer=checkpointer)
        async def parent_call_same_subgraph(inputs):
            return await call_same_subgraph(*inputs)

        config = {"configurable": {"thread_id": "1"}}
        assert await parent_call_same_subgraph.ainvoke([2, 3], config) == 15

        # Test calling multiple subgraphs
        @task
        async def call_multiple_subgraphs(a, b):
            add_result = (await add_subgraph.ainvoke({"a": a, "b": b}))["result"]
            multiply_result = (await multiply_subgraph.ainvoke({"a": a, "b": b}))[
                "result"
            ]
            return [add_result, multiply_result]

        @entrypoint(checkpointer=checkpointer)
        async def parent_call_multiple_subgraphs(inputs):
            return await call_multiple_subgraphs(*inputs)

        config = {"configurable": {"thread_id": "2"}}
        assert await parent_call_multiple_subgraphs.ainvoke([2, 3], config) == [5, 6]


@NEEDS_CONTEXTVARS
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_multiple_subgraphs_mixed_state_graph(
    request: pytest.FixtureRequest, checkpointer_name: str
) -> None:
    """Test calling multiple entrypoint "subgraphs" from a StateGraph."""
    async with awith_checkpointer(checkpointer_name) as checkpointer:

        class State(TypedDict):
            a: int
            b: int

        class Output(TypedDict):
            result: int

        # Define addition subgraph
        @entrypoint()
        async def add(inputs):
            a, b = inputs
            return a + b

        # Define multiplication subgraph using tasks
        @task
        async def multiply_task(a, b):
            return a * b

        @entrypoint()
        async def multiply(inputs):
            return await multiply_task(*inputs)

        # Test calling the same subgraph multiple times
        async def call_same_subgraph(state):
            result = await add.ainvoke([state["a"], state["b"]])
            another_result = await add.ainvoke([result, 10])
            return {"result": another_result}

        parent_call_same_subgraph = (
            StateGraph(State, output=Output)
            .add_node(call_same_subgraph)
            .add_edge(START, "call_same_subgraph")
            .compile(checkpointer=checkpointer)
        )
        config = {"configurable": {"thread_id": "1"}}
        assert await parent_call_same_subgraph.ainvoke({"a": 2, "b": 3}, config) == {
            "result": 15
        }

        # Test calling multiple subgraphs
        class Output(TypedDict):
            add_result: int
            multiply_result: int

        async def call_multiple_subgraphs(state):
            add_result = await add.ainvoke([state["a"], state["b"]])
            multiply_result = await multiply.ainvoke([state["a"], state["b"]])
            return {
                "add_result": add_result,
                "multiply_result": multiply_result,
            }

        parent_call_multiple_subgraphs = (
            StateGraph(State, output=Output)
            .add_node(call_multiple_subgraphs)
            .add_edge(START, "call_multiple_subgraphs")
            .compile(checkpointer=checkpointer)
        )
        config = {"configurable": {"thread_id": "2"}}
        assert await parent_call_multiple_subgraphs.ainvoke(
            {"a": 2, "b": 3}, config
        ) == {
            "add_result": 5,
            "multiply_result": 6,
        }


@NEEDS_CONTEXTVARS
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_multiple_subgraphs_checkpointer(checkpointer_name: str) -> None:
    async with awith_checkpointer(checkpointer_name) as checkpointer:

        class SubgraphState(TypedDict):
            sub_counter: Annotated[int, operator.add]

        async def subgraph_node(state):
            return {"sub_counter": 2}

        sub_graph_1 = (
            StateGraph(SubgraphState)
            .add_node(subgraph_node)
            .add_edge(START, "subgraph_node")
            .compile(checkpointer=True)
        )

        class OtherSubgraphState(TypedDict):
            other_sub_counter: Annotated[int, operator.add]

        async def other_subgraph_node(state):
            return {"other_sub_counter": 3}

        sub_graph_2 = (
            StateGraph(OtherSubgraphState)
            .add_node(other_subgraph_node)
            .add_edge(START, "other_subgraph_node")
            .compile()
        )

        class ParentState(TypedDict):
            parent_counter: int

        async def parent_node(state):
            result = await sub_graph_1.ainvoke({"sub_counter": state["parent_counter"]})
            other_result = await sub_graph_2.ainvoke(
                {"other_sub_counter": result["sub_counter"]}
            )
            return {"parent_counter": other_result["other_sub_counter"]}

        parent_graph = (
            StateGraph(ParentState)
            .add_node(parent_node)
            .add_edge(START, "parent_node")
            .compile(checkpointer=checkpointer)
        )

        config = {"configurable": {"thread_id": "1"}}
        assert await parent_graph.ainvoke({"parent_counter": 0}, config) == {
            "parent_counter": 5
        }
        assert await parent_graph.ainvoke({"parent_counter": 0}, config) == {
            "parent_counter": 7
        }
        config = {"configurable": {"thread_id": "2"}}
        assert [
            c
            async for c in parent_graph.astream(
                {"parent_counter": 0}, config, subgraphs=True, stream_mode="updates"
            )
        ] == [
            (("parent_node",), {"subgraph_node": {"sub_counter": 2}}),
            (
                (AnyStr("parent_node:"), "1"),
                {"other_subgraph_node": {"other_sub_counter": 3}},
            ),
            ((), {"parent_node": {"parent_counter": 5}}),
        ]
        assert [
            c
            async for c in parent_graph.astream(
                {"parent_counter": 0}, config, subgraphs=True, stream_mode="updates"
            )
        ] == [
            (("parent_node",), {"subgraph_node": {"sub_counter": 2}}),
            (
                (AnyStr("parent_node:"), "1"),
                {"other_subgraph_node": {"other_sub_counter": 3}},
            ),
            ((), {"parent_node": {"parent_counter": 7}}),
        ]


@NEEDS_CONTEXTVARS
async def test_async_entrypoint_without_checkpointer() -> None:
    """Test no checkpointer."""
    states = []
    config = {"configurable": {"thread_id": "1"}}

    # Test without previous
    @entrypoint()
    async def foo(inputs: Any) -> Any:
        states.append(inputs)
        return inputs

    assert (await foo.ainvoke({"a": "1"}, config)) == {"a": "1"}

    @entrypoint()
    async def foo(inputs: Any, *, previous: Any) -> Any:
        states.append(previous)
        return {"previous": previous, "current": inputs}

    assert (await foo.ainvoke({"a": "1"}, config)) == {
        "current": {"a": "1"},
        "previous": None,
    }
    assert (await foo.ainvoke({"a": "1"}, config)) == {
        "current": {"a": "1"},
        "previous": None,
    }


async def test_entrypoint_from_async_generator() -> None:
    """@entrypoint does not support sync generators."""
    with pytest.raises(NotImplementedError):

        @entrypoint(checkpointer=MemorySaver())
        async def foo(inputs) -> Any:
            yield "a"
            yield "b"


@NEEDS_CONTEXTVARS
async def test_named_tasks_functional() -> None:
    class Foo:
        async def foo(self, value: str) -> dict:
            return value + "foo"

    f = Foo()

    # class method task
    foo = task(f.foo, name="custom_foo")
    other_foo = task(f.foo, name="other_foo")

    # regular function task
    @task(name="custom_bar")
    async def bar(value: str) -> dict:
        return value + "|bar"

    async def baz(update: str, value: str) -> dict:
        return value + f"|{update}"

    # partial function task (unnamed)
    baz_task = task(functools.partial(baz, "baz"))
    # partial function task (named_)
    custom_baz_task = task(functools.partial(baz, "custom_baz"), name="custom_baz")

    class Qux:
        def __call__(self, value: str) -> dict:
            return value + "|qux"

    qux_task = task(Qux(), name="qux")

    @entrypoint()
    async def workflow(inputs: dict) -> dict:
        foo_result = await foo(inputs)
        await other_foo(inputs)
        bar_result = await bar(foo_result)
        baz_result = await baz_task(bar_result)
        custom_baz_result = await custom_baz_task(baz_result)
        qux_result = await qux_task(custom_baz_result)
        return qux_result

    assert [c async for c in workflow.astream("", stream_mode="updates")] == [
        {"custom_foo": "foo"},
        {"other_foo": "foo"},
        {"custom_bar": "foo|bar"},
        {"baz": "foo|bar|baz"},
        {"custom_baz": "foo|bar|baz|custom_baz"},
        {"qux": "foo|bar|baz|custom_baz|qux"},
        {"workflow": "foo|bar|baz|custom_baz|qux"},
    ]


@NEEDS_CONTEXTVARS
async def test_overriding_injectable_args_with_async_task() -> None:
    """Test overriding injectable args in tasks."""
    from langgraph.store.memory import InMemoryStore

    @task
    async def foo(store: BaseStore, writer: StreamWriter, value: Any) -> None:
        assert store is value
        assert writer is value

    @entrypoint(store=InMemoryStore())
    async def main(inputs, store: BaseStore) -> str:
        assert store is not None
        await foo(store=None, writer=None, value=None)
        await foo(store="hello", writer="hello", value="hello")
        return "OK"

    assert await main.ainvoke({}) == "OK"


async def test_tags_stream_mode_messages() -> None:
    model = GenericFakeChatModel(messages=iter(["foo"]), tags=["meow"])

    async def call_model(state, config):
        return {"messages": await model.ainvoke(state["messages"], config)}

    graph = (
        StateGraph(MessagesState)
        .add_node(call_model)
        .add_edge(START, "call_model")
        .compile()
    )
    assert [
        c
        async for c in graph.astream(
            {
                "messages": "hi",
            },
            stream_mode="messages",
        )
    ] == [
        (
            _AnyIdAIMessageChunk(content="foo"),
            {
                "langgraph_step": 1,
                "langgraph_node": "call_model",
                "langgraph_triggers": ("branch:to:call_model",),
                "langgraph_path": ("__pregel_pull", "call_model"),
                "langgraph_checkpoint_ns": AnyStr("call_model:"),
                "checkpoint_ns": AnyStr("call_model:"),
                "ls_provider": "genericfakechatmodel",
                "ls_model_type": "chat",
                "tags": ["meow"],
            },
        )
    ]


async def test_stream_mode_messages_command() -> None:
    from langchain_core.messages import HumanMessage

    async def my_node(state):
        return {"messages": HumanMessage(content="foo")}

    async def my_other_node(state):
        return Command(update={"messages": HumanMessage(content="bar")})

    graph = (
        StateGraph(MessagesState)
        .add_sequence([my_node, my_other_node])
        .add_edge(START, "my_node")
        .compile()
    )
    assert [
        c
        async for c in graph.astream(
            {
                "messages": [],
            },
            stream_mode="messages",
        )
    ] == [
        (
            _AnyIdHumanMessage(content="foo"),
            {
                "langgraph_step": 1,
                "langgraph_node": "my_node",
                "langgraph_triggers": ("branch:to:my_node",),
                "langgraph_path": ("__pregel_pull", "my_node"),
                "langgraph_checkpoint_ns": AnyStr("my_node:"),
            },
        ),
        (
            _AnyIdHumanMessage(content="bar"),
            {
                "langgraph_step": 2,
                "langgraph_node": "my_other_node",
                "langgraph_triggers": ("branch:to:my_other_node",),
                "langgraph_path": ("__pregel_pull", "my_other_node"),
                "langgraph_checkpoint_ns": AnyStr("my_other_node:"),
            },
        ),
    ]


async def test_stream_messages_dedupe_inputs() -> None:
    from langchain_core.messages import AIMessage

    async def call_model(state):
        return {"messages": AIMessage("hi", id="1")}

    async def route(state):
        return Command(goto="node_2", graph=Command.PARENT)

    subgraph = (
        StateGraph(MessagesState)
        .add_node(call_model)
        .add_node(route)
        .add_edge(START, "call_model")
        .add_edge("call_model", "route")
        .compile()
    )

    graph = (
        StateGraph(MessagesState)
        .add_node("node_1", subgraph)
        .add_node("node_2", lambda state: state)
        .add_edge(START, "node_1")
        .compile()
    )

    chunks = [
        chunk
        async for ns, chunk in graph.astream(
            {"messages": "hi"}, stream_mode="messages", subgraphs=True
        )
    ]

    assert len(chunks) == 1
    assert chunks[0][0] == AIMessage("hi", id="1")
    assert chunks[0][1]["langgraph_node"] == "call_model"


@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_stream_messages_dedupe_state(checkpointer_name: str) -> None:
    async with awith_checkpointer(checkpointer_name) as checkpointer:
        from langchain_core.messages import AIMessage

        to_emit = [AIMessage("bye", id="1"), AIMessage("bye again", id="2")]

        async def call_model(state):
            return {"messages": to_emit.pop(0)}

        async def route(state):
            return Command(goto="node_2", graph=Command.PARENT)

        subgraph = (
            StateGraph(MessagesState)
            .add_node(call_model)
            .add_node(route)
            .add_edge(START, "call_model")
            .add_edge("call_model", "route")
            .compile()
        )

        graph = (
            StateGraph(MessagesState)
            .add_node("node_1", subgraph)
            .add_node("node_2", lambda state: state)
            .add_edge(START, "node_1")
            .compile(checkpointer=checkpointer)
        )

        thread1 = {"configurable": {"thread_id": "1"}}

        chunks = [
            chunk
            async for ns, chunk in graph.astream(
                {"messages": "hi"}, thread1, stream_mode="messages", subgraphs=True
            )
        ]

        assert len(chunks) == 1
        assert chunks[0][0] == AIMessage("bye", id="1")
        assert chunks[0][1]["langgraph_node"] == "call_model"

        chunks = [
            chunk
            async for ns, chunk in graph.astream(
                {"messages": "hi again"},
                thread1,
                stream_mode="messages",
                subgraphs=True,
            )
        ]

        assert len(chunks) == 1
        assert chunks[0][0] == AIMessage("bye again", id="2")
        assert chunks[0][1]["langgraph_node"] == "call_model"


@NEEDS_CONTEXTVARS
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_interrupt_subgraph_reenter_checkpointer_true(
    checkpointer_name: str,
) -> None:
    async with awith_checkpointer(checkpointer_name) as checkpointer:

        class SubgraphState(TypedDict):
            foo: str
            bar: str

        class ParentState(TypedDict):
            foo: str
            counter: int

        called = []
        bar_values = []

        async def subnode_1(state: SubgraphState):
            called.append("subnode_1")
            bar_values.append(state.get("bar"))
            return {"foo": "subgraph_1"}

        async def subnode_2(state: SubgraphState):
            called.append("subnode_2")
            value = interrupt("Provide value")
            value += "baz"
            return {"foo": "subgraph_2", "bar": value}

        subgraph = (
            StateGraph(SubgraphState)
            .add_node(subnode_1)
            .add_node(subnode_2)
            .add_edge(START, "subnode_1")
            .add_edge("subnode_1", "subnode_2")
            .compile(checkpointer=True)
        )

        async def call_subgraph(state: ParentState):
            called.append("call_subgraph")
            return await subgraph.ainvoke(state)

        async def node(state: ParentState):
            called.append("parent")
            if state["counter"] < 1:
                return Command(
                    goto="call_subgraph", update={"counter": state["counter"] + 1}
                )

            return {"foo": state["foo"] + "|" + "parent"}

        parent = (
            StateGraph(ParentState)
            .add_node(call_subgraph)
            .add_node(node)
            .add_edge(START, "call_subgraph")
            .add_edge("call_subgraph", "node")
            .compile(checkpointer=checkpointer)
        )

        config = {"configurable": {"thread_id": "1"}}
        assert await parent.ainvoke({"foo": "", "counter": 0}, config) == {
            "foo": "",
            "counter": 0,
        }
        assert await parent.ainvoke(Command(resume="bar"), config) == {
            "foo": "subgraph_2",
            "counter": 1,
        }
        assert await parent.ainvoke(Command(resume="qux"), config) == {
            "foo": "subgraph_2|parent",
            "counter": 1,
        }
        assert called == [
            "call_subgraph",
            "subnode_1",
            "subnode_2",
            "call_subgraph",
            "subnode_2",
            "parent",
            "call_subgraph",
            "subnode_1",
            "subnode_2",
            "call_subgraph",
            "subnode_2",
            "parent",
        ]

        # invoke parent again (new turn)
        assert await parent.ainvoke({"foo": "meow", "counter": 0}, config) == {
            "foo": "meow",
            "counter": 0,
        }
        # confirm that we preserve the state values from the previous invocation
        assert bar_values == [None, "barbaz", "quxbaz"]


@NEEDS_CONTEXTVARS
async def test_handles_multiple_interrupts_from_tasks() -> None:
    @task
    async def add_participant(name: str) -> str:
        feedback = interrupt(f"Hey do you want to add {name}?")

        if feedback is False:
            return f"The user changed their mind and doesn't want to add {name}!"

        if feedback is True:
            return f"Added {name}!"

        raise ValueError("Invalid feedback")

    @entrypoint(checkpointer=MemorySaver())
    async def program(_state: Any) -> list[str]:
        first = await add_participant("James")
        second = await add_participant("Will")
        return [first, second]

    config = {"configurable": {"thread_id": "1"}}

    result = await program.ainvoke("this is ignored", config=config)
    assert result is None

    state = await program.aget_state(config=config)
    assert len(state.tasks[0].interrupts) == 1
    task_interrupt = state.tasks[0].interrupts[0]
    assert task_interrupt.resumable is True
    assert len(task_interrupt.ns) == 2
    assert task_interrupt.ns[0].startswith("program:")
    assert task_interrupt.ns[1].startswith("add_participant:")
    assert task_interrupt.value == "Hey do you want to add James?"

    result = await program.ainvoke(Command(resume=True), config=config)
    assert result is None

    state = await program.aget_state(config=config)
    assert len(state.tasks[0].interrupts) == 1
    task_interrupt = state.tasks[0].interrupts[0]
    assert task_interrupt.resumable is True
    assert len(task_interrupt.ns) == 2
    assert task_interrupt.ns[0].startswith("program:")
    assert task_interrupt.ns[1].startswith("add_participant:")
    assert task_interrupt.value == "Hey do you want to add Will?"

    result = await program.ainvoke(Command(resume=True), config=config)
    assert result is not None
    assert len(result) == 2
    assert result[0] == "Added James!"
    assert result[1] == "Added Will!"


async def test_pregel_loop_refcount():
    gc.collect()
    try:
        gc.disable()

        class State(TypedDict):
            messages: Annotated[list, add_messages]

        graph_builder = StateGraph(State)

        async def chatbot(state: State):
            return {"messages": [("ai", "HIYA")]}

        graph_builder.add_node("chatbot", chatbot)
        graph_builder.set_entry_point("chatbot")
        graph_builder.set_finish_point("chatbot")
        graph = graph_builder.compile()

        for _ in range(5):
            await graph.ainvoke({"messages": [{"role": "user", "content": "hi"}]})
            assert (
                len(
                    [
                        obj
                        for obj in gc.get_objects()
                        if isinstance(obj, AsyncPregelLoop)
                    ]
                )
                == 0
            )
            assert (
                len([obj for obj in gc.get_objects() if isinstance(obj, PregelRunner)])
                == 0
            )
    finally:
        gc.enable()


@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC)
async def test_bulk_state_updates(checkpointer_name: str) -> None:
    async with awith_checkpointer(checkpointer_name) as checkpointer:

        class State(TypedDict):
            foo: str
            baz: str

        def node_a(state: State) -> State:
            return {"foo": "bar"}

        def node_b(state: State) -> State:
            return {"baz": "qux"}

        graph = (
            StateGraph(State)
            .add_node("node_a", node_a)
            .add_node("node_b", node_b)
            .add_edge(START, "node_a")
            .add_edge("node_a", "node_b")
            .compile(checkpointer=checkpointer)
        )

        config = {"configurable": {"thread_id": "1"}}

        # First update with node_a
        await graph.abulk_update_state(
            config,
            [
                [
                    StateUpdate({"foo": "bar"}, "node_a"),
                ]
            ],
        )

        # Then bulk update with both nodes
        await graph.abulk_update_state(
            config,
            [
                [
                    StateUpdate({"foo": "updated"}, "node_a"),
                    StateUpdate({"baz": "new"}, "node_b"),
                ]
            ],
        )

        state = await graph.aget_state(config)
        assert state.values == {"foo": "updated", "baz": "new"}

        # Check if there are only two checkpoints
        checkpoints = [
            c async for c in checkpointer.alist({"configurable": {"thread_id": "1"}})
        ]
        assert len(checkpoints) == 2
        assert checkpoints[0].metadata["writes"] == {
            "node_a": {"foo": "updated"},
            "node_b": {"baz": "new"},
        }
        assert checkpoints[1].metadata["writes"] == {"node_a": {"foo": "bar"}}

        # perform multiple steps at the same time
        config = {"configurable": {"thread_id": "2"}}

        await graph.abulk_update_state(
            config,
            [
                [
                    StateUpdate({"foo": "bar"}, "node_a"),
                ],
                [
                    StateUpdate({"foo": "updated"}, "node_a"),
                    StateUpdate({"baz": "new"}, "node_b"),
                ],
            ],
        )

        state = await graph.aget_state(config)
        assert state.values == {"foo": "updated", "baz": "new"}

        checkpoints = [
            c async for c in checkpointer.alist({"configurable": {"thread_id": "1"}})
        ]
        assert len(checkpoints) == 2
        assert checkpoints[0].metadata["writes"] == {
            "node_a": {"foo": "updated"},
            "node_b": {"baz": "new"},
        }
        assert checkpoints[1].metadata["writes"] == {"node_a": {"foo": "bar"}}

        # Should raise error if updating without as_node
        with pytest.raises(InvalidUpdateError):
            await graph.abulk_update_state(
                config,
                [
                    [
                        StateUpdate(values={"foo": "error"}, as_node=None),
                        StateUpdate(values={"bar": "error"}, as_node=None),
                    ]
                ],
            )

        # Should raise if no updates are provided
        with pytest.raises(ValueError, match="No supersteps provided"):
            await graph.abulk_update_state(config, [])

        # Should raise if no updates are provided
        with pytest.raises(ValueError, match="No updates provided"):
            await graph.abulk_update_state(config, [[], []])

        # Should raise if __end__ or __copy__ update is applied in bulk
        with pytest.raises(InvalidUpdateError):
            await graph.abulk_update_state(
                config,
                [
                    [
                        StateUpdate(values=None, as_node="__end__"),
                        StateUpdate(values=None, as_node="__copy__"),
                    ],
                ],
            )


@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC)
async def test_update_as_input(checkpointer_name: str) -> None:
    async with awith_checkpointer(checkpointer_name) as checkpointer:

        class State(TypedDict):
            foo: str

        def agent(state: State) -> State:
            return {"foo": "agent"}

        def tool(state: State) -> State:
            return {"foo": "tool"}

        graph = (
            StateGraph(State)
            .add_node("agent", agent)
            .add_node("tool", tool)
            .add_edge(START, "agent")
            .add_edge("agent", "tool")
            .compile(checkpointer=checkpointer)
        )

        assert await graph.ainvoke(
            {"foo": "input"}, {"configurable": {"thread_id": "1"}}
        ) == {"foo": "tool"}

        assert await graph.ainvoke(
            {"foo": "input"}, {"configurable": {"thread_id": "1"}}
        ) == {"foo": "tool"}

        def map_snapshot(i: StateSnapshot) -> dict:
            return {
                "values": i.values,
                "next": i.next,
                "step": i.metadata.get("step"),
            }

        history = [
            map_snapshot(s)
            async for s in graph.aget_state_history(
                {"configurable": {"thread_id": "1"}}
            )
        ]

        await graph.abulk_update_state(
            {"configurable": {"thread_id": "2"}},
            [
                # First turn
                [StateUpdate({"foo": "input"}, "__input__")],
                [StateUpdate({"foo": "input"}, "__start__")],
                [StateUpdate({"foo": "agent"}, "agent")],
                [StateUpdate({"foo": "tool"}, "tool")],
                # Second turn
                [StateUpdate({"foo": "input"}, "__input__")],
                [StateUpdate({"foo": "input"}, "__start__")],
                [StateUpdate({"foo": "agent"}, "agent")],
                [StateUpdate({"foo": "tool"}, "tool")],
            ],
        )

        state = await graph.aget_state({"configurable": {"thread_id": "2"}})
        assert state.values == {"foo": "tool"}

        new_history = [
            map_snapshot(s)
            async for s in graph.aget_state_history(
                {"configurable": {"thread_id": "2"}}
            )
        ]

        assert new_history == history


@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC)
async def test_batch_update_as_input(checkpointer_name: str) -> None:
    async with awith_checkpointer(checkpointer_name) as checkpointer:

        class State(TypedDict):
            foo: str
            tasks: Annotated[list[int], operator.add]

        def agent(state: State) -> State:
            return {"foo": "agent"}

        def map(state: State) -> Command["task"]:
            return Command(
                goto=[
                    Send("task", {"index": 0}),
                    Send("task", {"index": 1}),
                    Send("task", {"index": 2}),
                ],
                update={"foo": "map"},
            )

        def task(state: dict) -> State:
            return {"tasks": [state["index"]]}

        graph = (
            StateGraph(State)
            .add_node("agent", agent)
            .add_node("map", map)
            .add_node("task", task)
            .add_edge(START, "agent")
            .add_edge("agent", "map")
            .compile(checkpointer=checkpointer)
        )

        assert await graph.ainvoke(
            {"foo": "input"}, {"configurable": {"thread_id": "1"}}
        ) == {"foo": "map", "tasks": [0, 1, 2]}

        def map_snapshot(i: StateSnapshot) -> dict:
            return {
                "values": i.values,
                "next": i.next,
                "step": i.metadata.get("step"),
                "tasks": [t.name for t in i.tasks],
            }

        history = [
            map_snapshot(s)
            async for s in graph.aget_state_history(
                {"configurable": {"thread_id": "1"}}
            )
        ]

        await graph.abulk_update_state(
            {"configurable": {"thread_id": "2"}},
            [
                [StateUpdate({"foo": "input"}, "__input__")],
                [StateUpdate({"foo": "input"}, "__start__")],
                [StateUpdate({"foo": "agent", "tasks": []}, "agent")],
                [
                    StateUpdate(
                        Command(
                            goto=[
                                Send("task", {"index": 0}),
                                Send("task", {"index": 1}),
                                Send("task", {"index": 2}),
                            ],
                            update={"foo": "map"},
                        ),
                        "map",
                    )
                ],
                [
                    StateUpdate({"tasks": [0]}, "task"),
                    StateUpdate({"tasks": [1]}, "task"),
                    StateUpdate({"tasks": [2]}, "task"),
                ],
            ],
        )

        state = await graph.aget_state({"configurable": {"thread_id": "2"}})
        assert state.values == {"foo": "map", "tasks": [0, 1, 2]}

        new_history = [
            map_snapshot(s)
            async for s in graph.aget_state_history(
                {"configurable": {"thread_id": "2"}}
            )
        ]

        assert new_history == history