From 9100f2c68203819ecd69ad71862b47336e8f725a Mon Sep 17 00:00:00 2001 From: Quanzheng Long Date: Wed, 17 Jun 2026 10:50:30 -0700 Subject: [PATCH] feat: add delta-channel-dump recovery script as examples to dump from Postgres for deltaChannel rollback (#8109) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a standalone recovery tool at `examples/delta-channel-dump/` for operators rolling back from langgraph >= 1.2 / deepagents 0.6.x to an older runtime that does not understand `EXT_DELTA_SNAPSHOT` msgpack blobs. Without this, channels like `messages` can appear empty after rollback because the old reducer does not decode delta snapshots. `dump.py` connects directly to Postgres, walks the checkpoint parent chain (mirroring `aget_delta_channel_history`), decodes msgpack blobs via `ormsgpack` with a hand-rolled EXT 0–7 hook (no langgraph imports), and emits JSON with per-channel `seed` + oldest-to-newest `writes`. The operator reduces and re-applies via `update_state` manually. ## Scope - Postgres checkpointer only - Deps: `psycopg[binary]`, `ormsgpack` (not added to repo `pyproject.toml` — operator installs at runtime) - No reducer application, no DB writes, no AES/custom encryption support (fails loudly on encrypted blob types) ## How verified - E2E against local Postgres: 15 runs of `delta_channel_messages_freq` (`snapshot_frequency=10`) - Asserted `delta_kind == "snapshot"`, `len(seed) == 10`, `len(writes) == 5`, message ids `ai-0`..`ai-14` - Confirmed portability: fresh venv with only `pip install "psycopg[binary]" ormsgpack` produced identical output --- examples/delta-channel-dump/README.md | 132 ++++++++ examples/delta-channel-dump/dump.py | 469 ++++++++++++++++++++++++++ 2 files changed, 601 insertions(+) create mode 100644 examples/delta-channel-dump/README.md create mode 100755 examples/delta-channel-dump/dump.py diff --git a/examples/delta-channel-dump/README.md b/examples/delta-channel-dump/README.md new file mode 100644 index 000000000..f760f0850 --- /dev/null +++ b/examples/delta-channel-dump/README.md @@ -0,0 +1,132 @@ +# delta-channel-dump + +Recover messages (and other channels) from a Postgres-backed LangGraph thread +written by **langgraph >= 1.2** (DeltaChannel format) — including **LangGraph +Server / langgraph-api** deployments on the Postgres runtime, **deepagents +0.6.x**, or any OSS app using `PostgresSaver` — before rolling back to an older +runtime such as **deepagents 0.5.x / langgraph < 1.2**. + +langgraph-api uses the same `checkpoints` / `checkpoint_blobs` / `checkpoint_writes` +schema as OSS `checkpoint-postgres`; this script reads those tables directly. + +On the older runtime, `add_messages` does not understand the `EXT_DELTA_SNAPSHOT` +msgpack ext code and silently returns an empty list for affected channels. This +tool reads the raw checkpoint blobs from Postgres and emits JSON you can inspect +and re-apply via `update_state` (LangGraph Server SDK) or `graph.update_state` +(OSS). + +## Install + +```bash +pip install "psycopg[binary]" ormsgpack +``` + +## Run + +```bash +export DATABASE_URI=postgres://user:pass@host:5432/dbname +python3 dump.py \ + --thread-id \ + --channel messages \ + [--channel files ...] \ + [--checkpoint-id ] \ + [--checkpoint-ns ""] \ + --output recovery.json +``` + +- `--thread-id` (required): thread UUID +- `--channel` (required, repeatable): channel names to recover +- `--checkpoint-id` (optional): target checkpoint; defaults to latest +- `--checkpoint-ns` (optional): namespace; defaults to `""` +- `--database-uri` (optional): Postgres URI; defaults to `DATABASE_URI` env var +- `--output` (optional): output file; defaults to stdout + +## Output + +```json +{ + "thread_id": "...", + "checkpoint_ns": "", + "target_checkpoint_id": "...", + "parent_checkpoint_id": "...", + "channels": { + "messages": { + "delta_kind": "snapshot", + "seed_checkpoint_id": "...", + "seed_version": "...", + "seed": [{ "type": "ai", "content": "...", "id": "ai-0" }], + "writes": [ + { + "checkpoint_id": "...", + "task_id": "...", + "idx": 0, + "value": [{ "type": "ai", "content": "...", "id": "ai-10" }] + } + ] + } + } +} +``` + +`delta_kind` is one of: + +- `snapshot` — DeltaChannel snapshot blob (`channel_values[ch] == true`) +- `legacy_plain` — pre-DeltaChannel inline or blob value +- `no_seed` — walked to root without finding a populated ancestor + +`writes` are ordered oldest-to-newest (the order a reducer would replay them). + +## Reducing back to a single list + +For deepagents-style messages, combine seed and writes, then deduplicate: + +```python +import json + +data = json.load(open("recovery.json")) +ch = data["channels"]["messages"] +messages = list(ch["seed"] or []) +for w in ch["writes"]: + messages.extend(w["value"] or []) + +# Dedup by id, keep last; drop RemoveMessage tombstones +by_id = {} +for m in messages: + if isinstance(m, dict) and m.get("type") == "remove": + by_id.pop(m.get("id"), None) + elif isinstance(m, dict) and m.get("id"): + by_id[m["id"]] = m + else: + by_id[id(m)] = m +reduced = list(by_id.values()) +``` + +This approximates `_messages_delta_reducer` semantics; adjust for your graph. + +## Re-applying + +```python +from langgraph_sdk import get_client + +client = get_client(url="http://localhost:8123") +await client.threads.update_state( + thread_id, + values={"messages": reduced}, +) +``` + +Review the recovered JSON before calling `update_state`. This tool is +read-only and intentionally does not mutate the database. + +## Scope / non-goals (v1) + +- **Postgres only** — OSS `PostgresSaver` or langgraph-api Postgres runtime; not + inmem, gRPC core, Mongo, or Redis checkpointer backends +- **No AES decryption** (`LANGGRAPH_AES_KEY`) or custom encryption +- **No reducer** — raw seed + writes only +- **No automatic `update_state`** — operator applies manually + +## Copying + +`dump.py` is self-contained. Copy it anywhere; only `psycopg[binary]` and +`ormsgpack` are required at runtime. No langgraph imports. diff --git a/examples/delta-channel-dump/dump.py b/examples/delta-channel-dump/dump.py new file mode 100755 index 000000000..bd23de0b0 --- /dev/null +++ b/examples/delta-channel-dump/dump.py @@ -0,0 +1,469 @@ +#!/usr/bin/env python3 +"""Recover delta-channel state from a Postgres-backed LangGraph thread. + +Works with OSS ``PostgresSaver`` and LangGraph Server / langgraph-api on the +Postgres runtime (same checkpoint schema). Use after rolling back from +langgraph >= 1.2 / deepagents 0.6.x to an older runtime that does not +understand ``EXT_DELTA_SNAPSHOT`` msgpack blobs. The script walks the checkpoint +parent chain, decodes msgpack blobs, and emits a JSON dump of per-channel +``seed`` plus oldest-to-newest ``writes``. Apply the recovered values manually +via ``client.threads.update_state(...)`` (Server) or ``graph.update_state`` +(OSS). + +Install:: + + pip install "psycopg[binary]" ormsgpack + +Run:: + + export DATABASE_URI=postgres://... + python3 dump.py --thread-id --channel messages --output recovery.json + +Scope (v1): Postgres only; no AES/custom encryption; no reducer application. +""" + +from __future__ import annotations + +import argparse +import base64 +import json +import os +import sys +import uuid +from typing import Any + +import ormsgpack +import psycopg + +# LangGraph msgpack EXT type codes (langgraph/checkpoint/serde/jsonplus.py). +EXT_CONSTRUCTOR_SINGLE_ARG = 0 +EXT_CONSTRUCTOR_POS_ARGS = 1 +EXT_CONSTRUCTOR_KW_ARGS = 2 +EXT_METHOD_SINGLE_ARG = 3 +EXT_PYDANTIC_V1 = 4 +EXT_PYDANTIC_V2 = 5 +EXT_NUMPY_ARRAY = 6 +EXT_DELTA_SNAPSHOT = 7 + +_MSGPACK_OPTION = ormsgpack.OPT_NON_STR_KEYS + + +def ext_hook(code: int, data: bytes) -> Any: + """Decode LangGraph msgpack EXT payloads to JSON-friendly Python values.""" + if code == EXT_DELTA_SNAPSHOT: + inner = ormsgpack.unpackb(data, ext_hook=ext_hook, option=_MSGPACK_OPTION) + return {"__delta_snapshot__": inner} + if code == EXT_CONSTRUCTOR_SINGLE_ARG: + try: + tup = ormsgpack.unpackb(data, ext_hook=ext_hook, option=_MSGPACK_OPTION) + if tup[0] == "uuid" and tup[1] == "UUID": + hex_ = tup[2] + return ( + f"{hex_[:8]}-{hex_[8:12]}-{hex_[12:16]}-" + f"{hex_[16:20]}-{hex_[20:]}" + ) + return tup[2] + except Exception: + return None + if code == EXT_CONSTRUCTOR_POS_ARGS: + try: + tup = ormsgpack.unpackb(data, ext_hook=ext_hook, option=_MSGPACK_OPTION) + if tup[0] == "langgraph.types" and tup[1] == "Send": + args = tup[2] + if len(args) == 2: + return {"__send__": {"node": args[0], "arg": args[1]}} + return { + "__send__": { + "node": args[0], + "arg": args[1], + "timeout": args[2], + } + } + return tup[2] + except Exception: + return None + if code in (EXT_CONSTRUCTOR_KW_ARGS, EXT_METHOD_SINGLE_ARG): + try: + tup = ormsgpack.unpackb(data, ext_hook=ext_hook, option=_MSGPACK_OPTION) + return tup[2] + except Exception: + return None + if code in (EXT_PYDANTIC_V1, EXT_PYDANTIC_V2): + try: + tup = ormsgpack.unpackb(data, ext_hook=ext_hook, option=_MSGPACK_OPTION) + return tup[2] + except Exception: + return None + if code == EXT_NUMPY_ARRAY: + try: + dtype_str, shape, order, buf = ormsgpack.unpackb( + data, ext_hook=ext_hook, option=_MSGPACK_OPTION + ) + return { + "__numpy_array__": { + "dtype": dtype_str, + "shape": shape, + "order": order, + "data_b64": base64.b64encode(buf).decode("ascii"), + } + } + except Exception: + return None + return None + + +def decode_blob(blob_type: str, blob_bytes: bytes | None) -> Any: + if blob_type in ("empty", "null") or blob_bytes is None: + return None + if blob_type == "msgpack": + return ormsgpack.unpackb( + blob_bytes, ext_hook=ext_hook, option=_MSGPACK_OPTION + ) + if blob_type in ("bytes", "bytearray"): + return base64.b64encode(blob_bytes).decode("ascii") + raise RuntimeError( + f"Unknown blob type {blob_type!r}. " + "AES/custom-encrypted deployments are out of scope for v1." + ) + + +def delta_unwrap(value: Any) -> Any: + if isinstance(value, dict) and "__delta_snapshot__" in value: + return value["__delta_snapshot__"] + return value + + +def json_default(obj: Any) -> Any: + if isinstance(obj, (bytes, bytearray)): + return base64.b64encode(bytes(obj)).decode("ascii") + if isinstance(obj, uuid.UUID): + return str(obj) + raise TypeError(f"Object of type {type(obj).__name__} is not JSON serializable") + + +def resolve_target_checkpoint_id( + conn: psycopg.Connection[Any], + thread_id: str, + checkpoint_ns: str, + checkpoint_id: str | None, +) -> str: + if checkpoint_id is not None: + return checkpoint_id + row = conn.execute( + """ + SELECT checkpoint_id::text + FROM checkpoints + WHERE thread_id = %s AND checkpoint_ns = %s + ORDER BY checkpoint_id DESC + LIMIT 1 + """, + (thread_id, checkpoint_ns), + ).fetchone() + if row is None: + raise SystemExit( + f"No checkpoints found for thread_id={thread_id!r} " + f"checkpoint_ns={checkpoint_ns!r}" + ) + return row[0] + + +def _load_checkpoint( + conn: psycopg.Connection[Any], + thread_id: str, + checkpoint_ns: str, + checkpoint_id: str, +) -> tuple[dict[str, Any], str | None] | None: + row = conn.execute( + """ + SELECT checkpoint, parent_checkpoint_id::text + FROM checkpoints + WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = %s + """, + (thread_id, checkpoint_ns, checkpoint_id), + ).fetchone() + if row is None: + return None + return row[0], row[1] + + +def _load_seed( + conn: psycopg.Connection[Any], + *, + thread_id: str, + checkpoint_ns: str, + channel: str, + checkpoint_id: str, + channel_values: dict[str, Any], + channel_versions: dict[str, str], +) -> dict[str, Any]: + cv = channel_values[channel] + version = channel_versions.get(channel) + if cv is True: + blob_row = conn.execute( + """ + SELECT type, blob + FROM checkpoint_blobs + WHERE thread_id = %s AND checkpoint_ns = %s + AND channel = %s AND version = %s + """, + (thread_id, checkpoint_ns, channel, version), + ).fetchone() + seed_value = None + if blob_row is not None and blob_row[0] != "empty": + seed_value = delta_unwrap(decode_blob(blob_row[0], blob_row[1])) + return { + "delta_kind": "snapshot", + "seed_checkpoint_id": checkpoint_id, + "seed_version": version, + "seed": seed_value, + } + if isinstance(cv, (int, float, str, bool)) or cv is None: + return { + "delta_kind": "legacy_plain", + "seed_checkpoint_id": checkpoint_id, + "seed_version": version, + "seed": cv, + } + blob_row = conn.execute( + """ + SELECT type, blob + FROM checkpoint_blobs + WHERE thread_id = %s AND checkpoint_ns = %s + AND channel = %s AND version = %s + """, + (thread_id, checkpoint_ns, channel, version), + ).fetchone() + seed_value = cv if blob_row is None else decode_blob(blob_row[0], blob_row[1]) + return { + "delta_kind": "legacy_plain", + "seed_checkpoint_id": checkpoint_id, + "seed_version": version, + "seed": seed_value, + } + + +def _load_writes_for_checkpoint( + conn: psycopg.Connection[Any], + *, + thread_id: str, + checkpoint_ns: str, + checkpoint_id: str, + channel: str, +) -> list[dict[str, Any]]: + """Load writes for one checkpoint, newest-first by ``(task_id, idx)``. + + ``walk_channel`` reverses the accumulated flat list before returning; + DESC here yields oldest-first within each checkpoint in the final output, + matching ``PostgresSaver._build_delta_channels_writes_history``. + """ + rows = conn.execute( + """ + SELECT task_id::text, idx, type, blob + FROM checkpoint_writes + WHERE thread_id = %s AND checkpoint_ns = %s + AND checkpoint_id = %s AND channel = %s + ORDER BY task_id DESC, idx DESC + """, + (thread_id, checkpoint_ns, checkpoint_id, channel), + ).fetchall() + return [ + { + "checkpoint_id": checkpoint_id, + "task_id": task_id, + "idx": idx, + "value": decode_blob(blob_type, blob_bytes), + } + for task_id, idx, blob_type, blob_bytes in rows + ] + + +def walk_channel( + conn: psycopg.Connection[Any], + *, + thread_id: str, + checkpoint_ns: str, + start_checkpoint_id: str | None, + channel: str, +) -> dict[str, Any]: + """Walk one channel's parent chain from target.parent backward to seed.""" + chain_writes_newest_first: list[dict[str, Any]] = [] + cur = start_checkpoint_id + while cur is not None: + loaded = _load_checkpoint(conn, thread_id, checkpoint_ns, cur) + if loaded is None: + break + checkpoint_json, parent_id = loaded + channel_values = checkpoint_json.get("channel_values") or {} + channel_versions = checkpoint_json.get("channel_versions") or {} + + chain_writes_newest_first.extend( + _load_writes_for_checkpoint( + conn, + thread_id=thread_id, + checkpoint_ns=checkpoint_ns, + checkpoint_id=cur, + channel=channel, + ) + ) + + if channel in channel_values: + seed = _load_seed( + conn, + thread_id=thread_id, + checkpoint_ns=checkpoint_ns, + channel=channel, + checkpoint_id=cur, + channel_values=channel_values, + channel_versions=channel_versions, + ) + seed["writes"] = list(reversed(chain_writes_newest_first)) + return seed + + cur = parent_id + + return { + "delta_kind": "no_seed", + "seed_checkpoint_id": None, + "seed_version": None, + "seed": None, + "writes": list(reversed(chain_writes_newest_first)), + } + + +def walk_parent_chain( + conn: psycopg.Connection[Any], + *, + thread_id: str, + checkpoint_ns: str, + target_checkpoint_id: str, + channels: list[str], +) -> tuple[str | None, dict[str, dict[str, Any]]]: + """Walk parent chain and return per-channel seed + writes (oldest-first).""" + parent_row = conn.execute( + """ + SELECT parent_checkpoint_id::text + FROM checkpoints + WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = %s + """, + (thread_id, checkpoint_ns, target_checkpoint_id), + ).fetchone() + if parent_row is None: + raise SystemExit( + f"Checkpoint {target_checkpoint_id!r} not found for thread {thread_id!r}" + ) + parent_checkpoint_id = parent_row[0] + + result = { + ch: walk_channel( + conn, + thread_id=thread_id, + checkpoint_ns=checkpoint_ns, + start_checkpoint_id=parent_checkpoint_id, + channel=ch, + ) + for ch in channels + } + return parent_checkpoint_id, result + + +def build_output( + *, + thread_id: str, + checkpoint_ns: str, + target_checkpoint_id: str, + parent_checkpoint_id: str | None, + channels: dict[str, dict[str, Any]], +) -> dict[str, Any]: + return { + "thread_id": thread_id, + "checkpoint_ns": checkpoint_ns, + "target_checkpoint_id": target_checkpoint_id, + "parent_checkpoint_id": parent_checkpoint_id, + "channels": channels, + } + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Recover delta-channel seed + writes from Postgres checkpoint data." + ), + ) + parser.add_argument("--thread-id", required=True, help="Thread UUID") + parser.add_argument( + "--channel", + action="append", + required=True, + dest="channels", + help="Channel name (repeatable)", + ) + parser.add_argument( + "--checkpoint-id", + default=None, + help="Target checkpoint UUID (default: latest for thread)", + ) + parser.add_argument( + "--checkpoint-ns", + default="", + help='Checkpoint namespace (default: "")', + ) + parser.add_argument( + "--database-uri", + default=os.environ.get("DATABASE_URI"), + help="Postgres URI (default: DATABASE_URI env var)", + ) + parser.add_argument( + "--output", + default="-", + help="Output JSON file path (default: stdout)", + ) + return parser.parse_args(argv) + + +def main(argv: list[str] | None = None) -> int: + args = parse_args(argv) + if not args.database_uri: + print( + "error: --database-uri or DATABASE_URI is required", + file=sys.stderr, + ) + return 2 + + thread_id = str(uuid.UUID(args.thread_id)) + channels = list(dict.fromkeys(args.channels)) + + with psycopg.connect(args.database_uri) as conn: + target_checkpoint_id = resolve_target_checkpoint_id( + conn, + thread_id, + args.checkpoint_ns, + args.checkpoint_id, + ) + parent_checkpoint_id, channel_data = walk_parent_chain( + conn, + thread_id=thread_id, + checkpoint_ns=args.checkpoint_ns, + target_checkpoint_id=target_checkpoint_id, + channels=channels, + ) + + output = build_output( + thread_id=thread_id, + checkpoint_ns=args.checkpoint_ns, + target_checkpoint_id=target_checkpoint_id, + parent_checkpoint_id=parent_checkpoint_id, + channels=channel_data, + ) + payload = json.dumps(output, indent=2, default=json_default) + + if args.output == "-": + print(payload) + else: + with open(args.output, "w", encoding="utf-8") as f: + f.write(payload) + f.write("\n") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())