diff --git a/cli/main.py b/cli/main.py index a3926e819..b64487367 100644 --- a/cli/main.py +++ b/cli/main.py @@ -1129,124 +1129,128 @@ def run_analysis(checkpoint: bool | None = None): # 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. + # checkpointing is disabled. Torn down in the finally below. 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. On resume, feed None so LangGraph continues the + # interrupted run instead of re-appending the initial state (#1249); the + # try/finally tears the checkpointer down even if the stream raises. trace = [] - for chunk in graph.graph.stream(init_agent_state, **args): - # Process all messages in chunk, deduplicating by message ID - for message in chunk.get("messages", []): - msg_id = getattr(message, "id", None) - if msg_id is not None: - if msg_id in message_buffer._processed_message_ids: - continue - message_buffer._processed_message_ids.add(msg_id) + try: + for chunk in graph.graph.stream(graph.checkpoint_input(init_agent_state), **args): + # Process all messages in chunk, deduplicating by message ID + for message in chunk.get("messages", []): + msg_id = getattr(message, "id", None) + if msg_id is not None: + if msg_id in message_buffer._processed_message_ids: + continue + message_buffer._processed_message_ids.add(msg_id) - msg_type, content = classify_message_type(message) - if content and content.strip(): - message_buffer.add_message(msg_type, content) + msg_type, content = classify_message_type(message) + if content and content.strip(): + message_buffer.add_message(msg_type, content) - if hasattr(message, "tool_calls") and message.tool_calls: - for tool_call in message.tool_calls: - if isinstance(tool_call, dict): - message_buffer.add_tool_call(tool_call["name"], tool_call["args"]) - else: - message_buffer.add_tool_call(tool_call.name, tool_call.args) + if hasattr(message, "tool_calls") and message.tool_calls: + for tool_call in message.tool_calls: + if isinstance(tool_call, dict): + message_buffer.add_tool_call(tool_call["name"], tool_call["args"]) + else: + message_buffer.add_tool_call(tool_call.name, tool_call.args) - # Update analyst statuses based on report state (runs on every chunk) - update_analyst_statuses( - message_buffer, - chunk, - wall_time_tracker=analyst_wall_time_tracker, - ) - - # Research Team - Handle Investment Debate State - if chunk.get("investment_debate_state"): - debate_state = chunk["investment_debate_state"] - bull_hist = debate_state.get("bull_history", "").strip() - bear_hist = debate_state.get("bear_history", "").strip() - judge = debate_state.get("judge_decision", "").strip() - - # Only update status when there's actual content - if bull_hist or bear_hist: - update_research_team_status("in_progress") - if bull_hist: - message_buffer.update_report_section( - "investment_plan", f"### Bull Researcher Analysis\n{bull_hist}" - ) - if bear_hist: - message_buffer.update_report_section( - "investment_plan", f"### Bear Researcher Analysis\n{bear_hist}" - ) - if judge: - message_buffer.update_report_section( - "investment_plan", f"### Research Manager Decision\n{judge}" - ) - update_research_team_status("completed") - message_buffer.update_agent_status("Trader", "in_progress") - - # Trading Team - if chunk.get("trader_investment_plan"): - message_buffer.update_report_section( - "trader_investment_plan", chunk["trader_investment_plan"] + # Update analyst statuses based on report state (runs on every chunk) + update_analyst_statuses( + message_buffer, + chunk, + wall_time_tracker=analyst_wall_time_tracker, ) - if message_buffer.agent_status.get("Trader") != "completed": - message_buffer.update_agent_status("Trader", "completed") - message_buffer.update_agent_status("Aggressive Analyst", "in_progress") - # Risk Management Team - Handle Risk Debate State - if chunk.get("risk_debate_state"): - risk_state = chunk["risk_debate_state"] - agg_hist = risk_state.get("aggressive_history", "").strip() - con_hist = risk_state.get("conservative_history", "").strip() - neu_hist = risk_state.get("neutral_history", "").strip() - judge = risk_state.get("judge_decision", "").strip() + # Research Team - Handle Investment Debate State + if chunk.get("investment_debate_state"): + debate_state = chunk["investment_debate_state"] + bull_hist = debate_state.get("bull_history", "").strip() + bear_hist = debate_state.get("bear_history", "").strip() + judge = debate_state.get("judge_decision", "").strip() - if agg_hist: - if message_buffer.agent_status.get("Aggressive Analyst") != "completed": + # Only update status when there's actual content + if bull_hist or bear_hist: + update_research_team_status("in_progress") + if bull_hist: + message_buffer.update_report_section( + "investment_plan", f"### Bull Researcher Analysis\n{bull_hist}" + ) + if bear_hist: + message_buffer.update_report_section( + "investment_plan", f"### Bear Researcher Analysis\n{bear_hist}" + ) + if judge: + message_buffer.update_report_section( + "investment_plan", f"### Research Manager Decision\n{judge}" + ) + update_research_team_status("completed") + message_buffer.update_agent_status("Trader", "in_progress") + + # Trading Team + if chunk.get("trader_investment_plan"): + message_buffer.update_report_section( + "trader_investment_plan", chunk["trader_investment_plan"] + ) + if message_buffer.agent_status.get("Trader") != "completed": + message_buffer.update_agent_status("Trader", "completed") message_buffer.update_agent_status("Aggressive Analyst", "in_progress") - message_buffer.update_report_section( - "final_trade_decision", f"### Aggressive Analyst Analysis\n{agg_hist}" - ) - if con_hist: - if message_buffer.agent_status.get("Conservative Analyst") != "completed": - message_buffer.update_agent_status("Conservative Analyst", "in_progress") - message_buffer.update_report_section( - "final_trade_decision", f"### Conservative Analyst Analysis\n{con_hist}" - ) - if neu_hist: - if message_buffer.agent_status.get("Neutral Analyst") != "completed": - message_buffer.update_agent_status("Neutral Analyst", "in_progress") - message_buffer.update_report_section( - "final_trade_decision", f"### Neutral Analyst Analysis\n{neu_hist}" - ) - if judge and message_buffer.agent_status.get("Portfolio Manager") != "completed": - message_buffer.update_agent_status("Portfolio Manager", "in_progress") - message_buffer.update_report_section( - "final_trade_decision", f"### Portfolio Manager Decision\n{judge}" - ) - message_buffer.update_agent_status("Aggressive Analyst", "completed") - message_buffer.update_agent_status("Conservative Analyst", "completed") - message_buffer.update_agent_status("Neutral Analyst", "completed") - message_buffer.update_agent_status("Portfolio Manager", "completed") - # Update the display - update_display(layout, stats_handler=stats_handler, start_time=start_time) + # Risk Management Team - Handle Risk Debate State + if chunk.get("risk_debate_state"): + risk_state = chunk["risk_debate_state"] + agg_hist = risk_state.get("aggressive_history", "").strip() + con_hist = risk_state.get("conservative_history", "").strip() + neu_hist = risk_state.get("neutral_history", "").strip() + judge = risk_state.get("judge_decision", "").strip() - trace.append(chunk) + if agg_hist: + if message_buffer.agent_status.get("Aggressive Analyst") != "completed": + message_buffer.update_agent_status("Aggressive Analyst", "in_progress") + message_buffer.update_report_section( + "final_trade_decision", f"### Aggressive Analyst Analysis\n{agg_hist}" + ) + if con_hist: + if message_buffer.agent_status.get("Conservative Analyst") != "completed": + message_buffer.update_agent_status("Conservative Analyst", "in_progress") + message_buffer.update_report_section( + "final_trade_decision", f"### Conservative Analyst Analysis\n{con_hist}" + ) + if neu_hist: + if message_buffer.agent_status.get("Neutral Analyst") != "completed": + message_buffer.update_agent_status("Neutral Analyst", "in_progress") + message_buffer.update_report_section( + "final_trade_decision", f"### Neutral Analyst Analysis\n{neu_hist}" + ) + if judge and message_buffer.agent_status.get("Portfolio Manager") != "completed": + message_buffer.update_agent_status("Portfolio Manager", "in_progress") + message_buffer.update_report_section( + "final_trade_decision", f"### Portfolio Manager Decision\n{judge}" + ) + message_buffer.update_agent_status("Aggressive Analyst", "completed") + message_buffer.update_agent_status("Conservative Analyst", "completed") + message_buffer.update_agent_status("Neutral Analyst", "completed") + message_buffer.update_agent_status("Portfolio Manager", "completed") - # 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() + # Update the display + update_display(layout, stats_handler=stats_handler, start_time=start_time) + + trace.append(chunk) + + # Clean run: drop this run's checkpoint so a later run starts fresh. + # A mid-stream failure skips this, keeping the checkpoint for resume. + graph.clear_checkpoint_on_success( + selections["ticker"], selections["analysis_date"], selections["asset_type"] + ) + finally: + # Always restore the plain uncheckpointed graph, even on failure. + graph.end_checkpoint() # Streamed chunks are per-node deltas, not full state. Merge them # so every report field populated across the run is present. diff --git a/tests/test_checkpoint_lifecycle.py b/tests/test_checkpoint_lifecycle.py index 5cdb247af..1e5e2adf5 100644 --- a/tests/test_checkpoint_lifecycle.py +++ b/tests/test_checkpoint_lifecycle.py @@ -84,6 +84,38 @@ def test_begin_returns_thread_id_and_recompiles(): assert g._checkpointer_ctx is None # restored +@pytest.mark.unit +def test_checkpoint_input_is_none_only_when_resuming(): + global _should_crash + with tempfile.TemporaryDirectory() as tmp: + init = {"count": 0} + args = ("AAPL", "2026-05-08", "stock") + # Fresh run: no checkpoint yet -> stream the initial state, then crash. + g1 = _bare_graph(tmp) + tid = g1.begin_checkpoint(*args) + try: + assert g1._resuming is False + assert g1.checkpoint_input(init) is init # not resuming -> initial state + _should_crash = True + with pytest.raises(RuntimeError): + for _ in g1.graph.stream(init, config={"configurable": {"thread_id": tid}}): + pass + finally: + g1.end_checkpoint() + assert g1.checkpoint_input(init) is init # reset after teardown + + # A later run finds the checkpoint -> resume by feeding None, not the + # initial state (re-passing it would duplicate messages, #1249). + _should_crash = False + g2 = _bare_graph(tmp) + g2.begin_checkpoint(*args) + try: + assert g2._resuming is True + assert g2.checkpoint_input(init) is None + finally: + g2.end_checkpoint() + + @pytest.mark.unit def test_cli_style_usage_saves_then_resumes(): global _should_crash diff --git a/tradingagents/graph/trading_graph.py b/tradingagents/graph/trading_graph.py index 2ae1e5bf9..c2a9418c1 100644 --- a/tradingagents/graph/trading_graph.py +++ b/tradingagents/graph/trading_graph.py @@ -163,6 +163,7 @@ class TradingAgentsGraph: self.workflow = self.graph_setup.setup_graph(selected_analysts) self.graph = self.workflow.compile() self._checkpointer_ctx = None + self._resuming = False def _get_provider_kwargs(self) -> dict[str, Any]: """Get provider-specific kwargs for LLM client creation.""" @@ -277,8 +278,10 @@ class TradingAgentsGraph: ``benchmark`` is the index used as the alpha baseline (resolved by the caller via ``_resolve_benchmark``). Returns ``(raw_return, alpha_return, - actual_holding_days)`` or ``(None, None, None)`` if price data is - unavailable (too recent, delisted, or network error). + actual_holding_days, resolution_date)`` — where ``resolution_date`` is + the date of the last price bar used, i.e. when the outcome became known + (#1251) — or ``(None, None, None, None)`` if price data is unavailable + (too recent, delisted, or network error). """ from tradingagents.dataflows.symbol_utils import normalize_symbol @@ -433,6 +436,7 @@ class TradingAgentsGraph: lived only inside ``propagate`` and the CLI streamed the checkpointer-less graph, making the flag a no-op. """ + self._resuming = False if not self.config.get("checkpoint_enabled"): return None signature = self._run_signature(asset_type) @@ -443,18 +447,30 @@ class TradingAgentsGraph: step = checkpoint_step( self.config["data_cache_dir"], company_name, str(trade_date), signature ) + self._resuming = step is not None 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 checkpoint_input(self, init_state): + """The value to stream/invoke: ``None`` to resume an existing checkpoint, + else the initial state for a fresh run. + + LangGraph resumes an interrupted thread when invoked with ``None``; + re-passing the initial state instead appends it through the message + reducer, duplicating messages in the resumed state (#1249). + """ + return None if self._resuming else init_state + 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() + self._resuming = False @contextmanager def checkpoint_scope(self, company_name, trade_date, asset_type: str = "stock"): @@ -512,10 +528,12 @@ class TradingAgentsGraph: if checkpoint_thread_id is not None: args.setdefault("config", {}).setdefault("configurable", {})["thread_id"] = checkpoint_thread_id + # None resumes an existing checkpoint; init_agent_state starts fresh (#1249). + graph_input = self.checkpoint_input(init_agent_state) if self.debug: trace = [] last_printed = None - for chunk in self.graph.stream(init_agent_state, **args): + for chunk in self.graph.stream(graph_input, **args): if chunk["messages"]: msg = chunk["messages"][-1] # Nodes after the trader don't append to messages, so the @@ -532,7 +550,7 @@ class TradingAgentsGraph: for chunk in trace: final_state.update(chunk) else: - final_state = self.graph.invoke(init_agent_state, **args) + final_state = self.graph.invoke(graph_input, **args) # Store current state for reflection. self.curr_state = final_state