feat!: Move graph_key to AIConfigTracker instantiation

packages/ai-providers/server-ai-langchain/src/ldai_langchain/langgraph_callback_handler.py

@@ -197,7 +197,6 @@ def flush(self, graph: AgentGraphDefinition, graph_tracker: Any) -> None:
         :param graph: The AgentGraphDefinition whose nodes hold the LD config trackers.
         :param graph_tracker: The AIGraphTracker for the overall graph (may be None).
         """
-        gk = graph_tracker.graph_key if graph_tracker is not None else None
         for node_key in self._path:
             node = graph.get_node(node_key)
             if not node:
@@ -208,13 +207,13 @@ def flush(self, graph: AgentGraphDefinition, graph_tracker: Any) -> None:
 
             usage = self._node_tokens.get(node_key)
             if usage:
-                config_tracker.track_tokens(usage, graph_key=gk)
+                config_tracker.track_tokens(usage)
 
             duration = self._node_duration_ms.get(node_key)
             if duration is not None:
-                config_tracker.track_duration(duration, graph_key=gk)
+                config_tracker.track_duration(duration)
 
-            config_tracker.track_success(graph_key=gk)
+            config_tracker.track_success()
 
             for tool_key in self._node_tool_calls.get(node_key, []):
-                config_tracker.track_tool_call(tool_key, graph_key=gk)
+                config_tracker.track_tool_call(tool_key)

packages/ai-providers/server-ai-langchain/tests/test_langgraph_callback_handler.py

@@ -35,6 +35,7 @@ def _make_graph(mock_ld_client: MagicMock, node_key: str = 'root-agent', graph_k
         model_name='gpt-4',
         provider_name='openai',
         context=context,
+        graph_key=graph_key,
     )
     graph_tracker = AIGraphTracker(
         ld_client=mock_ld_client,
@@ -389,16 +390,47 @@ def test_flush_includes_graph_key_in_node_events():
     assert token_data.get('graphKey') == 'my-graph'
 
 
-def test_flush_with_none_tracker_uses_no_graph_key():
-    """flush() with graph_tracker=None does not fail and omits graphKey."""
+def test_flush_with_no_graph_key_on_node_tracker():
+    """When node tracker has no graph_key, events omit graphKey."""
     mock_ld_client = MagicMock()
-    graph = _make_graph(mock_ld_client)
+    context = MagicMock()
+    node_tracker = LDAIConfigTracker(
+        ld_client=mock_ld_client,
+        variation_key='v1',
+        config_key='root-agent',
+        version=1,
+        model_name='gpt-4',
+        provider_name='openai',
+        context=context,
+    )
+    node_config = AIAgentConfig(
+        key='root-agent',
+        enabled=True,
+        model=ModelConfig(name='gpt-4', parameters={}),
+        provider=ProviderConfig(name='openai'),
+        instructions='Be helpful.',
+        tracker=node_tracker,
+    )
+    graph_config = AIAgentGraphConfig(
+        key='test-graph',
+        root_config_key='root-agent',
+        edges=[],
+        enabled=True,
+    )
+    nodes = AgentGraphDefinition.build_nodes(graph_config, {'root-agent': node_config})
+    graph = AgentGraphDefinition(
+        agent_graph=graph_config,
+        nodes=nodes,
+        context=context,
+        enabled=True,
+        tracker=None,
+    )
 
     handler = LDMetricsCallbackHandler({'root-agent'}, {})
     node_run_id = uuid4()
     handler.on_chain_start({}, {}, run_id=node_run_id, name='root-agent')
     handler.on_llm_end(_llm_result(5, 3, 2), run_id=uuid4(), parent_run_id=node_run_id)
-    handler.flush(graph, None)  # graph_tracker=None
+    handler.flush(graph, None)
 
     ev = _events(mock_ld_client)
     token_data = ev['$ld:ai:tokens:total'][0][0]

packages/ai-providers/server-ai-langchain/tests/test_tracking_langgraph.py

@@ -45,6 +45,7 @@ def _make_graph(
         model_name='gpt-4',
         provider_name='openai',
         context=context,
+        graph_key=graph_key,
     )
 
     graph_tracker = AIGraphTracker(
@@ -141,6 +142,7 @@ def _make_two_node_graph(mock_ld_client: MagicMock) -> 'AgentGraphDefinition':
         model_name='gpt-4',
         provider_name='openai',
         context=context,
+        graph_key='two-node-graph',
     )
     child_tracker = LDAIConfigTracker(
         ld_client=mock_ld_client,
@@ -150,6 +152,7 @@ def _make_two_node_graph(mock_ld_client: MagicMock) -> 'AgentGraphDefinition':
         model_name='gpt-4',
         provider_name='openai',
         context=context,
+        graph_key='two-node-graph',
     )
     graph_tracker = AIGraphTracker(
         ld_client=mock_ld_client,
@@ -517,6 +520,7 @@ def _node_tracker(key: str) -> LDAIConfigTracker:
             model_name='gpt-4',
             provider_name='openai',
             context=context,
+            graph_key='multi-child-graph',
         )
 
     graph_tracker = AIGraphTracker(
@@ -627,6 +631,7 @@ def _node_tracker(key: str) -> LDAIConfigTracker:
             model_name='gpt-4',
             provider_name='openai',
             context=context,
+            graph_key='multi-child-tools-graph',
         )
 
     graph_tracker = AIGraphTracker(

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

@@ -248,13 +248,12 @@ def _handle_handoff(
         except Exception:
             pass
 
-        gk = tracker.graph_key if tracker is not None else None
         if config_tracker is not None:
             if usage is not None:
-                config_tracker.track_tokens(usage, graph_key=gk)
+                config_tracker.track_tokens(usage)
             if duration_ms is not None:
-                config_tracker.track_duration(int(duration_ms), graph_key=gk)
-            config_tracker.track_success(graph_key=gk)
+                config_tracker.track_duration(int(duration_ms))
+            config_tracker.track_success()
 
     def _flush_final_segment(
         self,
@@ -283,15 +282,13 @@ def _flush_final_segment(
         except Exception:
             pass
 
-        gk = tracker.graph_key if tracker is not None else None
         if usage is not None:
-            config_tracker.track_tokens(usage, graph_key=gk)
-        config_tracker.track_duration(int(duration_ms), graph_key=gk)
-        config_tracker.track_success(graph_key=gk)
+            config_tracker.track_tokens(usage)
+        config_tracker.track_duration(int(duration_ms))
+        config_tracker.track_success()
 
     def _track_tool_calls(self, result: Any, tracker: Any) -> None:
         """Track all tool calls from the run result, attributed to the node that called them."""
-        gk = tracker.graph_key if tracker is not None else None
         for agent_name, tool_fn_name in get_tool_calls_from_run_items(result.new_items):
             agent_key = self._agent_name_map.get(agent_name, agent_name)
             tool_name = self._tool_name_map.get(tool_fn_name)
@@ -302,4 +299,4 @@ def _track_tool_calls(self, result: Any, tracker: Any) -> None:
                 continue
             config_tracker = node.get_config().tracker
             if config_tracker is not None:
-                config_tracker.track_tool_call(tool_name, graph_key=gk)
+                config_tracker.track_tool_call(tool_name)

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

@@ -40,6 +40,7 @@ def _make_graph(
         model_name='gpt-4',
         provider_name='openai',
         context=context,
+        graph_key=graph_key,
     )
 
     graph_tracker = AIGraphTracker(
@@ -178,6 +179,7 @@ def _make_two_node_graph(mock_ld_client: MagicMock) -> AgentGraphDefinition:
         model_name='gpt-4',
         provider_name='openai',
         context=context,
+        graph_key='two-node-graph',
     )
     child_tracker = LDAIConfigTracker(
         ld_client=mock_ld_client,
@@ -187,6 +189,7 @@ def _make_two_node_graph(mock_ld_client: MagicMock) -> AgentGraphDefinition:
         model_name='gpt-4',
         provider_name='openai',
         context=context,
+        graph_key='two-node-graph',
     )
     graph_tracker = AIGraphTracker(
         ld_client=mock_ld_client,

packages/sdk/server-ai/src/ldai/client.py

@@ -606,9 +606,12 @@ def agent_graph(
             for single_edge in variation.get("edges", {}).get(edge_key, []):
                 all_agent_keys.add(single_edge.get("key", ""))
 
+        graph_key_value = key
         agent_configs = {
-            key: self.agent_config(key, context, AIAgentConfigDefault(enabled=False))
-            for key in all_agent_keys
+            agent_key: self.__evaluate_agent(
+                agent_key, context, AIAgentConfigDefault(enabled=False), graph_key=graph_key_value
+            )
+            for agent_key in all_agent_keys
         }
 
         if not all(config.enabled for config in agent_configs.values()):
@@ -748,6 +751,7 @@ def __evaluate(
         context: Context,
         default_dict: Dict[str, Any],
         variables: Optional[Dict[str, Any]] = None,
+        graph_key: Optional[str] = None,
     ) -> Tuple[
         Optional[ModelConfig], Optional[ProviderConfig], Optional[List[LDMessage]],
         Optional[str], LDAIConfigTracker, bool, Optional[Any], Dict[str, Any]
@@ -759,6 +763,7 @@ def __evaluate(
         :param context: The evaluation context.
         :param default_dict: Default configuration as dictionary.
         :param variables: Variables for interpolation.
+        :param graph_key: When set, passed to the tracker so all events include ``graphKey``.
         :return: Tuple of (model, provider, messages, instructions, tracker, enabled, judge_configuration, variation).
         """
         variation = self._client.variation(key, context, default_dict)
@@ -809,6 +814,7 @@ def __evaluate(
             model.name if model else '',
             provider_config.name if provider_config else '',
             context,
+            graph_key=graph_key,
         )
 
         enabled = variation.get('_ldMeta', {}).get('enabled', False)
@@ -836,6 +842,7 @@ def __evaluate_agent(
         context: Context,
         default: AIAgentConfigDefault,
         variables: Optional[Dict[str, Any]] = None,
+        graph_key: Optional[str] = None,
     ) -> AIAgentConfig:
         """
         Internal method to evaluate an agent configuration.
@@ -844,10 +851,11 @@ def __evaluate_agent(
         :param context: The evaluation context.
         :param default: Default agent values.
         :param variables: Variables for interpolation.
+        :param graph_key: When set, passed to the tracker so all events include ``graphKey``.
         :return: Configured AIAgentConfig instance.
         """
         model, provider, messages, instructions, tracker, enabled, judge_configuration, _ = self.__evaluate(
-            key, context, default.to_dict(), variables
+            key, context, default.to_dict(), variables, graph_key=graph_key
         )
 
         # For agents, prioritize instructions over messages