fix: Cache node trackers per execution to preserve runId correlation

_flush_final_segment and _track_tool_calls were each calling
create_tracker() independently, generating new runIds that broke
per-execution event correlation. Now build_node creates one tracker
per node, cached in _node_trackers, and reused by all tracking methods.

Adds test_same_run_id_across_token_success_and_tool_call_events to
verify all node-level events for a single execution share one runId.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: jsonbailey

packages/ai-providers/server-ai-openai/src/ldai_openai/openai_agent_graph_runner.py

@@ -57,6 +57,7 @@ def __init__(self, graph: AgentGraphDefinition, tools: ToolRegistry):
         self._tools = tools
         self._agent_name_map: Dict[str, str] = {}
         self._tool_name_map: Dict[str, str] = {}
+        self._node_trackers: Dict[str, Any] = {}
 
     async def run(self, input: Any) -> AgentGraphResult:
         """
@@ -145,10 +146,12 @@ def _build_agents(self, path: List[str], state: _RunState) -> Any:
         tracker = self._graph.create_tracker()
         name_map: Dict[str, str] = {}
         tool_name_map: Dict[str, str] = {}
+        node_trackers: Dict[str, Any] = {}
 
         def build_node(node: AgentGraphNode, ctx: dict) -> Any:
             node_config = node.get_config()
             config_tracker = node_config.create_tracker()
+            node_trackers[node_config.key] = config_tracker
             model = node_config.model
 
             if not model:
@@ -204,6 +207,7 @@ def build_node(node: AgentGraphNode, ctx: dict) -> Any:
         root = self._graph.reverse_traverse(fn=build_node)
         self._agent_name_map = name_map
         self._tool_name_map = tool_name_map
+        self._node_trackers = node_trackers
         return root
 
     def _make_on_handoff(
@@ -263,10 +267,7 @@ def _flush_final_segment(
         """Record duration/tokens for the last active agent (no handoff after it)."""
         if not state.last_node_key:
             return
-        node = self._graph.get_node(state.last_node_key)
-        if node is None:
-            return
-        config_tracker = node.get_config().create_tracker()
+        config_tracker = self._node_trackers.get(state.last_node_key)
         if config_tracker is None:
             return
 
@@ -293,9 +294,6 @@ def _track_tool_calls(self, result: Any) -> None:
             tool_name = self._tool_name_map.get(tool_fn_name)
             if tool_name is None:
                 continue
-            node = self._graph.get_node(agent_key)
-            if node is None:
-                continue
-            config_tracker = node.get_config().create_tracker()
+            config_tracker = self._node_trackers.get(agent_key)
             if config_tracker is not None:
                 config_tracker.track_tool_call(tool_name)

packages/ai-providers/server-ai-openai/tests/test_openai_agent_graph_runner.py

@@ -138,7 +138,10 @@ async def test_openai_agent_graph_runner_run_success():
     tracker.track_path.assert_called_once()
     tracker.track_latency.assert_called_once()
 
-    root_tracker = graph.get_node('root-agent').get_config().create_tracker()
-    root_tracker.track_duration.assert_called_once()
-    root_tracker.track_tokens.assert_called_once()
-    root_tracker.track_success.assert_called_once()
+    # The runner caches one tracker per node — verify it is the same instance
+    # returned by create_tracker() and that all tracking calls hit it.
+    cached = runner._node_trackers['root-agent']
+    assert cached is graph.get_node('root-agent').get_config().create_tracker()
+    cached.track_duration.assert_called_once()
+    cached.track_tokens.assert_called_once()
+    cached.track_success.assert_called_once()

packages/ai-providers/server-ai-openai/tests/test_tracking_openai_agents.py

@@ -359,6 +359,40 @@ async def test_tracks_multiple_tool_calls():
     assert sorted(tool_keys) == ['search', 'summarize']
 
 
+@pytest.mark.asyncio
+async def test_same_run_id_across_token_success_and_tool_call_events():
+    """All node-level events for a single execution share the same runId."""
+    mock_ld_client = MagicMock()
+    graph = _make_graph(
+        mock_ld_client, node_key='root-agent', graph_key='g', tool_names=['search']
+    )
+
+    tool_item = _make_tool_call_item('root-agent', 'search')
+    run_result = _make_run_result(
+        output='ok', total_tokens=10, input_tokens=7, output_tokens=3,
+        tool_call_items=[tool_item],
+    )
+
+    with patch.dict('sys.modules', _make_agents_modules(run_result)):
+        runner = OpenAIAgentGraphRunner(graph, _tool_registry('search'))
+        await runner.run('go')
+
+    ev = _events(mock_ld_client)
+
+    # Collect runIds from node-level events
+    run_ids = set()
+    for event_name in (
+        '$ld:ai:tokens:total', '$ld:ai:tokens:input', '$ld:ai:tokens:output',
+        '$ld:ai:generation:success', '$ld:ai:generation:duration', '$ld:ai:tool_call',
+    ):
+        for data, _ in ev.get(event_name, []):
+            if data.get('configKey') == 'root-agent':
+                run_ids.add(data['runId'])
+
+    # All events must share a single runId
+    assert len(run_ids) == 1
+
+
 @pytest.mark.asyncio
 async def test_does_not_track_tool_calls_without_graph_and_registry_config():
     """RunResult tool items that are not backed by graph + registry tools are ignored."""