diff --git a/cli/main.py b/cli/main.py index 2e0c41d0d..a3926e819 100644 --- a/cli/main.py +++ b/cli/main.py @@ -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 = {} diff --git a/tests/test_checkpoint_lifecycle.py b/tests/test_checkpoint_lifecycle.py new file mode 100644 index 000000000..5cdb247af --- /dev/null +++ b/tests/test_checkpoint_lifecycle.py @@ -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 diff --git a/tradingagents/graph/trading_graph.py b/tradingagents/graph/trading_graph.py index 2ef8b9001..fefbd2342 100644 --- a/tradingagents/graph/trading_graph.py +++ b/tradingagents/graph/trading_graph.py @@ -3,6 +3,7 @@ import json import logging import os +from contextlib import contextmanager from datetime import datetime, timedelta from pathlib import Path from typing import Any @@ -400,32 +401,61 @@ class TradingAgentsGraph: # Resolve any pending memory-log entries for this ticker before the pipeline runs. self._resolve_pending_entries(company_name) - # Recompile with a checkpointer if the user opted in. - if self.config.get("checkpoint_enabled"): - self._checkpointer_ctx = get_checkpointer( - self.config["data_cache_dir"], company_name + with self.checkpoint_scope(company_name, trade_date, asset_type) as thread_id_value: + return self._run_graph( + company_name, trade_date, asset_type=asset_type, + 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._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: """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) - 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.""" # Initialize state — inject memory log context for PM and the # deterministically resolved instrument identity for all agents. @@ -457,11 +488,10 @@ class TradingAgentsGraph: ) args = self.propagator.get_graph_args() - # Inject thread_id so same ticker+date+graph-shape resumes; a different - # date or graph shape starts fresh (#1089). - if self.config.get("checkpoint_enabled"): - tid = thread_id(company_name, str(trade_date), self._run_signature(asset_type)) - args.setdefault("config", {}).setdefault("configurable", {})["thread_id"] = tid + # Inject the checkpoint thread_id (from checkpoint_scope) so the same + # ticker+date+graph-shape resumes; a different one starts fresh (#1089). + if checkpoint_thread_id is not None: + args.setdefault("config", {}).setdefault("configurable", {})["thread_id"] = checkpoint_thread_id if self.debug: trace = [] @@ -499,11 +529,7 @@ class TradingAgentsGraph: ) # Clear checkpoint on successful completion to avoid stale state. - if self.config.get("checkpoint_enabled"): - clear_checkpoint( - self.config["data_cache_dir"], company_name, str(trade_date), - self._run_signature(asset_type), - ) + self.clear_checkpoint_on_success(company_name, trade_date, asset_type) return final_state, self.process_signal(final_state["final_trade_decision"])