From a6318b31ccf221ca0241eb55516de8b8e3e65331 Mon Sep 17 00:00:00 2001 From: Yijia-Xiao Date: Fri, 25 Sep 2026 07:00:38 +0000 Subject: [PATCH] fix(graph): stream a run under its graph's own config - stream_run, which the CLI uses, scopes the graph's config like propagate() does, so a process-wide config changed elsewhere does not reach its tools --- tests/test_graph_end_to_end.py | 19 ++++++++++++++++++ tradingagents/dataflows/config.py | 13 ++++++++++++- tradingagents/graph/trading_graph.py | 29 ++++++++++++++++++---------- 3 files changed, 50 insertions(+), 11 deletions(-) diff --git a/tests/test_graph_end_to_end.py b/tests/test_graph_end_to_end.py index e51df46ad..7259108a8 100644 --- a/tests/test_graph_end_to_end.py +++ b/tests/test_graph_end_to_end.py @@ -213,3 +213,22 @@ def test_each_report_streams_as_soon_as_its_analyst_files_it(tmp_path, monkeypat if state and any(state.get(k) for k in reports)) assert sum(bool(first.get(k)) for k in reports) == 1 + + +@pytest.mark.unit +def test_a_streamed_run_reads_its_own_graph_config(tmp_path, monkeypatch, offline): + """The CLI streams the run; a process-wide config set elsewhere must not reach + its tools, and the graph's config must not reach the caller between steps.""" + from tradingagents.dataflows.config import get_config, set_config + + graph = _graph(tmp_path, monkeypatch, ScriptedModel(), output_language="French") + set_config({"output_language": "German"}) + seen = [] + monkeypatch.setitem(router.VENDOR_METHODS["get_stock_data"], "yfinance", + lambda *a, **k: seen.append(get_config()["output_language"]) or "prices") + + between = [get_config()["output_language"] + for _ in graph.stream_run(graph.create_run_state("NVDA", TRADE_DATE), **graph.propagator.get_graph_args())] + + assert seen and set(seen) == {"French"} + assert set(between) == {"German"} diff --git a/tradingagents/dataflows/config.py b/tradingagents/dataflows/config.py index 19a640a43..f7ae14f36 100644 --- a/tradingagents/dataflows/config.py +++ b/tradingagents/dataflows/config.py @@ -1,5 +1,5 @@ from contextlib import contextmanager -from contextvars import ContextVar +from contextvars import Context, ContextVar, copy_context from copy import deepcopy import tradingagents.default_config as default_config @@ -51,6 +51,17 @@ def run_config(config: dict): _run_config.reset(token) +def run_config_context(config: dict) -> Context: + """A context serving ``config``, over the defaults, to every read run in it. + + For a run that yields between steps, where a ``run_config`` block would + stay open in the caller's context while it holds the results. + """ + context = copy_context() + context.run(_run_config.set, _merge(deepcopy(default_config.DEFAULT_CONFIG), config)) + return context + + def get_config() -> dict: """Get the configuration of the run in progress, else the process-wide one.""" scoped = _run_config.get() diff --git a/tradingagents/graph/trading_graph.py b/tradingagents/graph/trading_graph.py index 3757ff5b2..c2872b5ff 100644 --- a/tradingagents/graph/trading_graph.py +++ b/tradingagents/graph/trading_graph.py @@ -9,7 +9,7 @@ from typing import Any import tradingagents from tradingagents.agents.context import build_instrument_context, resolve_instrument_identity from tradingagents.agents.rating import run_rating -from tradingagents.dataflows.config import run_config, set_config +from tradingagents.dataflows.config import run_config, run_config_context, set_config from tradingagents.dataflows.date_window import get_current_date from tradingagents.dataflows.symbols import safe_ticker_component from tradingagents.default_config import DEFAULT_CONFIG @@ -371,15 +371,24 @@ class TradingAgentsGraph: and reports come from their own finished steps ("tasks") as they happen. """ args = {**args, "stream_mode": ["values", "tasks"]} - for namespace, mode, chunk in self.graph.stream(graph_input, subgraphs=True, **args): - if namespace: - result = chunk.get("result") if mode == "tasks" else None - if isinstance(result, dict): - report = {k: v for k, v in result.items() if k != "messages" and v} - if result.get("messages") or report: - yield result.get("messages", []), report or None - elif mode == "values": - yield chunk.get("messages", []), chunk + # The graph's own config serves every tool call, as in propagate(), even + # when the process-wide config has changed since. Each step runs in the + # run's context, so the caller keeps its own between steps. + context = run_config_context(self.config) + stream = context.run(self.graph.stream, graph_input, subgraphs=True, **args) + try: + while (step := context.run(next, stream, None)) is not None: + namespace, mode, chunk = step + if namespace: + result = chunk.get("result") if mode == "tasks" else None + if isinstance(result, dict): + report = {k: v for k, v in result.items() if k != "messages" and v} + if result.get("messages") or report: + yield result.get("messages", []), report or None + elif mode == "values": + yield chunk.get("messages", []), chunk + finally: + context.run(stream.close) def _log_state(self, trade_date, final_state): """Write a run's final state to JSON under the run's own ticker."""