fix(langgraph): honor per-node cache_policy=None with set_node_defaults

Use a MISSING default on add_node so explicit cache_policy=None opts out
of graph-wide cache defaults instead of being treated as unset.
This commit is contained in:
Christian Bromann
2026-06-02 16:47:59 -07:00
parent 77a60e8101
commit 9ec5b16107
3 changed files with 105 additions and 3 deletions
+1
View File
@@ -88,6 +88,7 @@ class StateNodeSpec(Generic[NodeInputT, ContextT]):
input_schema: type[NodeInputT]
retry_policy: RetryPolicy | Sequence[RetryPolicy] | None
cache_policy: CachePolicy | None
cache_policy_opt_out: bool = False
is_error_handler: bool = False
error_handler_node: str | None = None
ends: tuple[str, ...] | dict[str, str] | None = EMPTY_SEQ
+17 -3
View File
@@ -292,7 +292,8 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
their own via `add_node(..., retry_policy=...)`. Also applies
to error-handler nodes.
cache_policy: Default cache policy for nodes that don't specify
their own via `add_node(..., cache_policy=...)`. Does **not**
their own via `add_node(..., cache_policy=...)`. Per-node
`cache_policy=None` opts out of this default. Does **not**
apply to error-handler nodes.
error_handler: Default error handler invoked when any regular node
raises and does not have its own `error_handler` set via
@@ -668,7 +669,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
metadata: dict[str, Any] | None = None,
input_schema: type[NodeInputT] | None = None,
retry_policy: RetryPolicy | Sequence[RetryPolicy] | None = None,
cache_policy: CachePolicy | None = None,
cache_policy: CachePolicy | None = MISSING,
error_handler: StateNode[Any, ContextT] | None = None,
destinations: dict[str, str] | tuple[str, ...] | None = None,
timeout: float | timedelta | TimeoutPolicy | None = None,
@@ -689,7 +690,10 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
retry_policy: The retry policy for the node.
If a sequence is provided, the first matching policy will be applied.
cache_policy: The cache policy for the node.
cache_policy: The cache policy for the node. Pass `None` explicitly to
disable caching for this node, including when graph defaults from
`set_node_defaults` would otherwise enable it. Omit the argument to
inherit graph defaults.
error_handler: Optional node-level error handler callable for this node.
destinations: Destinations that indicate where a node can route to.
@@ -765,6 +769,12 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
input_schema = cast(type[NodeInputT] | None, input_)
timeout = coerce_timeout_policy(timeout)
cache_policy_opt_out = False
if cache_policy is MISSING:
cache_policy = None
elif cache_policy is None:
cache_policy_opt_out = True
if not isinstance(node, str):
action = node
if isinstance(action, Runnable):
@@ -876,6 +886,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
input_schema=input_schema,
retry_policy=retry_policy,
cache_policy=cache_policy,
cache_policy_opt_out=cache_policy_opt_out,
error_handler_node=handler_node_name,
ends=ends,
defer=defer,
@@ -888,6 +899,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
input_schema=inferred_input_schema,
retry_policy=retry_policy,
cache_policy=cache_policy,
cache_policy_opt_out=cache_policy_opt_out,
error_handler_node=handler_node_name,
ends=ends,
defer=defer,
@@ -900,6 +912,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
input_schema=self.state_schema,
retry_policy=retry_policy,
cache_policy=cache_policy,
cache_policy_opt_out=cache_policy_opt_out,
error_handler_node=handler_node_name,
ends=ends,
defer=defer,
@@ -1317,6 +1330,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
not spec.is_error_handler
and defaults.cache_policy is not None
and spec.cache_policy is None
and not spec.cache_policy_opt_out
):
spec.cache_policy = defaults.cache_policy
# timeout: all nodes — a stuck handler should be cancelled the
+87
View File
@@ -58,6 +58,7 @@ from langgraph.pregel._retry import (
from langgraph.pregel.protocol import StreamProtocol
from langgraph.runtime import DEFAULT_RUNTIME, ExecutionInfo, Runtime
from langgraph.types import (
CachePolicy,
Command,
PregelExecutableTask,
RetryPolicy,
@@ -2570,6 +2571,92 @@ def test_set_node_defaults_retry_policy_per_node_wins():
assert attempts == 2
def test_set_node_defaults_cache_policy_applies_to_nodes_without_own() -> None:
from langgraph.cache.memory import InMemoryCache
class State(TypedDict):
foo: str
calls = 0
def cached_node(state: State) -> State:
nonlocal calls
calls += 1
return {"foo": "ok"}
graph = (
StateGraph(State)
.set_node_defaults(cache_policy=CachePolicy(ttl=60))
.add_node("cached", cached_node)
.add_edge(START, "cached")
.compile(cache=InMemoryCache())
)
assert graph.nodes["cached"].cache_policy == CachePolicy(ttl=60)
graph.invoke({"foo": ""})
graph.invoke({"foo": ""})
assert calls == 1
def test_set_node_defaults_cache_policy_per_node_wins() -> None:
from langgraph.cache.memory import InMemoryCache
class State(TypedDict):
foo: str
graph = (
StateGraph(State)
.set_node_defaults(cache_policy=CachePolicy(ttl=60))
.add_node("a", lambda state: {"foo": "a"}, cache_policy=CachePolicy(ttl=5))
.add_node("b", lambda state: {"foo": "b"})
.add_edge(START, "a")
.add_edge("a", "b")
.compile(cache=InMemoryCache())
)
assert graph.nodes["a"].cache_policy == CachePolicy(ttl=5)
assert graph.nodes["b"].cache_policy == CachePolicy(ttl=60)
def test_set_node_defaults_cache_policy_explicit_none_opts_out() -> None:
from langgraph.cache.memory import InMemoryCache
class State(TypedDict):
foo: str
cached_calls = 0
uncached_calls = 0
def cached_node(state: State) -> State:
nonlocal cached_calls
cached_calls += 1
return {"foo": "cached"}
def uncached_node(state: State) -> State:
nonlocal uncached_calls
uncached_calls += 1
return {"foo": "uncached"}
graph = (
StateGraph(State)
.set_node_defaults(cache_policy=CachePolicy())
.add_node("cached", cached_node)
.add_node("uncached", uncached_node, cache_policy=None)
.add_edge(START, "cached")
.add_edge("cached", "uncached")
.compile(cache=InMemoryCache())
)
assert graph.nodes["cached"].cache_policy == CachePolicy()
assert graph.nodes["uncached"].cache_policy is None
graph.invoke({"foo": ""})
graph.invoke({"foo": ""})
assert cached_calls == 1
assert uncached_calls == 2
@pytest.mark.anyio
async def test_set_node_defaults_timeout():
class State(TypedDict):