fix: Update provider packages for graph_key-on-init refactor

Remove graph_key kwarg from tracker method calls in langchain and
openai provider packages. Fix line length in client.py. Update test
helpers to pass graph_key at node tracker construction.

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

Author: jsonbailey

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

@@ -608,7 +608,9 @@ def agent_graph(
 
         graph_key_value = key
         agent_configs = {
-            agent_key: self.__evaluate_agent(agent_key, context, AIAgentConfigDefault(enabled=False), graph_key=graph_key_value)
+            agent_key: self.__evaluate_agent(
+                agent_key, context, AIAgentConfigDefault(enabled=False), graph_key=graph_key_value
+            )
             for agent_key in all_agent_keys
         }