How to Keep LangGraph Subgraphs Stateful with Typed Pydantic State and Checkpointing

Keep LangGraph subgraphs stateful with Pydantic BaseModel and StateGraph: use __pydantic_fields_set__ to preserve checkpoints when parents overwrite fields.

Keeping subgraphs stateful in LangGraph can be tricky when a parent flow manipulates shared fields. A common symptom is that the subgraph’s checkpoint appears to “forget” earlier values after the parent writes None into the same keys. The goal is to let the subgraph preserve its own memory across invocations, even if the parent overwrites those fields upstream.

Minimal example that shows the problem

The following snippet models a parent flow that sets a shared field to None, and a subgraph that reads from it with checkpointing turned on. The expectation is that the subgraph holds on to its own memory, but in practice it falls back to a default once the parent writes None.

from langgraph.graph import Graph

# Parent Flow
root_flow = Graph()
root_flow.add_node("root_stage", lambda ctx: {"alpha": None})  # Overrides alpha

# Subgraph with checkpointing
child_flow = Graph(checkpoint=True)
child_flow.add_node("child_stage", lambda ctx: {"beta": ctx.get("alpha", "default")})

# Wiring
root_flow.add_edge("root_stage", child_flow)

# Initial state (subgraph should remember alpha=5)
seed_payload = {"alpha": 5, "beta": "unset"}

# Run
outcome = root_flow.run(seed_payload)
print(outcome)  # Expect: {"alpha": None, "beta": 5} (but beta becomes "default")

What’s actually going on

The parent sets a shared key to None, and the subgraph consumes that state on entry. As a result, it no longer sees the earlier value and emits a default. Attempts to force checkpointing at call time, to rename fields to avoid collisions, or to push state into external storage can either fail to isolate the subgraph’s memory or add complexity that isn’t desirable.

The practical fix is to make the state model explicit and control which fields are considered “set” when the graph evaluates state. This lets you preserve subgraph-specific values across runs.

Solution: define a typed state and mark fields as set

By moving to a typed state via pydantic.BaseModel and StateGraph, and initializing the internal flag that tracks which fields are set, you give the graph a reliable basis for persisting and restoring the subgraph’s state. The key is to define your state schema and set __pydantic_fields_set__.

from langgraph.constants import START, END
from langgraph.graph.state import StateGraph
from langgraph.checkpoint.memory import MemorySaver
from pydantic import BaseModel
from typing import Optional


class FlowState(BaseModel):
    title: Optional[str] = None
    years: Optional[int] = None


def step_one(s: FlowState) -> FlowState:
    print(s)
    s.years = 10
    s.title = "step_one"
    return s

pipeline = StateGraph(FlowState)
pipeline.add_node("step_one", step_one)
pipeline.add_edge(START, "step_one")
pipeline.add_edge("step_one", END)
pipeline = pipeline.compile(checkpointer=MemorySaver())

pipeline.invoke(FlowState(title="Sam"), config={"configurable": {"thread_id": "1"}})
# output: title='Sam' years=None

blank = FlowState()

# Important: mark the fields as set on the model instance
blank.__pydantic_fields_set__ = set(FlowState.model_fields.keys())

pipeline.invoke(blank, config={"configurable": {"thread_id": "1"}})
# output: title=None years=None

This pattern gives you an explicit state contract and control over which keys are treated as present when the graph runs and checkpoints. In a parent–child setup, the subgraph can rely on its own typed state, preserving its values across invocations even when upstream modifies the same field names.

How to apply this to subgraphs

Model the subgraph’s state as a BaseModel and use StateGraph for the subgraph boundary. When creating or restoring the subgraph’s state instance, set __pydantic_fields_set__ to the fields you want the subgraph to treat as defined. With a checkpointer like MemorySaver in place, this ensures the subgraph stores and restores its own state independently, rather than inheriting None from the parent for those fields.

Why this matters

In real workflows, parent graphs may sanitize or reset fields that are meaningful to a subgraph’s local logic. Without a typed state boundary, the subgraph can unintentionally adopt upstream overrides and lose context it previously computed and persisted. Making the state schema explicit and marking fields as set avoids this class of issues and keeps the subgraph’s memory coherent.

Takeaways

If you need a subgraph to maintain its own state across runs, define a Pydantic-based state for the subgraph, run it via StateGraph, and initialize __pydantic_fields_set__. This lets the subgraph persist its own values and ignore upstream overrides for chosen keys, while keeping checkpointing simple and self-contained.