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) # (LLM tracking is handled separately via LLM constructor)
args = graph.propagator.get_graph_args(callbacks=[stats_handler]) 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 # Stream the analysis
trace = [] trace = []
for chunk in graph.graph.stream(init_agent_state, **args): for chunk in graph.graph.stream(init_agent_state, **args):
@@ -1231,6 +1240,14 @@ def run_analysis(checkpoint: bool | None = None):
trace.append(chunk) 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 # Streamed chunks are per-node deltas, not full state. Merge them
# so every report field populated across the run is present. # so every report field populated across the run is present.
final_state = {} final_state = {}

View File

@@ -0,0 +1,123 @@
"""The checkpoint lifecycle is reusable so --checkpoint works on the CLI path (#1249).
Checkpoint setup previously lived only inside ``propagate``; the CLI streamed the
checkpointer-less graph, so ``--checkpoint`` neither saved nor resumed. The
lifecycle is now ``begin_checkpoint`` / ``end_checkpoint`` /
``clear_checkpoint_on_success`` on TradingAgentsGraph, used by both paths. These
tests drive that lifecycle exactly as the CLI does (begin -> stream self.graph ->
clear/end) and prove state is saved and resumed.
"""
from __future__ import annotations
import tempfile
from typing import TypedDict
import pytest
from langgraph.graph import END, StateGraph
from tradingagents.graph.checkpointer import checkpoint_step
from tradingagents.graph.trading_graph import TradingAgentsGraph
_should_crash = False
class _State(TypedDict):
count: int
def _node_a(state: _State) -> dict:
return {"count": state["count"] + 1}
def _node_b(state: _State) -> dict:
if _should_crash:
raise RuntimeError("simulated mid-stream crash")
return {"count": state["count"] + 10}
def _workflow() -> StateGraph:
b = StateGraph(_State)
b.add_node("analyst", _node_a)
b.add_node("trader", _node_b)
b.set_entry_point("analyst")
b.add_edge("analyst", "trader")
b.add_edge("trader", END)
return b
def _bare_graph(tmpdir, *, enabled=True):
g = object.__new__(TradingAgentsGraph)
g.config = {
"checkpoint_enabled": enabled, "data_cache_dir": tmpdir,
"max_debate_rounds": 1, "max_risk_discuss_rounds": 1,
}
g.selected_analysts = ("market",)
g.workflow = _workflow()
g.graph = g.workflow.compile()
g._checkpointer_ctx = None
return g
@pytest.mark.unit
def test_disabled_is_a_noop():
with tempfile.TemporaryDirectory() as tmp:
g = _bare_graph(tmp, enabled=False)
plain = g.graph
assert g.begin_checkpoint("AAPL", "2026-05-08", "stock") is None
assert g.graph is plain # graph not recompiled
g.end_checkpoint() # safe no-op
@pytest.mark.unit
def test_begin_returns_thread_id_and_recompiles():
with tempfile.TemporaryDirectory() as tmp:
g = _bare_graph(tmp)
plain = g.graph
tid = g.begin_checkpoint("AAPL", "2026-05-08", "stock")
try:
assert tid # a real thread_id
assert g.graph is not plain # recompiled with a checkpointer
finally:
g.end_checkpoint()
assert g._checkpointer_ctx is None # restored
@pytest.mark.unit
def test_cli_style_usage_saves_then_resumes():
global _should_crash
with tempfile.TemporaryDirectory() as tmp:
cfg_args = ("AAPL", "2026-05-08", "stock")
# Run 1 (the CLI path): begin -> stream self.graph -> crash at 'trader'.
_should_crash = True
g1 = _bare_graph(tmp)
tid = g1.begin_checkpoint(*cfg_args)
args = {"config": {"configurable": {"thread_id": tid}}}
try:
with pytest.raises(RuntimeError):
for _ in g1.graph.stream({"count": 0}, **args):
pass
finally:
g1.end_checkpoint()
# A checkpoint was saved for this run signature (so --checkpoint works).
sig = g1._run_signature("stock")
assert checkpoint_step(tmp, "AAPL", "2026-05-08", sig) is not None
# Run 2 (fresh graph, as a new CLI invocation): resume and finish.
_should_crash = False
g2 = _bare_graph(tmp)
tid2 = g2.begin_checkpoint(*cfg_args)
assert tid2 == tid # stable id -> same thread resumes
try:
result = g2.graph.invoke(None, config={"configurable": {"thread_id": tid2}})
assert result["count"] == 11 # analyst(+1) resumed into trader(+10)
g2.clear_checkpoint_on_success(*cfg_args)
finally:
g2.end_checkpoint()
# Cleared on success -> a later run starts fresh.
assert checkpoint_step(tmp, "AAPL", "2026-05-08", sig) is None

View File

@@ -3,6 +3,7 @@
import json import json
import logging import logging
import os import os
from contextlib import contextmanager
from datetime import datetime, timedelta from datetime import datetime, timedelta
from pathlib import Path from pathlib import Path
from typing import Any from typing import Any
@@ -400,32 +401,61 @@ class TradingAgentsGraph:
# Resolve any pending memory-log entries for this ticker before the pipeline runs. # Resolve any pending memory-log entries for this ticker before the pipeline runs.
self._resolve_pending_entries(company_name) self._resolve_pending_entries(company_name)
# Recompile with a checkpointer if the user opted in. with self.checkpoint_scope(company_name, trade_date, asset_type) as thread_id_value:
if self.config.get("checkpoint_enabled"): return self._run_graph(
self._checkpointer_ctx = get_checkpointer( company_name, trade_date, asset_type=asset_type,
self.config["data_cache_dir"], company_name checkpoint_thread_id=thread_id_value,
) )
saver = self._checkpointer_ctx.__enter__()
self.graph = self.workflow.compile(checkpointer=saver)
step = checkpoint_step( def begin_checkpoint(self, company_name, trade_date, asset_type: str = "stock") -> str | None:
"""Recompile the graph with a per-ticker checkpointer and return the
``thread_id`` to inject into the stream/invoke ``config`` (or ``None``
when checkpointing is disabled).
Pair every call with :meth:`end_checkpoint` in a ``finally``. Both
``propagate`` (via :meth:`checkpoint_scope`) and the CLI stream path use
this so ``--checkpoint`` actually resumes (#1249); previously the setup
lived only inside ``propagate`` and the CLI streamed the checkpointer-less
graph, making the flag a no-op.
"""
if not self.config.get("checkpoint_enabled"):
return None
signature = self._run_signature(asset_type)
self._checkpointer_ctx = get_checkpointer(self.config["data_cache_dir"], company_name)
saver = self._checkpointer_ctx.__enter__()
self.graph = self.workflow.compile(checkpointer=saver)
step = checkpoint_step(
self.config["data_cache_dir"], company_name, str(trade_date), signature
)
if step is not None:
logger.info("Resuming from step %d for %s on %s", step, company_name, trade_date)
else:
logger.info("Starting fresh for %s on %s", company_name, trade_date)
return thread_id(company_name, str(trade_date), signature)
def end_checkpoint(self):
"""Restore the plain uncheckpointed graph after a checkpointed run."""
if self._checkpointer_ctx is not None:
self._checkpointer_ctx.__exit__(None, None, None)
self._checkpointer_ctx = None
self.graph = self.workflow.compile()
@contextmanager
def checkpoint_scope(self, company_name, trade_date, asset_type: str = "stock"):
"""Context-manager form of begin/end_checkpoint for the propagate path."""
try:
yield self.begin_checkpoint(company_name, trade_date, asset_type)
finally:
self.end_checkpoint()
def clear_checkpoint_on_success(self, company_name, trade_date, asset_type: str = "stock"):
"""Drop a completed run's checkpoint so a later run starts fresh (#1249)."""
if self.config.get("checkpoint_enabled"):
clear_checkpoint(
self.config["data_cache_dir"], company_name, str(trade_date), self.config["data_cache_dir"], company_name, str(trade_date),
self._run_signature(asset_type), self._run_signature(asset_type),
) )
if step is not None:
logger.info(
"Resuming from step %d for %s on %s", step, company_name, trade_date
)
else:
logger.info("Starting fresh for %s on %s", company_name, trade_date)
try:
return self._run_graph(company_name, trade_date, asset_type=asset_type)
finally:
if self._checkpointer_ctx is not None:
self._checkpointer_ctx.__exit__(None, None, None)
self._checkpointer_ctx = None
self.graph = self.workflow.compile()
def save_reports(self, final_state, ticker, save_path=None) -> Path: def save_reports(self, final_state, ticker, save_path=None) -> Path:
"""Write the markdown report tree for a completed run, like the CLI does. """Write the markdown report tree for a completed run, like the CLI does.
@@ -442,7 +472,8 @@ class TradingAgentsGraph:
) )
return write_report_tree(final_state, ticker, save_path) return write_report_tree(final_state, ticker, save_path)
def _run_graph(self, company_name, trade_date, asset_type: str = "stock"): def _run_graph(self, company_name, trade_date, asset_type: str = "stock",
checkpoint_thread_id: str | None = None):
"""Execute the graph and write the resulting state to disk and memory log.""" """Execute the graph and write the resulting state to disk and memory log."""
# Initialize state — inject memory log context for PM and the # Initialize state — inject memory log context for PM and the
# deterministically resolved instrument identity for all agents. # deterministically resolved instrument identity for all agents.
@@ -457,11 +488,10 @@ class TradingAgentsGraph:
) )
args = self.propagator.get_graph_args() args = self.propagator.get_graph_args()
# Inject thread_id so same ticker+date+graph-shape resumes; a different # Inject the checkpoint thread_id (from checkpoint_scope) so the same
# date or graph shape starts fresh (#1089). # ticker+date+graph-shape resumes; a different one starts fresh (#1089).
if self.config.get("checkpoint_enabled"): if checkpoint_thread_id is not None:
tid = thread_id(company_name, str(trade_date), self._run_signature(asset_type)) args.setdefault("config", {}).setdefault("configurable", {})["thread_id"] = checkpoint_thread_id
args.setdefault("config", {}).setdefault("configurable", {})["thread_id"] = tid
if self.debug: if self.debug:
trace = [] trace = []
@@ -499,11 +529,7 @@ class TradingAgentsGraph:
) )
# Clear checkpoint on successful completion to avoid stale state. # Clear checkpoint on successful completion to avoid stale state.
if self.config.get("checkpoint_enabled"): self.clear_checkpoint_on_success(company_name, trade_date, asset_type)
clear_checkpoint(
self.config["data_cache_dir"], company_name, str(trade_date),
self._run_signature(asset_type),
)
return final_state, self.process_signal(final_state["final_trade_decision"]) return final_state, self.process_signal(final_state["final_trade_decision"])