fix: honor resource auth actions

Co-authored-by: open-swe[bot] <open-swe@users.noreply.github.com>

libs/langgraph/langgraph/pregel/_checkpoint.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import uuid
-from collections.abc import Callable, Iterable, Mapping, Sequence
+from collections.abc import Callable, Mapping
 from datetime import datetime, timezone
 from typing import Any, cast
 
@@ -14,7 +14,6 @@ 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
@@ -71,26 +70,6 @@ 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,
@@ -123,6 +102,11 @@ 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

libs/langgraph/langgraph/pregel/main.py

@@ -108,6 +108,7 @@ 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
@@ -133,8 +134,6 @@ 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
@@ -1997,19 +1996,13 @@ class Pregel(
                         },
                     ),
                 )
-            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
-                        )
-
+            # 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
             apply_writes(
                 checkpoint,
                 channels,
@@ -2017,13 +2010,20 @@ 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,
-                updated_channels=updated_channels,
-                get_next_version=checkpointer.get_next_version,
-                channels_to_snapshot=channels_to_snapshot,
+                checkpoint, channels, step + 1, channels_to_snapshot=delta_snapshot
             )
             next_config = checkpointer.put(
                 checkpoint_config,
@@ -2038,11 +2038,7 @@ class Pregel(
                 ),
             )
             for task_id, task in zip(run_task_ids, run_tasks):
-                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)
+                # save push writes
                 if push_writes := [w for w in task.writes if w[0] == PUSH]:
                     checkpointer.put_writes(next_config, push_writes, task_id)
 
@@ -2459,19 +2455,15 @@ class Pregel(
                         },
                     ),
                 )
-            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
-                        )
-
+            # 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
             apply_writes(
                 checkpoint,
                 channels,
@@ -2479,14 +2471,22 @@ class Pregel(
                 checkpointer.get_next_version,
                 self.trigger_to_nodes,
             )
-            checkpoint = create_checkpoint(
-                checkpoint,
-                channels,
-                step + 1,
-                updated_channels=updated_channels,
-                get_next_version=checkpointer.get_next_version,
-                channels_to_snapshot=channels_to_snapshot,
+            # 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
+            )
+            # save checkpoint, after applying writes
             next_config = await checkpointer.aput(
                 checkpoint_config,
                 checkpoint,
@@ -2500,13 +2500,7 @@ class Pregel(
                 ),
             )
             for task_id, task in zip(run_task_ids, run_tasks):
-                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
-                        )
+                # save push writes
                 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)

libs/langgraph/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "langgraph"
-version = "1.2.7"
+version = "1.2.8"
 description = "Building stateful, multi-actor applications with LLMs"
 authors = []
 requires-python = ">=3.10"

libs/langgraph/tests/test_delta_channel_update_state.py

@@ -1,12 +1,14 @@
 """Tests for `update_state` / `aupdate_state` against `DeltaChannel`.
 
-Regression suite for deepagents#3774 and Postgres read-path compatibility:
-fresh-thread `update_state` must persist DeltaChannel state correctly.
+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.
 
-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.
+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.
 
 Coverage:
 
@@ -14,8 +16,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 shape on a fresh thread (single snapshotted head checkpoint)
-* head checkpoint snapshots updated DeltaChannels for Postgres read paths
+* state-history chain shape on a fresh thread (single self-contained update
+  checkpoint with the snapshot inline and no parent)
 """
 
 from typing import Annotated, Any
@@ -23,7 +25,6 @@ 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
@@ -95,7 +96,8 @@ async def test_aupdate_state_fresh_thread_delta_channel() -> None:
 
 
 def test_update_state_after_invoke_delta_channel() -> None:
-    """The non-fresh-thread path must keep working across snapshot changes."""
+    """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."""
     saver = InMemorySaver()
     graph = _build_graph(saver)
     config = {"configurable": {"thread_id": "after-invoke-sync"}}
@@ -134,8 +136,9 @@ async def test_aupdate_state_after_invoke_delta_channel() -> None:
 
 
 def test_consecutive_update_states_delta_channel() -> None:
-    """Two consecutive fresh-thread-style updates: the first creates a
-    snapshotted head; the second anchors on it. Both messages round-trip."""
+    """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."""
     saver = InMemorySaver()
     graph = _build_graph(saver)
     config = {"configurable": {"thread_id": "consecutive-sync"}}
@@ -209,8 +212,10 @@ 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 must accumulate
-    all N message writes in the snapshotted head state.
+    """`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.
 
     Explicit `task_id`s are required to disambiguate writes belonging to
     different `StateUpdate`s targeting the same node — otherwise both share
@@ -250,13 +255,14 @@ def test_bulk_update_state_multi_task_per_superstep_delta_channel() -> None:
 
 
 # ---------------------------------------------------------------------------
-# Public-API observation of fresh-thread checkpoint shape
+# Public-API observation of the forced-snapshot mechanism
 # ---------------------------------------------------------------------------
 
 
 def test_state_history_chain_after_fresh_update_state_delta_channel() -> None:
-    """Fresh-thread `update_state` on snapshotted DeltaChannels yields one
-    self-contained checkpoint (step=0, no parent, source='update')."""
+    """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."""
     saver = InMemorySaver()
     graph = _build_graph(saver)
     config = {"configurable": {"thread_id": "history-chain"}}
@@ -270,32 +276,9 @@ 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[0]
-
+    (update_snapshot,) = history
     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"
-    ]

libs/langgraph/uv.lock

@@ -1438,7 +1438,7 @@ wheels = [
 
 [[package]]
 name = "langgraph"
-version = "1.2.7"
+version = "1.2.8"
 source = { editable = "." }
 dependencies = [
     { name = "langchain-core" },

libs/prebuilt/uv.lock

@@ -285,7 +285,7 @@ wheels = [
 
 [[package]]
 name = "langgraph"
-version = "1.2.7"
+version = "1.2.8"
 source = { editable = "../langgraph" }
 dependencies = [
     { name = "langchain-core" },

libs/sdk-py/langgraph_sdk/auth/__init__.py

@@ -392,7 +392,7 @@ class _ResourceOn(typing.Generic[VCreate, VRead, VUpdate, VDelete, VSearch]):
     def __call__(
         self,
         *,
-        resources: str | Sequence[str],
+        resources: str | Sequence[str] | None = None,
         actions: str | Sequence[str] | None = None,
     ) -> Callable[
         [_ActionHandler[VCreate | VUpdate | VRead | VDelete | VSearch]],
@@ -416,25 +416,38 @@ class _ResourceOn(typing.Generic[VCreate, VRead, VUpdate, VDelete, VSearch]):
             _ActionHandler[VCreate | VUpdate | VRead | VDelete | VSearch],
         ]
     ):
-        if fn is not None:
-            _validate_handler(fn)
-            return typing.cast(
-                "_ActionHandler[VCreate | VUpdate | VRead | VDelete | VSearch]",
-                _register_handler(self.auth, self.resource, "*", fn),
-            )
-
-        def decorator(
+        def register(
             handler: _ActionHandler[VCreate | VUpdate | VRead | VDelete | VSearch],
         ) -> _ActionHandler[VCreate | VUpdate | VRead | VDelete | VSearch]:
             _validate_handler(handler)
-            return typing.cast(
-                "_ActionHandler[VCreate | VUpdate | VRead | VDelete | VSearch]",
-                _register_handler(self.auth, self.resource, "*", handler),
-            )
+            if isinstance(resources, str):
+                resource_list = [resources]
+            else:
+                resource_list = (
+                    list(resources) if resources is not None else [self.resource]
+                )
+            if resource_list != [self.resource]:
+                raise ValueError(
+                    f"Resource-specific decorator for {self.resource!r} cannot "
+                    f"register handlers for {resource_list!r}. Use @auth.on(...) "
+                    "for multiple resources."
+                )
+            if isinstance(actions, str):
+                action_list = [actions]
+            else:
+                action_list = list(actions) if actions is not None else ["*"]
+            for action in action_list:
+                _register_handler(self.auth, self.resource, action, handler)
+            return handler
 
-        # Accept keyword-only parameters for future filtering behavior; referenced to satisfy linters.
-        _ = resources, actions
-        return decorator
+        if fn is not None:
+            return register(
+                typing.cast(
+                    "_ActionHandler[VCreate | VUpdate | VRead | VDelete | VSearch]",
+                    fn,
+                )
+            )
+        return register
 
 
 class _AssistantsOn(

libs/sdk-py/tests/test_auth.py

@@ -0,0 +1,44 @@
+import pytest
+
+from langgraph_sdk import Auth
+
+
+def _handler():
+    async def handler(ctx, value):
+        return ctx is not None and value is not None
+
+    return handler
+
+
+def test_resource_decorator_registers_specific_actions():
+    auth = Auth()
+
+    handler = auth.on.threads(actions=["read", "search"])(_handler())
+
+    assert auth._handlers == {
+        ("threads", "read"): [handler],
+        ("threads", "search"): [handler],
+    }
+
+
+def test_resource_decorator_registers_single_action():
+    auth = Auth()
+
+    handler = auth.on.threads(actions="read")(_handler())
+
+    assert auth._handlers == {("threads", "read"): [handler]}
+
+
+def test_resource_decorator_without_actions_registers_resource_wildcard():
+    auth = Auth()
+
+    handler = auth.on.threads(_handler())
+
+    assert auth._handlers == {("threads", "*"): [handler]}
+
+
+def test_resource_decorator_rejects_mismatched_resources():
+    auth = Auth()
+
+    with pytest.raises(ValueError, match=r"Use @auth\.on"):
+        auth.on.threads(resources="assistants", actions="read")(_handler())

libs/sdk-py/uv.lock

@@ -298,7 +298,7 @@ wheels = [
 
 [[package]]
 name = "langgraph"
-version = "1.2.7"
+version = "1.2.8"
 source = { editable = "../langgraph" }
 dependencies = [
     { name = "langchain-core" },