mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 13:16:24 +00:00
d9e6341e5b
* fix(checkpoints): isolate cached history payloads from state readers Copy nested snapshot seeds and pending-write payloads at both cache boundaries. Keep user copy callbacks outside the registry lock so readers cannot poison later checkpoint materialization or deadlock by re-entering the cache. Cover sync/async state accessor reads, nested message payloads, disabled-cache controls and callback lock ownership. * docs(checkpoints): document synchronous history copy costs --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
117 lines
4.7 KiB
Python
117 lines
4.7 KiB
Python
"""Returned checkpoint snapshots must never mutate cached history payloads."""
|
|
|
|
from copy import deepcopy
|
|
from typing import Annotated, TypedDict
|
|
|
|
import pytest
|
|
from langgraph.channels import DeltaChannel
|
|
from langgraph.checkpoint.memory import InMemorySaver
|
|
from langgraph.graph import StateGraph
|
|
|
|
from deerflow.runtime.checkpoint_cache.memory import MemoryCheckpointHistoryCache
|
|
from deerflow.runtime.checkpoint_state import CheckpointStateAccessor
|
|
from deerflow.runtime.checkpointer.cached_saver import CachedHistorySaver
|
|
|
|
|
|
def _fold(state, writes):
|
|
result = list(state)
|
|
for batch in writes:
|
|
result.extend(batch)
|
|
return result
|
|
|
|
|
|
def _stack(frequency, enabled):
|
|
class State(TypedDict):
|
|
items: Annotated[list[dict], DeltaChannel(_fold, snapshot_frequency=frequency)]
|
|
|
|
inner = InMemorySaver()
|
|
cache = MemoryCheckpointHistoryCache(max_entries=32 if enabled else 0)
|
|
saver = CachedHistorySaver(inner, cache, key_prefix="payload-isolation")
|
|
builder = StateGraph(State)
|
|
for step in range(4):
|
|
builder.add_node(str(step), lambda state: {"items": [{"nested": {"value": "saved"}}]})
|
|
if step:
|
|
builder.add_edge(str(step - 1), str(step))
|
|
builder.set_entry_point("0")
|
|
builder.set_finish_point("3")
|
|
graph = builder.compile(checkpointer=saver)
|
|
return graph, CheckpointStateAccessor.bind(graph, saver, mode="delta"), cache
|
|
|
|
|
|
def _mutate(snapshot):
|
|
original = deepcopy(snapshot.values)
|
|
snapshot.values["items"][0]["nested"]["value"] = "tampered"
|
|
snapshot.values["items"].append({"nested": {"value": "injected"}})
|
|
return original
|
|
|
|
|
|
@pytest.mark.parametrize("frequency", [2, 100], ids=["snapshot-seed", "pending-writes"])
|
|
@pytest.mark.parametrize("enabled", [True, False], ids=["cached", "uncached-control"])
|
|
def test_sync_graph_snapshots_are_isolated(frequency, enabled):
|
|
graph, accessor, cache = _stack(frequency, enabled)
|
|
config = {"configurable": {"thread_id": "sync-payload-isolation"}}
|
|
graph.invoke({"items": [{"nested": {"value": "input"}}]}, config)
|
|
history = accessor.history(config)
|
|
checked = 0
|
|
for checkpoint in history:
|
|
snapshot = accessor.get(checkpoint.config)
|
|
if not snapshot.values.get("items"):
|
|
continue
|
|
original = _mutate(snapshot)
|
|
assert accessor.get(checkpoint.config).values == original
|
|
checked += 1
|
|
assert checked >= 4
|
|
assert cache.stats().hits > 0 if enabled else cache.stats().entries == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("frequency", [2, 100], ids=["snapshot-seed", "pending-writes"])
|
|
@pytest.mark.parametrize("enabled", [True, False], ids=["cached", "uncached-control"])
|
|
async def test_async_graph_snapshots_are_isolated(frequency, enabled):
|
|
graph, accessor, cache = _stack(frequency, enabled)
|
|
config = {"configurable": {"thread_id": "async-payload-isolation"}}
|
|
await graph.ainvoke({"items": [{"nested": {"value": "input"}}]}, config)
|
|
history = await accessor.ahistory(config)
|
|
checked = 0
|
|
for checkpoint in history:
|
|
snapshot = await accessor.aget(checkpoint.config)
|
|
if not snapshot.values.get("items"):
|
|
continue
|
|
original = _mutate(snapshot)
|
|
assert (await accessor.aget(checkpoint.config)).values == original
|
|
checked += 1
|
|
assert checked >= 4
|
|
assert cache.stats().hits > 0 if enabled else cache.stats().entries == 0
|
|
|
|
|
|
@pytest.mark.parametrize("direction", ["ingest", "read"])
|
|
def test_cache_copies_nested_write_payloads_and_snapshot_seeds(direction):
|
|
from langchain_core.messages import HumanMessage
|
|
from langgraph.checkpoint.serde.types import _DeltaSnapshot
|
|
|
|
cache = MemoryCheckpointHistoryCache()
|
|
seed = _DeltaSnapshot([HumanMessage(content="seed", id="seed", additional_kwargs={"nested": ["original"]})])
|
|
entry = {"seed": seed, "writes": [("task", "messages", [HumanMessage(content="write", id="write", additional_kwargs={"nested": ["original"]})])]}
|
|
original = deepcopy(entry)
|
|
cache.set_many({"key": entry})
|
|
exposed = entry if direction == "ingest" else cache.get_many(["key"])["key"]
|
|
exposed["seed"].value[0].additional_kwargs["nested"].append("corrupted")
|
|
exposed["writes"][0][2][0].content = "corrupted"
|
|
assert cache.get_many(["key"])["key"] == original
|
|
|
|
|
|
def test_payload_copying_runs_outside_registry_lock():
|
|
cache = MemoryCheckpointHistoryCache()
|
|
copies = []
|
|
|
|
class Probe:
|
|
def __deepcopy__(self, memo):
|
|
assert cache._lock.acquire(blocking=False), "copy callback ran under the cache lock"
|
|
cache._lock.release()
|
|
copies.append(True)
|
|
return Probe()
|
|
|
|
cache.set_many({"key": {"writes": [("task", "channel", Probe())]}})
|
|
cache.get_many(["key"])
|
|
assert len(copies) == 2
|