diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index db29d440c..5f3e04af7 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -4,7 +4,7 @@ import threading from collections import defaultdict from collections.abc import Iterator, Sequence from contextlib import contextmanager -from typing import Any +from typing import Any, cast from langchain_core.runnables import RunnableConfig from langgraph.checkpoint.base import ( @@ -28,6 +28,7 @@ from langgraph.checkpoint.postgres import _internal from langgraph.checkpoint.postgres.base import ( SELECT_DELTA_COMBINED_SQL, BasePostgresSaver, + _DeltaCombinedRow, ) from langgraph.checkpoint.postgres.shallow import ShallowPostgresSaver @@ -475,7 +476,7 @@ class PostgresSaver(BasePostgresSaver): return self._build_delta_channel_writes_history( channel=channel, target_id=checkpoint_id, - rows=rows, + rows=cast("list[_DeltaCombinedRow]", rows), ) def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple: diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index f2031258d..9b0fd2100 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -4,7 +4,7 @@ import asyncio from collections import defaultdict from collections.abc import AsyncIterator, Iterator, Sequence from contextlib import asynccontextmanager -from typing import Any +from typing import Any, cast from langchain_core.runnables import RunnableConfig from langgraph.checkpoint.base import ( @@ -28,6 +28,7 @@ from langgraph.checkpoint.postgres import _ainternal from langgraph.checkpoint.postgres.base import ( SELECT_DELTA_COMBINED_SQL, BasePostgresSaver, + _DeltaCombinedRow, ) from langgraph.checkpoint.postgres.shallow import AsyncShallowPostgresSaver @@ -434,7 +435,7 @@ class AsyncPostgresSaver(BasePostgresSaver): return self._build_delta_channel_writes_history( channel=channel, target_id=checkpoint_id, - rows=rows, + rows=cast("list[_DeltaCombinedRow]", rows), ) async def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple: diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py index 098ac3c7a..b0f806665 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py @@ -4,7 +4,7 @@ import random import warnings from collections.abc import Sequence from importlib.metadata import version as get_version -from typing import Any, cast +from typing import Any, TypedDict, cast from langchain_core.runnables import RunnableConfig from langgraph.checkpoint.base import ( @@ -155,6 +155,29 @@ INSERT_CHECKPOINT_WRITES_SQL = """ ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) DO NOTHING """ + +class _DeltaCombinedRow(TypedDict, total=False): + """One row from `SELECT_DELTA_COMBINED_SQL` (a UNION ALL of three tables). + + Every row carries `_kind` ("p" / "w" / "b") plus whichever columns are + relevant for that kind; irrelevant columns are NULL and typed as `None`. + """ + + _kind: str # always present: "p", "w", or "b" + # checkpoint row ("p") + checkpoint_id: str | None + parent_checkpoint_id: str | None + ver: str | None + # write / blob rows ("w", "b") + type: str | None + blob: bytes | None + # write row only ("w") + task_id: str | None + idx: int | None + # blob row only ("b") + version: str | None + + # DeltaChannel reconstruction: one combined CTE+UNION ALL query per channel. # Bench (notes/delta_channel_query_bench.md) showed the prior recursive CTE # carried a hidden O(ancestors x blobs_in_thread) join; plain SELECTs are @@ -241,7 +264,7 @@ class BasePostgresSaver(BaseCheckpointSaver[str]): *, channel: str, target_id: str, - rows: Sequence[Any], + rows: Sequence[_DeltaCombinedRow], ) -> _ChannelWritesHistory: """Reconstruct one delta channel's history from the combined UNION ALL rows. @@ -264,16 +287,21 @@ class BasePostgresSaver(BaseCheckpointSaver[str]): for r in rows: kind = r["_kind"] if kind == "p": - cid = r["checkpoint_id"] + cid = cast(str, r["checkpoint_id"]) parent_of[cid] = r["parent_checkpoint_id"] ver_of[cid] = r["ver"] elif kind == "w": - cid = r["checkpoint_id"] + cid = cast(str, r["checkpoint_id"]) writes_by_cid.setdefault(cid, []).append( - (r["type"], r["blob"], r["task_id"], r["idx"]) + cast( + "tuple[str, bytes, str, int]", + (r["type"], r["blob"], r["task_id"], r["idx"]), + ) ) else: # kind == "b" - blob_by_ver[r["version"]] = (r["type"], r["blob"]) + blob_by_ver[cast(str, r["version"])] = cast( + "tuple[str, bytes]", (r["type"], r["blob"]) + ) # Sort writes within each checkpoint (task_id DESC, idx DESC) to match # the prior CTE ordering — newest write first per ancestor. @@ -281,10 +309,10 @@ class BasePostgresSaver(BaseCheckpointSaver[str]): ws.sort(key=lambda w: (w[2], w[3]), reverse=True) ancestors: list[str] = [] - cid = parent_of.get(target_id) - while cid is not None: - ancestors.append(cid) - cid = parent_of.get(cid) + cur_cid: str | None = parent_of.get(target_id) + while cur_cid is not None: + ancestors.append(cur_cid) + cur_cid = parent_of.get(cur_cid) if not ancestors: return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[]) diff --git a/libs/langgraph/langgraph/channels/binop.py b/libs/langgraph/langgraph/channels/binop.py index 82f7e940e..5735ac65d 100644 --- a/libs/langgraph/langgraph/channels/binop.py +++ b/libs/langgraph/langgraph/channels/binop.py @@ -32,7 +32,7 @@ def _get_overwrite(value: Any) -> tuple[bool, Any]: """Inspects the given value and returns (is_overwrite, overwrite_value).""" if isinstance(value, Overwrite): return True, value.value - if isinstance(value, dict) and set(value.keys()) == {OVERWRITE}: + if isinstance(value, dict) and len(value) == 1 and OVERWRITE in value: return True, value[OVERWRITE] return False, None diff --git a/libs/langgraph/langgraph/channels/delta.py b/libs/langgraph/langgraph/channels/delta.py index a061e21a2..cacfd4d2f 100644 --- a/libs/langgraph/langgraph/channels/delta.py +++ b/libs/langgraph/langgraph/channels/delta.py @@ -58,7 +58,7 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]): super().__init__(list) self.operator = operator self.snapshot_frequency = snapshot_frequency - self.value: Any = [] + self.value: Any = MISSING def __eq__(self, other: object) -> bool: if not isinstance(other, DeltaChannel): @@ -78,16 +78,15 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]): def is_snapshot_step(self, step: int) -> bool: """True if pregel should write a snapshot blob at this step.""" return ( - self.snapshot_frequency is not None and step % self.snapshot_frequency == 0 + self.snapshot_frequency is not None + and step > 0 + and step % self.snapshot_frequency == 0 ) def _clone_empty(self) -> Self: - new = self.__class__.__new__(self.__class__) - new.typ = self.typ - new.key = self.key - new.operator = self.operator - new.snapshot_frequency = self.snapshot_frequency - new.value = MISSING + new = self.__class__(self.operator, snapshot_frequency=self.snapshot_frequency) + new.typ = self.typ # typ may differ from list when set via Annotated injection + new.key = self.key # key is injected externally by the graph builder return new def copy(self) -> Self: diff --git a/libs/langgraph/langgraph/pregel/_checkpoint.py b/libs/langgraph/langgraph/pregel/_checkpoint.py index aea3ed583..b58dddfa2 100644 --- a/libs/langgraph/langgraph/pregel/_checkpoint.py +++ b/libs/langgraph/langgraph/pregel/_checkpoint.py @@ -126,7 +126,10 @@ def channels_from_checkpoint( ch: BaseChannel stored = checkpoint["channel_values"].get(k, MISSING) if _needs_replay(spec, stored) and saver is not None and config is not None: - assert isinstance(spec, DeltaChannel) + if not isinstance(spec, DeltaChannel): + raise TypeError( + f"Expected DeltaChannel for channel requiring replay, got {type(spec)}" + ) history = saver._get_channel_writes_history(config, k) replay_ch = spec.from_checkpoint(history.seed) replay_ch.replay_writes(history.writes) @@ -158,7 +161,10 @@ async def achannels_from_checkpoint( ch: BaseChannel stored = checkpoint["channel_values"].get(k, MISSING) if _needs_replay(spec, stored) and saver is not None and config is not None: - assert isinstance(spec, DeltaChannel) + if not isinstance(spec, DeltaChannel): + raise TypeError( + f"Expected DeltaChannel for channel requiring replay, got {type(spec)}" + ) history = await saver._aget_channel_writes_history(config, k) replay_ch = spec.from_checkpoint(history.seed) replay_ch.replay_writes(history.writes)