From 1cc02825ea2cc47fc53ccc83761d6f94db2cf329 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Wed, 14 Aug 2024 15:53:40 -0700
Subject: [PATCH 01/30] Add ScopedValue

- state shared between threads
---
 libs/langgraph/langgraph/channels/base.py     |   5 +
 libs/langgraph/langgraph/constants.py         |   2 +
 libs/langgraph/langgraph/graph/state.py       |  30 ++++-
 libs/langgraph/langgraph/kv/__init__.py       |   0
 libs/langgraph/langgraph/kv/base.py           | 113 ++++++++++++++++++
 libs/langgraph/langgraph/kv/memory.py         |  33 +++++
 libs/langgraph/langgraph/managed/base.py      |  43 ++++++-
 .../langgraph/managed/scoped_value.py         |  99 +++++++++++++++
 libs/langgraph/langgraph/pregel/__init__.py   |   8 +-
 libs/langgraph/langgraph/pregel/algo.py       |  11 +-
 libs/langgraph/langgraph/pregel/loop.py       |  34 +++++-
 libs/langgraph/tests/test_kv.py               |  32 +++++
 libs/langgraph/tests/test_pregel.py           |  41 ++++++-
 13 files changed, 424 insertions(+), 27 deletions(-)
 create mode 100644 libs/langgraph/langgraph/kv/__init__.py
 create mode 100644 libs/langgraph/langgraph/kv/base.py
 create mode 100644 libs/langgraph/langgraph/kv/memory.py
 create mode 100644 libs/langgraph/langgraph/managed/scoped_value.py
 create mode 100644 libs/langgraph/tests/test_kv.py

diff --git a/libs/langgraph/langgraph/channels/base.py b/libs/langgraph/langgraph/channels/base.py
index fe47f0d8f..9c7794c46 100644
--- a/libs/langgraph/langgraph/channels/base.py
+++ b/libs/langgraph/langgraph/channels/base.py
@@ -33,6 +33,11 @@ class BaseChannel(Generic[Value, Update, C], ABC):
 
     # serialize/deserialize methods
 
+    def tap(self) -> Optional[C]:
+        """Return the current checkpoint of the channel, without consuming it.
+        By default, it just calls checkpoint()."""
+        return self.checkpoint()
+
     @abstractmethod
     def checkpoint(self) -> Optional[C]:
         """Return a serializable representation of the channel's current state.
diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py
index f85b33ba3..cf87be337 100644
--- a/libs/langgraph/langgraph/constants.py
+++ b/libs/langgraph/langgraph/constants.py
@@ -5,6 +5,7 @@ INPUT = "__input__"
 CONFIG_KEY_SEND = "__pregel_send"
 CONFIG_KEY_READ = "__pregel_read"
 CONFIG_KEY_CHECKPOINTER = "__pregel_checkpointer"
+CONFIG_KEY_KV = "__pregel_kv"
 CONFIG_KEY_RESUMING = "__pregel_resuming"
 CONFIG_KEY_TASK_ID = "__pregel_task_id"
 INTERRUPT = "__interrupt__"
@@ -17,6 +18,7 @@ RESERVED = {
     CONFIG_KEY_SEND,
     CONFIG_KEY_READ,
     CONFIG_KEY_CHECKPOINTER,
+    CONFIG_KEY_KV,
     CONFIG_KEY_RESUMING,
     CONFIG_KEY_TASK_ID,
     INPUT,
diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py
index 5810aaf29..cda357607 100644
--- a/libs/langgraph/langgraph/graph/state.py
+++ b/libs/langgraph/langgraph/graph/state.py
@@ -40,7 +40,14 @@ from langgraph.graph.graph import (
     Graph,
     Send,
 )
-from langgraph.managed.base import ManagedValue, is_managed_value
+from langgraph.kv.base import BaseKV
+from langgraph.managed.base import (
+    ChannelKeyPlaceholder,
+    ConfiguredManagedValue,
+    ManagedValue,
+    is_managed_value,
+    is_writable_managed_value,
+)
 from langgraph.pregel.read import ChannelRead, PregelNode
 from langgraph.pregel.types import All, RetryPolicy
 from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry
@@ -373,6 +380,8 @@ class StateGraph(Graph):
 
     def compile(
         self,
+        *,
+        kv: Optional[BaseKV] = None,
         checkpointer: Optional[BaseCheckpointSaver] = None,
         interrupt_before: Optional[Union[All, Sequence[str]]] = None,
         interrupt_after: Optional[Union[All, Sequence[str]]] = None,
@@ -442,6 +451,7 @@ class StateGraph(Graph):
             interrupt_after_nodes=interrupt_after,
             auto_validate=False,
             debug=debug,
+            kv=kv,
         )
 
         compiled.attach_node(START, None)
@@ -511,7 +521,11 @@ class CompiledStateGraph(CompiledGraph):
                 if not isinstance(v, Context) and not is_managed_value(v)
             ]
         else:
-            output_keys = list(self.builder.channels)
+            output_keys = list(self.builder.channels) + [
+                k
+                for k, v in self.builder.managed.items()
+                if is_writable_managed_value(v)
+            ]
 
         def _get_state_key(
             input: Union[None, dict, Any], config: RunnableConfig, *, key: str
@@ -684,7 +698,7 @@ def _get_channels(
         return {"__root__": _get_channel(schema, allow_managed=False)}, {}
 
     all_keys = {
-        name: _get_channel(typ)
+        name: _get_channel(name, typ)
         for name, typ in get_type_hints(schema, include_extras=True).items()
         if name != "__slots__"
     }
@@ -695,9 +709,9 @@ def _get_channels(
 
 
 def _get_channel(
-    annotation: Any, *, allow_managed: bool = True
+    name: str, annotation: Any, *, allow_managed: bool = True
 ) -> Union[BaseChannel, Type[ManagedValue]]:
-    if manager := _is_field_managed_value(annotation):
+    if manager := _is_field_managed_value(name, annotation):
         if allow_managed:
             return manager
         else:
@@ -736,12 +750,16 @@ def _is_field_binop(typ: Type[Any]) -> Optional[BinaryOperatorAggregate]:
     return None
 
 
-def _is_field_managed_value(typ: Type[Any]) -> Optional[Type[ManagedValue]]:
+def _is_field_managed_value(name: str, typ: Type[Any]) -> Optional[Type[ManagedValue]]:
     if hasattr(typ, "__metadata__"):
         meta = typ.__metadata__
         if len(meta) >= 1:
             decoration = get_origin(meta[-1]) or meta[-1]
             if is_managed_value(decoration):
+                if isinstance(decoration, ConfiguredManagedValue):
+                    for k, v in decoration.kwargs.items():
+                        if v is ChannelKeyPlaceholder:
+                            decoration.kwargs[k] = name
                 return decoration
 
     return None
diff --git a/libs/langgraph/langgraph/kv/__init__.py b/libs/langgraph/langgraph/kv/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/libs/langgraph/langgraph/kv/base.py b/libs/langgraph/langgraph/kv/base.py
new file mode 100644
index 000000000..570783ec1
--- /dev/null
+++ b/libs/langgraph/langgraph/kv/base.py
@@ -0,0 +1,113 @@
+import asyncio
+from typing import Any, List, NamedTuple, Optional, Union
+
+V = dict[str, Any]
+
+
+class BaseKV:
+    def get(self, pairs: List[tuple[str, str]]) -> dict[tuple[str, str], Optional[V]]:
+        # list[(namespace, key)] -> dict[(namespace, key), value | none]
+        raise NotImplementedError
+
+    def list(self, prefixes: List[str]) -> dict[str, dict[str, V]]:
+        # list[namespace] -> dict[namespace, list[value]]
+        raise NotImplementedError
+
+    def put(self, writes: List[tuple[str, str, Optional[V]]]) -> None:
+        # list[(namespace, key, value | none)] -> None
+        raise NotImplementedError
+
+    async def aget(
+        self, pairs: List[tuple[str, str]]
+    ) -> dict[tuple[str, str], Optional[V]]:
+        # list[(namespace, key)] -> dict[(namespace, key), value | none]
+        raise NotImplementedError
+
+    async def alist(self, prefixes: List[str]) -> dict[str, dict[str, V]]:
+        # list[namespace] -> dict[namespace, list[value]]
+        raise NotImplementedError
+
+    async def aput(self, writes: List[tuple[str, str, Optional[V]]]) -> None:
+        # list[(namespace, key, value | none)] -> None
+        raise NotImplementedError
+
+
+class GetOp(NamedTuple):
+    pairs: List[tuple[str, str]]
+
+
+class ListOp(NamedTuple):
+    prefixes: List[str]
+
+
+class PutOp(NamedTuple):
+    writes: List[tuple[str, str, Optional[V]]]
+
+
+class KeyValueStore(BaseKV):
+    def __init__(self, kv: BaseKV) -> None:
+        self.kv = kv
+        self.aqueue: dict[asyncio.Future, Union[GetOp, ListOp, PutOp]] = {}
+        self.task = asyncio.create_task(_run(self.aqueue, self.kv))
+
+    def __del__(self) -> None:
+        self.task.cancel()
+
+    async def aget(
+        self, pairs: List[tuple[str, str]]
+    ) -> dict[tuple[str, str], Optional[V]]:
+        fut = asyncio.get_running_loop().create_future()
+        self.aqueue[fut] = GetOp(pairs)
+        return await fut
+
+    async def alist(self, prefixes: List[str]) -> dict[str, dict[str, V]]:
+        fut = asyncio.get_running_loop().create_future()
+        self.aqueue[fut] = ListOp(prefixes)
+        return await fut
+
+    async def aput(self, writes: List[tuple[str, str, Optional[V]]]) -> None:
+        fut = asyncio.get_running_loop().create_future()
+        self.aqueue[fut] = PutOp(writes)
+        return await fut
+
+
+async def _run(
+    aqueue: dict[asyncio.Future, Union[GetOp, ListOp, PutOp]], kv: BaseKV
+) -> None:
+    while True:
+        await asyncio.sleep(0)
+        if not aqueue:
+            continue
+        # this could use a lock, if we want thread safety
+        taken = aqueue.copy()
+        aqueue.clear()
+        # action each operation
+        gets = {f: o for f, o in taken.items() if isinstance(o, GetOp)}
+        if gets:
+            try:
+                results = await kv.aget([p for op in gets.values() for p in op.pairs])
+                for fut, op in gets.items():
+                    fut.set_result({k: results.get(k) for k in op.pairs})
+            except Exception as e:
+                for fut in gets:
+                    fut.set_exception(e)
+        lists = {f: o for f, o in taken.items() if isinstance(o, ListOp)}
+        if lists:
+            try:
+                results = await kv.alist(
+                    [p for op in lists.values() for p in op.prefixes]
+                )
+                for fut, op in lists.items():
+                    fut.set_result({k: results.get(k) for k in op.prefixes})
+            except Exception as e:
+                for fut in lists:
+                    fut.set_exception(e)
+        puts = {f: o for f, o in taken.items() if isinstance(o, PutOp)}
+        if puts:
+            try:
+                await kv.aput([w for op in puts.values() for w in op.writes])
+                for fut in puts:
+                    fut.set_result(None)
+            except Exception as e:
+                for fut in puts:
+                    fut.set_exception(e)
diff --git a/libs/langgraph/langgraph/kv/memory.py b/libs/langgraph/langgraph/kv/memory.py
new file mode 100644
index 000000000..387ea1e7a
--- /dev/null
+++ b/libs/langgraph/langgraph/kv/memory.py
@@ -0,0 +1,33 @@
+from collections import defaultdict
+from typing import List
+
+from langgraph.kv.base import BaseKV, V
+
+
+class MemoryKV(BaseKV):
+    def __init__(self) -> None:
+        self.data: dict[str, dict[str, V]] = defaultdict(dict)
+
+    def get(self, pairs: List[tuple[str, str]]) -> dict[tuple[str, str], V | None]:
+        return {pair: self.data[pair[0]].get(pair[1]) for pair in pairs}
+
+    async def aget(
+        self, pairs: List[tuple[str, str]]
+    ) -> dict[tuple[str, str], V | None]:
+        return self.get(pairs)
+
+    def list(self, prefixes: List[str]) -> dict[str, dict[str, V]]:
+        return {prefix: self.data[prefix] for prefix in prefixes}
+
+    async def alist(self, prefixes: List[str]) -> dict[str, dict[str, V]]:
+        return self.list(prefixes)
+
+    def put(self, writes: List[tuple[str, str, V | None]]) -> None:
+        for namespace, key, value in writes:
+            if value is None:
+                self.data[namespace].pop(key, None)
+            else:
+                self.data[namespace][key] = value
+
+    async def aput(self, writes: List[tuple[str, str, V | None]]) -> None:
+        self.put(writes)
diff --git a/libs/langgraph/langgraph/managed/base.py b/libs/langgraph/langgraph/managed/base.py
index 0455ed58b..5820383f8 100644
--- a/libs/langgraph/langgraph/managed/base.py
+++ b/libs/langgraph/langgraph/managed/base.py
@@ -5,9 +5,12 @@ from inspect import isclass
 from typing import (
     Any,
     AsyncGenerator,
+    AsyncIterator,
     Generator,
     Generic,
+    Iterator,
     NamedTuple,
+    Sequence,
     Type,
     TypeVar,
     Union,
@@ -17,6 +20,7 @@ from langchain_core.runnables import RunnableConfig
 from typing_extensions import Self, TypeGuard
 
 V = TypeVar("V")
+U = TypeVar("U")
 
 
 class ManagedValue(ABC, Generic[V]):
@@ -25,9 +29,7 @@ class ManagedValue(ABC, Generic[V]):
 
     @classmethod
     @contextmanager
-    def enter(
-        cls, config: RunnableConfig, **kwargs: Any
-    ) -> Generator[Self, None, None]:
+    def enter(cls, config: RunnableConfig, **kwargs: Any) -> Iterator[Self]:
         try:
             value = cls(config, **kwargs)
             yield value
@@ -41,9 +43,7 @@ class ManagedValue(ABC, Generic[V]):
 
     @classmethod
     @asynccontextmanager
-    async def aenter(
-        cls, config: RunnableConfig, **kwargs: Any
-    ) -> AsyncGenerator[Self, None]:
+    async def aenter(cls, config: RunnableConfig, **kwargs: Any) -> AsyncIterator[Self]:
         try:
             value = cls(config, **kwargs)
             yield value
@@ -60,6 +60,16 @@ class ManagedValue(ABC, Generic[V]):
         ...
 
 
+class WritableManagedValue(Generic[V, U], ManagedValue[V], ABC):
+    @abstractmethod
+    def update(self, writes: Sequence[U]) -> None:
+        ...
+
+    @abstractmethod
+    async def aupdate(self, writes: Sequence[U]) -> None:
+        ...
+
+
 class ConfiguredManagedValue(NamedTuple):
     cls: Type[ManagedValue]
     kwargs: dict[str, Any]
@@ -76,6 +86,24 @@ def is_managed_value(value: Any) -> TypeGuard[ManagedValueSpec]:
     )
 
 
+def is_readonly_managed_value(value: Any) -> TypeGuard[Type[ManagedValue]]:
+    return (
+        isclass(value)
+        and issubclass(value, ManagedValue)
+        and not issubclass(value, WritableManagedValue)
+    ) or (
+        isinstance(value, ConfiguredManagedValue)
+        and not issubclass(value.cls, WritableManagedValue)
+    )
+
+
+def is_writable_managed_value(value: Any) -> TypeGuard[Type[WritableManagedValue]]:
+    return (isclass(value) and issubclass(value, WritableManagedValue)) or (
+        isinstance(value, ConfiguredManagedValue)
+        and issubclass(value.cls, WritableManagedValue)
+    )
+
+
 @contextmanager
 def ManagedValuesManager(
     values: dict[str, ManagedValueSpec],
@@ -119,3 +147,6 @@ async def AsyncManagedValuesManager(
             yield {tasks[task]: task.result() for task in done}
     else:
         yield {}
+
+
+ChannelKeyPlaceholder = object()
diff --git a/libs/langgraph/langgraph/managed/scoped_value.py b/libs/langgraph/langgraph/managed/scoped_value.py
new file mode 100644
index 000000000..adb16e1b4
--- /dev/null
+++ b/libs/langgraph/langgraph/managed/scoped_value.py
@@ -0,0 +1,99 @@
+from contextlib import asynccontextmanager, contextmanager
+from typing import (
+    Any,
+    AsyncIterator,
+    Iterator,
+    Optional,
+    Self,
+    Sequence,
+)
+
+from langchain_core.runnables import RunnableConfig
+
+from langgraph.constants import CONFIG_KEY_KV
+from langgraph.errors import InvalidUpdateError
+from langgraph.kv.base import BaseKV
+from langgraph.managed.base import (
+    ChannelKeyPlaceholder,
+    ConfiguredManagedValue,
+    WritableManagedValue,
+)
+from langgraph.pregel.types import PregelTaskDescription
+
+V = dict[str, Any]
+
+
+Value = dict[str, V]
+Update = dict[str, Optional[V]]
+
+
+class ScopedValue(WritableManagedValue[Value, Update]):
+    @staticmethod
+    def configure(scope: str) -> ConfiguredManagedValue:
+        return ConfiguredManagedValue(
+            ScopedValue, {"scope": scope, "key": ChannelKeyPlaceholder}
+        )
+
+    @classmethod
+    @contextmanager
+    def enter(cls, config: RunnableConfig, **kwargs: Any) -> Iterator[Self]:
+        with super().enter(config, **kwargs) as value:
+            if value.kv is not None:
+                saved = value.kv.list([value.ns])
+                value.value = saved[value.ns]
+            yield value
+
+    @classmethod
+    @asynccontextmanager
+    async def aenter(cls, config: RunnableConfig, **kwargs: Any) -> AsyncIterator[Self]:
+        async with super().aenter(config, **kwargs) as value:
+            if value.kv is not None:
+                saved = await value.kv.alist([value.ns])
+                value.value = saved[value.ns]
+            yield value
+
+    def __init__(self, config: RunnableConfig, *, scope: str, key: str) -> None:
+        self.scope = scope
+        self.config = config
+        self.value: Value = {}
+        self.kv: BaseKV = config["configurable"].get(CONFIG_KEY_KV)
+        if self.kv is None:
+            self.ns: Optional[str] = None
+        elif scope_value := config["configurable"].get(self.scope):
+            self.ns = f"scoped:{scope}:{key}:{scope_value}"
+        else:
+            raise ValueError(
+                f"Scope {scope} for shared state key not in config.configurable"
+            )
+
+    def __call__(self, step: int, task: PregelTaskDescription) -> Value:
+        return self.value.copy()
+
+    def _process_update(
+        self, values: Sequence[Update]
+    ) -> list[tuple[str, str, Optional[dict[str, Any]]]]:
+        writes = []
+        for vv in values:
+            for k, v in vv.items():
+                if v is None:
+                    if k in self.value:
+                        self.value[k] = None
+                        writes.append((self.ns, k, None))
+                elif not isinstance(v, dict):
+                    raise InvalidUpdateError("Received a non-dict value")
+                else:
+                    self.value[k] = v
+                    writes.append((self.ns, k, v))
+        return writes
+
+    def update(self, values: Sequence[Update]) -> None:
+        if self.kv is None:
+            self._process_update(values)
+        else:
+            return self.kv.put(self._process_update(values))
+
+    async def aupdate(self, writes: Sequence[Update]) -> None:
+        if self.kv is None:
+            self._process_update(writes)
+        else:
+            return await self.kv.aput(self._process_update(writes))
diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py
index d2f0790cc..833453197 100644
--- a/libs/langgraph/langgraph/pregel/__init__.py
+++ b/libs/langgraph/langgraph/pregel/__init__.py
@@ -73,6 +73,7 @@ from langgraph.constants import (
     Interrupt,
 )
 from langgraph.errors import GraphInterrupt, GraphRecursionError, InvalidUpdateError
+from langgraph.kv.base import BaseKV
 from langgraph.managed.base import (
     AsyncManagedValuesManager,
     ManagedValuesManager,
@@ -223,6 +224,9 @@ class Pregel(
     checkpointer: Optional[BaseCheckpointSaver] = None
     """Checkpointer used to save and load graph state. Defaults to None."""
 
+    kv: Optional[BaseKV] = None
+    """Key-value store to use. Defaults to None."""
+
     retry_policy: Optional[RetryPolicy] = None
     """Retry policy to use when running tasks. Set to None to disable."""
 
@@ -644,7 +648,7 @@ class Pregel(
                 ),
             )
             # apply to checkpoint and save
-            apply_writes(
+            assert not apply_writes(
                 checkpoint, channels, [task], self.checkpointer.get_next_version
             )
             checkpoint = create_checkpoint(checkpoint, channels, step + 1)
@@ -788,7 +792,7 @@ class Pregel(
                 ),
             )
             # apply to checkpoint and save
-            apply_writes(
+            assert not apply_writes(
                 checkpoint, channels, [task], self.checkpointer.get_next_version
             )
             checkpoint = create_checkpoint(checkpoint, channels, step + 1)
diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py
index 5922fd14e..994af018e 100644
--- a/libs/langgraph/langgraph/pregel/algo.py
+++ b/libs/langgraph/langgraph/pregel/algo.py
@@ -145,7 +145,7 @@ def apply_writes(
     channels: Mapping[str, BaseChannel],
     tasks: Sequence[WritesProtocol],
     get_next_version: Optional[Callable[[int, BaseChannel], int]],
-) -> None:
+) -> dict[str, list[Any]]:
     # update seen versions
     for task in tasks:
         checkpoint["versions_seen"].setdefault(task.name, {}).update(
@@ -161,6 +161,7 @@ def apply_writes(
         max_version = max(checkpoint["channel_versions"].values())
     else:
         max_version = None
+
     # Consume all channels that were read
     for chan in {
         chan for task in tasks for chan in task.triggers if chan not in RESERVED
@@ -177,12 +178,15 @@ def apply_writes(
 
     # Group writes by channel
     pending_writes_by_channel: dict[str, list[Any]] = defaultdict(list)
+    pending_writes_by_managed: dict[str, list[Any]] = defaultdict(list)
     for task in tasks:
         for chan, val in task.writes:
             if chan == TASKS:
                 checkpoint["pending_sends"].append(val)
-            else:
+            elif chan in channels:
                 pending_writes_by_channel[chan].append(val)
+            else:
+                pending_writes_by_managed[chan].append(val)
 
     # Find the highest version of all channels
     if checkpoint["channel_versions"]:
@@ -214,6 +218,9 @@ def apply_writes(
                     max_version, channels[chan]
                 )
 
+    # Return managed values writes to be applied externally
+    return pending_writes_by_managed
+
 
 @overload
 def prepare_next_tasks(
diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py
index f48f86fcd..aaa43be67 100644
--- a/libs/langgraph/langgraph/pregel/loop.py
+++ b/libs/langgraph/langgraph/pregel/loop.py
@@ -18,10 +18,11 @@ from typing import (
     Type,
     TypeVar,
     Union,
+    cast,
 )
 
 from langchain_core.callbacks import AsyncParentRunManager, ParentRunManager
-from langchain_core.runnables import RunnableConfig
+from langchain_core.runnables import RunnableConfig, patch_config
 from typing_extensions import Self
 
 from langgraph.channels.base import BaseChannel
@@ -40,6 +41,7 @@ from langgraph.checkpoint.base import (
     empty_checkpoint,
 )
 from langgraph.constants import (
+    CONFIG_KEY_KV,
     CONFIG_KEY_READ,
     CONFIG_KEY_RESUMING,
     ERROR,
@@ -52,6 +54,7 @@ from langgraph.managed.base import (
     AsyncManagedValuesManager,
     ManagedValueMapping,
     ManagedValuesManager,
+    WritableManagedValue,
 )
 from langgraph.pregel.algo import (
     PregelTaskWrites,
@@ -182,12 +185,15 @@ class PregelLoop:
         elif all(task.writes for task in self.tasks):
             writes = [w for t in self.tasks for w in t.writes]
             # all tasks have finished
-            apply_writes(
+            mv_writes = apply_writes(
                 self.checkpoint,
                 self.channels,
                 self.tasks,
                 self.checkpointer_get_next_version,
             )
+            # apply writes to managed values
+            for key, values in mv_writes.items():
+                self._update_mv(key, values)
             # produce values output
             self.stream.extend(
                 ("values", v)
@@ -324,12 +330,13 @@ class PregelLoop:
                 manager=None,
             )
             # apply input writes
-            apply_writes(
+            mv_writes = apply_writes(
                 self.checkpoint,
                 self.channels,
                 discard_tasks + [PregelTaskWrites(INPUT, input_writes, [])],
                 self.checkpointer_get_next_version,
             )
+            assert not mv_writes
             # save input checkpoint
             self._put_checkpoint({"source": "input", "writes": self.input})
         else:
@@ -395,6 +402,9 @@ class PregelLoop:
         # increment step
         self.step += 1
 
+    def _update_mv(self, key: str, values: Sequence[Any]) -> None:
+        raise NotImplementedError
+
     def _suppress_interrupt(
         self,
         exc_type: Optional[Type[BaseException]],
@@ -438,6 +448,9 @@ class SyncPregelLoop(PregelLoop, ContextManager):
         finally:
             self.checkpointer.put(config, checkpoint, metadata, new_versions)
 
+    def _update_mv(self, key: str, values: Sequence[Any]) -> None:
+        return self.submit(cast(WritableManagedValue, self.managed[key]).update, values)
+
     # context manager
 
     def __enter__(self) -> Self:
@@ -461,7 +474,10 @@ class SyncPregelLoop(PregelLoop, ContextManager):
             ChannelsManager(self.graph.channels, self.checkpoint, self.config)
         )
         self.managed = self.stack.enter_context(
-            ManagedValuesManager(self.graph.managed_values_dict, self.config)
+            ManagedValuesManager(
+                self.graph.managed_values_dict,
+                patch_config(self.config, configurable={CONFIG_KEY_KV: self.graph.kv}),
+            )
         )
         self.stack.push(self._suppress_interrupt)
         self.status = "pending"
@@ -515,6 +531,11 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
         finally:
             await self.checkpointer.aput(config, checkpoint, metadata, new_versions)
 
+    def _update_mv(self, key: str, values: Sequence[Any]) -> None:
+        return self.submit(
+            cast(WritableManagedValue, self.managed[key]).aupdate, values
+        )
+
     # context manager
 
     async def __aenter__(self) -> Self:
@@ -540,7 +561,10 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
             AsyncChannelsManager(self.graph.channels, self.checkpoint, self.config)
         )
         self.managed = await self.stack.enter_async_context(
-            AsyncManagedValuesManager(self.graph.managed_values_dict, self.config)
+            AsyncManagedValuesManager(
+                self.graph.managed_values_dict,
+                patch_config(self.config, configurable={CONFIG_KEY_KV: self.graph.kv}),
+            )
         )
         self.stack.push(self._suppress_interrupt)
         self.status = "pending"
diff --git a/libs/langgraph/tests/test_kv.py b/libs/langgraph/tests/test_kv.py
new file mode 100644
index 000000000..2e24ee07d
--- /dev/null
+++ b/libs/langgraph/tests/test_kv.py
@@ -0,0 +1,32 @@
+import asyncio
+from typing import Any, List
+
+from pytest_mock import MockerFixture
+
+from langgraph.kv.base import BaseKV, KeyValueStore
+
+
+async def test_kv_queue(mocker: MockerFixture) -> None:
+    aget = mocker.stub()
+
+    class MockKV(BaseKV):
+        async def aget(
+            self, pairs: List[tuple[str, str]]
+        ) -> dict[tuple[str, str], dict[str, Any] | None]:
+            aget(pairs)
+            return {pair: {0: pair[0], 1: pair[1]} for pair in pairs}
+
+    store = KeyValueStore(MockKV())
+
+    # concurrent calls are batched
+    results = await asyncio.gather(
+        store.aget([("a", "b")]),
+        store.aget([("c", "d")]),
+    )
+    assert results == [
+        {("a", "b"): {0: "a", 1: "b"}},
+        {("c", "d"): {0: "c", 1: "d"}},
+    ]
+    assert [c.args for c in aget.call_args_list] == [
+        ([("a", "b"), ("c", "d")],),
+    ]
diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py
index 138d4dba5..3cc1c4d23 100644
--- a/libs/langgraph/tests/test_pregel.py
+++ b/libs/langgraph/tests/test_pregel.py
@@ -58,6 +58,8 @@ from langgraph.graph import END, Graph
 from langgraph.graph.graph import START
 from langgraph.graph.message import MessageGraph, add_messages
 from langgraph.graph.state import StateGraph
+from langgraph.kv.memory import MemoryKV
+from langgraph.managed.scoped_value import ScopedValue
 from langgraph.prebuilt.chat_agent_executor import (
     create_tool_calling_executor,
 )
@@ -165,6 +167,7 @@ def test_graph_validation() -> None:
 
     class State(TypedDict):
         hello: str
+        shared_things: Annotated[dict[str, dict[str, Any]], ScopedValue("assistant_id")]
 
     def node_a(state: State) -> State:
         # typo
@@ -6202,10 +6205,34 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None:
     class State(TypedDict):
         my_key: Annotated[str, operator.add]
         market: str
+        shared: Annotated[
+            dict[str, dict[str, Any]], ScopedValue.configure("assistant_id")
+        ]
+
+    def assert_shared_value(data: State, config: RunnableConfig) -> State:
+        assert "shared" in data
+        if thread_id := config["configurable"].get("thread_id"):
+            if thread_id == "1":
+                # this is the first thread, so should not see a value
+                assert data["shared"] == {}
+                return {"shared": {"1": {"hello": "world"}}}
+            elif thread_id == "2":
+                # this should get value saved by thread 1
+                assert data["shared"] == {"1": {"hello": "world"}}
+            elif thread_id == "3":
+                # this is a different assistant, so should not see previous value
+                assert data["shared"] == {}
+        return {}
+
+    def tool_two_slow(data: State, config: RunnableConfig) -> State:
+        return {"my_key": " slow", **assert_shared_value(data, config)}
+
+    def tool_two_fast(data: State, config: RunnableConfig) -> State:
+        return {"my_key": " fast", **assert_shared_value(data, config)}
 
     tool_two_graph = StateGraph(State)
-    tool_two_graph.add_node("tool_two_slow", lambda s: {"my_key": " slow"})
-    tool_two_graph.add_node("tool_two_fast", lambda s: {"my_key": " fast"})
+    tool_two_graph.add_node("tool_two_slow", tool_two_slow)
+    tool_two_graph.add_node("tool_two_fast", tool_two_fast)
     tool_two_graph.set_conditional_entry_point(
         lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast", then=END
     )
@@ -6223,14 +6250,16 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None:
 
     with SqliteSaver.from_conn_string(":memory:") as saver:
         tool_two = tool_two_graph.compile(
-            checkpointer=saver, interrupt_before=["tool_two_fast", "tool_two_slow"]
+            kv=MemoryKV(),
+            checkpointer=saver,
+            interrupt_before=["tool_two_fast", "tool_two_slow"],
         )
 
         # missing thread_id
         with pytest.raises(ValueError, match="thread_id"):
             tool_two.invoke({"my_key": "value", "market": "DE"})
 
-        thread1 = {"configurable": {"thread_id": "1"}}
+        thread1 = {"configurable": {"thread_id": "1", "assistant_id": "a"}}
         # stop when about to enter node
         assert tool_two.invoke({"my_key": "value ⛰️", "market": "DE"}, thread1) == {
             "my_key": "value ⛰️",
@@ -6282,7 +6311,7 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None:
             parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config,
         )
 
-        thread2 = {"configurable": {"thread_id": "2"}}
+        thread2 = {"configurable": {"thread_id": "2", "assistant_id": "a"}}
         # stop when about to enter node
         assert tool_two.invoke({"my_key": "value", "market": "US"}, thread2) == {
             "my_key": "value",
@@ -6322,7 +6351,7 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None:
             parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config,
         )
 
-        thread3 = {"configurable": {"thread_id": "3"}}
+        thread3 = {"configurable": {"thread_id": "3", "assistant_id": "b"}}
         # stop when about to enter node
         assert tool_two.invoke({"my_key": "value", "market": "US"}, thread3) == {
             "my_key": "value",

From 1dc09dc45fd53ed9c7a795f344f14fd23c9c13d2 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Fri, 16 Aug 2024 09:58:49 -0700
Subject: [PATCH 02/30] Remove warning on write to managed channel

---
 libs/langgraph/langgraph/pregel/algo.py | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py
index 994af018e..a52d15f4d 100644
--- a/libs/langgraph/langgraph/pregel/algo.py
+++ b/libs/langgraph/langgraph/pregel/algo.py
@@ -122,6 +122,7 @@ def local_write(
     processes: Mapping[str, PregelNode],
     channels: Mapping[str, BaseChannel],
     writes: Sequence[tuple[str, Any]],
+    managed: ManagedValueMapping,
 ) -> None:
     for chan, value in writes:
         if chan == TASKS:
@@ -131,7 +132,7 @@ def local_write(
                 )
             if value.node not in processes:
                 raise InvalidUpdateError(f"Invalid node name {value.node} in packet")
-        elif chan not in channels:
+        elif chan not in channels and chan not in managed:
             logger.warning(f"Skipping write for channel '{chan}' which has no readers")
     commit(writes)
 
@@ -321,7 +322,11 @@ def prepare_next_tasks(
                                 CONFIG_KEY_TASK_ID: task_id,
                                 # deque.extend is thread-safe
                                 CONFIG_KEY_SEND: partial(
-                                    local_write, writes.extend, processes, channels
+                                    local_write,
+                                    writes.extend,
+                                    processes,
+                                    channels,
+                                    managed,
                                 ),
                                 CONFIG_KEY_READ: partial(
                                     local_read,
@@ -412,7 +417,11 @@ def prepare_next_tasks(
                                     CONFIG_KEY_TASK_ID: task_id,
                                     # deque.extend is thread-safe
                                     CONFIG_KEY_SEND: partial(
-                                        local_write, writes.extend, processes, channels
+                                        local_write,
+                                        writes.extend,
+                                        processes,
+                                        channels,
+                                        managed,
                                     ),
                                     CONFIG_KEY_READ: partial(
                                         local_read,

From 2bc0e2df4299259fc7e3187c3c288a0c5145dfc1 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Fri, 16 Aug 2024 15:53:05 -0700
Subject: [PATCH 03/30] Fix

---
 libs/langgraph/langgraph/channels/base.py        | 5 -----
 libs/langgraph/langgraph/managed/scoped_value.py | 3 +--
 2 files changed, 1 insertion(+), 7 deletions(-)

diff --git a/libs/langgraph/langgraph/channels/base.py b/libs/langgraph/langgraph/channels/base.py
index 9c7794c46..fe47f0d8f 100644
--- a/libs/langgraph/langgraph/channels/base.py
+++ b/libs/langgraph/langgraph/channels/base.py
@@ -33,11 +33,6 @@ class BaseChannel(Generic[Value, Update, C], ABC):
 
     # serialize/deserialize methods
 
-    def tap(self) -> Optional[C]:
-        """Return the current checkpoint of the channel, without consuming it.
-        By default, it just calls checkpoint()."""
-        return self.checkpoint()
-
     @abstractmethod
     def checkpoint(self) -> Optional[C]:
         """Return a serializable representation of the channel's current state.
diff --git a/libs/langgraph/langgraph/managed/scoped_value.py b/libs/langgraph/langgraph/managed/scoped_value.py
index adb16e1b4..997bc9469 100644
--- a/libs/langgraph/langgraph/managed/scoped_value.py
+++ b/libs/langgraph/langgraph/managed/scoped_value.py
@@ -18,7 +18,6 @@ from langgraph.managed.base import (
     ConfiguredManagedValue,
     WritableManagedValue,
 )
-from langgraph.pregel.types import PregelTaskDescription
 
 V = dict[str, Any]
 
@@ -66,7 +65,7 @@ class ScopedValue(WritableManagedValue[Value, Update]):
                 f"Scope {scope} for shared state key not in config.configurable"
             )
 
-    def __call__(self, step: int, task: PregelTaskDescription) -> Value:
+    def __call__(self, step: int) -> Value:
         return self.value.copy()
 
     def _process_update(

From 7ca37afc74095c86ad0e86f36ce948d0ed7e33b7 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Fri, 16 Aug 2024 16:03:51 -0700
Subject: [PATCH 04/30] Fix

---
 libs/langgraph/langgraph/graph/state.py          |  2 +-
 libs/langgraph/langgraph/managed/scoped_value.py |  6 +++---
 libs/langgraph/langgraph/pregel/algo.py          |  2 +-
 libs/langgraph/tests/test_pregel.py              | 10 +++++-----
 libs/langgraph/tests/test_pregel_async.py        |  2 +-
 5 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py
index cda357607..1c9242ea1 100644
--- a/libs/langgraph/langgraph/graph/state.py
+++ b/libs/langgraph/langgraph/graph/state.py
@@ -695,7 +695,7 @@ def _get_channels(
     schema: Type[dict],
 ) -> tuple[dict[str, BaseChannel], dict[str, Type[ManagedValue]]]:
     if not hasattr(schema, "__annotations__"):
-        return {"__root__": _get_channel(schema, allow_managed=False)}, {}
+        return {"__root__": _get_channel("__root__", schema, allow_managed=False)}, {}
 
     all_keys = {
         name: _get_channel(name, typ)
diff --git a/libs/langgraph/langgraph/managed/scoped_value.py b/libs/langgraph/langgraph/managed/scoped_value.py
index 997bc9469..38a666163 100644
--- a/libs/langgraph/langgraph/managed/scoped_value.py
+++ b/libs/langgraph/langgraph/managed/scoped_value.py
@@ -26,11 +26,11 @@ Value = dict[str, V]
 Update = dict[str, Optional[V]]
 
 
-class ScopedValue(WritableManagedValue[Value, Update]):
+class SharedValue(WritableManagedValue[Value, Update]):
     @staticmethod
-    def configure(scope: str) -> ConfiguredManagedValue:
+    def on(scope: str) -> ConfiguredManagedValue:
         return ConfiguredManagedValue(
-            ScopedValue, {"scope": scope, "key": ChannelKeyPlaceholder}
+            SharedValue, {"scope": scope, "key": ChannelKeyPlaceholder}
         )
 
     @classmethod
diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py
index a52d15f4d..37f168bf7 100644
--- a/libs/langgraph/langgraph/pregel/algo.py
+++ b/libs/langgraph/langgraph/pregel/algo.py
@@ -121,8 +121,8 @@ def local_write(
     commit: Callable[[Sequence[tuple[str, Any]]], None],
     processes: Mapping[str, PregelNode],
     channels: Mapping[str, BaseChannel],
-    writes: Sequence[tuple[str, Any]],
     managed: ManagedValueMapping,
+    writes: Sequence[tuple[str, Any]],
 ) -> None:
     for chan, value in writes:
         if chan == TASKS:
diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py
index 3cc1c4d23..4bcce7ccd 100644
--- a/libs/langgraph/tests/test_pregel.py
+++ b/libs/langgraph/tests/test_pregel.py
@@ -59,7 +59,7 @@ from langgraph.graph.graph import START
 from langgraph.graph.message import MessageGraph, add_messages
 from langgraph.graph.state import StateGraph
 from langgraph.kv.memory import MemoryKV
-from langgraph.managed.scoped_value import ScopedValue
+from langgraph.managed.scoped_value import SharedValue
 from langgraph.prebuilt.chat_agent_executor import (
     create_tool_calling_executor,
 )
@@ -167,7 +167,9 @@ def test_graph_validation() -> None:
 
     class State(TypedDict):
         hello: str
-        shared_things: Annotated[dict[str, dict[str, Any]], ScopedValue("assistant_id")]
+        shared_things: Annotated[
+            dict[str, dict[str, Any]], SharedValue.on("assistant_id")
+        ]
 
     def node_a(state: State) -> State:
         # typo
@@ -6205,9 +6207,7 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None:
     class State(TypedDict):
         my_key: Annotated[str, operator.add]
         market: str
-        shared: Annotated[
-            dict[str, dict[str, Any]], ScopedValue.configure("assistant_id")
-        ]
+        shared: Annotated[dict[str, dict[str, Any]], SharedValue.on("assistant_id")]
 
     def assert_shared_value(data: State, config: RunnableConfig) -> State:
         assert "shared" in data
diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py
index 7c5742aa0..3048ddbb4 100644
--- a/libs/langgraph/tests/test_pregel_async.py
+++ b/libs/langgraph/tests/test_pregel_async.py
@@ -563,7 +563,7 @@ async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
 
     assert app.input_schema.schema() == {"title": "LangGraphInput", "type": "integer"}
     assert app.output_schema.schema() == {"title": "LangGraphOutput", "type": "integer"}
-    assert await app.ainvoke(2) == 3
+    assert await app.ainvoke(2, debug=True) == 3
     assert await app.ainvoke(2, output_keys=["output"]) == {"output": 3}
 
     assert await gapp.ainvoke(2) == 3

From 77d7deb033820834dc431b735cd8b7c09b07aa98 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Fri, 16 Aug 2024 16:23:57 -0700
Subject: [PATCH 05/30] Rename

---
 .../langgraph/managed/{scoped_value.py => shared_value.py}      | 2 +-
 libs/langgraph/tests/test_pregel.py                             | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)
 rename libs/langgraph/langgraph/managed/{scoped_value.py => shared_value.py} (98%)

diff --git a/libs/langgraph/langgraph/managed/scoped_value.py b/libs/langgraph/langgraph/managed/shared_value.py
similarity index 98%
rename from libs/langgraph/langgraph/managed/scoped_value.py
rename to libs/langgraph/langgraph/managed/shared_value.py
index 38a666163..6853306f3 100644
--- a/libs/langgraph/langgraph/managed/scoped_value.py
+++ b/libs/langgraph/langgraph/managed/shared_value.py
@@ -4,11 +4,11 @@ from typing import (
     AsyncIterator,
     Iterator,
     Optional,
-    Self,
     Sequence,
 )
 
 from langchain_core.runnables import RunnableConfig
+from typing_extensions import Self
 
 from langgraph.constants import CONFIG_KEY_KV
 from langgraph.errors import InvalidUpdateError
diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py
index 4bcce7ccd..fef936c80 100644
--- a/libs/langgraph/tests/test_pregel.py
+++ b/libs/langgraph/tests/test_pregel.py
@@ -59,7 +59,7 @@ from langgraph.graph.graph import START
 from langgraph.graph.message import MessageGraph, add_messages
 from langgraph.graph.state import StateGraph
 from langgraph.kv.memory import MemoryKV
-from langgraph.managed.scoped_value import SharedValue
+from langgraph.managed.shared_value import SharedValue
 from langgraph.prebuilt.chat_agent_executor import (
     create_tool_calling_executor,
 )

From 9b90a24d9460ca70ee15a26541c7af706d82a886 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Fri, 16 Aug 2024 16:34:02 -0700
Subject: [PATCH 06/30] Lint

---
 libs/langgraph/langgraph/managed/base.py      |  1 +
 .../langgraph/managed/shared_value.py         | 26 +++++++++++++++++--
 2 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/libs/langgraph/langgraph/managed/base.py b/libs/langgraph/langgraph/managed/base.py
index 5820383f8..8dbf5c30c 100644
--- a/libs/langgraph/langgraph/managed/base.py
+++ b/libs/langgraph/langgraph/managed/base.py
@@ -150,3 +150,4 @@ async def AsyncManagedValuesManager(
 
 
 ChannelKeyPlaceholder = object()
+ChannelTypePlaceholder = object()
diff --git a/libs/langgraph/langgraph/managed/shared_value.py b/libs/langgraph/langgraph/managed/shared_value.py
index 6853306f3..f1b8d9516 100644
--- a/libs/langgraph/langgraph/managed/shared_value.py
+++ b/libs/langgraph/langgraph/managed/shared_value.py
@@ -1,3 +1,4 @@
+import collections.abc
 from contextlib import asynccontextmanager, contextmanager
 from typing import (
     Any,
@@ -5,10 +6,11 @@ from typing import (
     Iterator,
     Optional,
     Sequence,
+    Type,
 )
 
 from langchain_core.runnables import RunnableConfig
-from typing_extensions import Self
+from typing_extensions import NotRequired, Required, Self
 
 from langgraph.constants import CONFIG_KEY_KV
 from langgraph.errors import InvalidUpdateError
@@ -26,6 +28,17 @@ Value = dict[str, V]
 Update = dict[str, Optional[V]]
 
 
+# Adapted from typing_extensions
+def _strip_extras(t):
+    """Strips Annotated, Required and NotRequired from a given type."""
+    if hasattr(t, "__origin__"):
+        return _strip_extras(t.__origin__)
+    if hasattr(t, "__origin__") and t.__origin__ in (Required, NotRequired):
+        return _strip_extras(t.__args__[0])
+
+    return t
+
+
 class SharedValue(WritableManagedValue[Value, Update]):
     @staticmethod
     def on(scope: str) -> ConfiguredManagedValue:
@@ -51,7 +64,16 @@ class SharedValue(WritableManagedValue[Value, Update]):
                 value.value = saved[value.ns]
             yield value
 
-    def __init__(self, config: RunnableConfig, *, scope: str, key: str) -> None:
+    def __init__(
+        self, config: RunnableConfig, *, typ: Type[Any], scope: str, key: str
+    ) -> None:
+        if typ := _strip_extras(typ):
+            if typ not in (
+                dict,
+                collections.abc.Mapping,
+                collections.abc.MutableMapping,
+            ):
+                raise ValueError("SharedValue must be a dict")
         self.scope = scope
         self.config = config
         self.value: Value = {}

From 656f89e16ab7074c972d1920d774bf0852346996 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Fri, 16 Aug 2024 16:38:46 -0700
Subject: [PATCH 07/30] Lint

---
 libs/langgraph/langgraph/graph/state.py          | 3 +++
 libs/langgraph/langgraph/managed/shared_value.py | 9 ++++++++-
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py
index 1c9242ea1..261fbf1d5 100644
--- a/libs/langgraph/langgraph/graph/state.py
+++ b/libs/langgraph/langgraph/graph/state.py
@@ -43,6 +43,7 @@ from langgraph.graph.graph import (
 from langgraph.kv.base import BaseKV
 from langgraph.managed.base import (
     ChannelKeyPlaceholder,
+    ChannelTypePlaceholder,
     ConfiguredManagedValue,
     ManagedValue,
     is_managed_value,
@@ -760,6 +761,8 @@ def _is_field_managed_value(name: str, typ: Type[Any]) -> Optional[Type[ManagedV
                     for k, v in decoration.kwargs.items():
                         if v is ChannelKeyPlaceholder:
                             decoration.kwargs[k] = name
+                        if v is ChannelTypePlaceholder:
+                            decoration.kwargs[k] = typ.__origin__
                 return decoration
 
     return None
diff --git a/libs/langgraph/langgraph/managed/shared_value.py b/libs/langgraph/langgraph/managed/shared_value.py
index f1b8d9516..7c55fc858 100644
--- a/libs/langgraph/langgraph/managed/shared_value.py
+++ b/libs/langgraph/langgraph/managed/shared_value.py
@@ -17,6 +17,7 @@ from langgraph.errors import InvalidUpdateError
 from langgraph.kv.base import BaseKV
 from langgraph.managed.base import (
     ChannelKeyPlaceholder,
+    ChannelTypePlaceholder,
     ConfiguredManagedValue,
     WritableManagedValue,
 )
@@ -43,7 +44,12 @@ class SharedValue(WritableManagedValue[Value, Update]):
     @staticmethod
     def on(scope: str) -> ConfiguredManagedValue:
         return ConfiguredManagedValue(
-            SharedValue, {"scope": scope, "key": ChannelKeyPlaceholder}
+            SharedValue,
+            {
+                "scope": scope,
+                "key": ChannelKeyPlaceholder,
+                "typ": ChannelTypePlaceholder,
+            },
         )
 
     @classmethod
@@ -68,6 +74,7 @@ class SharedValue(WritableManagedValue[Value, Update]):
         self, config: RunnableConfig, *, typ: Type[Any], scope: str, key: str
     ) -> None:
         if typ := _strip_extras(typ):
+            print(typ)
             if typ not in (
                 dict,
                 collections.abc.Mapping,

From 8f8f3849fcfd55eefcdaf316ac0d0554e244d9bf Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Tue, 20 Aug 2024 13:30:37 -0700
Subject: [PATCH 08/30] Lint

---
 libs/langgraph/langgraph/kv/memory.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/libs/langgraph/langgraph/kv/memory.py b/libs/langgraph/langgraph/kv/memory.py
index 387ea1e7a..265dfe77a 100644
--- a/libs/langgraph/langgraph/kv/memory.py
+++ b/libs/langgraph/langgraph/kv/memory.py
@@ -1,5 +1,5 @@
 from collections import defaultdict
-from typing import List
+from typing import List, Optional
 
 from langgraph.kv.base import BaseKV, V
 
@@ -8,12 +8,12 @@ class MemoryKV(BaseKV):
     def __init__(self) -> None:
         self.data: dict[str, dict[str, V]] = defaultdict(dict)
 
-    def get(self, pairs: List[tuple[str, str]]) -> dict[tuple[str, str], V | None]:
+    def get(self, pairs: List[tuple[str, str]]) -> dict[tuple[str, str], Optional[V]]:
         return {pair: self.data[pair[0]].get(pair[1]) for pair in pairs}
 
     async def aget(
         self, pairs: List[tuple[str, str]]
-    ) -> dict[tuple[str, str], V | None]:
+    ) -> dict[tuple[str, str], Optional[V]]:
         return self.get(pairs)
 
     def list(self, prefixes: List[str]) -> dict[str, dict[str, V]]:
@@ -22,12 +22,12 @@ class MemoryKV(BaseKV):
     async def alist(self, prefixes: List[str]) -> dict[str, dict[str, V]]:
         return self.list(prefixes)
 
-    def put(self, writes: List[tuple[str, str, V | None]]) -> None:
+    def put(self, writes: List[tuple[str, str, Optional[V]]]) -> None:
         for namespace, key, value in writes:
             if value is None:
                 self.data[namespace].pop(key, None)
             else:
                 self.data[namespace][key] = value
 
-    async def aput(self, writes: List[tuple[str, str, V | None]]) -> None:
-        self.put(writes)
+    async def aput(self, writes: List[tuple[str, str, Optional[V]]]) -> None:
+        return self.put(writes)

From 630d9c79edca85ef0417afd78727487c973e25e2 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Tue, 20 Aug 2024 13:55:28 -0700
Subject: [PATCH 09/30] Split out async batch to sep file

---
 libs/langgraph/langgraph/graph/state.py |  2 +-
 libs/langgraph/langgraph/kv/base.py     | 84 +-----------------------
 libs/langgraph/langgraph/kv/batch.py    | 85 +++++++++++++++++++++++++
 libs/langgraph/tests/test_kv.py         | 31 ++++++---
 4 files changed, 110 insertions(+), 92 deletions(-)
 create mode 100644 libs/langgraph/langgraph/kv/batch.py

diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py
index 261fbf1d5..c116e6d47 100644
--- a/libs/langgraph/langgraph/graph/state.py
+++ b/libs/langgraph/langgraph/graph/state.py
@@ -381,9 +381,9 @@ class StateGraph(Graph):
 
     def compile(
         self,
+        checkpointer: Optional[BaseCheckpointSaver] = None,
         *,
         kv: Optional[BaseKV] = None,
-        checkpointer: Optional[BaseCheckpointSaver] = None,
         interrupt_before: Optional[Union[All, Sequence[str]]] = None,
         interrupt_after: Optional[Union[All, Sequence[str]]] = None,
         debug: bool = False,
diff --git a/libs/langgraph/langgraph/kv/base.py b/libs/langgraph/langgraph/kv/base.py
index 570783ec1..37e07ee84 100644
--- a/libs/langgraph/langgraph/kv/base.py
+++ b/libs/langgraph/langgraph/kv/base.py
@@ -1,5 +1,4 @@
-import asyncio
-from typing import Any, List, NamedTuple, Optional, Union
+from typing import Any, List, Optional
 
 V = dict[str, Any]
 
@@ -30,84 +29,3 @@ class BaseKV:
     async def aput(self, writes: List[tuple[str, str, Optional[V]]]) -> None:
         # list[(namespace, key, value | none)] -> None
         raise NotImplementedError
-
-
-class GetOp(NamedTuple):
-    pairs: List[tuple[str, str]]
-
-
-class ListOp(NamedTuple):
-    prefixes: List[str]
-
-
-class PutOp(NamedTuple):
-    writes: List[tuple[str, str, Optional[V]]]
-
-
-class KeyValueStore(BaseKV):
-    def __init__(self, kv: BaseKV) -> None:
-        self.kv = kv
-        self.aqueue: dict[asyncio.Future, Union[GetOp, ListOp, PutOp]] = {}
-        self.task = asyncio.create_task(_run(self.aqueue, self.kv))
-
-    def __del__(self) -> None:
-        self.task.cancel()
-
-    async def aget(
-        self, pairs: List[tuple[str, str]]
-    ) -> dict[tuple[str, str], Optional[V]]:
-        fut = asyncio.get_running_loop().create_future()
-        self.aqueue[fut] = GetOp(pairs)
-        return await fut
-
-    async def alist(self, prefixes: List[str]) -> dict[str, dict[str, V]]:
-        fut = asyncio.get_running_loop().create_future()
-        self.aqueue[fut] = ListOp(prefixes)
-        return await fut
-
-    async def aput(self, writes: List[tuple[str, str, Optional[V]]]) -> None:
-        fut = asyncio.get_running_loop().create_future()
-        self.aqueue[fut] = PutOp(writes)
-        return await fut
-
-
-async def _run(
-    aqueue: dict[asyncio.Future, Union[GetOp, ListOp, PutOp]], kv: BaseKV
-) -> None:
-    while True:
-        await asyncio.sleep(0)
-        if not aqueue:
-            continue
-        # this could use a lock, if we want thread safety
-        taken = aqueue.copy()
-        aqueue.clear()
-        # action each operation
-        gets = {f: o for f, o in taken.items() if isinstance(o, GetOp)}
-        if gets:
-            try:
-                results = await kv.aget([p for op in gets.values() for p in op.pairs])
-                for fut, op in gets.items():
-                    fut.set_result({k: results.get(k) for k in op.pairs})
-            except Exception as e:
-                for fut in gets:
-                    fut.set_exception(e)
-        lists = {f: o for f, o in taken.items() if isinstance(o, ListOp)}
-        if lists:
-            try:
-                results = await kv.alist(
-                    [p for op in lists.values() for p in op.prefixes]
-                )
-                for fut, op in lists.items():
-                    fut.set_result({k: results.get(k) for k in op.prefixes})
-            except Exception as e:
-                for fut in lists:
-                    fut.set_exception(e)
-        puts = {f: o for f, o in taken.items() if isinstance(o, PutOp)}
-        if puts:
-            try:
-                await kv.aput([w for op in puts.values() for w in op.writes])
-                for fut in puts:
-                    fut.set_result(None)
-            except Exception as e:
-                for fut in puts:
-                    fut.set_exception(e)
diff --git a/libs/langgraph/langgraph/kv/batch.py b/libs/langgraph/langgraph/kv/batch.py
new file mode 100644
index 000000000..fa6cfc59b
--- /dev/null
+++ b/libs/langgraph/langgraph/kv/batch.py
@@ -0,0 +1,85 @@
+import asyncio
+from typing import NamedTuple, Optional, Union
+
+from langgraph.kv.base import BaseKV, V
+
+
+class GetOp(NamedTuple):
+    pairs: list[tuple[str, str]]
+
+
+class ListOp(NamedTuple):
+    prefixes: list[str]
+
+
+class PutOp(NamedTuple):
+    writes: list[tuple[str, str, Optional[V]]]
+
+
+class AsyncBatchedKV(BaseKV):
+    def __init__(self, kv: BaseKV) -> None:
+        self.kv = kv
+        self.aqueue: dict[asyncio.Future, Union[GetOp, ListOp, PutOp]] = {}
+        self.task = asyncio.create_task(_run(self.aqueue, self.kv))
+
+    def __del__(self) -> None:
+        self.task.cancel()
+
+    async def aget(
+        self, pairs: list[tuple[str, str]]
+    ) -> dict[tuple[str, str], Optional[V]]:
+        fut = asyncio.get_running_loop().create_future()
+        self.aqueue[fut] = GetOp(pairs)
+        return await fut
+
+    async def alist(self, prefixes: list[str]) -> dict[str, dict[str, V]]:
+        fut = asyncio.get_running_loop().create_future()
+        self.aqueue[fut] = ListOp(prefixes)
+        return await fut
+
+    async def aput(self, writes: list[tuple[str, str, Optional[V]]]) -> None:
+        fut = asyncio.get_running_loop().create_future()
+        self.aqueue[fut] = PutOp(writes)
+        return await fut
+
+
+async def _run(
+    aqueue: dict[asyncio.Future, Union[GetOp, ListOp, PutOp]], kv: BaseKV
+) -> None:
+    while True:
+        await asyncio.sleep(0)
+        if not aqueue:
+            continue
+        # this could use a lock, if we want thread safety
+        taken = aqueue.copy()
+        aqueue.clear()
+        # action each operation
+        gets = {f: o for f, o in taken.items() if isinstance(o, GetOp)}
+        if gets:
+            try:
+                results = await kv.aget([p for op in gets.values() for p in op.pairs])
+                for fut, op in gets.items():
+                    fut.set_result({k: results.get(k) for k in op.pairs})
+            except Exception as e:
+                for fut in gets:
+                    fut.set_exception(e)
+        lists = {f: o for f, o in taken.items() if isinstance(o, ListOp)}
+        if lists:
+            try:
+                results = await kv.alist(
+                    [p for op in lists.values() for p in op.prefixes]
+                )
+                for fut, op in lists.items():
+                    fut.set_result({k: results.get(k) for k in op.prefixes})
+            except Exception as e:
+                for fut in lists:
+                    fut.set_exception(e)
+        puts = {f: o for f, o in taken.items() if isinstance(o, PutOp)}
+        if puts:
+            try:
+                await kv.aput([w for op in puts.values() for w in op.writes])
+                for fut in puts:
+                    fut.set_result(None)
+            except Exception as e:
+                for fut in puts:
+                    fut.set_exception(e)
diff --git a/libs/langgraph/tests/test_kv.py b/libs/langgraph/tests/test_kv.py
index 2e24ee07d..5af46ef9c 100644
--- a/libs/langgraph/tests/test_kv.py
+++ b/libs/langgraph/tests/test_kv.py
@@ -1,22 +1,28 @@
 import asyncio
-from typing import Any, List
+from typing import Any
 
 from pytest_mock import MockerFixture
 
-from langgraph.kv.base import BaseKV, KeyValueStore
+from langgraph.kv.base import BaseKV
+from langgraph.kv.batch import AsyncBatchedKV
 
 
-async def test_kv_queue(mocker: MockerFixture) -> None:
+async def test_kv_async_batch(mocker: MockerFixture) -> None:
     aget = mocker.stub()
+    alist = mocker.stub()
 
     class MockKV(BaseKV):
         async def aget(
-            self, pairs: List[tuple[str, str]]
+            self, pairs: list[tuple[str, str]]
         ) -> dict[tuple[str, str], dict[str, Any] | None]:
             aget(pairs)
-            return {pair: {0: pair[0], 1: pair[1]} for pair in pairs}
+            return {pair: 1 for pair in pairs}
 
-    store = KeyValueStore(MockKV())
+        async def alist(self, prefixes: list[str]) -> dict[str, dict[str, Any]]:
+            alist(prefixes)
+            return {prefix: {prefix: 1} for prefix in prefixes}
+
+    store = AsyncBatchedKV(MockKV())
 
     # concurrent calls are batched
     results = await asyncio.gather(
@@ -24,9 +30,18 @@ async def test_kv_queue(mocker: MockerFixture) -> None:
         store.aget([("c", "d")]),
     )
     assert results == [
-        {("a", "b"): {0: "a", 1: "b"}},
-        {("c", "d"): {0: "c", 1: "d"}},
+        {("a", "b"): 1},
+        {("c", "d"): 1},
     ]
     assert [c.args for c in aget.call_args_list] == [
         ([("a", "b"), ("c", "d")],),
     ]
+
+    results = await asyncio.gather(
+        store.alist(["a", "b"]),
+        store.alist(["c", "d"]),
+    )
+    assert results == [{"a": {"a": 1}, "b": {"b": 1}}, {"c": {"c": 1}, "d": {"d": 1}}]
+    assert [c.args for c in alist.call_args_list] == [
+        (["a", "b", "c", "d"],),
+    ]

From 4e1db854f6773758a06b0c288cc79f1cbfbd47c8 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Tue, 20 Aug 2024 14:55:23 -0700
Subject: [PATCH 10/30] Add async test

---
 .../langgraph/managed/shared_value.py         |  1 -
 libs/langgraph/tests/test_pregel_async.py     | 40 ++++++++++++++++---
 2 files changed, 34 insertions(+), 7 deletions(-)

diff --git a/libs/langgraph/langgraph/managed/shared_value.py b/libs/langgraph/langgraph/managed/shared_value.py
index 7c55fc858..c12482c5d 100644
--- a/libs/langgraph/langgraph/managed/shared_value.py
+++ b/libs/langgraph/langgraph/managed/shared_value.py
@@ -74,7 +74,6 @@ class SharedValue(WritableManagedValue[Value, Update]):
         self, config: RunnableConfig, *, typ: Type[Any], scope: str, key: str
     ) -> None:
         if typ := _strip_extras(typ):
-            print(typ)
             if typ not in (
                 dict,
                 collections.abc.Mapping,
diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py
index 3048ddbb4..02cd637ac 100644
--- a/libs/langgraph/tests/test_pregel_async.py
+++ b/libs/langgraph/tests/test_pregel_async.py
@@ -52,6 +52,9 @@ from langgraph.errors import InvalidUpdateError, NodeInterrupt
 from langgraph.graph import END, Graph, StateGraph
 from langgraph.graph.graph import START
 from langgraph.graph.message import MessageGraph, add_messages
+from langgraph.kv.batch import AsyncBatchedKV
+from langgraph.kv.memory import MemoryKV
+from langgraph.managed.shared_value import SharedValue
 from langgraph.prebuilt.chat_agent_executor import (
     create_tool_calling_executor,
 )
@@ -4778,10 +4781,33 @@ async def test_start_branch_then() -> None:
     class State(TypedDict):
         my_key: Annotated[str, operator.add]
         market: str
+        shared: Annotated[dict[str, dict[str, Any]], SharedValue.on("assistant_id")]
+        other: Annotated[dict[str, dict[str, Any]], SharedValue.on("assistant_id")]
+
+    def assert_shared_value(data: State, config: RunnableConfig) -> State:
+        assert "shared" in data
+        if thread_id := config["configurable"].get("thread_id"):
+            if thread_id == "1":
+                # this is the first thread, so should not see a value
+                assert data["shared"] == {}
+                return {"shared": {"1": {"hello": "world"}}, "other": {"2": {1: 2}}}
+            elif thread_id == "2":
+                # this should get value saved by thread 1
+                assert data["shared"] == {"1": {"hello": "world"}}
+            elif thread_id == "3":
+                # this is a different assistant, so should not see previous value
+                assert data["shared"] == {}
+        return {}
+
+    def tool_two_slow(data: State, config: RunnableConfig) -> State:
+        return {"my_key": " slow", **assert_shared_value(data, config)}
+
+    def tool_two_fast(data: State, config: RunnableConfig) -> State:
+        return {"my_key": " fast", **assert_shared_value(data, config)}
 
     tool_two_graph = StateGraph(State)
-    tool_two_graph.add_node("tool_two_slow", lambda s, config: {"my_key": " slow"})
-    tool_two_graph.add_node("tool_two_fast", lambda s: {"my_key": " fast"})
+    tool_two_graph.add_node("tool_two_slow", tool_two_slow)
+    tool_two_graph.add_node("tool_two_fast", tool_two_fast)
     tool_two_graph.set_conditional_entry_point(
         lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast", then=END
     )
@@ -4798,14 +4824,16 @@ async def test_start_branch_then() -> None:
 
     async with AsyncSqliteSaver.from_conn_string(":memory:") as saver:
         tool_two = tool_two_graph.compile(
-            checkpointer=saver, interrupt_before=["tool_two_fast", "tool_two_slow"]
+            kv=AsyncBatchedKV(MemoryKV()),
+            checkpointer=saver,
+            interrupt_before=["tool_two_fast", "tool_two_slow"],
         )
 
         # missing thread_id
         with pytest.raises(ValueError, match="thread_id"):
             await tool_two.ainvoke({"my_key": "value", "market": "DE"})
 
-        thread1 = {"configurable": {"thread_id": "1"}}
+        thread1 = {"configurable": {"thread_id": "1", "assistant_id": "a"}}
         # stop when about to enter node
         assert await tool_two.ainvoke({"my_key": "value", "market": "DE"}, thread1) == {
             "my_key": "value",
@@ -4865,7 +4893,7 @@ async def test_start_branch_then() -> None:
             ][-1].config,
         )
 
-        thread2 = {"configurable": {"thread_id": "2"}}
+        thread2 = {"configurable": {"thread_id": "2", "assistant_id": "a"}}
         # stop when about to enter node
         assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread2) == {
             "my_key": "value",
@@ -4913,7 +4941,7 @@ async def test_start_branch_then() -> None:
             ][-1].config,
         )
 
-        thread3 = {"configurable": {"thread_id": "3"}}
+        thread3 = {"configurable": {"thread_id": "3", "assistant_id": "b"}}
         # stop when about to enter node
         assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread3) == {
             "my_key": "value",

From c3794f1fd3211bbc55bfec225d6ae3aa64f4e0d5 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Tue, 20 Aug 2024 15:03:09 -0700
Subject: [PATCH 11/30] Use batched async kv inside loop

---
 libs/langgraph/langgraph/pregel/loop.py   | 10 +++++++---
 libs/langgraph/tests/test_pregel_async.py |  3 +--
 2 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py
index aaa43be67..903a3a644 100644
--- a/libs/langgraph/langgraph/pregel/loop.py
+++ b/libs/langgraph/langgraph/pregel/loop.py
@@ -50,6 +50,8 @@ from langgraph.constants import (
     Interrupt,
 )
 from langgraph.errors import EmptyInputError, GraphInterrupt
+from langgraph.kv.base import BaseKV
+from langgraph.kv.batch import AsyncBatchedKV
 from langgraph.managed.base import (
     AsyncManagedValuesManager,
     ManagedValueMapping,
@@ -103,7 +105,7 @@ class PregelLoop:
         ]
     ]
     graph: "Pregel"
-
+    kv: Optional[BaseKV]
     submit: Submit
     channels: Mapping[str, BaseChannel]
     managed: ManagedValueMapping
@@ -425,6 +427,7 @@ class SyncPregelLoop(PregelLoop, ContextManager):
         graph: "Pregel",
     ) -> None:
         super().__init__(input, config=config, checkpointer=checkpointer, graph=graph)
+        self.kv = graph.kv
         self.stack = ExitStack()
         if checkpointer:
             self.checkpointer_get_next_version = checkpointer.get_next_version
@@ -476,7 +479,7 @@ class SyncPregelLoop(PregelLoop, ContextManager):
         self.managed = self.stack.enter_context(
             ManagedValuesManager(
                 self.graph.managed_values_dict,
-                patch_config(self.config, configurable={CONFIG_KEY_KV: self.graph.kv}),
+                patch_config(self.config, configurable={CONFIG_KEY_KV: self.kv}),
             )
         )
         self.stack.push(self._suppress_interrupt)
@@ -508,6 +511,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
         graph: "Pregel",
     ) -> None:
         super().__init__(input, config=config, checkpointer=checkpointer, graph=graph)
+        self.kv = AsyncBatchedKV(graph.kv) if graph.kv else None
         self.stack = AsyncExitStack()
         if checkpointer:
             self.checkpointer_get_next_version = checkpointer.get_next_version
@@ -563,7 +567,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
         self.managed = await self.stack.enter_async_context(
             AsyncManagedValuesManager(
                 self.graph.managed_values_dict,
-                patch_config(self.config, configurable={CONFIG_KEY_KV: self.graph.kv}),
+                patch_config(self.config, configurable={CONFIG_KEY_KV: self.kv}),
             )
         )
         self.stack.push(self._suppress_interrupt)
diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py
index 02cd637ac..b383a4d3f 100644
--- a/libs/langgraph/tests/test_pregel_async.py
+++ b/libs/langgraph/tests/test_pregel_async.py
@@ -52,7 +52,6 @@ from langgraph.errors import InvalidUpdateError, NodeInterrupt
 from langgraph.graph import END, Graph, StateGraph
 from langgraph.graph.graph import START
 from langgraph.graph.message import MessageGraph, add_messages
-from langgraph.kv.batch import AsyncBatchedKV
 from langgraph.kv.memory import MemoryKV
 from langgraph.managed.shared_value import SharedValue
 from langgraph.prebuilt.chat_agent_executor import (
@@ -4824,7 +4823,7 @@ async def test_start_branch_then() -> None:
 
     async with AsyncSqliteSaver.from_conn_string(":memory:") as saver:
         tool_two = tool_two_graph.compile(
-            kv=AsyncBatchedKV(MemoryKV()),
+            kv=MemoryKV(),
             checkpointer=saver,
             interrupt_before=["tool_two_fast", "tool_two_slow"],
         )

From b228fc1a9bf75fcdf99de3694d12eb930a23293c Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Tue, 20 Aug 2024 16:36:35 -0700
Subject: [PATCH 12/30] Add serde

---
 libs/langgraph/langgraph/kv/base.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/libs/langgraph/langgraph/kv/base.py b/libs/langgraph/langgraph/kv/base.py
index 37e07ee84..bbd2e0899 100644
--- a/libs/langgraph/langgraph/kv/base.py
+++ b/libs/langgraph/langgraph/kv/base.py
@@ -1,9 +1,15 @@
 from typing import Any, List, Optional
 
+from langgraph.checkpoint.serde.base import SerializerProtocol
+from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
+
 V = dict[str, Any]
 
 
 class BaseKV:
+    def __init__(self, *, serde: SerializerProtocol = JsonPlusSerializer()) -> None:
+        self.serde = serde
+
     def get(self, pairs: List[tuple[str, str]]) -> dict[tuple[str, str], Optional[V]]:
         # list[(namespace, key)] -> dict[(namespace, key), value | none]
         raise NotImplementedError

From bd7b9cca21bb74868de6ff9b9407b9793215ba62 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Wed, 21 Aug 2024 09:04:43 -0700
Subject: [PATCH 13/30] WIP

---
 libs/langgraph/langgraph/graph/state.py       |  4 +--
 libs/langgraph/langgraph/kv/base.py           | 18 +----------
 libs/langgraph/langgraph/kv/batch.py          | 30 ++++---------------
 libs/langgraph/langgraph/kv/memory.py         | 12 ++------
 .../langgraph/managed/shared_value.py         |  4 +--
 libs/langgraph/langgraph/pregel/__init__.py   |  4 +--
 libs/langgraph/langgraph/pregel/loop.py       |  4 +--
 libs/langgraph/tests/test_kv.py               |  4 +--
 8 files changed, 18 insertions(+), 62 deletions(-)

diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py
index c116e6d47..26770bd20 100644
--- a/libs/langgraph/langgraph/graph/state.py
+++ b/libs/langgraph/langgraph/graph/state.py
@@ -40,7 +40,7 @@ from langgraph.graph.graph import (
     Graph,
     Send,
 )
-from langgraph.kv.base import BaseKV
+from langgraph.kv.base import BaseMemory
 from langgraph.managed.base import (
     ChannelKeyPlaceholder,
     ChannelTypePlaceholder,
@@ -383,7 +383,7 @@ class StateGraph(Graph):
         self,
         checkpointer: Optional[BaseCheckpointSaver] = None,
         *,
-        kv: Optional[BaseKV] = None,
+        kv: Optional[BaseMemory] = None,
         interrupt_before: Optional[Union[All, Sequence[str]]] = None,
         interrupt_after: Optional[Union[All, Sequence[str]]] = None,
         debug: bool = False,
diff --git a/libs/langgraph/langgraph/kv/base.py b/libs/langgraph/langgraph/kv/base.py
index bbd2e0899..0c5791eb6 100644
--- a/libs/langgraph/langgraph/kv/base.py
+++ b/libs/langgraph/langgraph/kv/base.py
@@ -1,19 +1,9 @@
 from typing import Any, List, Optional
 
-from langgraph.checkpoint.serde.base import SerializerProtocol
-from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
-
 V = dict[str, Any]
 
 
-class BaseKV:
-    def __init__(self, *, serde: SerializerProtocol = JsonPlusSerializer()) -> None:
-        self.serde = serde
-
-    def get(self, pairs: List[tuple[str, str]]) -> dict[tuple[str, str], Optional[V]]:
-        # list[(namespace, key)] -> dict[(namespace, key), value | none]
-        raise NotImplementedError
-
+class BaseMemory:
     def list(self, prefixes: List[str]) -> dict[str, dict[str, V]]:
         # list[namespace] -> dict[namespace, list[value]]
         raise NotImplementedError
@@ -22,12 +12,6 @@ class BaseKV:
         # list[(namespace, key, value | none)] -> None
         raise NotImplementedError
 
-    async def aget(
-        self, pairs: List[tuple[str, str]]
-    ) -> dict[tuple[str, str], Optional[V]]:
-        # list[(namespace, key)] -> dict[(namespace, key), value | none]
-        raise NotImplementedError
-
     async def alist(self, prefixes: List[str]) -> dict[str, dict[str, V]]:
         # list[namespace] -> dict[namespace, list[value]]
         raise NotImplementedError
diff --git a/libs/langgraph/langgraph/kv/batch.py b/libs/langgraph/langgraph/kv/batch.py
index fa6cfc59b..971d50df4 100644
--- a/libs/langgraph/langgraph/kv/batch.py
+++ b/libs/langgraph/langgraph/kv/batch.py
@@ -1,11 +1,7 @@
 import asyncio
 from typing import NamedTuple, Optional, Union
 
-from langgraph.kv.base import BaseKV, V
-
-
-class GetOp(NamedTuple):
-    pairs: list[tuple[str, str]]
+from langgraph.kv.base import BaseMemory, V
 
 
 class ListOp(NamedTuple):
@@ -16,22 +12,15 @@ class PutOp(NamedTuple):
     writes: list[tuple[str, str, Optional[V]]]
 
 
-class AsyncBatchedKV(BaseKV):
-    def __init__(self, kv: BaseKV) -> None:
+class AsyncBatchedKV(BaseMemory):
+    def __init__(self, kv: BaseMemory) -> None:
         self.kv = kv
-        self.aqueue: dict[asyncio.Future, Union[GetOp, ListOp, PutOp]] = {}
+        self.aqueue: dict[asyncio.Future, Union[ListOp, PutOp]] = {}
         self.task = asyncio.create_task(_run(self.aqueue, self.kv))
 
     def __del__(self) -> None:
         self.task.cancel()
 
-    async def aget(
-        self, pairs: list[tuple[str, str]]
-    ) -> dict[tuple[str, str], Optional[V]]:
-        fut = asyncio.get_running_loop().create_future()
-        self.aqueue[fut] = GetOp(pairs)
-        return await fut
-
     async def alist(self, prefixes: list[str]) -> dict[str, dict[str, V]]:
         fut = asyncio.get_running_loop().create_future()
         self.aqueue[fut] = ListOp(prefixes)
@@ -44,7 +33,7 @@ class AsyncBatchedKV(BaseKV):
 
 
 async def _run(
-    aqueue: dict[asyncio.Future, Union[GetOp, ListOp, PutOp]], kv: BaseKV
+    aqueue: dict[asyncio.Future, Union[ListOp, PutOp]], kv: BaseMemory
 ) -> None:
     while True:
         await asyncio.sleep(0)
@@ -54,15 +43,6 @@ async def _run(
         taken = aqueue.copy()
         aqueue.clear()
         # action each operation
-        gets = {f: o for f, o in taken.items() if isinstance(o, GetOp)}
-        if gets:
-            try:
-                results = await kv.aget([p for op in gets.values() for p in op.pairs])
-                for fut, op in gets.items():
-                    fut.set_result({k: results.get(k) for k in op.pairs})
-            except Exception as e:
-                for fut in gets:
-                    fut.set_exception(e)
         lists = {f: o for f, o in taken.items() if isinstance(o, ListOp)}
         if lists:
             try:
diff --git a/libs/langgraph/langgraph/kv/memory.py b/libs/langgraph/langgraph/kv/memory.py
index 265dfe77a..4b0e92a2d 100644
--- a/libs/langgraph/langgraph/kv/memory.py
+++ b/libs/langgraph/langgraph/kv/memory.py
@@ -1,21 +1,13 @@
 from collections import defaultdict
 from typing import List, Optional
 
-from langgraph.kv.base import BaseKV, V
+from langgraph.kv.base import BaseMemory, V
 
 
-class MemoryKV(BaseKV):
+class MemoryKV(BaseMemory):
     def __init__(self) -> None:
         self.data: dict[str, dict[str, V]] = defaultdict(dict)
 
-    def get(self, pairs: List[tuple[str, str]]) -> dict[tuple[str, str], Optional[V]]:
-        return {pair: self.data[pair[0]].get(pair[1]) for pair in pairs}
-
-    async def aget(
-        self, pairs: List[tuple[str, str]]
-    ) -> dict[tuple[str, str], Optional[V]]:
-        return self.get(pairs)
-
     def list(self, prefixes: List[str]) -> dict[str, dict[str, V]]:
         return {prefix: self.data[prefix] for prefix in prefixes}
 
diff --git a/libs/langgraph/langgraph/managed/shared_value.py b/libs/langgraph/langgraph/managed/shared_value.py
index c12482c5d..a1659bb8b 100644
--- a/libs/langgraph/langgraph/managed/shared_value.py
+++ b/libs/langgraph/langgraph/managed/shared_value.py
@@ -14,7 +14,7 @@ from typing_extensions import NotRequired, Required, Self
 
 from langgraph.constants import CONFIG_KEY_KV
 from langgraph.errors import InvalidUpdateError
-from langgraph.kv.base import BaseKV
+from langgraph.kv.base import BaseMemory
 from langgraph.managed.base import (
     ChannelKeyPlaceholder,
     ChannelTypePlaceholder,
@@ -83,7 +83,7 @@ class SharedValue(WritableManagedValue[Value, Update]):
         self.scope = scope
         self.config = config
         self.value: Value = {}
-        self.kv: BaseKV = config["configurable"].get(CONFIG_KEY_KV)
+        self.kv: BaseMemory = config["configurable"].get(CONFIG_KEY_KV)
         if self.kv is None:
             self.ns: Optional[str] = None
         elif scope_value := config["configurable"].get(self.scope):
diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py
index 833453197..a5049d36f 100644
--- a/libs/langgraph/langgraph/pregel/__init__.py
+++ b/libs/langgraph/langgraph/pregel/__init__.py
@@ -73,7 +73,7 @@ from langgraph.constants import (
     Interrupt,
 )
 from langgraph.errors import GraphInterrupt, GraphRecursionError, InvalidUpdateError
-from langgraph.kv.base import BaseKV
+from langgraph.kv.base import BaseMemory
 from langgraph.managed.base import (
     AsyncManagedValuesManager,
     ManagedValuesManager,
@@ -224,7 +224,7 @@ class Pregel(
     checkpointer: Optional[BaseCheckpointSaver] = None
     """Checkpointer used to save and load graph state. Defaults to None."""
 
-    kv: Optional[BaseKV] = None
+    kv: Optional[BaseMemory] = None
     """Key-value store to use. Defaults to None."""
 
     retry_policy: Optional[RetryPolicy] = None
diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py
index 903a3a644..41234e9c7 100644
--- a/libs/langgraph/langgraph/pregel/loop.py
+++ b/libs/langgraph/langgraph/pregel/loop.py
@@ -50,7 +50,7 @@ from langgraph.constants import (
     Interrupt,
 )
 from langgraph.errors import EmptyInputError, GraphInterrupt
-from langgraph.kv.base import BaseKV
+from langgraph.kv.base import BaseMemory
 from langgraph.kv.batch import AsyncBatchedKV
 from langgraph.managed.base import (
     AsyncManagedValuesManager,
@@ -105,7 +105,7 @@ class PregelLoop:
         ]
     ]
     graph: "Pregel"
-    kv: Optional[BaseKV]
+    kv: Optional[BaseMemory]
     submit: Submit
     channels: Mapping[str, BaseChannel]
     managed: ManagedValueMapping
diff --git a/libs/langgraph/tests/test_kv.py b/libs/langgraph/tests/test_kv.py
index 5af46ef9c..68294f42c 100644
--- a/libs/langgraph/tests/test_kv.py
+++ b/libs/langgraph/tests/test_kv.py
@@ -3,7 +3,7 @@ from typing import Any
 
 from pytest_mock import MockerFixture
 
-from langgraph.kv.base import BaseKV
+from langgraph.kv.base import BaseMemory
 from langgraph.kv.batch import AsyncBatchedKV
 
 
@@ -11,7 +11,7 @@ async def test_kv_async_batch(mocker: MockerFixture) -> None:
     aget = mocker.stub()
     alist = mocker.stub()
 
-    class MockKV(BaseKV):
+    class MockKV(BaseMemory):
         async def aget(
             self, pairs: list[tuple[str, str]]
         ) -> dict[tuple[str, str], dict[str, Any] | None]:

From c9e6ee6da70c3880df5ae54fd8f54e5c546a5666 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Wed, 21 Aug 2024 09:34:12 -0700
Subject: [PATCH 14/30] Rename

---
 libs/langgraph/langgraph/graph/state.py            |  6 +++---
 libs/langgraph/langgraph/managed/shared_value.py   |  4 ++--
 libs/langgraph/langgraph/pregel/__init__.py        |  4 ++--
 libs/langgraph/langgraph/pregel/loop.py            | 14 +++++++-------
 libs/langgraph/langgraph/{kv => store}/__init__.py |  0
 libs/langgraph/langgraph/{kv => store}/base.py     |  2 +-
 libs/langgraph/langgraph/{kv => store}/batch.py    |  8 ++++----
 libs/langgraph/langgraph/{kv => store}/memory.py   |  4 ++--
 libs/langgraph/tests/test_kv.py                    |  8 ++++----
 libs/langgraph/tests/test_pregel.py                |  4 ++--
 libs/langgraph/tests/test_pregel_async.py          |  4 ++--
 11 files changed, 29 insertions(+), 29 deletions(-)
 rename libs/langgraph/langgraph/{kv => store}/__init__.py (100%)
 rename libs/langgraph/langgraph/{kv => store}/base.py (97%)
 rename libs/langgraph/langgraph/{kv => store}/batch.py (93%)
 rename libs/langgraph/langgraph/{kv => store}/memory.py (91%)

diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py
index 26770bd20..8c85b7760 100644
--- a/libs/langgraph/langgraph/graph/state.py
+++ b/libs/langgraph/langgraph/graph/state.py
@@ -40,7 +40,6 @@ from langgraph.graph.graph import (
     Graph,
     Send,
 )
-from langgraph.kv.base import BaseMemory
 from langgraph.managed.base import (
     ChannelKeyPlaceholder,
     ChannelTypePlaceholder,
@@ -52,6 +51,7 @@ from langgraph.managed.base import (
 from langgraph.pregel.read import ChannelRead, PregelNode
 from langgraph.pregel.types import All, RetryPolicy
 from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry
+from langgraph.store.base import BaseStore
 from langgraph.utils import RunnableCallable, coerce_to_runnable
 
 logger = logging.getLogger(__name__)
@@ -383,7 +383,7 @@ class StateGraph(Graph):
         self,
         checkpointer: Optional[BaseCheckpointSaver] = None,
         *,
-        kv: Optional[BaseMemory] = None,
+        store: Optional[BaseStore] = None,
         interrupt_before: Optional[Union[All, Sequence[str]]] = None,
         interrupt_after: Optional[Union[All, Sequence[str]]] = None,
         debug: bool = False,
@@ -452,7 +452,7 @@ class StateGraph(Graph):
             interrupt_after_nodes=interrupt_after,
             auto_validate=False,
             debug=debug,
-            kv=kv,
+            store=store,
         )
 
         compiled.attach_node(START, None)
diff --git a/libs/langgraph/langgraph/managed/shared_value.py b/libs/langgraph/langgraph/managed/shared_value.py
index a1659bb8b..74a7c1ef4 100644
--- a/libs/langgraph/langgraph/managed/shared_value.py
+++ b/libs/langgraph/langgraph/managed/shared_value.py
@@ -14,13 +14,13 @@ from typing_extensions import NotRequired, Required, Self
 
 from langgraph.constants import CONFIG_KEY_KV
 from langgraph.errors import InvalidUpdateError
-from langgraph.kv.base import BaseMemory
 from langgraph.managed.base import (
     ChannelKeyPlaceholder,
     ChannelTypePlaceholder,
     ConfiguredManagedValue,
     WritableManagedValue,
 )
+from langgraph.store.base import BaseStore
 
 V = dict[str, Any]
 
@@ -83,7 +83,7 @@ class SharedValue(WritableManagedValue[Value, Update]):
         self.scope = scope
         self.config = config
         self.value: Value = {}
-        self.kv: BaseMemory = config["configurable"].get(CONFIG_KEY_KV)
+        self.kv: BaseStore = config["configurable"].get(CONFIG_KEY_KV)
         if self.kv is None:
             self.ns: Optional[str] = None
         elif scope_value := config["configurable"].get(self.scope):
diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py
index a5049d36f..5e058fd16 100644
--- a/libs/langgraph/langgraph/pregel/__init__.py
+++ b/libs/langgraph/langgraph/pregel/__init__.py
@@ -73,7 +73,6 @@ from langgraph.constants import (
     Interrupt,
 )
 from langgraph.errors import GraphInterrupt, GraphRecursionError, InvalidUpdateError
-from langgraph.kv.base import BaseMemory
 from langgraph.managed.base import (
     AsyncManagedValuesManager,
     ManagedValuesManager,
@@ -109,6 +108,7 @@ from langgraph.pregel.types import (
 from langgraph.pregel.utils import get_new_channel_versions
 from langgraph.pregel.validate import validate_graph, validate_keys
 from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
+from langgraph.store.base import BaseStore
 
 WriteValue = Union[
     Runnable[Input, Output],
@@ -224,7 +224,7 @@ class Pregel(
     checkpointer: Optional[BaseCheckpointSaver] = None
     """Checkpointer used to save and load graph state. Defaults to None."""
 
-    kv: Optional[BaseMemory] = None
+    store: Optional[BaseStore] = None
     """Key-value store to use. Defaults to None."""
 
     retry_policy: Optional[RetryPolicy] = None
diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py
index 41234e9c7..8079c0dc2 100644
--- a/libs/langgraph/langgraph/pregel/loop.py
+++ b/libs/langgraph/langgraph/pregel/loop.py
@@ -50,8 +50,6 @@ from langgraph.constants import (
     Interrupt,
 )
 from langgraph.errors import EmptyInputError, GraphInterrupt
-from langgraph.kv.base import BaseMemory
-from langgraph.kv.batch import AsyncBatchedKV
 from langgraph.managed.base import (
     AsyncManagedValuesManager,
     ManagedValueMapping,
@@ -74,6 +72,8 @@ from langgraph.pregel.executor import (
 from langgraph.pregel.io import map_input, map_output_updates, map_output_values, single
 from langgraph.pregel.types import PregelExecutableTask
 from langgraph.pregel.utils import get_new_channel_versions
+from langgraph.store.base import BaseStore
+from langgraph.store.batch import AsyncBatchedStore
 
 if TYPE_CHECKING:
     from langgraph.pregel import Pregel
@@ -105,7 +105,7 @@ class PregelLoop:
         ]
     ]
     graph: "Pregel"
-    kv: Optional[BaseMemory]
+    store: Optional[BaseStore]
     submit: Submit
     channels: Mapping[str, BaseChannel]
     managed: ManagedValueMapping
@@ -427,7 +427,7 @@ class SyncPregelLoop(PregelLoop, ContextManager):
         graph: "Pregel",
     ) -> None:
         super().__init__(input, config=config, checkpointer=checkpointer, graph=graph)
-        self.kv = graph.kv
+        self.store = graph.store
         self.stack = ExitStack()
         if checkpointer:
             self.checkpointer_get_next_version = checkpointer.get_next_version
@@ -479,7 +479,7 @@ class SyncPregelLoop(PregelLoop, ContextManager):
         self.managed = self.stack.enter_context(
             ManagedValuesManager(
                 self.graph.managed_values_dict,
-                patch_config(self.config, configurable={CONFIG_KEY_KV: self.kv}),
+                patch_config(self.config, configurable={CONFIG_KEY_KV: self.store}),
             )
         )
         self.stack.push(self._suppress_interrupt)
@@ -511,7 +511,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
         graph: "Pregel",
     ) -> None:
         super().__init__(input, config=config, checkpointer=checkpointer, graph=graph)
-        self.kv = AsyncBatchedKV(graph.kv) if graph.kv else None
+        self.store = AsyncBatchedStore(graph.store) if graph.store else None
         self.stack = AsyncExitStack()
         if checkpointer:
             self.checkpointer_get_next_version = checkpointer.get_next_version
@@ -567,7 +567,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
         self.managed = await self.stack.enter_async_context(
             AsyncManagedValuesManager(
                 self.graph.managed_values_dict,
-                patch_config(self.config, configurable={CONFIG_KEY_KV: self.kv}),
+                patch_config(self.config, configurable={CONFIG_KEY_KV: self.store}),
             )
         )
         self.stack.push(self._suppress_interrupt)
diff --git a/libs/langgraph/langgraph/kv/__init__.py b/libs/langgraph/langgraph/store/__init__.py
similarity index 100%
rename from libs/langgraph/langgraph/kv/__init__.py
rename to libs/langgraph/langgraph/store/__init__.py
diff --git a/libs/langgraph/langgraph/kv/base.py b/libs/langgraph/langgraph/store/base.py
similarity index 97%
rename from libs/langgraph/langgraph/kv/base.py
rename to libs/langgraph/langgraph/store/base.py
index 0c5791eb6..7f0030f56 100644
--- a/libs/langgraph/langgraph/kv/base.py
+++ b/libs/langgraph/langgraph/store/base.py
@@ -3,7 +3,7 @@ from typing import Any, List, Optional
 V = dict[str, Any]
 
 
-class BaseMemory:
+class BaseStore:
     def list(self, prefixes: List[str]) -> dict[str, dict[str, V]]:
         # list[namespace] -> dict[namespace, list[value]]
         raise NotImplementedError
diff --git a/libs/langgraph/langgraph/kv/batch.py b/libs/langgraph/langgraph/store/batch.py
similarity index 93%
rename from libs/langgraph/langgraph/kv/batch.py
rename to libs/langgraph/langgraph/store/batch.py
index 971d50df4..cfd05e548 100644
--- a/libs/langgraph/langgraph/kv/batch.py
+++ b/libs/langgraph/langgraph/store/batch.py
@@ -1,7 +1,7 @@
 import asyncio
 from typing import NamedTuple, Optional, Union
 
-from langgraph.kv.base import BaseMemory, V
+from langgraph.store.base import BaseStore, V
 
 
 class ListOp(NamedTuple):
@@ -12,8 +12,8 @@ class PutOp(NamedTuple):
     writes: list[tuple[str, str, Optional[V]]]
 
 
-class AsyncBatchedKV(BaseMemory):
-    def __init__(self, kv: BaseMemory) -> None:
+class AsyncBatchedStore(BaseStore):
+    def __init__(self, kv: BaseStore) -> None:
         self.kv = kv
         self.aqueue: dict[asyncio.Future, Union[ListOp, PutOp]] = {}
         self.task = asyncio.create_task(_run(self.aqueue, self.kv))
@@ -33,7 +33,7 @@ class AsyncBatchedKV(BaseMemory):
 
 
 async def _run(
-    aqueue: dict[asyncio.Future, Union[ListOp, PutOp]], kv: BaseMemory
+    aqueue: dict[asyncio.Future, Union[ListOp, PutOp]], kv: BaseStore
 ) -> None:
     while True:
         await asyncio.sleep(0)
diff --git a/libs/langgraph/langgraph/kv/memory.py b/libs/langgraph/langgraph/store/memory.py
similarity index 91%
rename from libs/langgraph/langgraph/kv/memory.py
rename to libs/langgraph/langgraph/store/memory.py
index 4b0e92a2d..48fa2884f 100644
--- a/libs/langgraph/langgraph/kv/memory.py
+++ b/libs/langgraph/langgraph/store/memory.py
@@ -1,10 +1,10 @@
 from collections import defaultdict
 from typing import List, Optional
 
-from langgraph.kv.base import BaseMemory, V
+from langgraph.store.base import BaseStore, V
 
 
-class MemoryKV(BaseMemory):
+class MemoryStore(BaseStore):
     def __init__(self) -> None:
         self.data: dict[str, dict[str, V]] = defaultdict(dict)
 
diff --git a/libs/langgraph/tests/test_kv.py b/libs/langgraph/tests/test_kv.py
index 68294f42c..0c2862b48 100644
--- a/libs/langgraph/tests/test_kv.py
+++ b/libs/langgraph/tests/test_kv.py
@@ -3,15 +3,15 @@ from typing import Any
 
 from pytest_mock import MockerFixture
 
-from langgraph.kv.base import BaseMemory
-from langgraph.kv.batch import AsyncBatchedKV
+from langgraph.store.base import BaseStore
+from langgraph.store.batch import AsyncBatchedStore
 
 
 async def test_kv_async_batch(mocker: MockerFixture) -> None:
     aget = mocker.stub()
     alist = mocker.stub()
 
-    class MockKV(BaseMemory):
+    class MockKV(BaseStore):
         async def aget(
             self, pairs: list[tuple[str, str]]
         ) -> dict[tuple[str, str], dict[str, Any] | None]:
@@ -22,7 +22,7 @@ async def test_kv_async_batch(mocker: MockerFixture) -> None:
             alist(prefixes)
             return {prefix: {prefix: 1} for prefix in prefixes}
 
-    store = AsyncBatchedKV(MockKV())
+    store = AsyncBatchedStore(MockKV())
 
     # concurrent calls are batched
     results = await asyncio.gather(
diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py
index fef936c80..8cf6aafa8 100644
--- a/libs/langgraph/tests/test_pregel.py
+++ b/libs/langgraph/tests/test_pregel.py
@@ -58,7 +58,6 @@ from langgraph.graph import END, Graph
 from langgraph.graph.graph import START
 from langgraph.graph.message import MessageGraph, add_messages
 from langgraph.graph.state import StateGraph
-from langgraph.kv.memory import MemoryKV
 from langgraph.managed.shared_value import SharedValue
 from langgraph.prebuilt.chat_agent_executor import (
     create_tool_calling_executor,
@@ -67,6 +66,7 @@ from langgraph.prebuilt.tool_node import ToolNode
 from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot
 from langgraph.pregel.retry import RetryPolicy
 from langgraph.pregel.types import PregelTask
+from langgraph.store.memory import MemoryStore
 from tests.any_str import AnyStr, ExceptionLike
 from tests.memory_assert import (
     MemorySaverAssertCheckpointMetadata,
@@ -6250,7 +6250,7 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None:
 
     with SqliteSaver.from_conn_string(":memory:") as saver:
         tool_two = tool_two_graph.compile(
-            kv=MemoryKV(),
+            store=MemoryStore(),
             checkpointer=saver,
             interrupt_before=["tool_two_fast", "tool_two_slow"],
         )
diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py
index b383a4d3f..d7a532489 100644
--- a/libs/langgraph/tests/test_pregel_async.py
+++ b/libs/langgraph/tests/test_pregel_async.py
@@ -52,7 +52,6 @@ from langgraph.errors import InvalidUpdateError, NodeInterrupt
 from langgraph.graph import END, Graph, StateGraph
 from langgraph.graph.graph import START
 from langgraph.graph.message import MessageGraph, add_messages
-from langgraph.kv.memory import MemoryKV
 from langgraph.managed.shared_value import SharedValue
 from langgraph.prebuilt.chat_agent_executor import (
     create_tool_calling_executor,
@@ -62,6 +61,7 @@ from langgraph.prebuilt.tool_node import ToolNode
 from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot
 from langgraph.pregel.retry import RetryPolicy
 from langgraph.pregel.types import PregelTask
+from langgraph.store.memory import MemoryStore
 from tests.any_str import AnyStr, ExceptionLike
 from tests.memory_assert import (
     MemorySaverAssertCheckpointMetadata,
@@ -4823,7 +4823,7 @@ async def test_start_branch_then() -> None:
 
     async with AsyncSqliteSaver.from_conn_string(":memory:") as saver:
         tool_two = tool_two_graph.compile(
-            kv=MemoryKV(),
+            store=MemoryStore(),
             checkpointer=saver,
             interrupt_before=["tool_two_fast", "tool_two_slow"],
         )

From aa1a6be160c624aa2f3e2eb233cb10852aa05169 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Wed, 21 Aug 2024 09:36:21 -0700
Subject: [PATCH 15/30] Rename more

---
 libs/langgraph/langgraph/constants.py         |  4 ++--
 .../langgraph/managed/shared_value.py         | 22 +++++++++----------
 libs/langgraph/langgraph/pregel/loop.py       |  6 ++---
 libs/langgraph/langgraph/store/batch.py       | 12 +++++-----
 .../tests/{test_kv.py => test_store.py}       | 18 +++------------
 5 files changed, 25 insertions(+), 37 deletions(-)
 rename libs/langgraph/tests/{test_kv.py => test_store.py} (68%)

diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py
index cf87be337..51b6e1437 100644
--- a/libs/langgraph/langgraph/constants.py
+++ b/libs/langgraph/langgraph/constants.py
@@ -5,7 +5,7 @@ INPUT = "__input__"
 CONFIG_KEY_SEND = "__pregel_send"
 CONFIG_KEY_READ = "__pregel_read"
 CONFIG_KEY_CHECKPOINTER = "__pregel_checkpointer"
-CONFIG_KEY_KV = "__pregel_kv"
+CONFIG_KEY_STORE = "__pregel_store"
 CONFIG_KEY_RESUMING = "__pregel_resuming"
 CONFIG_KEY_TASK_ID = "__pregel_task_id"
 INTERRUPT = "__interrupt__"
@@ -18,7 +18,7 @@ RESERVED = {
     CONFIG_KEY_SEND,
     CONFIG_KEY_READ,
     CONFIG_KEY_CHECKPOINTER,
-    CONFIG_KEY_KV,
+    CONFIG_KEY_STORE,
     CONFIG_KEY_RESUMING,
     CONFIG_KEY_TASK_ID,
     INPUT,
diff --git a/libs/langgraph/langgraph/managed/shared_value.py b/libs/langgraph/langgraph/managed/shared_value.py
index 74a7c1ef4..7f647d006 100644
--- a/libs/langgraph/langgraph/managed/shared_value.py
+++ b/libs/langgraph/langgraph/managed/shared_value.py
@@ -12,7 +12,7 @@ from typing import (
 from langchain_core.runnables import RunnableConfig
 from typing_extensions import NotRequired, Required, Self
 
-from langgraph.constants import CONFIG_KEY_KV
+from langgraph.constants import CONFIG_KEY_STORE
 from langgraph.errors import InvalidUpdateError
 from langgraph.managed.base import (
     ChannelKeyPlaceholder,
@@ -56,8 +56,8 @@ class SharedValue(WritableManagedValue[Value, Update]):
     @contextmanager
     def enter(cls, config: RunnableConfig, **kwargs: Any) -> Iterator[Self]:
         with super().enter(config, **kwargs) as value:
-            if value.kv is not None:
-                saved = value.kv.list([value.ns])
+            if value.store is not None:
+                saved = value.store.list([value.ns])
                 value.value = saved[value.ns]
             yield value
 
@@ -65,8 +65,8 @@ class SharedValue(WritableManagedValue[Value, Update]):
     @asynccontextmanager
     async def aenter(cls, config: RunnableConfig, **kwargs: Any) -> AsyncIterator[Self]:
         async with super().aenter(config, **kwargs) as value:
-            if value.kv is not None:
-                saved = await value.kv.alist([value.ns])
+            if value.store is not None:
+                saved = await value.store.alist([value.ns])
                 value.value = saved[value.ns]
             yield value
 
@@ -83,8 +83,8 @@ class SharedValue(WritableManagedValue[Value, Update]):
         self.scope = scope
         self.config = config
         self.value: Value = {}
-        self.kv: BaseStore = config["configurable"].get(CONFIG_KEY_KV)
-        if self.kv is None:
+        self.store: BaseStore = config["configurable"].get(CONFIG_KEY_STORE)
+        if self.store is None:
             self.ns: Optional[str] = None
         elif scope_value := config["configurable"].get(self.scope):
             self.ns = f"scoped:{scope}:{key}:{scope_value}"
@@ -114,13 +114,13 @@ class SharedValue(WritableManagedValue[Value, Update]):
         return writes
 
     def update(self, values: Sequence[Update]) -> None:
-        if self.kv is None:
+        if self.store is None:
             self._process_update(values)
         else:
-            return self.kv.put(self._process_update(values))
+            return self.store.put(self._process_update(values))
 
     async def aupdate(self, writes: Sequence[Update]) -> None:
-        if self.kv is None:
+        if self.store is None:
             self._process_update(writes)
         else:
-            return await self.kv.aput(self._process_update(writes))
+            return await self.store.aput(self._process_update(writes))
diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py
index 8079c0dc2..a1eadb697 100644
--- a/libs/langgraph/langgraph/pregel/loop.py
+++ b/libs/langgraph/langgraph/pregel/loop.py
@@ -41,9 +41,9 @@ from langgraph.checkpoint.base import (
     empty_checkpoint,
 )
 from langgraph.constants import (
-    CONFIG_KEY_KV,
     CONFIG_KEY_READ,
     CONFIG_KEY_RESUMING,
+    CONFIG_KEY_STORE,
     ERROR,
     INPUT,
     INTERRUPT,
@@ -479,7 +479,7 @@ class SyncPregelLoop(PregelLoop, ContextManager):
         self.managed = self.stack.enter_context(
             ManagedValuesManager(
                 self.graph.managed_values_dict,
-                patch_config(self.config, configurable={CONFIG_KEY_KV: self.store}),
+                patch_config(self.config, configurable={CONFIG_KEY_STORE: self.store}),
             )
         )
         self.stack.push(self._suppress_interrupt)
@@ -567,7 +567,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
         self.managed = await self.stack.enter_async_context(
             AsyncManagedValuesManager(
                 self.graph.managed_values_dict,
-                patch_config(self.config, configurable={CONFIG_KEY_KV: self.store}),
+                patch_config(self.config, configurable={CONFIG_KEY_STORE: self.store}),
             )
         )
         self.stack.push(self._suppress_interrupt)
diff --git a/libs/langgraph/langgraph/store/batch.py b/libs/langgraph/langgraph/store/batch.py
index cfd05e548..54eb20d47 100644
--- a/libs/langgraph/langgraph/store/batch.py
+++ b/libs/langgraph/langgraph/store/batch.py
@@ -13,10 +13,10 @@ class PutOp(NamedTuple):
 
 
 class AsyncBatchedStore(BaseStore):
-    def __init__(self, kv: BaseStore) -> None:
-        self.kv = kv
+    def __init__(self, store: BaseStore) -> None:
+        self.store = store
         self.aqueue: dict[asyncio.Future, Union[ListOp, PutOp]] = {}
-        self.task = asyncio.create_task(_run(self.aqueue, self.kv))
+        self.task = asyncio.create_task(_run(self.aqueue, self.store))
 
     def __del__(self) -> None:
         self.task.cancel()
@@ -33,7 +33,7 @@ class AsyncBatchedStore(BaseStore):
 
 
 async def _run(
-    aqueue: dict[asyncio.Future, Union[ListOp, PutOp]], kv: BaseStore
+    aqueue: dict[asyncio.Future, Union[ListOp, PutOp]], store: BaseStore
 ) -> None:
     while True:
         await asyncio.sleep(0)
@@ -46,7 +46,7 @@ async def _run(
         lists = {f: o for f, o in taken.items() if isinstance(o, ListOp)}
         if lists:
             try:
-                results = await kv.alist(
+                results = await store.alist(
                     [p for op in lists.values() for p in op.prefixes]
                 )
                 for fut, op in lists.items():
@@ -57,7 +57,7 @@ async def _run(
         puts = {f: o for f, o in taken.items() if isinstance(o, PutOp)}
         if puts:
             try:
-                await kv.aput([w for op in puts.values() for w in op.writes])
+                await store.aput([w for op in puts.values() for w in op.writes])
                 for fut in puts:
                     fut.set_result(None)
             except Exception as e:
diff --git a/libs/langgraph/tests/test_kv.py b/libs/langgraph/tests/test_store.py
similarity index 68%
rename from libs/langgraph/tests/test_kv.py
rename to libs/langgraph/tests/test_store.py
index 0c2862b48..a9fd0d1d4 100644
--- a/libs/langgraph/tests/test_kv.py
+++ b/libs/langgraph/tests/test_store.py
@@ -7,11 +7,11 @@ from langgraph.store.base import BaseStore
 from langgraph.store.batch import AsyncBatchedStore
 
 
-async def test_kv_async_batch(mocker: MockerFixture) -> None:
+async def test_async_batch_store(mocker: MockerFixture) -> None:
     aget = mocker.stub()
     alist = mocker.stub()
 
-    class MockKV(BaseStore):
+    class MockStore(BaseStore):
         async def aget(
             self, pairs: list[tuple[str, str]]
         ) -> dict[tuple[str, str], dict[str, Any] | None]:
@@ -22,21 +22,9 @@ async def test_kv_async_batch(mocker: MockerFixture) -> None:
             alist(prefixes)
             return {prefix: {prefix: 1} for prefix in prefixes}
 
-    store = AsyncBatchedStore(MockKV())
+    store = AsyncBatchedStore(MockStore())
 
     # concurrent calls are batched
-    results = await asyncio.gather(
-        store.aget([("a", "b")]),
-        store.aget([("c", "d")]),
-    )
-    assert results == [
-        {("a", "b"): 1},
-        {("c", "d"): 1},
-    ]
-    assert [c.args for c in aget.call_args_list] == [
-        ([("a", "b"), ("c", "d")],),
-    ]
-
     results = await asyncio.gather(
         store.alist(["a", "b"]),
         store.alist(["c", "d"]),

From f037a2e9cb6ae5dc75f22c2f0e1138521c8dca84 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Wed, 21 Aug 2024 09:39:58 -0700
Subject: [PATCH 16/30] Update docstring

---
 libs/langgraph/langgraph/pregel/__init__.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py
index 5e058fd16..436ed5ed0 100644
--- a/libs/langgraph/langgraph/pregel/__init__.py
+++ b/libs/langgraph/langgraph/pregel/__init__.py
@@ -225,7 +225,7 @@ class Pregel(
     """Checkpointer used to save and load graph state. Defaults to None."""
 
     store: Optional[BaseStore] = None
-    """Key-value store to use. Defaults to None."""
+    """Memory store to use for SharedValues. Defaults to None."""
 
     retry_policy: Optional[RetryPolicy] = None
     """Retry policy to use when running tasks. Set to None to disable."""

From 1af0367b3459beb19008dfd0463a81d6f5c37082 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Wed, 21 Aug 2024 09:41:53 -0700
Subject: [PATCH 17/30] Add error message

---
 libs/langgraph/langgraph/pregel/__init__.py | 4 ++--
 libs/langgraph/langgraph/pregel/loop.py     | 5 ++---
 2 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py
index 436ed5ed0..429226214 100644
--- a/libs/langgraph/langgraph/pregel/__init__.py
+++ b/libs/langgraph/langgraph/pregel/__init__.py
@@ -650,7 +650,7 @@ class Pregel(
             # apply to checkpoint and save
             assert not apply_writes(
                 checkpoint, channels, [task], self.checkpointer.get_next_version
-            )
+            ), "Can't write to SharedValues from update_state"
             checkpoint = create_checkpoint(checkpoint, channels, step + 1)
             # check interrupt before
             if tasks := should_interrupt(
@@ -794,7 +794,7 @@ class Pregel(
             # apply to checkpoint and save
             assert not apply_writes(
                 checkpoint, channels, [task], self.checkpointer.get_next_version
-            )
+            ), "Can't write to SharedValues from update_state"
             checkpoint = create_checkpoint(checkpoint, channels, step + 1)
             # check interrupt before
             if tasks := should_interrupt(
diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py
index a1eadb697..672d42049 100644
--- a/libs/langgraph/langgraph/pregel/loop.py
+++ b/libs/langgraph/langgraph/pregel/loop.py
@@ -332,13 +332,12 @@ class PregelLoop:
                 manager=None,
             )
             # apply input writes
-            mv_writes = apply_writes(
+            assert not apply_writes(
                 self.checkpoint,
                 self.channels,
                 discard_tasks + [PregelTaskWrites(INPUT, input_writes, [])],
                 self.checkpointer_get_next_version,
-            )
-            assert not mv_writes
+            ), "Can't write to SharedValues in graph input"
             # save input checkpoint
             self._put_checkpoint({"source": "input", "writes": self.input})
         else:

From 2106b5e4a6bed75c284146458a07217f5e5bbf06 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Wed, 21 Aug 2024 10:48:04 -0700
Subject: [PATCH 18/30] Lint

---
 libs/langgraph/tests/test_pregel.py | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py
index 8cf6aafa8..57df91165 100644
--- a/libs/langgraph/tests/test_pregel.py
+++ b/libs/langgraph/tests/test_pregel.py
@@ -167,9 +167,6 @@ def test_graph_validation() -> None:
 
     class State(TypedDict):
         hello: str
-        shared_things: Annotated[
-            dict[str, dict[str, Any]], SharedValue.on("assistant_id")
-        ]
 
     def node_a(state: State) -> State:
         # typo

From 7fd4a9ed300fac144e0c07b66e898d02d84bb069 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Wed, 21 Aug 2024 10:48:52 -0700
Subject: [PATCH 19/30] Lint

---
 libs/langgraph/tests/test_pregel_async.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py
index d7a532489..70ce65daa 100644
--- a/libs/langgraph/tests/test_pregel_async.py
+++ b/libs/langgraph/tests/test_pregel_async.py
@@ -565,7 +565,7 @@ async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
 
     assert app.input_schema.schema() == {"title": "LangGraphInput", "type": "integer"}
     assert app.output_schema.schema() == {"title": "LangGraphOutput", "type": "integer"}
-    assert await app.ainvoke(2, debug=True) == 3
+    assert await app.ainvoke(2) == 3
     assert await app.ainvoke(2, output_keys=["output"]) == {"output": 3}
 
     assert await gapp.ainvoke(2) == 3

From 4250ff92b8dbbaea3fb8fff31be10a25a9adaa00 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Wed, 21 Aug 2024 11:25:37 -0700
Subject: [PATCH 20/30] Fix

---
 libs/langgraph/langgraph/managed/shared_value.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/libs/langgraph/langgraph/managed/shared_value.py b/libs/langgraph/langgraph/managed/shared_value.py
index 7f647d006..7bb6e23b7 100644
--- a/libs/langgraph/langgraph/managed/shared_value.py
+++ b/libs/langgraph/langgraph/managed/shared_value.py
@@ -58,7 +58,7 @@ class SharedValue(WritableManagedValue[Value, Update]):
         with super().enter(config, **kwargs) as value:
             if value.store is not None:
                 saved = value.store.list([value.ns])
-                value.value = saved[value.ns]
+                value.value = saved[value.ns] or {}
             yield value
 
     @classmethod
@@ -67,7 +67,7 @@ class SharedValue(WritableManagedValue[Value, Update]):
         async with super().aenter(config, **kwargs) as value:
             if value.store is not None:
                 saved = await value.store.alist([value.ns])
-                value.value = saved[value.ns]
+                value.value = saved[value.ns] or {}
             yield value
 
     def __init__(

From 5fb2c2c6f804fc55961610483ed8ed95338d767d Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Wed, 21 Aug 2024 11:27:04 -0700
Subject: [PATCH 21/30] Lint

---
 libs/langgraph/tests/test_store.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/libs/langgraph/tests/test_store.py b/libs/langgraph/tests/test_store.py
index a9fd0d1d4..cd53ee407 100644
--- a/libs/langgraph/tests/test_store.py
+++ b/libs/langgraph/tests/test_store.py
@@ -1,5 +1,5 @@
 import asyncio
-from typing import Any
+from typing import Any, Optional
 
 from pytest_mock import MockerFixture
 
@@ -14,7 +14,7 @@ async def test_async_batch_store(mocker: MockerFixture) -> None:
     class MockStore(BaseStore):
         async def aget(
             self, pairs: list[tuple[str, str]]
-        ) -> dict[tuple[str, str], dict[str, Any] | None]:
+        ) -> dict[tuple[str, str], Optional[dict[str, Any]]]:
             aget(pairs)
             return {pair: 1 for pair in pairs}
 

From 2f41b2891b59f5729b5d604a828e8f143b3d7f36 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Wed, 21 Aug 2024 11:33:54 -0700
Subject: [PATCH 22/30] lib0.2.7

---
 libs/langgraph/pyproject.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/libs/langgraph/pyproject.toml b/libs/langgraph/pyproject.toml
index 4dff54572..9bf6d068b 100644
--- a/libs/langgraph/pyproject.toml
+++ b/libs/langgraph/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph"
-version = "0.2.6"
+version = "0.2.7"
 description = "Building stateful, multi-actor applications with LLMs"
 authors = []
 license = "MIT"

From 14976d4c56c6516004769831219099483612121c Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Wed, 21 Aug 2024 12:08:58 -0700
Subject: [PATCH 23/30] Interrupts shouldn't be retried (#1413)

* Interrupts shouldn't be retried

* Add async test

* Lint
---
 libs/langgraph/langgraph/pregel/retry.py  |  7 +++
 libs/langgraph/tests/test_pregel.py       |  7 ++-
 libs/langgraph/tests/test_pregel_async.py | 76 +++++++++++++++++++++++
 3 files changed, 89 insertions(+), 1 deletion(-)

diff --git a/libs/langgraph/langgraph/pregel/retry.py b/libs/langgraph/langgraph/pregel/retry.py
index 88940b584..486584809 100644
--- a/libs/langgraph/langgraph/pregel/retry.py
+++ b/libs/langgraph/langgraph/pregel/retry.py
@@ -4,6 +4,7 @@ import random
 import time
 from typing import Optional
 
+from langgraph.errors import GraphInterrupt
 from langgraph.pregel.types import PregelExecutableTask, RetryPolicy
 
 logger = logging.getLogger(__name__)
@@ -25,6 +26,9 @@ def run_with_retry(
             task.proc.invoke(task.input, task.config)
             # if successful, end
             break
+        except GraphInterrupt:
+            # if interrupted, end
+            raise
         except Exception as exc:
             if retry_policy is None:
                 raise
@@ -75,6 +79,9 @@ async def arun_with_retry(
                 await task.proc.ainvoke(task.input, task.config)
             # if successful, end
             break
+        except GraphInterrupt:
+            # if interrupted, end
+            raise
         except Exception as exc:
             if retry_policy is None:
                 raise
diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py
index 57df91165..070c59476 100644
--- a/libs/langgraph/tests/test_pregel.py
+++ b/libs/langgraph/tests/test_pregel.py
@@ -6139,13 +6139,17 @@ def test_dynamic_interrupt(snapshot: SnapshotAssertion) -> None:
         my_key: Annotated[str, operator.add]
         market: str
 
+    tool_two_node_count = 0
+
     def tool_two_node(s: State) -> State:
+        nonlocal tool_two_node_count
+        tool_two_node_count += 1
         if s["market"] == "DE":
             raise NodeInterrupt("Just because...")
         return {"my_key": " all good"}
 
     tool_two_graph = StateGraph(State)
-    tool_two_graph.add_node("tool_two", tool_two_node)
+    tool_two_graph.add_node("tool_two", tool_two_node, retry=RetryPolicy())
     tool_two_graph.add_edge(START, "tool_two")
     tool_two = tool_two_graph.compile()
 
@@ -6153,6 +6157,7 @@ def test_dynamic_interrupt(snapshot: SnapshotAssertion) -> None:
         "my_key": "value",
         "market": "DE",
     }
+    assert tool_two_node_count == 1, "interrupts aren't retried"
     assert tool_two.invoke({"my_key": "value", "market": "US"}) == {
         "my_key": "value all good",
         "market": "US",
diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py
index 70ce65daa..ebbfc7b29 100644
--- a/libs/langgraph/tests/test_pregel_async.py
+++ b/libs/langgraph/tests/test_pregel_async.py
@@ -206,6 +206,82 @@ async def test_node_cancellation_on_other_node_exception() -> None:
     assert inner_task_cancelled
 
 
+async def test_dynamic_interrupt(snapshot: SnapshotAssertion) -> 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":
+            raise NodeInterrupt("Just because...")
+        return {"my_key": " all good"}
+
+    tool_two_graph = StateGraph(State)
+    tool_two_graph.add_node("tool_two", tool_two_node, retry=RetryPolicy())
+    tool_two_graph.add_edge(START, "tool_two")
+    tool_two = tool_two_graph.compile()
+
+    assert await tool_two.ainvoke({"my_key": "value", "market": "DE"}) == {
+        "my_key": "value",
+        "market": "DE",
+    }
+    assert tool_two_node_count == 1, "interrupts aren't retried"
+    assert await tool_two.ainvoke({"my_key": "value", "market": "US"}) == {
+        "my_key": "value all good",
+        "market": "US",
+    }
+
+    async with AsyncSqliteSaver.from_conn_string(":memory:") as saver:
+        tool_two = tool_two_graph.compile(checkpointer=saver)
+
+        # missing thread_id
+        with pytest.raises(ValueError, match="thread_id"):
+            await tool_two.ainvoke({"my_key": "value", "market": "DE"})
+
+        thread1 = {"configurable": {"thread_id": "1"}}
+        # stop when about to enter node
+        assert await tool_two.ainvoke(
+            {"my_key": "value ⛰️", "market": "DE"}, thread1
+        ) == {
+            "my_key": "value ⛰️",
+            "market": "DE",
+        }
+        assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [
+            {
+                "source": "loop",
+                "step": 0,
+                "writes": None,
+            },
+            {
+                "source": "input",
+                "step": -1,
+                "writes": {"my_key": "value ⛰️", "market": "DE"},
+            },
+        ]
+        tup = await tool_two.checkpointer.aget_tuple(thread1)
+        assert await tool_two.aget_state(thread1) == StateSnapshot(
+            values={"my_key": "value ⛰️", "market": "DE"},
+            next=("tool_two",),
+            tasks=(
+                PregelTask(
+                    AnyStr(),
+                    "tool_two",
+                    interrupts=(Interrupt("during", "Just because..."),),
+                ),
+            ),
+            config=tup.config,
+            created_at=tup.checkpoint["ts"],
+            metadata={"source": "loop", "step": 0, "writes": None},
+            parent_config=[
+                c async for c in tool_two.checkpointer.alist(thread1, limit=2)
+            ][-1].config,
+        )
+
+
 @pytest.mark.parametrize(
     "checkpointer_name",
     ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"],

From 14ec51601c174a466f5d044e480961869f28c2a0 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Wed, 21 Aug 2024 12:09:11 -0700
Subject: [PATCH 24/30] lib0.2.8

---
 libs/langgraph/pyproject.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/libs/langgraph/pyproject.toml b/libs/langgraph/pyproject.toml
index 9bf6d068b..5ba045061 100644
--- a/libs/langgraph/pyproject.toml
+++ b/libs/langgraph/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph"
-version = "0.2.7"
+version = "0.2.8"
 description = "Building stateful, multi-actor applications with LLMs"
 authors = []
 license = "MIT"

From e76f4cc434f0d3a631222db2c4bdc00a77a9af03 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Wed, 21 Aug 2024 13:30:18 -0700
Subject: [PATCH 25/30] Combine channel and managed values manager

---
 libs/langgraph/langgraph/channels/manager.py |  39 -------
 libs/langgraph/langgraph/graph/state.py      |  25 ++---
 libs/langgraph/langgraph/managed/base.py     |  50 +--------
 libs/langgraph/langgraph/pregel/__init__.py  | 108 +++++++------------
 libs/langgraph/langgraph/pregel/algo.py      |  23 ++--
 libs/langgraph/langgraph/pregel/loop.py      |  79 ++++++++------
 libs/langgraph/langgraph/pregel/manager.py   | 104 ++++++++++++++++++
 libs/langgraph/langgraph/pregel/read.py      |   3 +-
 libs/langgraph/pyproject.toml                |   2 +-
 libs/langgraph/tests/test_algo.py            |   7 +-
 10 files changed, 213 insertions(+), 227 deletions(-)
 delete mode 100644 libs/langgraph/langgraph/channels/manager.py
 create mode 100644 libs/langgraph/langgraph/pregel/manager.py

diff --git a/libs/langgraph/langgraph/channels/manager.py b/libs/langgraph/langgraph/channels/manager.py
deleted file mode 100644
index 9a7c8d87e..000000000
--- a/libs/langgraph/langgraph/channels/manager.py
+++ /dev/null
@@ -1,39 +0,0 @@
-from contextlib import AsyncExitStack, ExitStack, asynccontextmanager, contextmanager
-from typing import AsyncGenerator, Generator, Mapping
-
-from langchain_core.runnables import RunnableConfig
-
-from langgraph.channels.base import BaseChannel
-from langgraph.checkpoint.base import Checkpoint
-
-
-@contextmanager
-def ChannelsManager(
-    channels: Mapping[str, BaseChannel],
-    checkpoint: Checkpoint,
-    config: RunnableConfig,
-) -> Generator[Mapping[str, BaseChannel], None, None]:
-    """Manage channels for the lifetime of a Pregel invocation (multiple steps)."""
-    with ExitStack() as stack:
-        yield {
-            k: stack.enter_context(
-                v.from_checkpoint(checkpoint["channel_values"].get(k), config)
-            )
-            for k, v in channels.items()
-        }
-
-
-@asynccontextmanager
-async def AsyncChannelsManager(
-    channels: Mapping[str, BaseChannel],
-    checkpoint: Checkpoint,
-    config: RunnableConfig,
-) -> AsyncGenerator[Mapping[str, BaseChannel], None]:
-    """Manage channels for the lifetime of a Pregel invocation (multiple steps)."""
-    async with AsyncExitStack() as stack:
-        yield {
-            k: await stack.enter_async_context(
-                v.afrom_checkpoint(checkpoint["channel_values"].get(k), config)
-            )
-            for k, v in channels.items()
-        }
diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py
index 8c85b7760..d846bc8d3 100644
--- a/libs/langgraph/langgraph/graph/state.py
+++ b/libs/langgraph/langgraph/graph/state.py
@@ -44,7 +44,7 @@ from langgraph.managed.base import (
     ChannelKeyPlaceholder,
     ChannelTypePlaceholder,
     ConfiguredManagedValue,
-    ManagedValue,
+    ManagedValueSpec,
     is_managed_value,
     is_writable_managed_value,
 )
@@ -129,8 +129,8 @@ class StateGraph(Graph):
 
     nodes: dict[str, StateNodeSpec]
     channels: dict[str, BaseChannel]
-    managed: dict[str, Type[ManagedValue]]
-    schemas: dict[Type[Any], dict[str, Union[BaseChannel, Type[ManagedValue]]]]
+    managed: dict[str, ManagedValueSpec]
+    schemas: dict[Type[Any], dict[str, Union[BaseChannel, ManagedValueSpec]]]
 
     def __init__(
         self,
@@ -442,7 +442,11 @@ class StateGraph(Graph):
             builder=self,
             config_type=self.config_schema,
             nodes={},
-            channels={**self.channels, START: EphemeralValue(self.input)},
+            channels={
+                **self.channels,
+                **self.managed,
+                START: EphemeralValue(self.input),
+            },
             input_channels=START,
             stream_mode="updates",
             output_channels=output_channels,
@@ -497,7 +501,7 @@ class CompiledStateGraph(CompiledGraph):
                     **{
                         k: (self.channels[k].UpdateType, None)
                         for k in self.builder.schemas[self.builder.input]
-                        if k in self.channels
+                        if isinstance(self.channels[k], BaseChannel)
                         and not isinstance(self.channels[k], Context)
                     },
                 )
@@ -572,10 +576,7 @@ class CompiledStateGraph(CompiledGraph):
             )
         else:
             input_schema = node.input if node else self.builder.schema
-            input_values = {
-                k: v if is_managed_value(v) else k
-                for k, v in self.builder.schemas[input_schema].items()
-            }
+            input_values = {k: k for k in self.builder.schemas[input_schema]}
             is_single_input = len(input_values) == 1 and "__root__" in input_values
 
             self.channels[key] = EphemeralValue(Any, guard=False)
@@ -694,7 +695,7 @@ def _coerce_state(schema: Type[Any], input: dict[str, Any]) -> dict[str, Any]:
 
 def _get_channels(
     schema: Type[dict],
-) -> tuple[dict[str, BaseChannel], dict[str, Type[ManagedValue]]]:
+) -> tuple[dict[str, BaseChannel], dict[str, ManagedValueSpec]]:
     if not hasattr(schema, "__annotations__"):
         return {"__root__": _get_channel("__root__", schema, allow_managed=False)}, {}
 
@@ -711,7 +712,7 @@ def _get_channels(
 
 def _get_channel(
     name: str, annotation: Any, *, allow_managed: bool = True
-) -> Union[BaseChannel, Type[ManagedValue]]:
+) -> Union[BaseChannel, ManagedValueSpec]:
     if manager := _is_field_managed_value(name, annotation):
         if allow_managed:
             return manager
@@ -751,7 +752,7 @@ def _is_field_binop(typ: Type[Any]) -> Optional[BinaryOperatorAggregate]:
     return None
 
 
-def _is_field_managed_value(name: str, typ: Type[Any]) -> Optional[Type[ManagedValue]]:
+def _is_field_managed_value(name: str, typ: Type[Any]) -> Optional[ManagedValueSpec]:
     if hasattr(typ, "__metadata__"):
         meta = typ.__metadata__
         if len(meta) >= 1:
diff --git a/libs/langgraph/langgraph/managed/base.py b/libs/langgraph/langgraph/managed/base.py
index 8dbf5c30c..bebca1be2 100644
--- a/libs/langgraph/langgraph/managed/base.py
+++ b/libs/langgraph/langgraph/managed/base.py
@@ -1,12 +1,9 @@
-import asyncio
 from abc import ABC, abstractmethod
-from contextlib import AsyncExitStack, ExitStack, asynccontextmanager, contextmanager
+from contextlib import asynccontextmanager, contextmanager
 from inspect import isclass
 from typing import (
     Any,
-    AsyncGenerator,
     AsyncIterator,
-    Generator,
     Generic,
     Iterator,
     NamedTuple,
@@ -104,50 +101,5 @@ def is_writable_managed_value(value: Any) -> TypeGuard[Type[WritableManagedValue
     )
 
 
-@contextmanager
-def ManagedValuesManager(
-    values: dict[str, ManagedValueSpec],
-    config: RunnableConfig,
-) -> Generator[ManagedValueMapping, None, None]:
-    if values:
-        with ExitStack() as stack:
-            yield {
-                key: stack.enter_context(
-                    value.cls.enter(config, **value.kwargs)
-                    if isinstance(value, ConfiguredManagedValue)
-                    else value.enter(config)
-                )
-                for key, value in values.items()
-            }
-    else:
-        yield {}
-
-
-@asynccontextmanager
-async def AsyncManagedValuesManager(
-    values: dict[str, ManagedValueSpec],
-    config: RunnableConfig,
-) -> AsyncGenerator[ManagedValueMapping, None]:
-    if values:
-        async with AsyncExitStack() as stack:
-            # create enter tasks with reference to spec
-            tasks = {
-                asyncio.create_task(
-                    stack.enter_async_context(
-                        value.cls.aenter(config, **value.kwargs)
-                        if isinstance(value, ConfiguredManagedValue)
-                        else value.aenter(config)
-                    )
-                ): key
-                for key, value in values.items()
-            }
-            # wait for all enter tasks
-            done, _ = await asyncio.wait(tasks, return_when=asyncio.ALL_COMPLETED)
-            # build mapping from spec to result
-            yield {tasks[task]: task.result() for task in done}
-    else:
-        yield {}
-
-
 ChannelKeyPlaceholder = object()
 ChannelTypePlaceholder = object()
diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py
index 429226214..be04d187d 100644
--- a/libs/langgraph/langgraph/pregel/__init__.py
+++ b/libs/langgraph/langgraph/pregel/__init__.py
@@ -52,11 +52,6 @@ from langgraph.channels.base import (
     BaseChannel,
 )
 from langgraph.channels.context import Context
-from langgraph.channels.last_value import LastValue
-from langgraph.channels.manager import (
-    AsyncChannelsManager,
-    ChannelsManager,
-)
 from langgraph.checkpoint.base import (
     BaseCheckpointSaver,
     copy_checkpoint,
@@ -73,12 +68,7 @@ from langgraph.constants import (
     Interrupt,
 )
 from langgraph.errors import GraphInterrupt, GraphRecursionError, InvalidUpdateError
-from langgraph.managed.base import (
-    AsyncManagedValuesManager,
-    ManagedValuesManager,
-    ManagedValueSpec,
-    is_managed_value,
-)
+from langgraph.managed.base import ManagedValueSpec
 from langgraph.pregel.algo import (
     apply_writes,
     local_read,
@@ -97,6 +87,7 @@ from langgraph.pregel.io import (
     read_channels,
 )
 from langgraph.pregel.loop import AsyncPregelLoop, SyncPregelLoop
+from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager
 from langgraph.pregel.read import PregelNode
 from langgraph.pregel.retry import RetryPolicy, arun_with_retry, run_with_retry
 from langgraph.pregel.types import (
@@ -197,7 +188,9 @@ class Pregel(
 ):
     nodes: Mapping[str, PregelNode]
 
-    channels: Mapping[str, BaseChannel] = Field(default_factory=dict)
+    channels: Mapping[str, Union[BaseChannel, ManagedValueSpec]] = Field(
+        default_factory=dict
+    )
 
     auto_validate: bool = True
 
@@ -350,16 +343,6 @@ class Pregel(
             k for k in self.channels if not isinstance(self.channels[k], Context)
         ]
 
-    @property
-    def managed_values_dict(self) -> dict[str, ManagedValueSpec]:
-        return {
-            k: v
-            for node in self.nodes.values()
-            if isinstance(node.channels, dict)
-            for k, v in node.channels.items()
-            if is_managed_value(v)
-        }
-
     def get_state(self, config: RunnableConfig) -> StateSnapshot:
         """Get the current state of the graph."""
         if not self.checkpointer:
@@ -368,16 +351,10 @@ class Pregel(
         saved = self.checkpointer.get_tuple(config)
         checkpoint = saved.checkpoint if saved else empty_checkpoint()
         config = saved.config if saved else config
-        with ChannelsManager(
-            {
-                k: LastValue(None) if isinstance(c, Context) else c
-                for k, c in self.channels.items()
-            },
-            checkpoint,
-            config,
-        ) as channels, ManagedValuesManager(
-            self.managed_values_dict, ensure_config(config)
-        ) as managed:
+        with ChannelsManager(self.channels, checkpoint, config, skip_context=True) as (
+            channels,
+            managed,
+        ):
             next_tasks = prepare_next_tasks(
                 checkpoint,
                 self.nodes,
@@ -408,15 +385,8 @@ class Pregel(
 
         config = saved.config if saved else config
         async with AsyncChannelsManager(
-            {
-                k: LastValue(None) if isinstance(c, Context) else c
-                for k, c in self.channels.items()
-            },
-            checkpoint,
-            config,
-        ) as channels, AsyncManagedValuesManager(
-            self.managed_values_dict, ensure_config(config)
-        ) as managed:
+            self.channels, checkpoint, config, skip_context=True
+        ) as (channels, managed):
             next_tasks = prepare_next_tasks(
                 checkpoint,
                 self.nodes,
@@ -460,15 +430,8 @@ class Pregel(
             pending_writes,
         ) in self.checkpointer.list(config, before=before, limit=limit, filter=filter):
             with ChannelsManager(
-                {
-                    k: LastValue(None) if isinstance(c, Context) else c
-                    for k, c in self.channels.items()
-                },
-                checkpoint,
-                config,
-            ) as channels, ManagedValuesManager(
-                self.managed_values_dict, ensure_config(config)
-            ) as managed:
+                self.channels, checkpoint, config, skip_context=True
+            ) as (channels, managed):
                 next_tasks = prepare_next_tasks(
                     checkpoint,
                     self.nodes,
@@ -512,15 +475,8 @@ class Pregel(
             pending_writes,
         ) in self.checkpointer.alist(config, before=before, limit=limit, filter=filter):
             async with AsyncChannelsManager(
-                {
-                    k: LastValue(None) if isinstance(c, Context) else c
-                    for k, c in self.channels.items()
-                },
-                checkpoint,
-                config,
-            ) as channels, AsyncManagedValuesManager(
-                self.managed_values_dict, ensure_config(config)
-            ) as managed:
+                self.channels, checkpoint, config, skip_context=True
+            ) as (channels, managed):
                 next_tasks = prepare_next_tasks(
                     checkpoint,
                     self.nodes,
@@ -613,11 +569,10 @@ class Pregel(
         if as_node not in self.nodes:
             raise InvalidUpdateError(f"Node {as_node} does not exist")
         # update channels
-        with ChannelsManager(
-            self.channels, checkpoint, config
-        ) as channels, ManagedValuesManager(
-            self.managed_values_dict, ensure_config(config)
-        ) as managed:
+        with ChannelsManager(self.channels, checkpoint, config) as (
+            channels,
+            managed,
+        ):
             # create task to run all writers of the chosen node
             writers = self.nodes[as_node].get_writers()
             if not writers:
@@ -757,11 +712,10 @@ class Pregel(
         if as_node not in self.nodes:
             raise InvalidUpdateError(f"Node {as_node} does not exist")
         # update channels, acting as the chosen node
-        async with AsyncChannelsManager(
-            self.channels, checkpoint, config
-        ) as channels, AsyncManagedValuesManager(
-            self.managed_values_dict, ensure_config(config)
-        ) as managed:
+        async with AsyncChannelsManager(self.channels, checkpoint, config) as (
+            channels,
+            managed,
+        ):
             # create task to run all writers of the chosen node
             writers = self.nodes[as_node].get_writers()
             if not writers:
@@ -998,7 +952,13 @@ class Pregel(
             )
 
             with SyncPregelLoop(
-                input, config=config, checkpointer=checkpointer, graph=self
+                input,
+                config=config,
+                store=self.store,
+                checkpointer=checkpointer,
+                graph=self,
+                nodes=self.nodes,
+                specs=self.channels,
             ) as loop:
                 # Similarly to Bulk Synchronous Parallel / Pregel model
                 # computation proceeds in steps, while there are channel updates
@@ -1248,7 +1208,13 @@ class Pregel(
                 debug=debug,
             )
             async with AsyncPregelLoop(
-                input, config=config, checkpointer=checkpointer, graph=self
+                input,
+                config=config,
+                store=self.store,
+                checkpointer=checkpointer,
+                graph=self,
+                nodes=self.nodes,
+                specs=self.channels,
             ) as loop:
                 aioloop = asyncio.get_event_loop()
                 # Similarly to Bulk Synchronous Parallel / Pregel model
diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py
index 37f168bf7..28a8dcd5d 100644
--- a/libs/langgraph/langgraph/pregel/algo.py
+++ b/libs/langgraph/langgraph/pregel/algo.py
@@ -25,7 +25,6 @@ from langchain_core.runnables.config import (
 
 from langgraph.channels.base import BaseChannel
 from langgraph.channels.context import Context
-from langgraph.channels.manager import ChannelsManager
 from langgraph.checkpoint.base import (
     BaseCheckpointSaver,
     Checkpoint,
@@ -46,9 +45,10 @@ from langgraph.constants import (
     Send,
 )
 from langgraph.errors import EmptyChannelError, InvalidUpdateError
-from langgraph.managed.base import ManagedValueMapping, is_managed_value
+from langgraph.managed.base import ManagedValueMapping
 from langgraph.pregel.io import read_channel, read_channels
 from langgraph.pregel.log import logger
+from langgraph.pregel.manager import ChannelsManager
 from langgraph.pregel.read import PregelNode
 from langgraph.pregel.types import All, PregelExecutableTask, PregelTask
 
@@ -105,11 +105,10 @@ def local_read(
     if fresh:
         new_checkpoint = create_checkpoint(copy_checkpoint(checkpoint), channels, -1)
         context_channels = {k: v for k, v in channels.items() if isinstance(v, Context)}
-        with ChannelsManager(
-            {k: v for k, v in channels.items() if k not in context_channels},
-            new_checkpoint,
-            config,
-        ) as channels:
+        with ChannelsManager(channels, new_checkpoint, config, skip_context=True) as (
+            channels,
+            _,
+        ):
             all_channels = {**channels, **context_channels}
             apply_writes(new_checkpoint, all_channels, [task], None)
             return read_channels(all_channels, select)
@@ -470,16 +469,10 @@ def _proc_input(
                     chan,
                     catch=chan not in proc.triggers,
                 )
+                if chan in channels
+                else managed[k](step)
                 for k, chan in proc.channels.items()
-                if isinstance(chan, str)
             }
-
-            managed_values = {}
-            for key, chan in proc.channels.items():
-                if is_managed_value(chan):
-                    managed_values[key] = managed[key](step)
-
-            val.update(managed_values)
         except EmptyChannelError:
             return
     elif isinstance(proc.channels, list):
diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py
index 672d42049..bc19b39cf 100644
--- a/libs/langgraph/langgraph/pregel/loop.py
+++ b/libs/langgraph/langgraph/pregel/loop.py
@@ -22,14 +22,10 @@ from typing import (
 )
 
 from langchain_core.callbacks import AsyncParentRunManager, ParentRunManager
-from langchain_core.runnables import RunnableConfig, patch_config
+from langchain_core.runnables import RunnableConfig
 from typing_extensions import Self
 
 from langgraph.channels.base import BaseChannel
-from langgraph.channels.manager import (
-    AsyncChannelsManager,
-    ChannelsManager,
-)
 from langgraph.checkpoint.base import (
     BaseCheckpointSaver,
     Checkpoint,
@@ -43,7 +39,6 @@ from langgraph.checkpoint.base import (
 from langgraph.constants import (
     CONFIG_KEY_READ,
     CONFIG_KEY_RESUMING,
-    CONFIG_KEY_STORE,
     ERROR,
     INPUT,
     INTERRUPT,
@@ -51,9 +46,8 @@ from langgraph.constants import (
 )
 from langgraph.errors import EmptyInputError, GraphInterrupt
 from langgraph.managed.base import (
-    AsyncManagedValuesManager,
     ManagedValueMapping,
-    ManagedValuesManager,
+    ManagedValueSpec,
     WritableManagedValue,
 )
 from langgraph.pregel.algo import (
@@ -70,6 +64,8 @@ from langgraph.pregel.executor import (
     Submit,
 )
 from langgraph.pregel.io import map_input, map_output_updates, map_output_values, single
+from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager
+from langgraph.pregel.read import PregelNode
 from langgraph.pregel.types import PregelExecutableTask
 from langgraph.pregel.utils import get_new_channel_versions
 from langgraph.store.base import BaseStore
@@ -88,7 +84,12 @@ EMPTY_SEQ = ()
 class PregelLoop:
     input: Optional[Any]
     config: RunnableConfig
+    store: Optional[BaseStore]
     checkpointer: Optional[BaseCheckpointSaver]
+    nodes: Mapping[str, PregelNode]
+    specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]]
+    is_nested: bool
+
     checkpointer_get_next_version: Callable[[Optional[V]], V]
     checkpointer_put_writes: Optional[
         Callable[[RunnableConfig, Sequence[tuple[str, Any]], str], Any]
@@ -123,7 +124,6 @@ class PregelLoop:
     ]
     tasks: Sequence[PregelExecutableTask]
     stream: deque[Tuple[str, Any]]
-    is_nested: bool
 
     # public
 
@@ -132,16 +132,20 @@ class PregelLoop:
         input: Optional[Any],
         *,
         config: RunnableConfig,
+        store: Optional[BaseStore],
         checkpointer: Optional[BaseCheckpointSaver],
+        nodes: Mapping[str, PregelNode],
+        specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
         graph: "Pregel",
     ) -> None:
         self.stream = deque()
         self.input = input
         self.config = config
+        self.store = store
         self.checkpointer = checkpointer
         self.graph = graph
-        # TODO if managed values no longer needs graph we can replace with
-        # managed_specs, channel_specs
+        self.nodes = nodes
+        self.specs = specs
         self.is_nested = CONFIG_KEY_READ in self.config.get("configurable", {})
 
     def mark_tasks_scheduled(self, tasks: Sequence[PregelExecutableTask]) -> None:
@@ -235,7 +239,7 @@ class PregelLoop:
         # prepare next tasks
         self.tasks = prepare_next_tasks(
             self.checkpoint,
-            self.graph.nodes,
+            self.nodes,
             self.channels,
             self.managed,
             self.config,
@@ -323,7 +327,7 @@ class PregelLoop:
             # discard any unfinished tasks from previous checkpoint
             discard_tasks = prepare_next_tasks(
                 self.checkpoint,
-                self.graph.nodes,
+                self.nodes,
                 self.channels,
                 self.managed,
                 self.config,
@@ -422,11 +426,21 @@ class SyncPregelLoop(PregelLoop, ContextManager):
         input: Optional[Any],
         *,
         config: RunnableConfig,
+        store: Optional[BaseStore],
         checkpointer: Optional[BaseCheckpointSaver],
+        nodes: Mapping[str, PregelNode],
+        specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
         graph: "Pregel",
     ) -> None:
-        super().__init__(input, config=config, checkpointer=checkpointer, graph=graph)
-        self.store = graph.store
+        super().__init__(
+            input,
+            config=config,
+            checkpointer=checkpointer,
+            graph=graph,
+            store=store,
+            nodes=nodes,
+            specs=specs,
+        )
         self.stack = ExitStack()
         if checkpointer:
             self.checkpointer_get_next_version = checkpointer.get_next_version
@@ -472,14 +486,8 @@ class SyncPregelLoop(PregelLoop, ContextManager):
         self.checkpoint_pending_writes = saved.pending_writes or []
 
         self.submit = self.stack.enter_context(BackgroundExecutor(self.config))
-        self.channels = self.stack.enter_context(
-            ChannelsManager(self.graph.channels, self.checkpoint, self.config)
-        )
-        self.managed = self.stack.enter_context(
-            ManagedValuesManager(
-                self.graph.managed_values_dict,
-                patch_config(self.config, configurable={CONFIG_KEY_STORE: self.store}),
-            )
+        self.channels, self.managed = self.stack.enter_context(
+            ChannelsManager(self.specs, self.checkpoint, self.config, self.store)
         )
         self.stack.push(self._suppress_interrupt)
         self.status = "pending"
@@ -506,11 +514,22 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
         input: Optional[Any],
         *,
         config: RunnableConfig,
+        store: Optional[BaseStore],
         checkpointer: Optional[BaseCheckpointSaver],
+        nodes: Mapping[str, PregelNode],
+        specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
         graph: "Pregel",
     ) -> None:
-        super().__init__(input, config=config, checkpointer=checkpointer, graph=graph)
-        self.store = AsyncBatchedStore(graph.store) if graph.store else None
+        super().__init__(
+            input,
+            config=config,
+            checkpointer=checkpointer,
+            graph=graph,
+            store=store,
+            nodes=nodes,
+            specs=specs,
+        )
+        self.store = AsyncBatchedStore(self.store) if self.store else None
         self.stack = AsyncExitStack()
         if checkpointer:
             self.checkpointer_get_next_version = checkpointer.get_next_version
@@ -560,14 +579,8 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
         self.checkpoint_pending_writes = saved.pending_writes or []
 
         self.submit = await self.stack.enter_async_context(AsyncBackgroundExecutor())
-        self.channels = await self.stack.enter_async_context(
-            AsyncChannelsManager(self.graph.channels, self.checkpoint, self.config)
-        )
-        self.managed = await self.stack.enter_async_context(
-            AsyncManagedValuesManager(
-                self.graph.managed_values_dict,
-                patch_config(self.config, configurable={CONFIG_KEY_STORE: self.store}),
-            )
+        self.channels, self.managed = await self.stack.enter_async_context(
+            AsyncChannelsManager(self.specs, self.checkpoint, self.config, self.store)
         )
         self.stack.push(self._suppress_interrupt)
         self.status = "pending"
diff --git a/libs/langgraph/langgraph/pregel/manager.py b/libs/langgraph/langgraph/pregel/manager.py
new file mode 100644
index 000000000..437019113
--- /dev/null
+++ b/libs/langgraph/langgraph/pregel/manager.py
@@ -0,0 +1,104 @@
+import asyncio
+from contextlib import AsyncExitStack, ExitStack, asynccontextmanager, contextmanager
+from typing import AsyncIterator, Iterator, Mapping, Optional, Union
+
+from langchain_core.runnables import RunnableConfig, patch_config
+
+from langgraph.channels.base import BaseChannel
+from langgraph.channels.context import Context
+from langgraph.channels.last_value import LastValue
+from langgraph.checkpoint.base import Checkpoint
+from langgraph.constants import CONFIG_KEY_STORE
+from langgraph.managed.base import (
+    ConfiguredManagedValue,
+    ManagedValueMapping,
+    ManagedValueSpec,
+)
+from langgraph.store.base import BaseStore
+
+
+@contextmanager
+def ChannelsManager(
+    specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
+    checkpoint: Checkpoint,
+    config: RunnableConfig,
+    store: Optional[BaseStore] = None,
+    *,
+    skip_context: bool = False,
+) -> Iterator[tuple[Mapping[str, BaseChannel], ManagedValueMapping]]:
+    """Manage channels for the lifetime of a Pregel invocation (multiple steps)."""
+    config_for_managed = patch_config(config, configurable={CONFIG_KEY_STORE: store})
+    channel_specs: Mapping[str, BaseChannel] = {}
+    managed_specs: Mapping[str, ManagedValueSpec] = {}
+    for k, v in specs.items():
+        if skip_context and isinstance(v, Context):
+            channel_specs[k] = LastValue(None)
+        elif isinstance(v, BaseChannel):
+            channel_specs[k] = v
+        else:
+            managed_specs[k] = v
+    with ExitStack() as stack:
+        yield (
+            {
+                k: stack.enter_context(
+                    v.from_checkpoint(checkpoint["channel_values"].get(k), config)
+                )
+                for k, v in channel_specs.items()
+            },
+            {
+                key: stack.enter_context(
+                    value.cls.enter(config_for_managed, **value.kwargs)
+                    if isinstance(value, ConfiguredManagedValue)
+                    else value.enter(config_for_managed)
+                )
+                for key, value in managed_specs.items()
+            },
+        )
+
+
+@asynccontextmanager
+async def AsyncChannelsManager(
+    specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
+    checkpoint: Checkpoint,
+    config: RunnableConfig,
+    store: Optional[BaseStore] = None,
+    *,
+    skip_context: bool = False,
+) -> AsyncIterator[Mapping[str, BaseChannel]]:
+    """Manage channels for the lifetime of a Pregel invocation (multiple steps)."""
+    config_for_managed = patch_config(config, configurable={CONFIG_KEY_STORE: store})
+    channel_specs: Mapping[str, BaseChannel] = {}
+    managed_specs: Mapping[str, ManagedValueSpec] = {}
+    for k, v in specs.items():
+        if skip_context and isinstance(v, Context):
+            channel_specs[k] = LastValue(None)
+        elif isinstance(v, BaseChannel):
+            channel_specs[k] = v
+        else:
+            managed_specs[k] = v
+    async with AsyncExitStack() as stack:
+        # managed: create enter tasks with reference to spec, await them
+        if tasks := {
+            asyncio.create_task(
+                stack.enter_async_context(
+                    value.cls.aenter(config_for_managed, **value.kwargs)
+                    if isinstance(value, ConfiguredManagedValue)
+                    else value.aenter(config_for_managed)
+                )
+            ): key
+            for key, value in managed_specs.items()
+        }:
+            done, _ = await asyncio.wait(tasks, return_when=asyncio.ALL_COMPLETED)
+        else:
+            done = set()
+        yield (
+            # channels: enter each channel with checkpoint
+            {
+                k: await stack.enter_async_context(
+                    v.afrom_checkpoint(checkpoint["channel_values"].get(k), config)
+                )
+                for k, v in channel_specs.items()
+            },
+            # managed: build mapping from spec to result
+            {tasks[task]: task.result() for task in done},
+        )
diff --git a/libs/langgraph/langgraph/pregel/read.py b/libs/langgraph/langgraph/pregel/read.py
index 6e492e2ac..b5c971e4a 100644
--- a/libs/langgraph/langgraph/pregel/read.py
+++ b/libs/langgraph/langgraph/pregel/read.py
@@ -15,7 +15,6 @@ from langchain_core.runnables.config import merge_configs
 from langchain_core.runnables.utils import ConfigurableFieldSpec
 
 from langgraph.constants import CONFIG_KEY_READ
-from langgraph.managed.base import ManagedValueSpec
 from langgraph.pregel.retry import RetryPolicy
 from langgraph.pregel.write import ChannelWrite
 from langgraph.utils import RunnableCallable
@@ -101,7 +100,7 @@ DEFAULT_BOUND: RunnablePassthrough = RunnablePassthrough()
 
 
 class PregelNode(RunnableBindingBase):
-    channels: Union[list[str], Mapping[str, Union[str, ManagedValueSpec]]]
+    channels: Union[list[str], Mapping[str, str]]
 
     triggers: list[str] = Field(default_factory=list)
 
diff --git a/libs/langgraph/pyproject.toml b/libs/langgraph/pyproject.toml
index 5ba045061..93dc203f2 100644
--- a/libs/langgraph/pyproject.toml
+++ b/libs/langgraph/pyproject.toml
@@ -65,7 +65,7 @@ omit = ["tests/*"]
 [tool.pytest-watcher]
 now = true
 delay = 0.1
-runner_args = ["--ff", "-v", "-n", "auto", "--dist", "worksteal", "--snapshot-update", "--tb", "short"]
+runner_args = ["--ff", "-v", "-x", "-n", "auto", "--dist", "worksteal", "--snapshot-update", "--tb", "short"]
 patterns = ["*.py"]
 
 [build-system]
diff --git a/libs/langgraph/tests/test_algo.py b/libs/langgraph/tests/test_algo.py
index cdf179c24..2102261ae 100644
--- a/libs/langgraph/tests/test_algo.py
+++ b/libs/langgraph/tests/test_algo.py
@@ -1,7 +1,6 @@
-from langgraph.channels.manager import ChannelsManager
 from langgraph.checkpoint.base import empty_checkpoint
-from langgraph.managed.base import ManagedValuesManager
 from langgraph.pregel.algo import prepare_next_tasks
+from langgraph.pregel.manager import ChannelsManager
 
 
 def test_prepare_next_tasks() -> None:
@@ -9,9 +8,7 @@ def test_prepare_next_tasks() -> None:
     processes = {}
     checkpoint = empty_checkpoint()
 
-    with ManagedValuesManager({}, config) as managed, ChannelsManager(
-        {}, checkpoint, config
-    ) as channels:
+    with ChannelsManager({}, checkpoint, config) as (channels, managed):
         assert (
             prepare_next_tasks(
                 checkpoint, processes, channels, managed, config, 0, for_execution=False

From 02697b5712c59d40ad39a8878d880043b0639a88 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Wed, 21 Aug 2024 13:45:42 -0700
Subject: [PATCH 26/30] Remove graph arg to PregelLoop

---
 libs/langgraph/langgraph/pregel/__init__.py | 11 +++--
 libs/langgraph/langgraph/pregel/loop.py     | 34 ++++----------
 libs/langgraph/tests/test_pregel.py         | 52 ++++++++++++---------
 libs/langgraph/tests/test_pregel_async.py   | 52 ++++++++++++---------
 4 files changed, 75 insertions(+), 74 deletions(-)

diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py
index be04d187d..6ac6c737a 100644
--- a/libs/langgraph/langgraph/pregel/__init__.py
+++ b/libs/langgraph/langgraph/pregel/__init__.py
@@ -340,7 +340,10 @@ class Pregel(
     @property
     def stream_channels_asis(self) -> Union[str, Sequence[str]]:
         return self.stream_channels or [
-            k for k in self.channels if not isinstance(self.channels[k], Context)
+            k
+            for k in self.channels
+            if isinstance(self.channels[k], BaseChannel)
+            and not isinstance(self.channels[k], Context)
         ]
 
     def get_state(self, config: RunnableConfig) -> StateSnapshot:
@@ -956,7 +959,6 @@ class Pregel(
                 config=config,
                 store=self.store,
                 checkpointer=checkpointer,
-                graph=self,
                 nodes=self.nodes,
                 specs=self.channels,
             ) as loop:
@@ -966,7 +968,9 @@ class Pregel(
                 # channels are guaranteed to be immutable for the duration of the step,
                 # with channel updates applied only at the transition between steps
                 while loop.tick(
+                    input_keys=self.input_channels,
                     output_keys=output_keys,
+                    stream_keys=self.stream_channels_asis,
                     interrupt_before=interrupt_before,
                     interrupt_after=interrupt_after,
                     manager=run_manager,
@@ -1212,7 +1216,6 @@ class Pregel(
                 config=config,
                 store=self.store,
                 checkpointer=checkpointer,
-                graph=self,
                 nodes=self.nodes,
                 specs=self.channels,
             ) as loop:
@@ -1223,7 +1226,9 @@ class Pregel(
                 # channels are guaranteed to be immutable for the duration of the step,
                 # with channel updates applied only at the transition between steps
                 while loop.tick(
+                    input_keys=self.input_channels,
                     output_keys=output_keys,
+                    stream_keys=self.stream_channels_asis,
                     interrupt_before=interrupt_before,
                     interrupt_after=interrupt_after,
                     manager=run_manager,
diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py
index bc19b39cf..db165513b 100644
--- a/libs/langgraph/langgraph/pregel/loop.py
+++ b/libs/langgraph/langgraph/pregel/loop.py
@@ -4,7 +4,6 @@ from collections import deque
 from contextlib import AsyncExitStack, ExitStack
 from types import TracebackType
 from typing import (
-    TYPE_CHECKING,
     Any,
     AsyncContextManager,
     Callable,
@@ -71,10 +70,6 @@ from langgraph.pregel.utils import get_new_channel_versions
 from langgraph.store.base import BaseStore
 from langgraph.store.batch import AsyncBatchedStore
 
-if TYPE_CHECKING:
-    from langgraph.pregel import Pregel
-
-
 V = TypeVar("V")
 INPUT_DONE = object()
 INPUT_RESUMING = object()
@@ -105,8 +100,6 @@ class PregelLoop:
             Any,
         ]
     ]
-    graph: "Pregel"
-    store: Optional[BaseStore]
     submit: Submit
     channels: Mapping[str, BaseChannel]
     managed: ManagedValueMapping
@@ -136,14 +129,12 @@ class PregelLoop:
         checkpointer: Optional[BaseCheckpointSaver],
         nodes: Mapping[str, PregelNode],
         specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
-        graph: "Pregel",
     ) -> None:
         self.stream = deque()
         self.input = input
         self.config = config
         self.store = store
         self.checkpointer = checkpointer
-        self.graph = graph
         self.nodes = nodes
         self.specs = specs
         self.is_nested = CONFIG_KEY_READ in self.config.get("configurable", {})
@@ -175,7 +166,9 @@ class PregelLoop:
     def tick(
         self,
         *,
+        input_keys: Union[str, Sequence[str]],
         output_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
+        stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
         interrupt_after: Sequence[str] = EMPTY_SEQ,
         interrupt_before: Sequence[str] = EMPTY_SEQ,
         manager: Union[None, AsyncParentRunManager, ParentRunManager] = None,
@@ -187,7 +180,7 @@ class PregelLoop:
             raise RuntimeError("Cannot tick when status is no longer 'pending'")
 
         if self.input not in (INPUT_DONE, INPUT_RESUMING):
-            self._first()
+            self._first(input_keys=input_keys)
         elif all(task.writes for task in self.tasks):
             writes = [w for t in self.tasks for w in t.writes]
             # all tasks have finished
@@ -211,11 +204,7 @@ class PregelLoop:
             self._put_checkpoint(
                 {
                     "source": "loop",
-                    "writes": single(
-                        map_output_updates(output_keys, self.tasks)
-                        if self.graph.stream_mode == "updates"
-                        else map_output_values(output_keys, writes, self.channels)
-                    ),
+                    "writes": single(map_output_updates(output_keys, self.tasks)),
                 }
             )
             # after execution, check if we should interrupt
@@ -258,7 +247,7 @@ class PregelLoop:
                     self.step - 1,  # printing checkpoint for previous step
                     self.checkpoint_config,
                     self.channels,
-                    self.graph.stream_channels_asis,
+                    stream_keys,
                     self.checkpoint_metadata,
                     self.checkpoint,
                     self.tasks,
@@ -282,6 +271,7 @@ class PregelLoop:
         # if all tasks have finished, re-tick
         if all(task.writes for task in self.tasks):
             return self.tick(
+                input_keys=input_keys,
                 output_keys=output_keys,
                 interrupt_after=interrupt_after,
                 interrupt_before=interrupt_before,
@@ -306,7 +296,7 @@ class PregelLoop:
 
     # private
 
-    def _first(self) -> None:
+    def _first(self, *, input_keys: Union[str, Sequence[str]]) -> None:
         # resuming from previous checkpoint requires
         # - finding a previous checkpoint
         # - receiving None input (outer graph) or RESUMING flag (subgraph)
@@ -323,7 +313,7 @@ class PregelLoop:
                     version = self.checkpoint["channel_versions"][k]
                     self.checkpoint["versions_seen"][INTERRUPT][k] = version
         # map inputs to channel updates
-        elif input_writes := deque(map_input(self.graph.input_channels, self.input)):
+        elif input_writes := deque(map_input(input_keys, self.input)):
             # discard any unfinished tasks from previous checkpoint
             discard_tasks = prepare_next_tasks(
                 self.checkpoint,
@@ -345,7 +335,7 @@ class PregelLoop:
             # save input checkpoint
             self._put_checkpoint({"source": "input", "writes": self.input})
         else:
-            raise EmptyInputError(f"Received no input for {self.graph.input_channels}")
+            raise EmptyInputError(f"Received no input for {input_keys}")
         # done with input
         self.input = INPUT_RESUMING if is_resuming else INPUT_DONE
 
@@ -430,13 +420,11 @@ class SyncPregelLoop(PregelLoop, ContextManager):
         checkpointer: Optional[BaseCheckpointSaver],
         nodes: Mapping[str, PregelNode],
         specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
-        graph: "Pregel",
     ) -> None:
         super().__init__(
             input,
             config=config,
             checkpointer=checkpointer,
-            graph=graph,
             store=store,
             nodes=nodes,
             specs=specs,
@@ -504,7 +492,6 @@ class SyncPregelLoop(PregelLoop, ContextManager):
         traceback: Optional[TracebackType],
     ) -> Optional[bool]:
         # unwind stack
-        del self.graph
         return self.stack.__exit__(exc_type, exc_value, traceback)
 
 
@@ -518,13 +505,11 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
         checkpointer: Optional[BaseCheckpointSaver],
         nodes: Mapping[str, PregelNode],
         specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
-        graph: "Pregel",
     ) -> None:
         super().__init__(
             input,
             config=config,
             checkpointer=checkpointer,
-            graph=graph,
             store=store,
             nodes=nodes,
             specs=specs,
@@ -598,7 +583,6 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
         traceback: Optional[TracebackType],
     ) -> Optional[bool]:
         # unwind stack
-        del self.graph
         return await asyncio.shield(
             self.stack.__aexit__(exc_type, exc_value, traceback)
         )
diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py
index 070c59476..197dc6c61 100644
--- a/libs/langgraph/tests/test_pregel.py
+++ b/libs/langgraph/tests/test_pregel.py
@@ -657,7 +657,7 @@ def test_invoke_two_processes_in_out_interrupt(
                     "checkpoint_id": AnyStr(),
                 }
             },
-            metadata={"source": "loop", "step": 6, "writes": 5},
+            metadata={"source": "loop", "step": 6, "writes": {"two": 5}},
             created_at=AnyStr(),
             parent_config=history[1].config,
         ),
@@ -672,7 +672,7 @@ def test_invoke_two_processes_in_out_interrupt(
                     "checkpoint_id": AnyStr(),
                 }
             },
-            metadata={"source": "loop", "step": 5, "writes": None},
+            metadata={"source": "loop", "step": 5, "writes": {"one": None}},
             created_at=AnyStr(),
             parent_config=history[2].config,
         ),
@@ -702,7 +702,7 @@ def test_invoke_two_processes_in_out_interrupt(
                     "checkpoint_id": AnyStr(),
                 }
             },
-            metadata={"source": "loop", "step": 3, "writes": None},
+            metadata={"source": "loop", "step": 3, "writes": {"one": None}},
             created_at=AnyStr(),
             parent_config=history[4].config,
         ),
@@ -732,7 +732,7 @@ def test_invoke_two_processes_in_out_interrupt(
                     "checkpoint_id": AnyStr(),
                 }
             },
-            metadata={"source": "loop", "step": 1, "writes": 4},
+            metadata={"source": "loop", "step": 1, "writes": {"two": 4}},
             created_at=AnyStr(),
             parent_config=history[6].config,
         ),
@@ -747,7 +747,7 @@ def test_invoke_two_processes_in_out_interrupt(
                     "checkpoint_id": AnyStr(),
                 }
             },
-            metadata={"source": "loop", "step": 0, "writes": None},
+            metadata={"source": "loop", "step": 0, "writes": {"one": None}},
             created_at=AnyStr(),
             parent_config=history[7].config,
         ),
@@ -1995,12 +1995,14 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
             "step": 0,
             "writes": {
                 "agent": {
-                    "input": "what is weather in sf",
-                    "agent_outcome": AgentAction(
-                        tool="search_api",
-                        tool_input="query",
-                        log="tool:search_api:query",
-                    ),
+                    "agent": {
+                        "input": "what is weather in sf",
+                        "agent_outcome": AgentAction(
+                            tool="search_api",
+                            tool_input="query",
+                            log="tool:search_api:query",
+                        ),
+                    }
                 },
             },
         },
@@ -2206,12 +2208,14 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
             "step": 0,
             "writes": {
                 "agent": {
-                    "input": "what is weather in sf",
-                    "agent_outcome": AgentAction(
-                        tool="search_api",
-                        tool_input="query",
-                        log="tool:search_api:query",
-                    ),
+                    "agent": {
+                        "input": "what is weather in sf",
+                        "agent_outcome": AgentAction(
+                            tool="search_api",
+                            tool_input="query",
+                            log="tool:search_api:query",
+                        ),
+                    }
                 }
             },
         },
@@ -2411,12 +2415,14 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
             "step": 0,
             "writes": {
                 "agent": {
-                    "input": "what is weather in sf",
-                    "agent_outcome": AgentAction(
-                        tool="search_api",
-                        tool_input="query",
-                        log="tool:search_api:query",
-                    ),
+                    "agent": {
+                        "input": "what is weather in sf",
+                        "agent_outcome": AgentAction(
+                            tool="search_api",
+                            tool_input="query",
+                            log="tool:search_api:query",
+                        ),
+                    }
                 }
             },
         },
diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py
index ebbfc7b29..e9987c601 100644
--- a/libs/langgraph/tests/test_pregel_async.py
+++ b/libs/langgraph/tests/test_pregel_async.py
@@ -876,7 +876,7 @@ async def test_invoke_two_processes_in_out_interrupt(
                     "checkpoint_id": AnyStr(),
                 }
             },
-            metadata={"source": "loop", "step": 6, "writes": 5},
+            metadata={"source": "loop", "step": 6, "writes": {"two": 5}},
             created_at=AnyStr(),
             parent_config=history[1].config,
         ),
@@ -891,7 +891,7 @@ async def test_invoke_two_processes_in_out_interrupt(
                     "checkpoint_id": AnyStr(),
                 }
             },
-            metadata={"source": "loop", "step": 5, "writes": None},
+            metadata={"source": "loop", "step": 5, "writes": {"one": None}},
             created_at=AnyStr(),
             parent_config=history[2].config,
         ),
@@ -921,7 +921,7 @@ async def test_invoke_two_processes_in_out_interrupt(
                     "checkpoint_id": AnyStr(),
                 }
             },
-            metadata={"source": "loop", "step": 3, "writes": None},
+            metadata={"source": "loop", "step": 3, "writes": {"one": None}},
             created_at=AnyStr(),
             parent_config=history[4].config,
         ),
@@ -951,7 +951,7 @@ async def test_invoke_two_processes_in_out_interrupt(
                     "checkpoint_id": AnyStr(),
                 }
             },
-            metadata={"source": "loop", "step": 1, "writes": 4},
+            metadata={"source": "loop", "step": 1, "writes": {"two": 4}},
             created_at=AnyStr(),
             parent_config=history[6].config,
         ),
@@ -966,7 +966,7 @@ async def test_invoke_two_processes_in_out_interrupt(
                     "checkpoint_id": AnyStr(),
                 }
             },
-            metadata={"source": "loop", "step": 0, "writes": None},
+            metadata={"source": "loop", "step": 0, "writes": {"one": None}},
             created_at=AnyStr(),
             parent_config=history[7].config,
         ),
@@ -2290,12 +2290,14 @@ async def test_conditional_graph() -> None:
             "step": 0,
             "writes": {
                 "agent": {
-                    "input": "what is weather in sf",
-                    "agent_outcome": AgentAction(
-                        tool="search_api",
-                        tool_input="query",
-                        log="tool:search_api:query",
-                    ),
+                    "agent": {
+                        "input": "what is weather in sf",
+                        "agent_outcome": AgentAction(
+                            tool="search_api",
+                            tool_input="query",
+                            log="tool:search_api:query",
+                        ),
+                    }
                 }
             },
         },
@@ -2516,12 +2518,14 @@ async def test_conditional_graph() -> None:
             "step": 0,
             "writes": {
                 "agent": {
-                    "input": "what is weather in sf",
-                    "agent_outcome": AgentAction(
-                        tool="search_api",
-                        tool_input="query",
-                        log="tool:search_api:query",
-                    ),
+                    "agent": {
+                        "input": "what is weather in sf",
+                        "agent_outcome": AgentAction(
+                            tool="search_api",
+                            tool_input="query",
+                            log="tool:search_api:query",
+                        ),
+                    }
                 }
             },
         },
@@ -2748,12 +2752,14 @@ async def test_conditional_graph() -> None:
             "step": 0,
             "writes": {
                 "agent": {
-                    "input": "what is weather in sf",
-                    "agent_outcome": AgentAction(
-                        tool="search_api",
-                        tool_input="query",
-                        log="tool:search_api:query",
-                    ),
+                    "agent": {
+                        "input": "what is weather in sf",
+                        "agent_outcome": AgentAction(
+                            tool="search_api",
+                            tool_input="query",
+                            log="tool:search_api:query",
+                        ),
+                    }
                 }
             },
         },

From 3a65f83ae1f08a66a84addb4e0a2f980acfd6bc9 Mon Sep 17 00:00:00 2001
From: Vadym Barda <vadym@langchain.dev>
Date: Wed, 21 Aug 2024 17:05:03 -0400
Subject: [PATCH 27/30] sdk-py: allow passing custom headers (#1416)

---
 libs/sdk-py/langgraph_sdk/client.py | 38 +++++++++++++++++++++++------
 1 file changed, 30 insertions(+), 8 deletions(-)

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

From 1907646bd472f2bf1ac99eb82ef9197e8269dbc8 Mon Sep 17 00:00:00 2001
From: Vadym Barda <vadym@langchain.dev>
Date: Wed, 21 Aug 2024 17:10:18 -0400
Subject: [PATCH 28/30] sdk-py: release 0.1.28 (#1417)

---
 libs/sdk-py/pyproject.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/libs/sdk-py/pyproject.toml b/libs/sdk-py/pyproject.toml
index 3c50594db..2755c3fc6 100644
--- a/libs/sdk-py/pyproject.toml
+++ b/libs/sdk-py/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph-sdk"
-version = "0.1.27"
+version = "0.1.28"
 description = "SDK for interacting with LangGraph API"
 authors = []
 license = "MIT"

From 47ed3d97e9819e7c353dea2817521256ef225ca1 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Wed, 21 Aug 2024 14:25:00 -0700
Subject: [PATCH 29/30] Fix pending run when interrupt exception is used

---
 libs/langgraph/langgraph/pregel/__init__.py | 16 ++--
 libs/langgraph/langgraph/pregel/loop.py     | 30 +++++--
 libs/langgraph/tests/fake_tracer.py         | 91 +++++++++++++++++++++
 libs/langgraph/tests/test_pregel.py         | 12 ++-
 libs/langgraph/tests/test_pregel_async.py   | 12 ++-
 5 files changed, 143 insertions(+), 18 deletions(-)
 create mode 100644 libs/langgraph/tests/fake_tracer.py

diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py
index 6ac6c737a..6e4d3e002 100644
--- a/libs/langgraph/langgraph/pregel/__init__.py
+++ b/libs/langgraph/langgraph/pregel/__init__.py
@@ -961,6 +961,7 @@ class Pregel(
                 checkpointer=checkpointer,
                 nodes=self.nodes,
                 specs=self.channels,
+                output_keys=output_keys,
             ) as loop:
                 # Similarly to Bulk Synchronous Parallel / Pregel model
                 # computation proceeds in steps, while there are channel updates
@@ -969,7 +970,6 @@ class Pregel(
                 # with channel updates applied only at the transition between steps
                 while loop.tick(
                     input_keys=self.input_channels,
-                    output_keys=output_keys,
                     stream_keys=self.stream_channels_asis,
                     interrupt_before=interrupt_before,
                     interrupt_after=interrupt_after,
@@ -1088,8 +1088,8 @@ class Pregel(
                         "without hitting a stop condition. You can increase the "
                         "limit by setting the `recursion_limit` config key."
                     )
-                # set final channel values as run output
-                run_manager.on_chain_end(read_channels(loop.channels, output_keys))
+            # set final channel values as run output
+            run_manager.on_chain_end(loop.output)
         except BaseException as e:
             run_manager.on_chain_error(e)
             raise
@@ -1218,6 +1218,7 @@ class Pregel(
                 checkpointer=checkpointer,
                 nodes=self.nodes,
                 specs=self.channels,
+                output_keys=output_keys,
             ) as loop:
                 aioloop = asyncio.get_event_loop()
                 # Similarly to Bulk Synchronous Parallel / Pregel model
@@ -1227,7 +1228,6 @@ class Pregel(
                 # with channel updates applied only at the transition between steps
                 while loop.tick(
                     input_keys=self.input_channels,
-                    output_keys=output_keys,
                     stream_keys=self.stream_channels_asis,
                     interrupt_before=interrupt_before,
                     interrupt_after=interrupt_after,
@@ -1349,13 +1349,9 @@ class Pregel(
                         "without hitting a stop condition. You can increase the "
                         "limit by setting the `recursion_limit` config key."
                     )
-
-                # set final channel values as run output
-                await run_manager.on_chain_end(
-                    read_channels(loop.channels, output_keys)
-                )
+            # set final channel values as run output
+            await run_manager.on_chain_end(loop.output)
         except BaseException as e:
-            # TODO use on_chain_end if exc is GraphInterrupt
             await asyncio.shield(run_manager.on_chain_error(e))
             raise
 
diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py
index db165513b..aa256b2f7 100644
--- a/libs/langgraph/langgraph/pregel/loop.py
+++ b/libs/langgraph/langgraph/pregel/loop.py
@@ -62,7 +62,13 @@ from langgraph.pregel.executor import (
     BackgroundExecutor,
     Submit,
 )
-from langgraph.pregel.io import map_input, map_output_updates, map_output_values, single
+from langgraph.pregel.io import (
+    map_input,
+    map_output_updates,
+    map_output_values,
+    read_channels,
+    single,
+)
 from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager
 from langgraph.pregel.read import PregelNode
 from langgraph.pregel.types import PregelExecutableTask
@@ -83,6 +89,7 @@ class PregelLoop:
     checkpointer: Optional[BaseCheckpointSaver]
     nodes: Mapping[str, PregelNode]
     specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]]
+    output_keys: Union[str, Sequence[str]]
     is_nested: bool
 
     checkpointer_get_next_version: Callable[[Optional[V]], V]
@@ -117,6 +124,7 @@ class PregelLoop:
     ]
     tasks: Sequence[PregelExecutableTask]
     stream: deque[Tuple[str, Any]]
+    output: Union[None, dict[str, Any], Any] = None
 
     # public
 
@@ -129,6 +137,7 @@ class PregelLoop:
         checkpointer: Optional[BaseCheckpointSaver],
         nodes: Mapping[str, PregelNode],
         specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
+        output_keys: Union[str, Sequence[str]],
     ) -> None:
         self.stream = deque()
         self.input = input
@@ -137,6 +146,7 @@ class PregelLoop:
         self.checkpointer = checkpointer
         self.nodes = nodes
         self.specs = specs
+        self.output_keys = output_keys
         self.is_nested = CONFIG_KEY_READ in self.config.get("configurable", {})
 
     def mark_tasks_scheduled(self, tasks: Sequence[PregelExecutableTask]) -> None:
@@ -167,7 +177,6 @@ class PregelLoop:
         self,
         *,
         input_keys: Union[str, Sequence[str]],
-        output_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
         stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
         interrupt_after: Sequence[str] = EMPTY_SEQ,
         interrupt_before: Sequence[str] = EMPTY_SEQ,
@@ -196,7 +205,7 @@ class PregelLoop:
             # produce values output
             self.stream.extend(
                 ("values", v)
-                for v in map_output_values(output_keys, writes, self.channels)
+                for v in map_output_values(self.output_keys, writes, self.channels)
             )
             # clear pending writes
             self.checkpoint_pending_writes.clear()
@@ -204,7 +213,7 @@ class PregelLoop:
             self._put_checkpoint(
                 {
                     "source": "loop",
-                    "writes": single(map_output_updates(output_keys, self.tasks)),
+                    "writes": single(map_output_updates(self.output_keys, self.tasks)),
                 }
             )
             # after execution, check if we should interrupt
@@ -272,7 +281,7 @@ class PregelLoop:
         if all(task.writes for task in self.tasks):
             return self.tick(
                 input_keys=input_keys,
-                output_keys=output_keys,
+                stream_keys=stream_keys,
                 interrupt_after=interrupt_after,
                 interrupt_before=interrupt_before,
                 manager=manager,
@@ -406,7 +415,12 @@ class PregelLoop:
         exc_value: Optional[BaseException],
         traceback: Optional[TracebackType],
     ) -> Optional[bool]:
-        if isinstance(exc_value, GraphInterrupt) and not self.is_nested:
+        suppress = isinstance(exc_value, GraphInterrupt) and not self.is_nested
+        if suppress or exc_type is None:
+            # save final output
+            self.output = read_channels(self.channels, self.output_keys)
+        if suppress:
+            # suppress interrupt
             return True
 
 
@@ -420,6 +434,7 @@ class SyncPregelLoop(PregelLoop, ContextManager):
         checkpointer: Optional[BaseCheckpointSaver],
         nodes: Mapping[str, PregelNode],
         specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
+        output_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
     ) -> None:
         super().__init__(
             input,
@@ -428,6 +443,7 @@ class SyncPregelLoop(PregelLoop, ContextManager):
             store=store,
             nodes=nodes,
             specs=specs,
+            output_keys=output_keys,
         )
         self.stack = ExitStack()
         if checkpointer:
@@ -505,6 +521,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
         checkpointer: Optional[BaseCheckpointSaver],
         nodes: Mapping[str, PregelNode],
         specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
+        output_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
     ) -> None:
         super().__init__(
             input,
@@ -513,6 +530,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
             store=store,
             nodes=nodes,
             specs=specs,
+            output_keys=output_keys,
         )
         self.store = AsyncBatchedStore(self.store) if self.store else None
         self.stack = AsyncExitStack()
diff --git a/libs/langgraph/tests/fake_tracer.py b/libs/langgraph/tests/fake_tracer.py
new file mode 100644
index 000000000..28ecc88db
--- /dev/null
+++ b/libs/langgraph/tests/fake_tracer.py
@@ -0,0 +1,91 @@
+from typing import Any, Optional
+from uuid import UUID
+
+from langchain_core.messages.base import BaseMessage
+from langchain_core.outputs.chat_generation import ChatGeneration
+from langchain_core.outputs.llm_result import LLMResult
+from langchain_core.tracers import BaseTracer, Run
+
+
+class FakeTracer(BaseTracer):
+    """Fake tracer that records LangChain execution.
+    It replaces run ids with deterministic UUIDs for snapshotting."""
+
+    def __init__(self) -> None:
+        """Initialize the tracer."""
+        super().__init__()
+        self.runs: list[Run] = []
+        self.uuids_map: dict[UUID, UUID] = {}
+        self.uuids_generator = (
+            UUID(f"00000000-0000-4000-8000-{i:012}", version=4) for i in range(10000)
+        )
+
+    def _replace_uuid(self, uuid: UUID) -> UUID:
+        if uuid not in self.uuids_map:
+            self.uuids_map[uuid] = next(self.uuids_generator)
+        return self.uuids_map[uuid]
+
+    def _replace_message_id(self, maybe_message: Any) -> Any:
+        if isinstance(maybe_message, BaseMessage):
+            maybe_message.id = str(next(self.uuids_generator))
+        if isinstance(maybe_message, ChatGeneration):
+            maybe_message.message.id = str(next(self.uuids_generator))
+        if isinstance(maybe_message, LLMResult):
+            for i, gen_list in enumerate(maybe_message.generations):
+                for j, gen in enumerate(gen_list):
+                    maybe_message.generations[i][j] = self._replace_message_id(gen)
+        if isinstance(maybe_message, dict):
+            for k, v in maybe_message.items():
+                maybe_message[k] = self._replace_message_id(v)
+        if isinstance(maybe_message, list):
+            for i, v in enumerate(maybe_message):
+                maybe_message[i] = self._replace_message_id(v)
+
+        return maybe_message
+
+    def _copy_run(self, run: Run) -> Run:
+        if run.dotted_order:
+            levels = run.dotted_order.split(".")
+            processed_levels = []
+            for level in levels:
+                timestamp, run_id = level.split("Z")
+                new_run_id = self._replace_uuid(UUID(run_id))
+                processed_level = f"{timestamp}Z{new_run_id}"
+                processed_levels.append(processed_level)
+            new_dotted_order = ".".join(processed_levels)
+        else:
+            new_dotted_order = None
+        return run.copy(
+            update={
+                "id": self._replace_uuid(run.id),
+                "parent_run_id": (
+                    self.uuids_map[run.parent_run_id] if run.parent_run_id else None
+                ),
+                "child_runs": [self._copy_run(child) for child in run.child_runs],
+                "trace_id": self._replace_uuid(run.trace_id) if run.trace_id else None,
+                "dotted_order": new_dotted_order,
+                "inputs": self._replace_message_id(run.inputs),
+                "outputs": self._replace_message_id(run.outputs),
+            }
+        )
+
+    def _persist_run(self, run: Run) -> None:
+        """Persist a run."""
+
+        self.runs.append(self._copy_run(run))
+
+    def flattened_runs(self) -> list[Run]:
+        q = [] + self.runs
+        result = []
+        while q:
+            parent = q.pop()
+            result.append(parent)
+            if parent.child_runs:
+                q.extend(parent.child_runs)
+        return result
+
+    @property
+    def run_ids(self) -> list[Optional[UUID]]:
+        runs = self.flattened_runs()
+        uuids_map = {v: k for k, v in self.uuids_map.items()}
+        return [uuids_map.get(r.id) for r in runs]
diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py
index 197dc6c61..74897efcc 100644
--- a/libs/langgraph/tests/test_pregel.py
+++ b/libs/langgraph/tests/test_pregel.py
@@ -68,6 +68,7 @@ from langgraph.pregel.retry import RetryPolicy
 from langgraph.pregel.types import PregelTask
 from langgraph.store.memory import MemoryStore
 from tests.any_str import AnyStr, ExceptionLike
+from tests.fake_tracer import FakeTracer
 from tests.memory_assert import (
     MemorySaverAssertCheckpointMetadata,
     MemorySaverAssertImmutable,
@@ -6159,11 +6160,20 @@ def test_dynamic_interrupt(snapshot: SnapshotAssertion) -> None:
     tool_two_graph.add_edge(START, "tool_two")
     tool_two = tool_two_graph.compile()
 
-    assert tool_two.invoke({"my_key": "value", "market": "DE"}) == {
+    tracer = FakeTracer()
+    assert tool_two.invoke(
+        {"my_key": "value", "market": "DE"}, {"callbacks": [tracer]}
+    ) == {
         "my_key": "value",
         "market": "DE",
     }
     assert tool_two_node_count == 1, "interrupts aren't retried"
+    assert len(tracer.runs) == 1
+    run = tracer.runs[0]
+    assert run.end_time is not None
+    assert run.error is None
+    assert run.outputs == {"market": "DE", "my_key": "value"}
+
     assert tool_two.invoke({"my_key": "value", "market": "US"}) == {
         "my_key": "value all good",
         "market": "US",
diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py
index e9987c601..ef78c8bdb 100644
--- a/libs/langgraph/tests/test_pregel_async.py
+++ b/libs/langgraph/tests/test_pregel_async.py
@@ -63,6 +63,7 @@ from langgraph.pregel.retry import RetryPolicy
 from langgraph.pregel.types import PregelTask
 from langgraph.store.memory import MemoryStore
 from tests.any_str import AnyStr, ExceptionLike
+from tests.fake_tracer import FakeTracer
 from tests.memory_assert import (
     MemorySaverAssertCheckpointMetadata,
     MemorySaverAssertImmutable,
@@ -225,11 +226,20 @@ async def test_dynamic_interrupt(snapshot: SnapshotAssertion) -> None:
     tool_two_graph.add_edge(START, "tool_two")
     tool_two = tool_two_graph.compile()
 
-    assert await tool_two.ainvoke({"my_key": "value", "market": "DE"}) == {
+    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",

From 9147d05cc47b8380cd8f6d76ad5f7351bb538d60 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Wed, 21 Aug 2024 14:32:15 -0700
Subject: [PATCH 30/30] lib0.2.9

---
 libs/cli/langgraph_cli/docker.py | 17 +++++++++++++++--
 libs/langgraph/pyproject.toml    |  2 +-
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/libs/cli/langgraph_cli/docker.py b/libs/cli/langgraph_cli/docker.py
index 79e0dd25a..d2cffc26b 100644
--- a/libs/cli/langgraph_cli/docker.py
+++ b/libs/cli/langgraph_cli/docker.py
@@ -12,6 +12,15 @@ DEFAULT_POSTGRES_URI = (
     "postgres://postgres:postgres@langgraph-postgres:5432/postgres?sslmode=disable"
 )
 
+REDIS = """
+    langgraph-redis:
+        image: redis:6
+        healthcheck:
+            test: redis-cli ping
+            interval: 5s
+            timeout: 1s
+            retries: 5
+"""
 
 DB = """
     langgraph-postgres:
@@ -166,18 +175,22 @@ def compose(
 
     compose_str = f"""{volumes}services:
 {db}
+{REDIS}
 {debugger_compose(port=debugger_port, base_url=debugger_base_url)}
     langgraph-api:
         ports:
-            - "{port}:8000\""""
+            - "{port}:8000\"
+        depends_on:
+            langgraph-redis:
+                condition: service_healthy"""
     if include_db:
         compose_str += """
-        depends_on:
             langgraph-postgres:
                 condition: service_healthy"""
     compose_str += f"""
         environment:
             POSTGRES_URI: {postgres_uri}
+            REDIS_URI: redis://langgraph-redis:6379
 """
     if capabilities.healthcheck_start_interval:
         compose_str += """        healthcheck:
diff --git a/libs/langgraph/pyproject.toml b/libs/langgraph/pyproject.toml
index 93dc203f2..1df951ee4 100644
--- a/libs/langgraph/pyproject.toml
+++ b/libs/langgraph/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph"
-version = "0.2.8"
+version = "0.2.9"
 description = "Building stateful, multi-actor applications with LLMs"
 authors = []
 license = "MIT"