run-llama / llama-agents / 38063514549

85%

push

github

Add state checkpoints with JSON Patch diffs to BasicRuntime (#786)

## What changed

`BasicRuntime.state_checkpoint(run_id)` returns a `StateCheckpoint`
holding a reference to the run's committed state model. Taking one
copies nothing. `checkpoint.diff(since)` returns JSON Patch ops from an
earlier checkpoint to this one, `checkpoint.to_dict()` is a full payload
that works as a snapshot's `state`, and `apply_state_patch(payload,
ops)` applies ops to a plain dict with no model classes loaded. A run
resumes from a state rebuilt that way by putting it in the broker
snapshot the journal restores.

```python
base = basic_runtime.state_checkpoint(run_id)
put("state/base", base.to_dict())                       # once
...
cur = basic_runtime.state_checkpoint(run_id)
put(f"state/{seq}", cur.diff(base))                     # per flush, alongside the journal batch
# [{"op": "add", "path": "/state_data/_data/messages/4", "value": {...}}]
base = cur

# resume, any process: broker from the journal, state from base plus patches
state = get("state/base")
for patch in patches:
    state = apply_state_patch(state, patch)
snapshot = basic_runtime.restore(wf, None, records).to_dict(include_state=False)
snapshot["state"] = state
await wf.run(ctx=Context.from_dict(wf, snapshot))
```

`to_dict()` writes a new tree-shaped `state_data` form marked with
`_tree`, which `decode_state` reads by the marker. The two existing
forms load as before. Nothing on the store's write path changes.

## Why

The tick journal from https://github.com/run-llama/llama-agents/pull/784
covers broker state and nothing else. `ctx.store` is written by steps
directly, and it is persisted today by serializing the whole store on
every write. A step that appends one message to a long conversation
re-serializes the whole conversation.

`ctx.store.set(path, value)` already rebuilds only the containers on the
written path and shares everything else with the previous state. That
sharing is what makes a cheap ... (continued)

173 of 194 new or added lines in 2 files covered. (89.18%)

18906 of 22215 relevant lines covered (85.1%)

1.15 hits per line