Merge branch 'main' into wfh/_/pgsearch

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

@@ -6,11 +6,15 @@ scoped to user IDs, assistant IDs, or other arbitrary namespaces.
 
 from abc import ABC, abstractmethod
 from datetime import datetime
-from typing import ( Any, Iterable, Literal, NamedTuple,
-                    Optional, TypedDict, Union, cast)
+from typing import Any, Iterable, Literal, NamedTuple, Optional, TypedDict, Union, cast
 
 from langchain_core.embeddings import Embeddings
-from langgraph.store.base._embed import AEmbeddingsFunc, EmbeddingsFunc, ensure_embeddings
+
+from langgraph.store.base._embed import (
+    AEmbeddingsFunc,
+    EmbeddingsFunc,
+    ensure_embeddings,
+)
 
 
 class Item:
@@ -528,6 +532,7 @@ class BaseStore(ABC):
         )
         return (await self.abatch([op]))[0]
 
+
 __all__ = [
     "BaseStore",
     "Item",
@@ -541,4 +546,4 @@ __all__ = [
     "NamespaceMatchType",
     "Embeddings",
     "ensure_embeddings",
-]
+]

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

@@ -7,8 +7,7 @@ asynchronous operations.
 """
 
 import asyncio
-from typing import (Any, Awaitable, Callable, List, Optional, Sequence,
-                    TypeGuard, Union)
+from typing import Any, Awaitable, Callable, List, Optional, Sequence, Union
 
 from langchain_core.embeddings import Embeddings
 
@@ -30,7 +29,7 @@ Similar to EmbeddingsFunc, but returns an awaitable that resolves to the embeddi
 def ensure_embeddings(
     embed: Union[Embeddings, EmbeddingsFunc, AEmbeddingsFunc, None],
     *,
-    aembed: Optional[AEmbeddingsFunc] = None
+    aembed: Optional[AEmbeddingsFunc] = None,
 ) -> Embeddings:
     """Ensure that an embedding function conforms to LangChain's Embeddings interface.
 
@@ -180,14 +179,14 @@ class EmbeddingsLambda(Embeddings):
 
 def _is_async_callable(
     func: Any,
-) -> TypeGuard[Callable[..., Awaitable]]:
+) -> bool:
     """Check if a function is async.
-    
+
     This includes both async def functions and classes with async __call__ methods.
-    
+
     Args:
         func: Function or callable object to check.
-        
+
     Returns:
         True if the function is async, False otherwise.
     """

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

@@ -6,6 +6,9 @@ from langgraph.store.base import (
     BaseStore,
     GetOp,
     Item,
+    ListNamespacesOp,
+    MatchCondition,
+    NameSpacePath,
     Op,
     PutOp,
     SearchOp,
@@ -69,6 +72,74 @@ class AsyncBatchedBaseStore(BaseStore):
         self._aqueue[fut] = PutOp(namespace, key, None)
         return await fut
 
+    async def alist_namespaces(
+        self,
+        *,
+        prefix: Optional[NameSpacePath] = None,
+        suffix: Optional[NameSpacePath] = None,
+        max_depth: Optional[int] = None,
+        limit: int = 100,
+        offset: int = 0,
+    ) -> list[tuple[str, ...]]:
+        fut = self._loop.create_future()
+        match_conditions = []
+        if prefix:
+            match_conditions.append(MatchCondition(match_type="prefix", path=prefix))
+        if suffix:
+            match_conditions.append(MatchCondition(match_type="suffix", path=suffix))
+
+        op = ListNamespacesOp(
+            match_conditions=tuple(match_conditions),
+            max_depth=max_depth,
+            limit=limit,
+            offset=offset,
+        )
+        self._aqueue[fut] = op
+        return await fut
+
+
+def _dedupe_ops(values: list[Op]) -> tuple[Optional[list[int]], list[Op]]:
+    """Dedupe operations while preserving order for results.
+
+    Args:
+        values: List of operations to dedupe
+
+    Returns:
+        Tuple of (listen indices, deduped operations)
+        where listen indices map deduped operation results back to original positions
+    """
+    if len(values) <= 1:
+        return None, list(values)
+
+    dedupped: list[Op] = []
+    listen: list[int] = []
+    puts: dict[tuple[tuple[str, ...], str], int] = {}
+
+    for op in values:
+        if isinstance(op, (GetOp, SearchOp, ListNamespacesOp)):
+            try:
+                listen.append(dedupped.index(op))
+            except ValueError:
+                listen.append(len(dedupped))
+                dedupped.append(op)
+        elif isinstance(op, PutOp):
+            putkey = (op.namespace, op.key)
+            if putkey in puts:
+                # Overwrite previous put
+                ix = puts[putkey]
+                dedupped[ix] = op
+                listen.append(ix)
+            else:
+                puts[putkey] = len(dedupped)
+                listen.append(len(dedupped))
+                dedupped.append(op)
+
+        else:  # Any new ops will be treated regularly
+            listen.append(len(dedupped))
+            dedupped.append(op)
+
+    return listen, dedupped
+
 
 async def _run(
     aqueue: dict[asyncio.Future, Op], store: weakref.ReferenceType[BaseStore]
@@ -82,7 +153,12 @@ async def _run(
             taken = aqueue.copy()
             # action each operation
             try:
-                results = await s.abatch(taken.values())
+                values = list(taken.values())
+                listen, dedupped = _dedupe_ops(values)
+                results = await s.abatch(dedupped)
+                if listen is not None:
+                    results = [results[ix] for ix in listen]
+
                 # set the results of each operation
                 for fut, result in zip(taken, results):
                     fut.set_result(result)

libs/checkpoint/tests/test_store.py

@@ -10,6 +10,18 @@ from langgraph.store.base.batch import AsyncBatchedBaseStore
 from langgraph.store.memory import InMemoryStore
 
 
+class MockAsyncBatchedStore(AsyncBatchedBaseStore):
+    def __init__(self) -> None:
+        super().__init__()
+        self._store = InMemoryStore()
+
+    def batch(self, ops: Iterable[Op]) -> list[Result]:
+        return self._store.batch(ops)
+
+    async def abatch(self, ops: Iterable[Op]) -> list[Result]:
+        return self._store.batch(ops)
+
+
 async def test_async_batch_store(mocker: MockerFixture) -> None:
     abatch = mocker.stub()
 
@@ -313,17 +325,6 @@ async def test_cannot_put_empty_namespace() -> None:
     store.delete(("langgraph", "foo"), "bar")
     assert store.get(("langgraph", "foo"), "bar") is None
 
-    class MockAsyncBatchedStore(AsyncBatchedBaseStore):
-        def __init__(self) -> None:
-            super().__init__()
-            self._store = InMemoryStore()
-
-        def batch(self, ops: Iterable[Op]) -> list[Result]:
-            return self._store.batch(ops)
-
-        async def abatch(self, ops: Iterable[Op]) -> list[Result]:
-            return self._store.batch(ops)
-
     async_store = MockAsyncBatchedStore()
     doc = {"foo": "bar"}
 
@@ -354,3 +355,68 @@ async def test_cannot_put_empty_namespace() -> None:
     assert (await async_store.asearch(("valid", "namespace")))[0].value == doc
     await async_store.adelete(("valid", "namespace"), "key")
     assert (await async_store.aget(("valid", "namespace"), "key")) is None
+
+
+async def test_async_batch_store_deduplication(mocker: MockerFixture) -> None:
+    abatch = mocker.spy(InMemoryStore, "batch")
+    store = MockAsyncBatchedStore()
+
+    same_doc = {"value": "same"}
+    diff_doc = {"value": "different"}
+    await asyncio.gather(
+        store.aput(namespace=("test",), key="same", value=same_doc),
+        store.aput(namespace=("test",), key="different", value=diff_doc),
+    )
+    abatch.reset_mock()
+
+    results = await asyncio.gather(
+        store.aget(namespace=("test",), key="same"),
+        store.aget(namespace=("test",), key="same"),
+        store.aget(namespace=("test",), key="different"),
+    )
+
+    assert len(results) == 3
+    assert results[0] == results[1]
+    assert results[0] != results[2]
+    assert results[0].value == same_doc  # type: ignore
+    assert results[2].value == diff_doc  # type: ignore
+    assert len(abatch.call_args_list) == 1
+    ops = list(abatch.call_args_list[0].args[1])
+    assert len(ops) == 2
+    assert GetOp(("test",), "same") in ops
+    assert GetOp(("test",), "different") in ops
+
+    abatch.reset_mock()
+
+    doc1 = {"value": 1}
+    doc2 = {"value": 2}
+    results = await asyncio.gather(
+        store.aput(namespace=("test",), key="key", value=doc1),
+        store.aput(namespace=("test",), key="key", value=doc2),
+    )
+    assert len(abatch.call_args_list) == 1
+    ops = list(abatch.call_args_list[0].args[1])
+    assert len(ops) == 1
+    assert ops[0] == PutOp(("test",), "key", doc2)
+    assert len(results) == 2
+    assert all(result is None for result in results)
+
+    result = await store.aget(namespace=("test",), key="key")
+    assert result is not None
+    assert result.value == doc2
+
+    abatch.reset_mock()
+
+    results = await asyncio.gather(
+        store.asearch(("test",), filter={"value": 2}),
+        store.asearch(("test",), filter={"value": 2}),
+    )
+    assert len(abatch.call_args_list) == 1
+    ops = list(abatch.call_args_list[0].args[1])
+    assert len(ops) == 1
+    assert len(results) == 2
+    assert results[0] == results[1]
+    assert len(results[0]) == 1
+    assert results[0][0].value == doc2
+
+    abatch.reset_mock()

libs/langgraph/langgraph/errors.py

@@ -2,7 +2,7 @@ from enum import Enum
 from typing import Any, Sequence
 
 from langgraph.checkpoint.base import EmptyChannelError  # noqa: F401
-from langgraph.types import Interrupt
+from langgraph.types import Command, Interrupt
 
 # EmptyChannelError re-exported for backwards compatibility
 
@@ -58,7 +58,11 @@ class InvalidUpdateError(Exception):
     pass
 
 
-class GraphInterrupt(Exception):
+class GraphBubbleUp(Exception):
+    pass
+
+
+class GraphInterrupt(GraphBubbleUp):
     """Raised when a subgraph is interrupted, suppressed by the root graph.
     Never raised directly, or surfaced to the user."""
 
@@ -73,13 +77,20 @@ class NodeInterrupt(GraphInterrupt):
         super().__init__([Interrupt(value=value)])
 
 
-class GraphDelegate(Exception):
+class GraphDelegate(GraphBubbleUp):
     """Raised when a graph is delegated (for distributed mode)."""
 
     def __init__(self, *args: dict[str, Any]) -> None:
         super().__init__(*args)
 
 
+class ParentCommand(GraphBubbleUp):
+    args: tuple[Command]
+
+    def __init__(self, command: Command) -> None:
+        super().__init__(command)
+
+
 class EmptyInputError(Exception):
     """Raised when graph receives an empty input."""
 

libs/langgraph/langgraph/graph/state.py

@@ -37,7 +37,12 @@ from langgraph.channels.ephemeral_value import EphemeralValue
 from langgraph.channels.last_value import LastValue
 from langgraph.channels.named_barrier_value import NamedBarrierValue
 from langgraph.constants import EMPTY_SEQ, NS_END, NS_SEP, SELF, TAG_HIDDEN
-from langgraph.errors import ErrorCode, InvalidUpdateError, create_error_message
+from langgraph.errors import (
+    ErrorCode,
+    InvalidUpdateError,
+    ParentCommand,
+    create_error_message,
+)
 from langgraph.graph.graph import END, START, Branch, CompiledGraph, Graph, Send
 from langgraph.managed.base import (
     ChannelKeyPlaceholder,
@@ -623,6 +628,8 @@ class CompiledStateGraph(CompiledGraph):
 
         def _get_root(input: Any) -> Any:
             if isinstance(input, Command):
+                if input.graph == Command.PARENT:
+                    return SKIP_WRITE
                 return input.update
             else:
                 return input
@@ -640,6 +647,8 @@ class CompiledStateGraph(CompiledGraph):
                     )
                 return input.get(key, SKIP_WRITE)
             elif isinstance(input, Command):
+                if input.graph == Command.PARENT:
+                    return SKIP_WRITE
                 return _get_state_key(input.update, key=key)
             elif get_type_hints(type(input)):
                 value = getattr(input, key, SKIP_WRITE)
@@ -822,6 +831,8 @@ def _control_branch(value: Any) -> Sequence[Union[str, Send]]:
         return [value]
     if not isinstance(value, GraphCommand):
         return EMPTY_SEQ
+    if value.graph == Command.PARENT:
+        raise ParentCommand(value)
     rtn: list[Union[str, Send]] = []
     if isinstance(value.goto, str):
         rtn.append(value.goto)
@@ -839,6 +850,8 @@ async def _acontrol_branch(value: Any) -> Sequence[Union[str, Send]]:
         return [value]
     if not isinstance(value, GraphCommand):
         return EMPTY_SEQ
+    if value.graph == Command.PARENT:
+        raise ParentCommand(value)
     rtn: list[Union[str, Send]] = []
     if isinstance(value.goto, str):
         rtn.append(value.goto)

libs/langgraph/langgraph/prebuilt/tool_node.py

@@ -37,7 +37,7 @@ from langchain_core.tools import tool as create_tool
 from langchain_core.tools.base import get_all_basemodel_annotations
 from typing_extensions import Annotated, get_args, get_origin
 
-from langgraph.errors import GraphInterrupt
+from langgraph.errors import GraphBubbleUp
 from langgraph.store.base import BaseStore
 from langgraph.utils.runnable import RunnableCallable
 
@@ -275,7 +275,7 @@ class ToolNode(RunnableCallable):
         # (2) a NodeInterrupt is raised inside a graph node for a graph called as a tool
         # (3) a GraphInterrupt is raised when a subgraph is interrupted inside a graph called as a tool
         # (2 and 3 can happen in a "supervisor w/ tools" multi-agent architecture)
-        except GraphInterrupt as e:
+        except GraphBubbleUp as e:
             raise e
         except Exception as e:
             if isinstance(self.handle_tool_errors, tuple):
@@ -316,7 +316,7 @@ class ToolNode(RunnableCallable):
         # (2) a NodeInterrupt is raised inside a graph node for a graph called as a tool
         # (3) a GraphInterrupt is raised when a subgraph is interrupted inside a graph called as a tool
         # (2 and 3 can happen in a "supervisor w/ tools" multi-agent architecture)
-        except GraphInterrupt as e:
+        except GraphBubbleUp as e:
             raise e
         except Exception as e:
             if isinstance(self.handle_tool_errors, tuple):

libs/langgraph/langgraph/pregel/algo.py

@@ -602,6 +602,7 @@ def prepare_single_task(
                     None,
                     task_id,
                     task_path,
+                    writers=proc.flat_writers,
                 )
 
         else:
@@ -720,6 +721,7 @@ def prepare_single_task(
                         None,
                         task_id,
                         task_path,
+                        writers=proc.flat_writers,
                     )
             else:
                 return PregelTask(task_id, name, task_path)

libs/langgraph/langgraph/pregel/executor.py

@@ -20,7 +20,7 @@ from langchain_core.runnables import RunnableConfig
 from langchain_core.runnables.config import get_executor_for_config
 from typing_extensions import ParamSpec
 
-from langgraph.errors import GraphInterrupt
+from langgraph.errors import GraphBubbleUp
 
 P = ParamSpec("P")
 T = TypeVar("T")
@@ -68,7 +68,7 @@ class BackgroundExecutor(ContextManager):
     def done(self, task: concurrent.futures.Future) -> None:
         try:
             task.result()
-        except GraphInterrupt:
+        except GraphBubbleUp:
             # This exception is an interruption signal, not an error
             # so we don't want to re-raise it on exit
             self.tasks.pop(task)
@@ -155,7 +155,7 @@ class AsyncBackgroundExecutor(AsyncContextManager):
             if exc := task.exception():
                 # This exception is an interruption signal, not an error
                 # so we don't want to re-raise it on exit
-                if isinstance(exc, GraphInterrupt):
+                if isinstance(exc, GraphBubbleUp):
                     self.tasks.pop(task)
             else:
                 self.tasks.pop(task)

libs/langgraph/langgraph/pregel/io.py

@@ -15,6 +15,7 @@ from langgraph.constants import (
     TAG_HIDDEN,
     TASKS,
 )
+from langgraph.errors import InvalidUpdateError
 from langgraph.pregel.log import logger
 from langgraph.types import Command, PregelExecutableTask, Send
 
@@ -68,6 +69,8 @@ def map_command(
     cmd: Command,
 ) -> Iterator[tuple[str, str, Any]]:
     """Map input chunk to a sequence of pending writes in the form (channel, value)."""
+    if cmd.graph == Command.PARENT:
+        raise InvalidUpdateError("There is not parent graph")
     if cmd.send:
         if isinstance(cmd.send, (tuple, list)):
             sends = cmd.send

libs/langgraph/langgraph/pregel/retry.py

@@ -2,6 +2,7 @@ import asyncio
 import logging
 import random
 import time
+from dataclasses import replace
 from functools import partial
 from typing import Any, Callable, Optional, Sequence
 
@@ -10,9 +11,10 @@ from langgraph.constants import (
     CONFIG_KEY_CHECKPOINT_NS,
     CONFIG_KEY_RESUMING,
     CONFIG_KEY_SEND,
+    NS_SEP,
 )
-from langgraph.errors import _SEEN_CHECKPOINT_NS, GraphInterrupt
-from langgraph.types import PregelExecutableTask, RetryPolicy
+from langgraph.errors import _SEEN_CHECKPOINT_NS, GraphBubbleUp, ParentCommand
+from langgraph.types import Command, PregelExecutableTask, RetryPolicy
 from langgraph.utils.config import patch_configurable
 
 logger = logging.getLogger(__name__)
@@ -40,7 +42,21 @@ def run_with_retry(
             task.proc.invoke(task.input, config)
             # if successful, end
             break
-        except GraphInterrupt:
+        except ParentCommand as exc:
+            ns: str = config[CONF][CONFIG_KEY_CHECKPOINT_NS]
+            cmd = exc.args[0]
+            if cmd.graph == ns:
+                # this command is for the current graph, handle it
+                for w in task.writers:
+                    w.invoke(cmd, config)
+                break
+            elif cmd.graph == Command.PARENT:
+                # this command is for the parent graph, assign it to the parent
+                parent_ns = NS_SEP.join(ns.split(NS_SEP)[:-1])
+                exc.args = (replace(cmd, graph=parent_ns),)
+            # bubble up
+            raise
+        except GraphBubbleUp:
             # if interrupted, end
             raise
         except Exception as exc:
@@ -118,7 +134,21 @@ async def arun_with_retry(
                 await task.proc.ainvoke(task.input, config)
             # if successful, end
             break
-        except GraphInterrupt:
+        except ParentCommand as exc:
+            ns: str = config[CONF][CONFIG_KEY_CHECKPOINT_NS]
+            cmd = exc.args[0]
+            if cmd.graph == ns:
+                # this command is for the current graph, handle it
+                for w in task.writers:
+                    w.invoke(cmd, config)
+                break
+            elif cmd.graph == Command.PARENT:
+                # this command is for the parent graph, assign it to the parent
+                parent_ns = NS_SEP.join(ns.split(NS_SEP)[:-1])
+                exc.args = (replace(cmd, graph=parent_ns),)
+            # bubble up
+            raise
+        except GraphBubbleUp:
             # if interrupted, end
             raise
         except Exception as exc:

libs/langgraph/langgraph/pregel/runner.py

@@ -23,7 +23,7 @@ from langgraph.constants import (
     PUSH,
     TAG_HIDDEN,
 )
-from langgraph.errors import GraphDelegate, GraphInterrupt
+from langgraph.errors import GraphBubbleUp, GraphInterrupt
 from langgraph.pregel.executor import Submit
 from langgraph.pregel.retry import arun_with_retry, run_with_retry
 from langgraph.types import PregelExecutableTask, RetryPolicy
@@ -298,7 +298,7 @@ class PregelRunner:
                 # save interrupt to checkpointer
                 if interrupts := [(INTERRUPT, i) for i in exception.args[0]]:
                     self.put_writes(task.id, interrupts)
-            elif isinstance(exception, GraphDelegate):
+            elif isinstance(exception, GraphBubbleUp):
                 raise exception
             else:
                 # save error to checkpointer
@@ -324,7 +324,7 @@ def _should_stop_others(
         if fut.cancelled():
             return True
         if exc := fut.exception():
-            return not isinstance(exc, GraphInterrupt)
+            return not isinstance(exc, GraphBubbleUp)
     else:
         return False
 

libs/langgraph/langgraph/types.py

@@ -5,6 +5,7 @@ from typing import (
     TYPE_CHECKING,
     Any,
     Callable,
+    ClassVar,
     Generic,
     Hashable,
     Literal,
@@ -140,6 +141,7 @@ class PregelExecutableTask(NamedTuple):
     id: str
     path: tuple[Union[str, int, tuple], ...]
     scheduled: bool = False
+    writers: Sequence[Runnable] = ()
 
 
 class StateSnapshot(NamedTuple):
@@ -239,6 +241,7 @@ N = TypeVar("N", bound=Hashable)
 class Command(Generic[N]):
     """One or more commands to update the graph's state and send messages to nodes."""
 
+    graph: Optional[str] = None
     update: Optional[dict[str, Any]] = None
     send: Union[Send, Sequence[Send]] = ()
     resume: Optional[Union[Any, dict[str, Any]]] = None
@@ -252,6 +255,8 @@ class Command(Generic[N]):
         )
         return f"Command({contents})"
 
+    PARENT: ClassVar[Literal["__parent__"]] = "__parent__"
+
 
 StreamChunk = tuple[tuple[str, ...], str, Any]
 

libs/langgraph/tests/test_pregel.py

@@ -14395,3 +14395,79 @@ def test_runnable_passthrough_node_graph() -> None:
     graph = graph_builder.compile()
 
     assert graph.get_graph(xray=True).to_json() == graph.get_graph(xray=False).to_json()
+
+
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
+def test_parent_command(request: pytest.FixtureRequest, checkpointer_name: str) -> None:
+    from langchain_core.messages import BaseMessage
+    from langchain_core.tools import tool
+
+    @tool(return_direct=True)
+    def get_user_name() -> GraphCommand:
+        """Retrieve user name"""
+        return GraphCommand(update={"user_name": "Meow"}, graph=GraphCommand.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")
+    checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
+    graph = builder.compile(checkpointer=checkpointer)
+
+    config = {"configurable": {"thread_id": "1"}}
+
+    assert graph.invoke({"messages": [("user", "get user name")]}, config) == {
+        "messages": [
+            _AnyIdHumanMessage(
+                content="get user name", additional_kwargs={}, response_metadata={}
+            ),
+        ],
+        "user_name": "Meow",
+    }
+    assert graph.get_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={
+            "configurable": {
+                "thread_id": "1",
+                "checkpoint_ns": "",
+                "checkpoint_id": AnyStr(),
+            }
+        },
+        tasks=(),
+    )

libs/langgraph/tests/test_pregel_async.py

@@ -12597,3 +12597,83 @@ async def test_debug_nested_subgraphs():
                 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() -> GraphCommand:
+        """Retrieve user name"""
+        return GraphCommand(update={"user_name": "Meow"}, graph=GraphCommand.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={
+                "configurable": {
+                    "thread_id": "1",
+                    "checkpoint_ns": "",
+                    "checkpoint_id": AnyStr(),
+                }
+            },
+            tasks=(),
+        )