langchain-ai--langgraph/libs/langgraph/tests/test_delta_channel_supersteps_bound.py

"""Tests for the supersteps-since-last-snapshot bound on DeltaChannel.

Validates that a delta channel which stops receiving writes is still
force-snapshotted after DELTA_MAX_SUPERSTEPS_SINCE_SNAPSHOT supersteps,
preventing unbounded ancestor walks.
"""

from typing import Annotated, Any
from unittest.mock import patch

import pytest
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.checkpoint.serde.types import _DeltaSnapshot
from typing_extensions import TypedDict

from langgraph.channels.delta import DeltaChannel
from langgraph.graph import END, START, StateGraph
from langgraph.pregel._checkpoint import delta_channels_to_snapshot

pytestmark = pytest.mark.anyio


def _simple_reducer(current: list, updates: list) -> list:
    """Flatten updates into current list (each update is itself a list)."""
    result = list(current)
    for u in updates:
        if isinstance(u, list):
            result.extend(u)
        else:
            result.append(u)
    return result


def _build_two_channel_graph(
    checkpointer: InMemorySaver,
    *,
    freq_a: int = 10_000,
    freq_b: int = 10_000,
    n_loops: int = 1,
) -> Any:
    """Graph with two delta channels A and B.

    The node only writes to channel A; B is never written by the node.
    `n_loops` controls how many supersteps the graph runs (via chained nodes).
    """
    ch_a = DeltaChannel(_simple_reducer, list, snapshot_frequency=freq_a)
    ch_b = DeltaChannel(_simple_reducer, list, snapshot_frequency=freq_b)
    State = TypedDict(  # noqa: UP013
        "State",
        {"a": Annotated[list, ch_a], "b": Annotated[list, ch_b]},
    )  # type: ignore[call-overload]

    builder = StateGraph(State)

    for i in range(n_loops):
        name = f"step_{i}"

        def node_fn(state: dict, _i: int = i) -> dict:
            return {"a": [f"a-val-{_i}"]}

        builder.add_node(name, node_fn)
        if i == 0:
            builder.add_edge(START, name)
        else:
            builder.add_edge(f"step_{i - 1}", name)
        if i == n_loops - 1:
            builder.add_edge(name, END)

    return builder.compile(checkpointer=checkpointer)


async def test_forced_snapshot_single_run() -> None:
    """A single invoke with enough supersteps triggers snapshot on the
    unwritten channel B via the supersteps bound."""
    max_ss = 3
    with patch(
        "langgraph.pregel._checkpoint.DELTA_MAX_SUPERSTEPS_SINCE_SNAPSHOT",
        max_ss,
    ):
        saver = InMemorySaver()
        graph = _build_two_channel_graph(saver, n_loops=4)
        config = {"configurable": {"thread_id": "single-run-ss"}}

        graph.invoke({"a": ["seed-a"], "b": ["seed-b"]}, config)

        head = saver.get_tuple(config)
        assert head is not None
        assert isinstance(head.checkpoint["channel_values"].get("b"), _DeltaSnapshot), (
            "Channel B should have been force-snapshotted via supersteps bound"
        )

        state = graph.get_state(config)
        assert state.values["b"] == ["seed-b"]
        assert "seed-a" in state.values["a"]


async def test_forced_snapshot_accumulates_across_runs() -> None:
    """Supersteps counter for an unwritten channel persists across separate
    invoke() calls. After enough runs, the channel is force-snapshotted."""
    max_ss = 5
    with patch(
        "langgraph.pregel._checkpoint.DELTA_MAX_SUPERSTEPS_SINCE_SNAPSHOT",
        max_ss,
    ):
        saver = InMemorySaver()
        graph = _build_two_channel_graph(saver, n_loops=1)
        config = {"configurable": {"thread_id": "multi-run-ss"}}

        graph.invoke({"a": ["init-a"], "b": ["init-b"]}, config)

        for i in range(1, 6):
            graph.invoke({"a": [f"run-{i}"]}, config)

            head = saver.get_tuple(config)
            assert head is not None
            counters = head.metadata.get("counters_since_delta_snapshot", {})
            b_counters = counters.get("b", (0, 0))

            if b_counters == (0, 0):
                assert isinstance(
                    head.checkpoint["channel_values"].get("b"), _DeltaSnapshot
                ), f"Run {i}: counter reset but no snapshot blob for B"
                break
        else:
            pytest.fail("Channel B was never force-snapshotted after multiple runs")

        state = graph.get_state(config)
        assert state.values["b"] == ["init-b"]
        assert "init-a" in state.values["a"]


async def test_predicate_fires_on_supersteps_overflow() -> None:
    """Unit test: delta_channels_to_snapshot fires when supersteps >= MAX
    even when updates == 0."""
    ch = DeltaChannel(_simple_reducer, list, snapshot_frequency=10_000)
    ch.key = "x"
    ch_instance = ch.from_checkpoint(None)

    channels = {"x": ch_instance}
    counters: dict[str, tuple[int, int]] = {"x": (0, 5000)}

    result = delta_channels_to_snapshot(channels, counters)
    assert "x" in result

    counters_below: dict[str, tuple[int, int]] = {"x": (0, 4999)}
    result2 = delta_channels_to_snapshot(channels, counters_below)
    assert "x" not in result2


async def test_counter_reset_after_supersteps_snapshot() -> None:
    """After the supersteps bound triggers a snapshot, the counters for
    that channel reset. Verify by using a bound higher than one run's
    supersteps so we can see the counter in an intermediate state."""
    max_ss = 15
    with patch(
        "langgraph.pregel._checkpoint.DELTA_MAX_SUPERSTEPS_SINCE_SNAPSHOT",
        max_ss,
    ):
        saver = InMemorySaver()
        graph = _build_two_channel_graph(saver, n_loops=4)
        config = {"configurable": {"thread_id": "counter-reset"}}

        graph.invoke({"a": ["seed-a"], "b": ["seed-b"]}, config)

        head = saver.get_tuple(config)
        assert head is not None
        counters = head.metadata.get("counters_since_delta_snapshot", {})
        b_counters = counters.get("b", (0, 0))
        run1_supersteps = b_counters[1]
        assert run1_supersteps > 0, "Should have some supersteps"
        assert b_counters[0] == 1, "B written once (input step)"

        graph.invoke({"a": ["more-a"]}, config)
        head2 = saver.get_tuple(config)
        assert head2 is not None
        counters2 = head2.metadata.get("counters_since_delta_snapshot", {})
        b_counters2 = counters2.get("b", (0, 0))
        run2_supersteps = b_counters2[1]
        assert run2_supersteps > run1_supersteps, "Supersteps should accumulate"
        assert b_counters2[0] == 1, "B written once total (only original input)"

        graph.invoke({"a": ["even-more"]}, config)
        head3 = saver.get_tuple(config)
        assert head3 is not None
        assert isinstance(
            head3.checkpoint["channel_values"].get("b"), _DeltaSnapshot
        ), "B should have snapshotted at supersteps >= max_ss"
        counters3 = head3.metadata.get("counters_since_delta_snapshot", {})
        b_counters3 = counters3.get("b", (0, 0))
        assert b_counters3[1] < max_ss, (
            f"After snapshot, supersteps should have reset, got {b_counters3}"
        )

        state = graph.get_state(config)
        assert state.values["b"] == ["seed-b"]