diff --git a/tmp_verify_send_blob.py b/tmp_verify_send_blob.py new file mode 100644 index 000000000..4b30434bf --- /dev/null +++ b/tmp_verify_send_blob.py @@ -0,0 +1,217 @@ +"""Verify that SEND values land in checkpoint_blobs (TASKS channel), +not just in checkpoint_writes. + +Strategy: + 1. Build a tiny graph: node "1" emits a Send to node "2" via + conditional edges. interrupt_before=["2"] freezes the loop after + super-step 1 completes (cp_1 persisted) but before "2" runs. + 2. Run graph.invoke against PostgresSaver. + 3. Read checkpoint_blobs raw via psycopg. + 4. Find the row with channel = '__pregel_tasks' and decode the blob + using the saver's serde. Assert it contains a non-empty list of + Send objects. + 5. As a control, also dump checkpoint_writes to confirm the *same* + value is independently present in the writes table. +""" + +from __future__ import annotations + +import operator +from typing import Annotated +from typing_extensions import TypedDict +from uuid import uuid4 + +from psycopg import Connection +from psycopg.rows import dict_row + +from langgraph.checkpoint.postgres import PostgresSaver +from langgraph.checkpoint.serde.types import TASKS +from langgraph.graph import START, StateGraph +from langgraph.types import Send + +DEFAULT_POSTGRES_URI = "postgres://postgres:postgres@localhost:5441/" + + +class State(TypedDict): + history: Annotated[list[str], operator.add] + + +def make_graph(checkpointer: PostgresSaver, *, with_interrupt: bool = True): + def node_one(state: State) -> State: + return {"history": ["1"]} + + def node_two(state: dict) -> State: + return {"history": [f"2:{state['payload']}"]} + + def fanout(state: State): + return [Send("two", {"payload": x}) for x in ("a", "b", "c")] + + builder = StateGraph(State) + builder.add_node("one", node_one) + builder.add_node("two", node_two) + builder.add_edge(START, "one") + builder.add_conditional_edges("one", fanout, ["two"]) + return builder.compile( + checkpointer=checkpointer, + interrupt_before=["two"] if with_interrupt else [], + ) + + +def main() -> None: + database = f"verify_{uuid4().hex[:12]}" + with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: + conn.execute(f"CREATE DATABASE {database}") + try: + uri = DEFAULT_POSTGRES_URI + database + with Connection.connect( + uri, autocommit=True, prepare_threshold=0, row_factory=dict_row + ) as conn: + saver = PostgresSaver(conn) + saver.setup() + + graph = make_graph(saver) + thread_id = "t1" + cfg = {"configurable": {"thread_id": thread_id}} + result = graph.invoke({"history": []}, cfg) + print("=== invoke result (interrupted before 'two') ===") + print(result) + print() + + # Inspect raw checkpoint_blobs. + rows = conn.execute( + """ + SELECT thread_id, checkpoint_ns, channel, version, type, + octet_length(blob) AS nbytes, blob + FROM checkpoint_blobs + WHERE thread_id = %s + ORDER BY channel, version + """, + (thread_id,), + ).fetchall() + + print(f"=== checkpoint_blobs ({len(rows)} rows) ===") + for r in rows: + preview = ( + "(empty)" + if r["blob"] is None + else f"{r['blob'][:60]!r}{'...' if len(r['blob']) > 60 else ''}" + ) + print( + f" channel={r['channel']!r:30s} v={r['version']!s:18s} " + f"type={r['type']:8s} nbytes={r['nbytes']!s:>5} blob={preview}" + ) + print() + + tasks_rows = [r for r in rows if r["channel"] == TASKS] + assert tasks_rows, ( + f"expected at least one row in checkpoint_blobs for {TASKS!r}, " + f"got channels={[r['channel'] for r in rows]}" + ) + + # Decode the latest TASKS blob using the saver's serde. + tasks_row = tasks_rows[-1] + decoded = saver.serde.loads_typed((tasks_row["type"], tasks_row["blob"])) + print(f"=== decoded {TASKS!r} blob ({tasks_row['type']}) ===") + print(f" type(decoded) = {type(decoded).__name__}") + print(f" len(decoded) = {len(decoded)}") + for i, v in enumerate(decoded): + print(f" [{i}] type={type(v).__name__} value={v!r}") + print() + + assert isinstance(decoded, list), ( + f"expected list of Sends, got {type(decoded).__name__}" + ) + assert len(decoded) == 3, f"expected 3 Sends (a, b, c), got {len(decoded)}" + for v in decoded: + assert isinstance(v, Send), ( + f"expected Send instance in blob, got {type(v).__name__}" + ) + assert v.node == "two" + print("✅ ASSERTION PASSED: checkpoint_blobs contains the Sends.") + print() + + # Control: confirm the same Sends are also in checkpoint_writes. + wrows = conn.execute( + """ + SELECT task_id, idx, channel, type, octet_length(blob) AS nbytes, blob + FROM checkpoint_writes + WHERE thread_id = %s AND channel = %s + ORDER BY task_id, idx + """, + (thread_id, TASKS), + ).fetchall() + print( + f"=== checkpoint_writes (channel={TASKS!r}, {len(wrows)} rows) ===" + ) + for r in wrows: + v = saver.serde.loads_typed((r["type"], r["blob"])) + print( + f" task={r['task_id']} idx={r['idx']} type={r['type']} " + f"nbytes={r['nbytes']} -> {v!r}" + ) + + assert len(wrows) == 3, ( + f"expected 3 writes-table rows for {TASKS!r}, got {len(wrows)}" + ) + print( + "✅ ASSERTION PASSED: checkpoint_writes also contains the Sends." + ) + print() + print( + "Conclusion (cp_1): SEND is double-stored. blobs row carries " + "the authoritative Send list; writes rows carry the same " + "payload as a redundant per-task record." + ) + + # ---------------------------------------------------------------- + # Now resume: let 'two' actually run, so the Sends get consumed. + # We expect the next checkpoint's TASKS blob to be cleared, + # proving the accumulate=False lifecycle. + # ---------------------------------------------------------------- + graph2 = make_graph(saver, with_interrupt=False) + result2 = graph2.invoke(None, cfg) + print() + print("=== invoke result (resumed; 'two' fanout completed) ===") + print(result2) + print() + + rows2 = conn.execute( + """ + SELECT channel, version, type, octet_length(blob) AS nbytes, blob + FROM checkpoint_blobs + WHERE thread_id = %s AND channel = %s + ORDER BY version + """, + (thread_id, TASKS), + ).fetchall() + print(f"=== checkpoint_blobs rows for {TASKS!r} across all cps ===") + for r in rows2: + decoded_v = saver.serde.loads_typed((r["type"], r["blob"])) + print( + f" v={r['version']} type={r['type']} nbytes={r['nbytes']} " + f"-> {decoded_v!r}" + ) + + assert len(rows2) >= 2, ( + "expected at least 2 versions of TASKS blob across the run" + ) + latest = saver.serde.loads_typed( + (rows2[-1]["type"], rows2[-1]["blob"]) + ) + assert latest == [], ( + f"expected latest TASKS blob to be cleared (==[]) " + f"after Sends were consumed, got {latest!r}" + ) + print() + print( + "✅ ASSERTION PASSED: after Sends were consumed, the latest " + "TASKS blob is [] (Topic with accumulate=False clears itself " + "via update(EMPTY_SEQ) in apply_writes line 304-311)." + ) + finally: + with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: + conn.execute(f"DROP DATABASE {database}") + + +if __name__ == "__main__": + main()