From 9c701273c6e3c0d3f57ea1210a36ec21ea9f4555 Mon Sep 17 00:00:00 2001
From: Hector Hernandez <39923391+hectorhdzg@users.noreply.github.com>
Date: Mon, 8 Jun 2026 13:40:44 -0700
Subject: [PATCH] Fix token count mismatch

---
 .../_genai/_langchain/_tracer.py              |  4 ++--
 tests/langchain/test_tracer.py                | 22 +++++++++++++++----
 2 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/src/microsoft/opentelemetry/_genai/_langchain/_tracer.py b/src/microsoft/opentelemetry/_genai/_langchain/_tracer.py
index 213afb33..6e4d1074 100644
--- a/src/microsoft/opentelemetry/_genai/_langchain/_tracer.py
+++ b/src/microsoft/opentelemetry/_genai/_langchain/_tracer.py
@@ -394,8 +394,8 @@ def _is_agent_run(self, run: Run) -> bool:
         """Detect whether a LangChain run should be the top-level agent span."""
         if not self._is_agent_like_chain(run):
             return False
-        # Don't nest agents — if a parent is already an agent, this is internal
-        if run.parent_run_id and run.parent_run_id in self._agent_run_ids:
+        # Don't nest agents — if any ancestor is already an agent, this is internal
+        if self._find_agent_ancestor(run) is not None:
             return False
         return True
 
diff --git a/tests/langchain/test_tracer.py b/tests/langchain/test_tracer.py
index 05430aca..115ef29e 100644
--- a/tests/langchain/test_tracer.py
+++ b/tests/langchain/test_tracer.py
@@ -160,6 +160,22 @@ def test_non_agent_returns_false(self):
         run = _make_run(run_type="chain", name="RunnableSequence")
         self.assertFalse(tracer._is_agent_run(run))
 
+    def test_deeply_nested_agent_returns_false(self):
+        """Agent-like chain nested under a non-agent chain whose ancestor is
+        already an agent must NOT be treated as a separate agent.  This
+        prevents token aggregation from splitting across multiple agent
+        spans (see token-count mismatch in multi-agent LangGraph setups)."""
+        tracer, _, _ = _make_tracer()
+        agent_id = uuid4()
+        chain_id = uuid4()
+        tracer._agent_run_ids.add(agent_id)
+        # Intermediate non-agent chain between top-level agent and sub-graph.
+        chain_run = _make_run(run_type="chain", name="node_step", id=chain_id, parent_run_id=agent_id)
+        tracer.run_map[str(chain_id)] = chain_run
+        # Sub-graph whose direct parent is the intermediate chain, not the agent.
+        sub_graph = _make_run(run_type="chain", name="SubAgentGraph", parent_run_id=chain_id)
+        self.assertFalse(tracer._is_agent_run(sub_graph))
+
 
 # ---- Agent name resolution ---------------------------------------------------
 
@@ -760,6 +776,7 @@ def test_attach_uses_inner_span_for_agent(self, mock_ctx):
             self.assertIs(attach_call_args[0][0], inner_span)
         mock_ctx.attach.assert_called_once()
 
+
 # ---- invoke_agent aggregation fixes (issue #172) -----------------------------
 
 
@@ -1070,9 +1087,7 @@ def test_tool_role_message_becomes_tool_call_response(self):
                     "id": ["langchain", "schema", "messages", "AIMessage"],
                     "kwargs": {
                         "content": "",
-                        "tool_calls": [
-                            {"name": "get_weather", "args": {"location": "Paris"}, "id": "tc1"}
-                        ],
+                        "tool_calls": [{"name": "get_weather", "args": {"location": "Paris"}, "id": "tc1"}],
                         "type": "ai",
                     },
                 },
@@ -1086,4 +1101,3 @@ def test_tool_role_message_becomes_tool_call_response(self):
         self.assertEqual(tool_part.type, "tool_call_response")
         self.assertEqual(tool_part.id, "tc1")
         self.assertEqual(tool_part.response, "rainy")
-