github-mirrors/tradingagents
1
0
Fork 0
mirror of https://github.com/TauricResearch/TradingAgents.git synced 2026-09-19 19:25:24 +03:00
Code Issues Packages Projects Releases Wiki Activity

fix(cli): resume a checkpoint without duplicating messages or leaking the saver

Browse Source 
- on resume, the CLI and propagate re-passed the initial state to a thread with
  an existing checkpoint; nodes do not re-run, but the message reducer appended
  the initial messages again, duplicating them in the resumed state
- feed None on resume (checkpoint_input) so LangGraph continues the interrupted
  run, and wrap the CLI stream in try/finally so the checkpointer tears down even
  if the stream raises
- correct the _fetch_returns docstring to the 4-tuple return #1249
This commit is contained in:
Yijia-Xiao
2026-08-31 01:29:14 +00:00
parent 8db41f6bca
commit b43bc31479
3 changed files with 159 additions and 105 deletions
Download Patch File Download Diff File Expand all files Collapse all files

206
cli/main.py
View File

@@ -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.

32
tests/test_checkpoint_lifecycle.py
View File

@@ -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

26
tradingagents/graph/trading_graph.py
View File

@@ -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