fix: ensure LLM stats tracking is accurate by including completed subagents (#441)

tests/telemetry/test_tracer.py

@@ -10,6 +10,7 @@ from opentelemetry.sdk.trace.export import SimpleSpanProcessor, SpanExportResult
 from strix.telemetry import tracer as tracer_module
 from strix.telemetry import utils as telemetry_utils
 from strix.telemetry.tracer import Tracer, set_global_tracer
+from strix.tools.agents_graph import agents_graph_actions
 
 
 def _load_events(events_path: Path) -> list[dict[str, Any]]:
@@ -255,6 +256,75 @@ def test_events_with_agent_id_include_agent_name(monkeypatch, tmp_path) -> None:
     assert chat_event["actor"]["agent_name"] == "Root Agent"
 
 
+def test_get_total_llm_stats_includes_completed_subagents(monkeypatch, tmp_path) -> None:
+    monkeypatch.chdir(tmp_path)
+
+    class DummyStats:
+        def __init__(
+            self,
+            *,
+            input_tokens: int,
+            output_tokens: int,
+            cached_tokens: int,
+            cost: float,
+            requests: int,
+        ) -> None:
+            self.input_tokens = input_tokens
+            self.output_tokens = output_tokens
+            self.cached_tokens = cached_tokens
+            self.cost = cost
+            self.requests = requests
+
+    class DummyLLM:
+        def __init__(self, stats: DummyStats) -> None:
+            self._total_stats = stats
+
+    class DummyAgent:
+        def __init__(self, stats: DummyStats) -> None:
+            self.llm = DummyLLM(stats)
+
+    tracer = Tracer("cost-rollup")
+    set_global_tracer(tracer)
+
+    monkeypatch.setattr(
+        agents_graph_actions,
+        "_agent_instances",
+        {
+            "root-agent": DummyAgent(
+                DummyStats(
+                    input_tokens=1_000,
+                    output_tokens=250,
+                    cached_tokens=100,
+                    cost=0.12345,
+                    requests=2,
+                )
+            )
+        },
+    )
+    monkeypatch.setattr(
+        agents_graph_actions,
+        "_completed_agent_llm_totals",
+        {
+            "input_tokens": 2_000,
+            "output_tokens": 500,
+            "cached_tokens": 400,
+            "cost": 0.54321,
+            "requests": 3,
+        },
+    )
+
+    stats = tracer.get_total_llm_stats()
+
+    assert stats["total"] == {
+        "input_tokens": 3_000,
+        "output_tokens": 750,
+        "cached_tokens": 500,
+        "cost": 0.6667,
+        "requests": 5,
+    }
+    assert stats["total_tokens"] == 3_750
+
+
 def test_run_metadata_is_only_on_run_lifecycle_events(monkeypatch, tmp_path) -> None:
     monkeypatch.chdir(tmp_path)
 

tests/tools/test_agents_graph_whitebox.py

@@ -5,16 +5,24 @@ from strix.llm.config import LLMConfig
 from strix.tools.agents_graph import agents_graph_actions
 
 
-def test_create_agent_inherits_parent_whitebox_flag(monkeypatch) -> None:
-    monkeypatch.setenv("STRIX_LLM", "openai/gpt-5")
-
+def _reset_agent_graph_state() -> None:
     agents_graph_actions._agent_graph["nodes"].clear()
     agents_graph_actions._agent_graph["edges"].clear()
     agents_graph_actions._agent_messages.clear()
     agents_graph_actions._running_agents.clear()
     agents_graph_actions._agent_instances.clear()
+    agents_graph_actions._completed_agent_llm_totals.clear()
+    agents_graph_actions._completed_agent_llm_totals.update(
+        agents_graph_actions._empty_llm_stats_totals()
+    )
     agents_graph_actions._agent_states.clear()
 
+
+def test_create_agent_inherits_parent_whitebox_flag(monkeypatch) -> None:
+    monkeypatch.setenv("STRIX_LLM", "openai/gpt-5")
+
+    _reset_agent_graph_state()
+
     parent_id = "parent-agent"
     parent_llm = LLMConfig(timeout=123, scan_mode="standard", is_whitebox=True)
     agents_graph_actions._agent_instances[parent_id] = SimpleNamespace(
@@ -66,12 +74,7 @@ def test_create_agent_inherits_parent_whitebox_flag(monkeypatch) -> None:
 def test_delegation_prompt_includes_wiki_memory_instruction_in_whitebox(monkeypatch) -> None:
     monkeypatch.setenv("STRIX_LLM", "openai/gpt-5")
 
-    agents_graph_actions._agent_graph["nodes"].clear()
-    agents_graph_actions._agent_graph["edges"].clear()
-    agents_graph_actions._agent_messages.clear()
-    agents_graph_actions._running_agents.clear()
-    agents_graph_actions._agent_instances.clear()
-    agents_graph_actions._agent_states.clear()
+    _reset_agent_graph_state()
 
     parent_id = "parent-1"
     child_id = "child-1"
@@ -116,12 +119,7 @@ def test_delegation_prompt_includes_wiki_memory_instruction_in_whitebox(monkeypa
 def test_agent_finish_appends_wiki_update_for_whitebox(monkeypatch) -> None:
     monkeypatch.setenv("STRIX_LLM", "openai/gpt-5")
 
-    agents_graph_actions._agent_graph["nodes"].clear()
-    agents_graph_actions._agent_graph["edges"].clear()
-    agents_graph_actions._agent_messages.clear()
-    agents_graph_actions._running_agents.clear()
-    agents_graph_actions._agent_instances.clear()
-    agents_graph_actions._agent_states.clear()
+    _reset_agent_graph_state()
 
     parent_id = "parent-2"
     child_id = "child-2"
@@ -192,12 +190,7 @@ def test_agent_finish_appends_wiki_update_for_whitebox(monkeypatch) -> None:
 def test_run_agent_in_thread_injects_shared_wiki_context_in_whitebox(monkeypatch) -> None:
     monkeypatch.setenv("STRIX_LLM", "openai/gpt-5")
 
-    agents_graph_actions._agent_graph["nodes"].clear()
-    agents_graph_actions._agent_graph["edges"].clear()
-    agents_graph_actions._agent_messages.clear()
-    agents_graph_actions._running_agents.clear()
-    agents_graph_actions._agent_instances.clear()
-    agents_graph_actions._agent_states.clear()
+    _reset_agent_graph_state()
 
     parent_id = "parent-3"
     child_id = "child-3"