Compare commits

..
Author SHA1 Message Date
Quanzheng Long d2848b19c9 fix 2026-07-06 10:35:55 -07:00
11 changed files with 125 additions and 247 deletions
+22 -6
View File
@@ -1,7 +1,7 @@
from __future__ import annotations
import uuid
from collections.abc import Callable, Mapping
from collections.abc import Callable, Iterable, Mapping, Sequence
from datetime import datetime, timezone
from typing import Any, cast
@@ -14,6 +14,7 @@ from langgraph.checkpoint.base.id import uuid6
from langgraph.checkpoint.serde.types import _DeltaSnapshot
from langgraph._internal._config import DELTA_MAX_SUPERSTEPS_SINCE_SNAPSHOT
from langgraph._internal._constants import PUSH
from langgraph._internal._typing import MISSING
from langgraph.channels.base import BaseChannel
from langgraph.channels.delta import DeltaChannel
@@ -70,6 +71,26 @@ def delta_channels_to_snapshot(
return result
def update_state_channels_plan(
run_tasks: Iterable[Any],
channels: Mapping[str, BaseChannel],
) -> tuple[set[str], set[str]]:
"""Return channels written and DeltaChannels to snapshot on update_state."""
updated_channels = {c for task in run_tasks for c, _ in task.writes if c != PUSH}
channels_to_snapshot = {
c for c in updated_channels if isinstance(channels.get(c), DeltaChannel)
}
return updated_channels, channels_to_snapshot
def update_state_channel_writes(
writes: Sequence[tuple[str, Any]],
channels_to_snapshot: set[str],
) -> list[tuple[str, Any]]:
"""Channel writes to persist separately from a head snapshot."""
return [w for w in writes if w[0] != PUSH and w[0] not in channels_to_snapshot]
def create_checkpoint(
checkpoint: Checkpoint,
channels: Mapping[str, BaseChannel] | None,
@@ -102,11 +123,6 @@ def create_checkpoint(
continue
ch = channels[k]
if k in channels_to_snapshot:
# Callers force a full snapshot blob here: exit mode when a
# delta channel reaches its snapshot cadence, and update_state
# on a fresh thread (no ancestor to replay writes from). The
# manual version-bump below only applies to the exit-mode case.
#
# In exit mode, the snapshot decision is deferred to exit
# time (intermediate steps have do_checkpoint=False). The
# channel's count may have reached snapshot_frequency over
+52 -46
View File
@@ -108,7 +108,6 @@ from langgraph.callbacks import (
get_sync_graph_callback_manager_for_config,
)
from langgraph.channels.base import BaseChannel
from langgraph.channels.delta import DeltaChannel
from langgraph.channels.topic import Topic
from langgraph.config import get_config
from langgraph.constants import END
@@ -134,6 +133,8 @@ from langgraph.pregel._checkpoint import (
copy_checkpoint,
create_checkpoint,
empty_checkpoint,
update_state_channel_writes,
update_state_channels_plan,
)
from langgraph.pregel._draw import draw_graph
from langgraph.pregel._io import map_input, read_channels
@@ -1996,13 +1997,19 @@ class Pregel(
},
),
)
# save task writes
for task_id, task in zip(run_task_ids, run_tasks):
# channel writes are saved to current checkpoint
channel_writes = [w for w in task.writes if w[0] != PUSH]
if saved and channel_writes:
checkpointer.put_writes(checkpoint_config, channel_writes, task_id)
# apply to checkpoint and save
updated_channels, channels_to_snapshot = update_state_channels_plan(
run_tasks, channels
)
if saved is not None:
for task_id, task in zip(run_task_ids, run_tasks):
if channel_writes := update_state_channel_writes(
task.writes, channels_to_snapshot
):
checkpointer.put_writes(
checkpoint_config, channel_writes, task_id
)
apply_writes(
checkpoint,
channels,
@@ -2010,20 +2017,13 @@ class Pregel(
checkpointer.get_next_version,
self.trigger_to_nodes,
)
# On a fresh thread there is no ancestor to replay DeltaChannel
# writes from, so force a self-contained snapshot in the first
# checkpoint instead of relying on ancestor write-replay.
delta_snapshot = (
{
k
for k, ch in channels.items()
if isinstance(ch, DeltaChannel) and ch.is_available()
}
if saved is None
else None
)
checkpoint = create_checkpoint(
checkpoint, channels, step + 1, channels_to_snapshot=delta_snapshot
checkpoint,
channels,
step + 1,
updated_channels=updated_channels,
get_next_version=checkpointer.get_next_version,
channels_to_snapshot=channels_to_snapshot,
)
next_config = checkpointer.put(
checkpoint_config,
@@ -2038,7 +2038,11 @@ class Pregel(
),
)
for task_id, task in zip(run_task_ids, run_tasks):
# save push writes
if saved is None:
if channel_writes := update_state_channel_writes(
task.writes, channels_to_snapshot
):
checkpointer.put_writes(next_config, channel_writes, task_id)
if push_writes := [w for w in task.writes if w[0] == PUSH]:
checkpointer.put_writes(next_config, push_writes, task_id)
@@ -2455,15 +2459,19 @@ class Pregel(
},
),
)
# save task writes
for task_id, task in zip(run_task_ids, run_tasks):
# channel writes are saved to current checkpoint
channel_writes = [w for w in task.writes if w[0] != PUSH]
if saved and channel_writes:
await checkpointer.aput_writes(
checkpoint_config, channel_writes, task_id
)
# apply to checkpoint and save
updated_channels, channels_to_snapshot = update_state_channels_plan(
run_tasks, channels
)
if saved is not None:
for task_id, task in zip(run_task_ids, run_tasks):
if channel_writes := update_state_channel_writes(
task.writes, channels_to_snapshot
):
await checkpointer.aput_writes(
checkpoint_config, channel_writes, task_id
)
apply_writes(
checkpoint,
channels,
@@ -2471,22 +2479,14 @@ class Pregel(
checkpointer.get_next_version,
self.trigger_to_nodes,
)
# On a fresh thread there is no ancestor to replay DeltaChannel
# writes from, so force a self-contained snapshot in the first
# checkpoint instead of relying on ancestor write-replay.
delta_snapshot = (
{
k
for k, ch in channels.items()
if isinstance(ch, DeltaChannel) and ch.is_available()
}
if saved is None
else None
)
checkpoint = create_checkpoint(
checkpoint, channels, step + 1, channels_to_snapshot=delta_snapshot
checkpoint,
channels,
step + 1,
updated_channels=updated_channels,
get_next_version=checkpointer.get_next_version,
channels_to_snapshot=channels_to_snapshot,
)
# save checkpoint, after applying writes
next_config = await checkpointer.aput(
checkpoint_config,
checkpoint,
@@ -2500,7 +2500,13 @@ class Pregel(
),
)
for task_id, task in zip(run_task_ids, run_tasks):
# save push writes
if saved is None:
if channel_writes := update_state_channel_writes(
task.writes, channels_to_snapshot
):
await checkpointer.aput_writes(
next_config, channel_writes, task_id
)
if push_writes := [w for w in task.writes if w[0] == PUSH]:
await checkpointer.aput_writes(next_config, push_writes, task_id)
return patch_checkpoint_map(next_config, saved.metadata if saved else None)
+1 -1
View File
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
[project]
name = "langgraph"
version = "1.2.8"
version = "1.2.7"
description = "Building stateful, multi-actor applications with LLMs"
authors = []
requires-python = ">=3.10"
@@ -1,14 +1,12 @@
"""Tests for `update_state` / `aupdate_state` against `DeltaChannel`.
Originally a regression suite for deepagents#3774 — `update_state` on a *fresh*
thread silently dropped the first write to a `DeltaChannel`-backed channel
because channel writes were only persisted when a previous checkpoint existed
and no snapshot was written either, so the checkpoint reconstructed to empty.
Regression suite for deepagents#3774 and Postgres read-path compatibility:
fresh-thread `update_state` must persist DeltaChannel state correctly.
Fixed by forcing a self-contained `_DeltaSnapshot` blob into the first
checkpoint on a fresh thread (`saved is None`), so the value is stored inline
and no ancestor write-replay is required. This keeps the read/replay path
untouched.
Fresh-thread updates snapshot updated DeltaChannels on the head checkpoint
(self-contained for Postgres readers). Delta writes are not persisted via
`put_writes`; non-delta channel writes on a fresh thread are attached to
the head after it is saved.
Coverage:
@@ -16,8 +14,8 @@ Coverage:
* non-fresh thread: `update_state` after `invoke`, after another `update_state`,
and `bulk_update_state` with multiple per-superstep updates
* update-by-id end-to-end via `update_state` (DeltaChannel reducer semantics)
* state-history chain shape on a fresh thread (single self-contained update
checkpoint with the snapshot inline and no parent)
* state-history shape on a fresh thread (single snapshotted head checkpoint)
* head checkpoint snapshots updated DeltaChannels for Postgres read paths
"""
from typing import Annotated, Any
@@ -25,6 +23,7 @@ from typing import Annotated, Any
import pytest
from langchain_core.messages import HumanMessage
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.checkpoint.serde.types import _DeltaSnapshot
from typing_extensions import TypedDict
from langgraph.channels.delta import DeltaChannel
@@ -96,8 +95,7 @@ async def test_aupdate_state_fresh_thread_delta_channel() -> None:
def test_update_state_after_invoke_delta_channel() -> None:
"""The non-fresh-thread path was already working before the fix; pin it
down so the forced-snapshot change for fresh threads doesn't regress it."""
"""The non-fresh-thread path must keep working across snapshot changes."""
saver = InMemorySaver()
graph = _build_graph(saver)
config = {"configurable": {"thread_id": "after-invoke-sync"}}
@@ -136,9 +134,8 @@ async def test_aupdate_state_after_invoke_delta_channel() -> None:
def test_consecutive_update_states_delta_channel() -> None:
"""First update_state forces a self-contained snapshot seed; the second
sees a real parent (`saved is not None`) and anchors its writes under that
seed. Both messages must round-trip in chronological order."""
"""Two consecutive fresh-thread-style updates: the first creates a
snapshotted head; the second anchors on it. Both messages round-trip."""
saver = InMemorySaver()
graph = _build_graph(saver)
config = {"configurable": {"thread_id": "consecutive-sync"}}
@@ -212,10 +209,8 @@ def test_update_state_replaces_message_by_id_delta_channel() -> None:
def test_bulk_update_state_multi_task_per_superstep_delta_channel() -> None:
"""`bulk_update_state` with N updates in one superstep produces N tasks
that each call `put_writes`. Guards the regression where moving
`put_writes` outside the per-task loop would persist only the last
task's writes.
"""`bulk_update_state` with N updates in one superstep must accumulate
all N message writes in the snapshotted head state.
Explicit `task_id`s are required to disambiguate writes belonging to
different `StateUpdate`s targeting the same node — otherwise both share
@@ -255,14 +250,13 @@ def test_bulk_update_state_multi_task_per_superstep_delta_channel() -> None:
# ---------------------------------------------------------------------------
# Public-API observation of the forced-snapshot mechanism
# Public-API observation of fresh-thread checkpoint shape
# ---------------------------------------------------------------------------
def test_state_history_chain_after_fresh_update_state_delta_channel() -> None:
"""A fresh-thread `update_state` should produce a single self-contained
checkpoint visible via `get_state_history`: step=0, `source='update'`,
no parent, with the DeltaChannel value snapshotted inline."""
"""Fresh-thread `update_state` on snapshotted DeltaChannels yields one
self-contained checkpoint (step=0, no parent, source='update')."""
saver = InMemorySaver()
graph = _build_graph(saver)
config = {"configurable": {"thread_id": "history-chain"}}
@@ -276,9 +270,32 @@ def test_state_history_chain_after_fresh_update_state_delta_channel() -> None:
history = list(graph.get_state_history(config))
assert len(history) == 1
(update_snapshot,) = history
update_snapshot = history[0]
assert update_snapshot.metadata is not None
assert update_snapshot.metadata["source"] == "update"
assert update_snapshot.metadata["step"] == 0
assert update_snapshot.parent_config is None
assert [m.content for m in update_snapshot.values["messages"]] == ["hello"]
def test_fresh_update_state_head_snapshots_delta_channel() -> None:
"""Postgres checkpointers skip the ancestor walk when the head checkpoint
has no `counters_since_delta_snapshot` entry. Force-snapshot updated
DeltaChannels on the update checkpoint so the head is self-contained."""
saver = InMemorySaver()
graph = _build_graph(saver)
config = {"configurable": {"thread_id": "head-snapshot"}}
graph.update_state(
config,
{"messages": [HumanMessage(content="hello", id="m1")]},
as_node="model",
)
head = saver.get_tuple(config)
assert head is not None
assert isinstance(head.checkpoint["channel_values"].get("messages"), _DeltaSnapshot)
assert [m.content for m in head.checkpoint["channel_values"]["messages"].value] == [
"hello"
]
+1 -1
View File
@@ -1438,7 +1438,7 @@ wheels = [
[[package]]
name = "langgraph"
version = "1.2.8"
version = "1.2.7"
source = { editable = "." }
dependencies = [
{ name = "langchain-core" },
+1 -1
View File
@@ -285,7 +285,7 @@ wheels = [
[[package]]
name = "langgraph"
version = "1.2.8"
version = "1.2.7"
source = { editable = "../langgraph" }
dependencies = [
{ name = "langchain-core" },
+3 -32
View File
@@ -911,34 +911,6 @@ class _SubgraphsProjection:
def __aiter__(self) -> AsyncIterator[ScopedStreamHandle]:
return self._subgraphs_iter()
@staticmethod
def _put_root_message(root_inbox: asyncio.Queue[Event | None], item: Event) -> None:
if root_inbox.maxsize > 0 and root_inbox.qsize() >= root_inbox.maxsize - 1:
raise RuntimeError(
"Root messages inbox exceeded max_queue_size while buffering "
"root-scope messages. Iterate thread.messages concurrently "
"or increase max_queue_size."
)
try:
root_inbox.put_nowait(item)
except asyncio.QueueFull as exc:
raise RuntimeError(
"Root messages inbox exceeded max_queue_size while buffering "
"root-scope messages. Iterate thread.messages concurrently "
"or increase max_queue_size."
) from exc
@staticmethod
def _signal_root_inbox_closed(root_inbox: asyncio.Queue[Event | None]) -> None:
try:
root_inbox.put_nowait(None)
except asyncio.QueueFull as exc:
raise RuntimeError(
"Root messages inbox exceeded max_queue_size while closing "
"root-scope messages. Iterate thread.messages concurrently "
"or increase max_queue_size."
) from exc
async def _subgraphs_iter(self) -> AsyncGenerator[ScopedStreamHandle, None]:
if self._thread._transport is None:
raise RuntimeError("AsyncThreadStream not entered - use `async with`.")
@@ -971,7 +943,7 @@ class _SubgraphsProjection:
and item.get("method") == "messages"
and tuple(_event_namespace(params_field)) == self._scope
):
self._put_root_message(root_inbox, item)
root_inbox.put_nowait(item)
for handle in decoder.feed(item):
yield handle
finally:
@@ -989,7 +961,7 @@ class _SubgraphsProjection:
handle._finish(terminal_status)
self._thread._unregister_subscription(sub.id)
if root_inbox is not None:
self._signal_root_inbox_closed(root_inbox)
root_inbox.put_nowait(None)
class ToolCallHandle:
@@ -1363,8 +1335,7 @@ class AsyncThreadStream:
that arrive at namespace `[]` before `thread.messages` has subscribed.
"""
if self._root_messages_inbox is None:
maxsize = self._max_queue_size + 1 if self._max_queue_size > 0 else 0
self._root_messages_inbox = asyncio.Queue(maxsize=maxsize)
self._root_messages_inbox = asyncio.Queue()
return self._root_messages_inbox
def _register_active_message_stream(self, stream: AsyncChatModelStream) -> None:
+3 -31
View File
@@ -954,34 +954,6 @@ class _SyncSubgraphsProjection:
def __iter__(self) -> Iterator[SyncScopedStreamHandle]:
return self._subgraphs_iter()
@staticmethod
def _put_root_message(root_inbox: queue.Queue[Event | None], item: Event) -> None:
if root_inbox.maxsize > 0 and root_inbox.qsize() >= root_inbox.maxsize - 1:
raise RuntimeError(
"Root messages inbox exceeded max_queue_size while buffering "
"root-scope messages. Iterate thread.messages concurrently "
"or increase max_queue_size."
)
try:
root_inbox.put_nowait(item)
except queue.Full as exc:
raise RuntimeError(
"Root messages inbox exceeded max_queue_size while buffering "
"root-scope messages. Iterate thread.messages concurrently "
"or increase max_queue_size."
) from exc
@staticmethod
def _signal_root_inbox_closed(root_inbox: queue.Queue[Event | None]) -> None:
try:
root_inbox.put_nowait(None)
except queue.Full as exc:
raise RuntimeError(
"Root messages inbox exceeded max_queue_size while closing "
"root-scope messages. Iterate thread.messages concurrently "
"or increase max_queue_size."
) from exc
def _subgraphs_iter(self) -> Iterator[SyncScopedStreamHandle]:
if self._thread._transport is None:
raise RuntimeError("SyncThreadStream not entered — use `with`.")
@@ -1014,7 +986,7 @@ class _SyncSubgraphsProjection:
and item.get("method") == "messages"
and tuple(_event_namespace(params_field)) == self._scope
):
self._put_root_message(root_inbox, item)
root_inbox.put_nowait(item)
for handle in decoder.feed(cast(dict[str, Any], item)):
yield handle
finally:
@@ -1035,7 +1007,7 @@ class _SyncSubgraphsProjection:
handle._finish(terminal_status)
self._thread._unregister_subscription(sub.id)
if root_inbox is not None:
self._signal_root_inbox_closed(root_inbox)
root_inbox.put_nowait(None)
class _SyncExtensionsProjection:
@@ -1250,7 +1222,7 @@ class SyncThreadStream:
def _activate_root_messages_inbox(self) -> queue.Queue[Event | None]:
if self._root_messages_inbox is None:
self._root_messages_inbox = queue.Queue(maxsize=1025)
self._root_messages_inbox = queue.Queue()
return self._root_messages_inbox
def _register_active_message_stream(self, stream: ChatModelStream) -> None:
@@ -2,13 +2,7 @@
from __future__ import annotations
import asyncio
from typing import cast
from unittest.mock import MagicMock
import httpx
import pytest
from langchain_protocol import Event
from langgraph_sdk._async.http import HttpClient
from langgraph_sdk._async.threads import ThreadsClient
@@ -453,56 +447,6 @@ def test_scoped_handle_inboxes_bounded_by_max_queue_size():
assert handle._tasks_inbox.maxsize == 16
def test_root_messages_inbox_bounded_by_max_queue_size():
"""Root messages inbox must reserve one terminal sentinel slot."""
from langgraph_sdk._async.stream import AsyncThreadStream
thread = AsyncThreadStream(
http=MagicMock(),
thread_id="t-1",
assistant_id="agent",
max_queue_size=16,
)
inbox = thread._activate_root_messages_inbox()
assert inbox.maxsize == 17
def test_subgraphs_root_message_overflow_raises_runtime_error():
"""Overflowing the root messages inbox must fail explicitly."""
from langgraph_sdk._async.stream import _SubgraphsProjection
inbox: asyncio.Queue[Event | None] = asyncio.Queue(maxsize=2)
inbox.put_nowait(cast(Event, message_start_event(seq=1, message_id="msg-1")))
with pytest.raises(RuntimeError, match="Root messages inbox exceeded"):
_SubgraphsProjection._put_root_message(
inbox,
cast(
Event,
message_text_delta_event(seq=2, text="overflow", message_id="msg-1"),
),
)
def test_subgraphs_root_message_close_preserves_full_inbox():
"""Closing a full root messages inbox must not evict buffered events."""
from langgraph_sdk._async.stream import _SubgraphsProjection
inbox: asyncio.Queue[Event | None] = asyncio.Queue(maxsize=3)
first = cast(Event, message_start_event(seq=1, message_id="msg-1"))
second = cast(Event, message_finish_event(seq=2, message_id="msg-1"))
inbox.put_nowait(first)
inbox.put_nowait(second)
_SubgraphsProjection._signal_root_inbox_closed(inbox)
assert inbox.get_nowait() is first
assert inbox.get_nowait() is second
assert inbox.get_nowait() is None
async def test_child_handle_inherits_max_queue_size_from_parent():
"""Grandchild ScopedStreamHandles created by _HandleSubgraphsProjection
inherit the parent's max_queue_size so all queues are consistently bounded."""
@@ -2,7 +2,6 @@
from __future__ import annotations
import queue
from typing import Any, cast
import httpx
@@ -399,53 +398,6 @@ def test_sync_tool_calls_run_error_fails_active_handle():
# ---------------------------------------------------------------------------
def test_sync_root_messages_inbox_is_bounded():
"""Root messages inbox must reserve one terminal sentinel slot."""
fake = SyncFakeServer()
fake.script([lifecycle_completed_event(seq=1)])
fake.set_state({})
with httpx.Client(transport=fake.transport, base_url="http://test") as raw:
threads = SyncThreadsClient(SyncHttpClient(raw))
with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
inbox = thread._activate_root_messages_inbox()
assert inbox.maxsize == 1025
def test_sync_subgraphs_root_message_overflow_raises_runtime_error():
"""Overflowing the root messages inbox must fail explicitly."""
from langgraph_sdk._sync.stream import _SyncSubgraphsProjection
inbox: queue.Queue[Event | None] = queue.Queue(maxsize=2)
inbox.put_nowait(cast(Event, message_start_event(seq=1, message_id="msg-1")))
with pytest.raises(RuntimeError, match="Root messages inbox exceeded"):
_SyncSubgraphsProjection._put_root_message(
inbox,
cast(
Event,
message_text_delta_event(seq=2, text="overflow", message_id="msg-1"),
),
)
def test_sync_subgraphs_root_message_close_preserves_full_inbox():
"""Closing a full root messages inbox must not evict buffered events."""
from langgraph_sdk._sync.stream import _SyncSubgraphsProjection
inbox: queue.Queue[Event | None] = queue.Queue(maxsize=3)
first = cast(Event, message_start_event(seq=1, message_id="msg-1"))
second = cast(Event, message_finish_event(seq=2, message_id="msg-1"))
inbox.put_nowait(first)
inbox.put_nowait(second)
_SyncSubgraphsProjection._signal_root_inbox_closed(inbox)
assert inbox.get_nowait() is first
assert inbox.get_nowait() is second
assert inbox.get_nowait() is None
def test_sync_drain_messages_inbox_pre_dispatches_before_yield():
"""When draining the root inbox, str(message.text) must work immediately on yield."""
fake = SyncFakeServer()
+1 -1
View File
@@ -298,7 +298,7 @@ wheels = [
[[package]]
name = "langgraph"
version = "1.2.8"
version = "1.2.7"
source = { editable = "../langgraph" }
dependencies = [
{ name = "langchain-core" },