fix(cli): make --checkpoint actually resume on the CLI path

- checkpoint setup lived only inside propagate(); the CLI streamed the
  checkpointer-less graph with no thread_id, so --checkpoint neither saved nor
  resumed a run
- extract the lifecycle into reusable begin_checkpoint / end_checkpoint /
  clear_checkpoint_on_success (checkpoint_scope wraps them for propagate) and use
  them around the CLI stream #1249
This commit is contained in:
Yijia-Xiao
2026-08-30 06:47:00 +00:00
parent 43fc275b36
commit 51a245dbe1
3 changed files with 198 additions and 32 deletions

View File

@@ -1127,6 +1127,15 @@ def run_analysis(checkpoint: bool | None = None):
# (LLM tracking is handled separately via LLM constructor)
args = graph.propagator.get_graph_args(callbacks=[stats_handler])
# Recompile with a checkpointer and inject the thread_id so --checkpoint
# actually saves and resumes on the CLI path (#1249); a no-op when
# checkpointing is disabled. Paired with end_checkpoint after the stream.
checkpoint_tid = graph.begin_checkpoint(
selections["ticker"], selections["analysis_date"], selections["asset_type"]
)
if checkpoint_tid is not None:
args.setdefault("config", {}).setdefault("configurable", {})["thread_id"] = checkpoint_tid
# Stream the analysis
trace = []
for chunk in graph.graph.stream(init_agent_state, **args):
@@ -1231,6 +1240,14 @@ def run_analysis(checkpoint: bool | None = None):
trace.append(chunk)
# The stream completed: drop this run's checkpoint and restore the plain
# graph (#1249). A mid-stream failure skips this, leaving the checkpoint
# in place so the next run resumes.
graph.clear_checkpoint_on_success(
selections["ticker"], selections["analysis_date"], selections["asset_type"]
)
graph.end_checkpoint()
# Streamed chunks are per-node deltas, not full state. Merge them
# so every report field populated across the run is present.
final_state = {}