[Pydantic Graph] Unexpected behaviour in FullStatePersistence

[TheFirstMe]

Question

In count_down_from_persistence.py, I replaced FileStatePersistence with FullStatePersistence. The persistence history after the end is not the same in both the cases.

History for FileStatePersistence:

[
  {
    "state": {
      "counter": 5
    },
    "node": {
      "node_id": "CountDown"
    },
    "start_ts": "2025-05-10T22:30:08.463174Z",
    "duration": 4.887999978109292e-6,
    "status": "success",
    "kind": "node",
    "id": "CountDown:8b39bead1166477392f419152160cd2f"
  },
  {
    "state": {
      "counter": 4
    },
    "node": {
      "node_id": "CountDown"
    },
    "start_ts": "2025-05-10T22:30:08.464833Z",
    "duration": 3.698999989865115e-6,
    "status": "success",
    "kind": "node",
    "id": "CountDown:777ca6d4a2b84f2abf17d49844923a6a"
  },
  {
    "state": {
      "counter": 3
    },
    "node": {
      "node_id": "CountDown"
    },
    "start_ts": "2025-05-10T22:30:08.466064Z",
    "duration": 3.1830000466470665e-6,
    "status": "success",
    "kind": "node",
    "id": "CountDown:268719f7e56c480389569a90a07b0ccc"
  },
  {
    "state": {
      "counter": 2
    },
    "node": {
      "node_id": "CountDown"
    },
    "start_ts": "2025-05-10T22:30:08.467180Z",
    "duration": 2.833999985796254e-6,
    "status": "success",
    "kind": "node",
    "id": "CountDown:bd395c57d13e466b8c4f950c2ede80d6"
  },
  {
    "state": {
      "counter": 1
    },
    "node": {
      "node_id": "CountDown"
    },
    "start_ts": "2025-05-10T22:30:08.468323Z",
    "duration": 2.8519999659692985e-6,
    "status": "success",
    "kind": "node",
    "id": "CountDown:3e88aeab47884cceb8070699be977d95"
  },
  {
    "state": {
      "counter": 0
    },
    "node": {
      "node_id": "CountDown"
    },
    "start_ts": "2025-05-10T22:30:08.469617Z",
    "duration": 4.914999976790568e-6,
    "status": "success",
    "kind": "node",
    "id": "CountDown:4f7f714ba6a24addba516b4a473d25db"
  },
  {
    "state": {
      "counter": 0
    },
    "result": {
      "data": 0
    },
    "ts": "2025-05-10T22:30:08.470297Z",
    "kind": "end",
    "id": "end:0fccfb59d15b4c258200ae66a6540a1e"
  }
]

History for FullStatePersistence (states):

Here's the code that uses FullStatePersistence:

from count_down import CountDown, CountDownState, count_down_graph
from pydantic_graph import End
from pydantic_graph.persistence.in_mem import FullStatePersistence


async def main():
    persistence = FullStatePersistence()
    state = CountDownState(counter=5)
    await count_down_graph.initialize(CountDown(), state=state, persistence=persistence)

    done = False
    while not done:
        done = await run_node(persistence)
    print("History: ")
    for step in persistence.history:
        print(step.state)


async def run_node(persistence) -> bool:
    async with count_down_graph.iter_from_persistence(persistence) as run:
        node_or_end = await run.next()

    print("Node:", node_or_end)
    # > Node: CountDown()
    # > Node: CountDown()
    # > Node: CountDown()
    # > Node: CountDown()
    # > Node: CountDown()
    # > Node: End(data=0)
    return isinstance(node_or_end, End)

Expected behaviour: Both FullStatePersistence and FileStatePersistence should have the same state for all nodes.

Additional Context

Pydantic Graph Version: 0.1.10
Python Version: 3.11.0

[DouweM]

@TheFirstMe Good catch, the issue appears to be that iter_from_persistence passes snapshot.state directly to the new graph run, without deep-copying it, meaning that the node will be passed the same state that was stored in-memory in persistence.history, and be allowed to modify it in place:

pydantic-ai/pydantic_graph/pydantic_graph/graph.py

Line 315 in 63af922

state=snapshot.state,

That means the snapshot with state counter=5 ends up getting updated to have counter=4. The same thing doesn't happen with file persistence, because it always reads the snapshots from disk so the in-memory update doesn't stick.

We'd likely need to update either iter_from_persistence or FullStatePersistence.load_next or to deep-copy state before passing it to the new graph run.

Would you be up for submitting a PR to that effect? :)

[lienminhquang]

@DouweM I would like to submit a PR !