perf(checkpoint-postgres): run a lone store read outside a pipeline (#9635)

`PostgresStore` and `AsyncPostgresStore` always ran a batch inside a
pipeline. For a read that costs two round trips: fetching the result
forces one, and leaving the pipeline sends a Sync and waits for a
second. `store.get()`, `store.search()` and `store.list_namespaces()`
are batches of one read, so each paid two round trips, which is why a
`get` took longer than a `put`.

A batch of exactly one read now runs outside the pipeline, in one round
trip. Batches with several ops or any write still use the pipeline, and
so does a store opened with `pipeline=True`. In the sync store the
non-pipeline cursor now takes the store's lock, as every batch already
did through the pipeline branch, so a lone read never uses a shared
connection while another thread has it in a pipeline. That also covers
`sweep_ttl`, which went through that branch without the lock.

Median of four alternating rounds on main and on this branch, Python
3.14, Postgres 15 in local Docker:

| workload | main | this PR | speedup |
|---|---|---|---|
| round trips for 200 `get` calls, counted with a TCP proxy | 400 | 200
| 2x |
| 200 `PostgresStore.get` | 85.2 ms | 47.8 ms | 1.78x |
| 50 `PostgresStore.search` | 25.7 ms | 16.6 ms | 1.55x |
| 200 `AsyncPostgresStore.aget`, one after another | 103 ms | 59.5 ms |
1.73x |
| the same 200 `get` with 0.5 ms added per round trip | 433 ms | 228 ms
| 1.90x |
| graph node doing one `get` and one `put` per step, 50 steps | 46.8 ms
| 37.2 ms | 1.26x |
| `put`, and many concurrent `aget` batched together | | | about 1x |

## Test plan
- [x] The existing store suites pass for the sync and async stores in
their default, `pipeline=True` and pool setups, which cover the
lone-read path and the pipeline path.

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

@@ -3,7 +3,7 @@ from __future__ import annotations
 import asyncio
 from collections import defaultdict
 from collections.abc import AsyncIterator, Iterator, Mapping, Sequence
-from contextlib import asynccontextmanager
+from contextlib import asynccontextmanager, nullcontext
 from typing import Any, cast
 
 from langchain_core.runnables import RunnableConfig
@@ -425,7 +425,12 @@ class AsyncPostgresSaver(BasePostgresSaver):
                 Will be applied regardless of whether the AsyncPostgresSaver instance was initialized with a pipeline.
                 If pipeline mode is not supported, will fall back to using transaction context manager.
         """
-        async with self.lock, _ainternal.get_connection(self.conn) as conn:
+        # A pool hands each call its own connection, so only a shared one needs
+        # the lock.
+        lock = (
+            nullcontext() if isinstance(self.conn, AsyncConnectionPool) else self.lock
+        )
+        async with lock, _ainternal.get_connection(self.conn) as conn:
             if self.pipe:
                 # a connection in pipeline mode can be used concurrently
                 # in multiple threads/coroutines, but only one cursor can be
@@ -594,7 +599,11 @@ class AsyncPostgresSaver(BasePostgresSaver):
                 if value["parent_checkpoint_id"]
                 else None
             ),
-            await asyncio.to_thread(self._load_writes, value["pending_writes"]),
+            (
+                await asyncio.to_thread(self._load_writes, value["pending_writes"])
+                if value["pending_writes"]
+                else []
+            ),
         )
 
     def list(

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

@@ -413,7 +413,14 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
     ) -> None:
         # Keep `conn` for compatibility with subclasses overriding this private hook.
         # All database I/O goes through `_cursor()`, which owns connection acquisition.
-        async with self._cursor(pipeline=True) as cur:
+        # A lone read is one statement: outside a pipeline it costs one round
+        # trip instead of two (its fetch, then the Sync on leaving the pipeline).
+        lone_read = (
+            not self.pipe
+            and PutOp not in grouped_ops
+            and sum(map(len, grouped_ops.values())) == 1
+        )
+        async with self._cursor(pipeline=not lone_read) as cur:
             if GetOp in grouped_ops:
                 await self._batch_get_ops(
                     cast(Sequence[tuple[int, GetOp]], grouped_ops[GetOp]),

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

@@ -977,14 +977,19 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
                     ):
                         yield cur
             else:
-                with conn.cursor(binary=True, row_factory=dict_row) as cur:
+                # Lone reads land here, so take the lock like the pipeline branch:
+                # another thread may be holding this connection in a pipeline.
+                with self.lock, conn.cursor(binary=True, row_factory=dict_row) as cur:
                     yield cur
 
     def batch(self, ops: Iterable[Op]) -> list[Result]:
         grouped_ops, num_ops = _group_ops(ops)
         results: list[Result] = [None] * num_ops
 
-        with self._cursor(pipeline=True) as cur:
+        # A lone read is one statement: outside a pipeline it costs one round
+        # trip instead of two (its fetch, then the Sync on leaving the pipeline).
+        lone_read = num_ops == 1 and PutOp not in grouped_ops and not self.pipe
+        with self._cursor(pipeline=not lone_read) as cur:
             if GetOp in grouped_ops:
                 self._batch_get_ops(
                     cast(Sequence[tuple[int, GetOp]], grouped_ops[GetOp]), results, cur

libs/checkpoint-postgres/tests/test_async.py

@@ -1,5 +1,6 @@
 # type: ignore
 
+import asyncio
 from contextlib import asynccontextmanager
 from typing import Any
 from uuid import uuid4
@@ -417,3 +418,33 @@ async def test_delta_channel_chain_reconstruction(saver_name: str) -> None:
         assert msgs[1].content == "reply-1"
         assert msgs[2].content == "there"
         assert msgs[3].content == "reply-3"
+
+
+async def test_pool_saver_reads_while_a_list_is_open(test_data) -> None:
+    async with _pool_saver() as saver:
+        config = await saver.aput(
+            test_data["configs"][0],
+            test_data["checkpoints"][0],
+            test_data["metadata"][0],
+            {},
+        )
+        listing = saver.alist(None)
+        try:
+            assert await anext(listing) is not None
+            assert await asyncio.wait_for(saver.aget_tuple(config), timeout=5)
+        finally:
+            await listing.aclose()
+
+
+async def test_base_saver_serializes_concurrent_calls(test_data) -> None:
+    async with _base_saver() as saver:
+        config = await saver.aput(
+            test_data["configs"][0],
+            test_data["checkpoints"][0],
+            test_data["metadata"][0],
+            {},
+        )
+        await asyncio.gather(
+            *(saver.aput_writes(config, [("x", i)], str(uuid4())) for i in range(20)),
+            *(saver.aget_tuple(config) for _ in range(20)),
+        )

libs/checkpoint-postgres/tests/test_store.py

@@ -2,6 +2,7 @@
 from __future__ import annotations
 
 import re
+import threading
 import time
 from contextlib import contextmanager
 from typing import Any
@@ -289,6 +290,30 @@ def test_basic_store_ops(store) -> None:
     assert deleted_item is None
 
 
+def test_lone_reads_and_puts_share_a_connection(store: PostgresStore) -> None:
+    errors: list[BaseException] = []
+
+    def run(op: Any) -> None:
+        try:
+            for i in range(100):
+                op(i)
+        except BaseException as e:
+            errors.append(e)
+
+    def read(i: int) -> None:
+        store.get(("race",), f"k{i % 10}")
+
+    def write(i: int) -> None:
+        store.put(("race",), f"k{i % 10}", {"i": i})
+
+    threads = [threading.Thread(target=run, args=(op,)) for op in [read, write] * 4]
+    for t in threads:
+        t.start()
+    for t in threads:
+        t.join()
+    assert errors == []
+
+
 def test_list_namespaces(store) -> None:
     # Create test data with various namespaces
     test_namespaces = [

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

@@ -302,14 +302,20 @@ def _dedupe_ops(values: list[Op]) -> tuple[list[int] | None, list[Op]]:
     dedupped: list[Op] = []
     listen: list[int] = []
     puts: dict[tuple[tuple[str, ...], str], int] = {}
+    reads: dict[Op, int] = {}
 
     for op in values:
         if isinstance(op, (GetOp, SearchOp, ListNamespacesOp)):
             try:
-                listen.append(dedupped.index(op))
-            except ValueError:
-                listen.append(len(dedupped))
+                ix = reads.setdefault(op, len(dedupped))
+            except TypeError:  # an op holding a dict or a list is unhashable
+                try:
+                    ix = dedupped.index(op)
+                except ValueError:
+                    ix = len(dedupped)
+            if ix == len(dedupped):
                 dedupped.append(op)
+            listen.append(ix)
         elif isinstance(op, PutOp):
             putkey = (op.namespace, op.key)
             if putkey in puts:

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

@@ -248,13 +248,20 @@ class InMemoryStore(BaseStore):
                 for key, filter_value in op.filter.items()
             )
 
+        # Without a filter or a query embedding, the caller keeps only
+        # filtered[offset:offset + limit], so stop collecting there. Negative
+        # bounds keep the full scan, since a slice counts them from the end.
+        stop = (
+            op.offset + op.limit
+            if not op.filter
+            and not (op.query and self.embeddings)
+            and min(op.offset, op.limit) >= 0
+            else None
+        )
+        n = len(namespace_prefix)
         filtered = []
         for namespace in self._data:
-            if not (
-                namespace[: len(namespace_prefix)] == namespace_prefix
-                if len(namespace) >= len(namespace_prefix)
-                else False
-            ):
+            if namespace[:n] != namespace_prefix:
                 continue
 
             for key, item in self._data[namespace].items():
@@ -263,6 +270,8 @@ class InMemoryStore(BaseStore):
                         filtered.append((item, list(embeddings.values())))
                     else:
                         filtered.append((item, []))
+                    if len(filtered) == stop:
+                        return filtered
         return filtered
 
     def _embed_search_queries(

libs/checkpoint/tests/test_store.py

@@ -22,7 +22,7 @@ from langgraph.store.base import (
     get_text_at_path,
     validate_op_namespace,
 )
-from langgraph.store.base.batch import AsyncBatchedBaseStore
+from langgraph.store.base.batch import AsyncBatchedBaseStore, _dedupe_ops
 from langgraph.store.memory import InMemoryStore
 from tests.embed_test_utils import CharacterEmbeddings
 
@@ -75,6 +75,29 @@ async def test_async_batch_store_resilience() -> None:
     assert result2.value == doc2
 
 
+def test_dedupe_ops_maps_repeated_reads_and_puts_to_one_op() -> None:
+    ops = [
+        GetOp(("a",), "k"),
+        PutOp(("a",), "k", {"v": 1}),
+        GetOp(("a",), "k"),
+        SearchOp(("a",), filter={"x": 1}),
+        SearchOp(("a",), filter={"x": 1}),
+        PutOp(("a",), "k", {"v": 2}),
+        GetOp(("b",), "k"),
+        ListNamespacesOp(match_conditions=(MatchCondition("prefix", ("a",)),)),
+        ListNamespacesOp(match_conditions=(MatchCondition("prefix", ("a",)),)),
+    ]
+    listen, dedupped = _dedupe_ops(ops)
+    assert dedupped == [
+        GetOp(("a",), "k"),
+        PutOp(("a",), "k", {"v": 2}),
+        SearchOp(("a",), filter={"x": 1}),
+        GetOp(("b",), "k"),
+        ListNamespacesOp(match_conditions=(MatchCondition("prefix", ("a",)),)),
+    ]
+    assert listen == [0, 1, 0, 2, 2, 1, 3, 4, 4]
+
+
 def test_get_text_at_path() -> None:
     nested_data = {
         "name": "test",
@@ -437,6 +460,43 @@ def test_list_namespaces_no_conditions() -> None:
     assert sorted(result) == sorted(expected)
 
 
+def test_search_pages_match_full_listing() -> None:
+    store = InMemoryStore()
+    items_per_namespace = 7
+    for ns in [("a",), ("a", "b"), ("ab",), ("a", "c", "d")]:
+        for i in range(items_per_namespace):
+            store.put(ns, f"k{i}", {"i": i, "even": i % 2 == 0})
+
+    def keys(items: list) -> list[tuple[tuple[str, ...], str]]:
+        return [(item.namespace, item.key) for item in items]
+
+    full = keys(store.search(("a",), limit=100))
+    assert {ns for ns, _ in full} == {("a",), ("a", "b"), ("a", "c", "d")}
+    assert len(full) == 3 * items_per_namespace
+    paged = [
+        key
+        for offset in range(0, len(full), 4)
+        for key in keys(store.search(("a",), limit=4, offset=offset))
+    ]
+    assert paged == full
+    even = keys(store.search(("a",), filter={"even": True}, limit=100))
+    assert (
+        keys(store.search(("a",), filter={"even": True}, limit=3, offset=2))
+        == even[2:5]
+    )
+
+
+def test_vector_search_scores_every_item(fake_embeddings: CharacterEmbeddings) -> None:
+    store = InMemoryStore(
+        index={"dims": fake_embeddings.dims, "embed": fake_embeddings}
+    )
+    for i in range(4):
+        store.put(("test",), f"other{i}", {"text": f"zzz qqq {i}"})
+    store.put(("test",), "match", {"text": "apple pie"})
+    results = store.search(("test",), query="apple pie", limit=1)
+    assert [r.key for r in results] == ["match"]
+
+
 def test_list_namespaces_empty_store() -> None:
     store = InMemoryStore()
 

libs/langgraph/langgraph/graph/message.py

@@ -194,14 +194,8 @@ def add_messages(
     if not isinstance(right, list):
         right = [right]  # type: ignore[assignment]
     # coerce to message
-    left = [
-        message_chunk_to_message(cast(BaseMessageChunk, m))
-        for m in convert_to_messages(left)
-    ]
-    right = [
-        message_chunk_to_message(cast(BaseMessageChunk, m))
-        for m in convert_to_messages(right)
-    ]
+    left = _coerce_to_messages(left)
+    right = _coerce_to_messages(right)
     # assign missing ids
     for m in left:
         if m.id is None:
@@ -234,7 +228,8 @@ def add_messages(
 
             merged_by_id[m.id] = len(merged)
             merged.append(m)
-    merged = [m for m in merged if m.id not in ids_to_remove]
+    if ids_to_remove:
+        merged = [m for m in merged if m.id not in ids_to_remove]
 
     if format == "langchain-openai":
         merged = _format_messages(merged)
@@ -247,6 +242,20 @@ def add_messages(
     return merged
 
 
+def _coerce_to_messages(messages: list) -> list[Any]:
+    # Both conversions return a plain message unchanged, and a long history is
+    # all plain messages, so check the few distinct types instead of each item.
+    if all(
+        issubclass(t, BaseMessage) and not issubclass(t, BaseMessageChunk)
+        for t in set(map(type, messages))
+    ):
+        return messages
+    return [
+        message_chunk_to_message(cast(BaseMessageChunk, m))
+        for m in convert_to_messages(messages)
+    ]
+
+
 @deprecated(
     "MessageGraph is deprecated in langgraph 1.0.0, to be removed in 2.0.0. Please use StateGraph with a `messages` key instead.",
     category=None,

libs/langgraph/langgraph/graph/state.py

@@ -1452,6 +1452,8 @@ class CompiledStateGraph(
             output_keys = list(self.builder.channels) + [
                 k for k, v in self.builder.managed.items()
             ]
+        # for membership checks, which run on every key of every node update
+        output_key_set = frozenset(output_keys)
 
         def _get_updates(
             input: None | dict | Any,
@@ -1459,12 +1461,12 @@ class CompiledStateGraph(
             if input is None:
                 return None
             elif isinstance(input, dict):
-                return [(k, v) for k, v in input.items() if k in output_keys]
+                return [(k, v) for k, v in input.items() if k in output_key_set]
             elif isinstance(input, Command):
                 if input.graph == Command.PARENT:
                     return None
                 return [
-                    (k, v) for k, v in input._update_as_tuples() if k in output_keys
+                    (k, v) for k, v in input._update_as_tuples() if k in output_key_set
                 ]
             elif (
                 isinstance(input, (list, tuple))
@@ -1477,7 +1479,9 @@ class CompiledStateGraph(
                         if i.graph == Command.PARENT:
                             continue
                         updates.extend(
-                            (k, v) for k, v in i._update_as_tuples() if k in output_keys
+                            (k, v)
+                            for k, v in i._update_as_tuples()
+                            if k in output_key_set
                         )
                     else:
                         updates.extend(_get_updates(i) or ())

libs/langgraph/langgraph/pregel/_io.py

@@ -130,35 +130,35 @@ def map_output_updates(
     ]
     if not output_tasks:
         return
+    # a set: a wide state would otherwise scan every key for each write
+    output_keys = (
+        frozenset(output_channels) if not isinstance(output_channels, str) else None
+    )
     updated: list[tuple[str, Any]] = []
     for task, writes in output_tasks:
         rtn = next((value for chan, value in writes if chan == RETURN), MISSING)
         if rtn is not MISSING:
             updated.append((task.name, rtn))
-        elif isinstance(output_channels, str):
+        elif output_keys is None:
             updated.extend(
                 (task.name, value) for chan, value in writes if chan == output_channels
             )
-        elif any(chan in output_channels for chan, _ in writes):
+        elif any(chan in output_keys for chan, _ in writes):
             counts = Counter(chan for chan, _ in writes)
-            if any(counts[chan] > 1 for chan in output_channels):
+            if any(n > 1 for chan, n in counts.items() if chan in output_keys):
                 updated.extend(
                     (
                         task.name,
                         {chan: value},
                     )
                     for chan, value in writes
-                    if chan in output_channels
+                    if chan in output_keys
                 )
             else:
                 updated.append(
                     (
                         task.name,
-                        {
-                            chan: value
-                            for chan, value in writes
-                            if chan in output_channels
-                        },
+                        {chan: value for chan, value in writes if chan in output_keys},
                     )
                 )
     grouped: dict[str, Any] = {t.name: [] for t, _ in output_tasks}

libs/langgraph/langgraph/pregel/_loop.py

@@ -358,6 +358,14 @@ class PregelLoop:
         self.retry_policy = retry_policy
         self.cache_policy = cache_policy
         self.durability = durability
+        # Channel types are fixed for the run, so per-step code reads these
+        # instead of scanning every channel.
+        self._has_untracked_channels = any(
+            isinstance(spec, UntrackedValue) for spec in specs.values()
+        )
+        self._delta_channel_names = [
+            k for k, spec in specs.items() if isinstance(spec, DeltaChannel)
+        ]
         self._has_graph_lifecycle_callbacks = has_graph_lifecycle_callbacks
         self._graph_lifecycle_events = deque()
         if self.stream is not None and CONFIG_KEY_STREAM in config[CONF]:
@@ -496,9 +504,7 @@ class PregelLoop:
             writes_to_save = writes
 
         # check if any writes are to an UntrackedValue channel
-        if any(
-            isinstance(channel, UntrackedValue) for channel in self.channels.values()
-        ):
+        if self._has_untracked_channels:
             # we do not persist untracked values in checkpoints
             writes_to_save = [
                 # sanitize UntrackedValues that are nested within Send packets
@@ -840,10 +846,12 @@ class PregelLoop:
         Unfinished (failed or interrupted) tasks keep empty writes, so the
         runner re-executes them or routes them to error handlers.
         """
-        for tid, status in read_task_statuses(self.checkpoint_pending_writes).items():
-            if task := tasks.get(tid):
-                task.writes.extend(status.output)
-                self._reapplied_task_ids.add(tid)
+        # Filter first: this runs once per task call, and reading the status of
+        # every task that wrote would make runs with many task calls quadratic.
+        writes = (w for w in self.checkpoint_pending_writes if w[0] in tasks)
+        for tid, status in read_task_statuses(writes).items():
+            tasks[tid].writes.extend(status.output)
+            self._reapplied_task_ids.add(tid)
 
     def _resume_error_handlers_if_applicable(self) -> None:
         """On resume, schedule error handlers for tasks that failed in a prior run.
@@ -1300,9 +1308,7 @@ class PregelLoop:
             )
             new_counters: dict[str, tuple[int, int]] = {}
             updated = self.updated_channels or set()
-            for ch_name, ch in self.channels.items():
-                if not isinstance(ch, DeltaChannel):
-                    continue
+            for ch_name in self._delta_channel_names:
                 u, s = prev_counters.get(ch_name, (0, 0))
                 s += 1
                 if ch_name in updated:
@@ -1322,7 +1328,9 @@ class PregelLoop:
         # create new checkpoint
         channels_to_snapshot = (
             delta_channels_to_snapshot(
-                self.channels, new_counters, self.checkpoint["channel_versions"]
+                {k: self.channels[k] for k in self._delta_channel_names},
+                new_counters,
+                self.checkpoint["channel_versions"],
             )
             | self._delta_channels_forced_snapshot
             if do_checkpoint
@@ -1351,9 +1359,7 @@ class PregelLoop:
         elif "counters_since_delta_snapshot" in self.checkpoint_metadata:
             del self.checkpoint_metadata["counters_since_delta_snapshot"]
         # sanitize TASK channel in the checkpoint before saving (durability=="exit")
-        if TASKS in self.checkpoint["channel_values"] and any(
-            isinstance(channel, UntrackedValue) for channel in self.channels.values()
-        ):
+        if TASKS in self.checkpoint["channel_values"] and self._has_untracked_channels:
             sanitized_tasks = [
                 sanitize_untracked_values_in_send(value, self.channels)
                 if isinstance(value, Send)

libs/langgraph/langgraph/pregel/_runner.py

@@ -78,6 +78,7 @@ class FuturesDict(Generic[F, E], dict[F, PregelExecutableTask | None]):
     should_stop: Callable[[set[F]], bool]
     counter: int
     done: set[F]
+    failed: set[F]
     lock: threading.Lock
 
     def __init__(
@@ -97,6 +98,7 @@ class FuturesDict(Generic[F, E], dict[F, PregelExecutableTask | None]):
         self.should_stop = should_stop
         self.counter = 0
         self.done: set[F] = set()
+        self.failed: set[F] = set()
 
     def __setitem__(
         self,
@@ -123,9 +125,13 @@ class FuturesDict(Generic[F, E], dict[F, PregelExecutableTask | None]):
             with self.lock:
                 self.done.add(fut)
                 self.counter -= 1
+                # Only a future that raised can meet the stop condition, so check
+                # those instead of every done future (quadratic in wide steps).
+                if not fut.cancelled() and fut.exception() is not None:
+                    self.failed.add(fut)
                 # Wake waiter when all tracked futures are done, or when runner-level
                 # stop condition is met (for example, a non-handled fatal exception).
-                if self.counter == 0 or self.should_stop(self.done):
+                if self.counter == 0 or self.should_stop(self.failed):
                     self.event.set()
 
 
@@ -234,6 +240,7 @@ class PregelRunner:
                         fut: concurrent.futures.Future = concurrent.futures.Future()
                         fut.set_exception(exc)
                         futures.done.add(fut)
+                        futures.failed.add(fut)
                 elif reraise and id(exc) not in self._handled_exception_ids:
                     # let the caller stream this task's result before it raises
                     yield
@@ -427,6 +434,7 @@ class PregelRunner:
                         fut: asyncio.Future = loop.create_future()
                         fut.set_exception(exc)
                         futures.done.add(fut)
+                        futures.failed.add(fut)
                 elif reraise and id(exc) not in self._handled_exception_ids:
                     # let the caller stream this task's result before it raises
                     yield

libs/langgraph/tests/test_messages_state.py

@@ -5,6 +5,7 @@ import langchain_core
 import pytest
 from langchain_core.messages import (
     AIMessage,
+    AIMessageChunk,
     AnyMessage,
     HumanMessage,
     RemoveMessage,
@@ -168,6 +169,41 @@ def test_non_list_inputs():
     assert result == expected_result
 
 
+def test_message_likes_and_chunks_are_converted_on_both_sides():
+    left = [
+        HumanMessage(content="Hello", id="1"),
+        {"role": "assistant", "content": "Hi", "id": "2"},
+    ]
+    right = [AIMessageChunk(content="chunk", id="3"), ("user", "Thanks")]
+    result = add_messages(left, right)
+    assert [type(m) for m in result] == [
+        HumanMessage,
+        AIMessage,
+        AIMessage,
+        HumanMessage,
+    ]
+    assert result[:3] == [
+        HumanMessage(content="Hello", id="1"),
+        AIMessage(content="Hi", id="2"),
+        AIMessage(content="chunk", id="3"),
+    ]
+    assert result[3].content == "Thanks"
+
+
+def test_chunk_only_lists_are_converted():
+    left = [AIMessageChunk(content="a", id="1")]
+    right = [AIMessageChunk(content="b", id="2")]
+    result = add_messages(left, right)
+    assert [type(m) for m in result] == [AIMessage, AIMessage]
+
+
+def test_result_is_a_new_list():
+    left = [HumanMessage(content="Hello", id="1")]
+    result = add_messages(left, [])
+    assert result == left
+    assert result is not left
+
+
 def test_delete_all():
     left = [
         HumanMessage(content="Hello", id="1"),

libs/langgraph/tests/test_runner_futures.py

@@ -0,0 +1,50 @@
+import concurrent.futures
+import threading
+import weakref
+from functools import partial
+
+from langgraph.pregel._runner import FuturesDict, _should_stop_others
+
+
+def _on_done(task: object, exc: BaseException | None) -> None:
+    pass
+
+
+def _futures(handled_exception_ids: set[int]) -> FuturesDict:
+    return FuturesDict(
+        event=threading.Event(),
+        callback=weakref.ref(_on_done),
+        should_stop=partial(
+            _should_stop_others, handled_exception_ids=handled_exception_ids
+        ),
+        future_type=concurrent.futures.Future,
+    )
+
+
+def test_failed_task_wakes_waiter_before_others_finish() -> None:
+    futures = _futures(set())
+    ok, bad, pending = (concurrent.futures.Future() for _ in range(3))
+    for fut in (ok, bad, pending):
+        futures[fut] = object()
+
+    ok.set_result(None)
+    assert not futures.event.is_set()
+
+    bad.set_exception(ValueError("boom"))
+    assert futures.event.is_set()
+
+
+def test_handled_failure_waits_for_all_tasks() -> None:
+    handled: set[int] = set()
+    futures = _futures(handled)
+    bad, pending = (concurrent.futures.Future() for _ in range(2))
+    for fut in (bad, pending):
+        futures[fut] = object()
+
+    exc = ValueError("handled by an error handler")
+    handled.add(id(exc))
+    bad.set_exception(exc)
+    assert not futures.event.is_set()
+
+    pending.set_result(None)
+    assert futures.event.is_set()

libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py

@@ -638,15 +638,15 @@ def create_react_agent(
             messages = (
                 _get_state_value(state, "llm_input_messages")
             ) or _get_state_value(state, "messages")
-            error_msg = f"Expected input to call_model to have 'llm_input_messages' or 'messages' key, but got {state}"
+            expected_keys = "'llm_input_messages' or 'messages' key"
         else:
             messages = _get_state_value(state, "messages")
-            error_msg = (
-                f"Expected input to call_model to have 'messages' key, but got {state}"
-            )
+            expected_keys = "'messages' key"
 
         if messages is None:
-            raise ValueError(error_msg)
+            raise ValueError(
+                f"Expected input to call_model to have {expected_keys}, but got {state}"
+            )
 
         _validate_chat_history(messages)
         # we're passing messages under `messages` key, as this is expected by the prompt