feat: migrate ManagedAgentGraph to new Runner protocol, update AI providers

ManagedAgentGraph now expects AgentGraphRunnerResult (with GraphMetrics) from
all runners and drives graph-level and per-node LD tracking from result.metrics.
OpenAIAgentGraphRunner and LangGraphAgentGraphRunner are migrated to return
AgentGraphRunnerResult; all direct tracker calls are removed from the runners.

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

Author: jsonbailey

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

@@ -7,8 +7,8 @@
 
 from ldai import log
 from ldai.agent_graph import AgentGraphDefinition, AgentGraphNode
-from ldai.providers import AgentGraphResult, AgentGraphRunner, ToolRegistry
-from ldai.providers.types import LDAIMetrics
+from ldai.providers import AgentGraphRunner, ToolRegistry
+from ldai.providers.types import AgentGraphRunnerResult, GraphMetrics, LDAIMetrics
 
 from ldai_langchain.langchain_helper import (
     build_structured_tools,
@@ -65,9 +65,10 @@ class LangGraphAgentGraphRunner(AgentGraphRunner):
 
     AgentGraphRunner implementation for LangGraph.
 
-    Compiles and runs the agent graph with LangGraph and automatically records
-    graph- and node-level AI metric data to the LaunchDarkly trackers on the
-    graph definition and each node.
+    Compiles and runs the agent graph with LangGraph and collects per-node AI
+    metric data via a LangChain callback handler.  All LaunchDarkly tracking is
+    driven by the managed layer (``ManagedAgentGraph``) from the returned
+    :class:`~ldai.providers.types.AgentGraphRunnerResult`.
 
     Requires ``langgraph`` to be installed.
     """
@@ -298,20 +299,21 @@ def route(state: WorkflowState) -> str:
         compiled = agent_builder.compile()
         return compiled, fn_name_to_config_key, node_keys
 
-    async def run(self, input: Any) -> AgentGraphResult:
+    async def run(self, input: Any) -> AgentGraphRunnerResult:
         """
         Run the agent graph with the given input.
 
         Builds a LangGraph StateGraph from the AgentGraphDefinition, compiles
         it, and invokes it. Uses a LangChain callback handler to collect
-        per-node metrics, then flushes them to LaunchDarkly trackers.
+        per-node metrics, then returns them as an
+        :class:`~ldai.providers.types.AgentGraphRunnerResult` with
+        :class:`~ldai.providers.types.GraphMetrics`.
 
         :param input: The string prompt to send to the agent graph
-        :return: AgentGraphResult with the final output and metrics
+        :return: AgentGraphRunnerResult with the final content and GraphMetrics
         """
         pending_eval_tasks: Dict[str, List[asyncio.Task]] = {}
         token = _run_eval_tasks.set(pending_eval_tasks)
-        tracker = self._graph.create_tracker()
         start_ns = time.perf_counter_ns()
 
         try:
@@ -329,20 +331,26 @@ async def run(self, input: Any) -> AgentGraphResult:
             messages = result.get('messages', [])
             output = extract_last_message_content(messages)
 
-            # Flush per-node metrics to LD trackers; eval results are tracked
-            # internally and intentionally not exposed on AgentGraphResult here
-            # — judge dispatch is the managed layer's responsibility.
-            await handler.flush(self._graph, pending_eval_tasks)
+            # Await eval tasks collected during node traversal (judge results are
+            # returned in node_metrics but not tracked here — the managed layer
+            # is responsible for driving tracker calls).
+            await handler.flush_eval_tasks(pending_eval_tasks)
 
-            tracker.track_path(handler.path)
-            tracker.track_duration(duration)
-            tracker.track_invocation_success()
-            tracker.track_total_tokens(sum_token_usage_from_messages(messages))
+            total_usage = sum_token_usage_from_messages(messages)
 
-            return AgentGraphResult(
-                output=output,
+            # Build per-node LDAIMetrics from the callback handler's collected data.
+            node_metrics = self._build_node_metrics(handler)
+
+            return AgentGraphRunnerResult(
+                content=output,
                 raw=result,
-                metrics=LDAIMetrics(success=True),
+                metrics=GraphMetrics(
+                    success=True,
+                    path=handler.path,
+                    duration_ms=duration,
+                    usage=total_usage,
+                    node_metrics=node_metrics,
+                ),
             )
 
         except Exception as exc:
@@ -354,12 +362,28 @@ async def run(self, input: Any) -> AgentGraphResult:
             else:
                 log.warning(f'LangGraphAgentGraphRunner run failed: {exc}')
             duration = (time.perf_counter_ns() - start_ns) // 1_000_000
-            tracker.track_duration(duration)
-            tracker.track_invocation_failure()
-            return AgentGraphResult(
-                output='',
+            return AgentGraphRunnerResult(
+                content='',
                 raw=None,
-                metrics=LDAIMetrics(success=False),
+                metrics=GraphMetrics(
+                    success=False,
+                    duration_ms=duration,
+                ),
             )
         finally:
             _run_eval_tasks.reset(token)
+
+    def _build_node_metrics(self, handler: LDMetricsCallbackHandler) -> Dict[str, LDAIMetrics]:
+        """Build per-node LDAIMetrics from callback handler collected data."""
+        node_metrics: Dict[str, LDAIMetrics] = {}
+        for node_key in handler.path:
+            usage = handler.node_tokens.get(node_key)
+            duration_ms = handler.node_durations_ms.get(node_key)
+            tool_calls = handler.node_tool_calls.get(node_key) or []
+            node_metrics[node_key] = LDAIMetrics(
+                success=True,
+                usage=usage,
+                duration_ms=duration_ms,
+                tool_calls=tool_calls if tool_calls else None,
+            )
+        return node_metrics

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

@@ -189,12 +189,35 @@ def on_tool_end(
     # Flush
     # ------------------------------------------------------------------
 
+    async def flush_eval_tasks(self, eval_tasks: Optional[Dict[str, List[Any]]] = None) -> None:
+        """
+        Await all pending judge evaluation tasks collected during the graph run.
+
+        Tracker calls are driven by the managed layer (``ManagedAgentGraph``)
+        from the returned ``GraphMetrics``.  This method only ensures eval
+        coroutines are awaited so they complete and do not leak.
+
+        :param eval_tasks: Dict mapping node key to a list of awaitables that
+            return judge evaluation results.
+        """
+        if not eval_tasks:
+            return
+        for node_key in self._path:
+            for eval_task in eval_tasks.get(node_key, []):
+                await eval_task
+
     async def flush(
         self, graph: AgentGraphDefinition, eval_tasks=None
     ) -> List[JudgeResult]:
         """
         Emit all collected per-node metrics to the LaunchDarkly trackers.
 
+        .. deprecated::
+            Use ``flush_eval_tasks()`` instead.  Direct tracker calls from
+            the runner are no longer supported; the managed layer
+            (``ManagedAgentGraph``) drives all tracking from the returned
+            ``GraphMetrics``.
+
         Call this once after the graph run completes.
 
         :param graph: The AgentGraphDefinition whose nodes hold the LD config trackers.

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

@@ -6,7 +6,7 @@
 from ldai.agent_graph import AgentGraphDefinition
 from ldai.evaluator import Evaluator
 from ldai.models import AIAgentGraphConfig, AIAgentConfig, ModelConfig, ProviderConfig
-from ldai.providers import AgentGraphResult, ToolRegistry
+from ldai.providers import AgentGraphRunnerResult, ToolRegistry
 from ldai_langchain.langgraph_agent_graph_runner import LangGraphAgentGraphRunner
 from ldai_langchain.langchain_runner_factory import LangChainRunnerFactory
 
@@ -75,28 +75,36 @@ async def test_langgraph_runner_run_raises_when_langgraph_not_installed():
 
     with patch.dict('sys.modules', {'langgraph': None, 'langgraph.graph': None}):
         result = await runner.run("test")
-        assert isinstance(result, AgentGraphResult)
+        assert isinstance(result, AgentGraphRunnerResult)
         assert result.metrics.success is False
 
 
 @pytest.mark.asyncio
 async def test_langgraph_runner_run_tracks_failure_on_exception():
     graph = _make_graph()
-    tracker = graph.create_tracker()
     runner = LangGraphAgentGraphRunner(graph, {})
 
     with patch.dict('sys.modules', {'langgraph': None, 'langgraph.graph': None}):
         result = await runner.run("fail")
 
     assert result.metrics.success is False
-    tracker.track_invocation_failure.assert_called_once()
-    tracker.track_duration.assert_called_once()
+    assert result.metrics.duration_ms is not None
+
+    # Runner no longer calls graph-level tracker directly.
+    graph_tracker = graph.create_tracker()
+    graph_tracker.track_invocation_failure.assert_not_called()
 
 
 @pytest.mark.asyncio
 async def test_langgraph_runner_run_success():
+    """
+    Tests successful run by injecting a pre-compiled graph mock.
+
+    Bypasses ``_build_graph()`` (which requires a real LangChain/LangGraph
+    installation) by setting ``_compiled``, ``_node_keys``, and
+    ``_fn_name_to_config_key`` directly on the runner before calling ``run()``.
+    """
     graph = _make_graph()
-    tracker = graph.create_tracker()
 
     mock_message = MagicMock()
     mock_message.content = "langgraph answer"
@@ -106,50 +114,33 @@ async def test_langgraph_runner_run_success():
     mock_compiled = MagicMock()
     mock_compiled.ainvoke = AsyncMock(return_value={'messages': [mock_message]})
 
-    mock_state_graph_instance = MagicMock()
-    mock_state_graph_instance.add_node = MagicMock()
-    mock_state_graph_instance.add_edge = MagicMock()
-    mock_state_graph_instance.compile = MagicMock(return_value=mock_compiled)
-
-    mock_langgraph_graph = MagicMock()
-    mock_langgraph_graph.END = 'END'
-    mock_langgraph_graph.START = 'START'
-    mock_langgraph_graph.StateGraph = MagicMock(return_value=mock_state_graph_instance)
-
-    mock_human_message = MagicMock()
     mock_lc_core_messages = MagicMock()
+    mock_human_message = MagicMock()
     mock_lc_core_messages.HumanMessage = MagicMock(return_value=mock_human_message)
-    mock_lc_core_messages.AnyMessage = MagicMock()
-
-    mock_model_response = MagicMock()
-    mock_model_response.content = 'langgraph answer'
-    mock_model_response.usage_metadata = None
-    mock_model_response.response_metadata = None
-    mock_model_response.tool_calls = None
 
-    mock_llm = MagicMock()
-    mock_llm.ainvoke = AsyncMock(return_value=mock_model_response)
-
-    mock_init_model = MagicMock()
-    mock_init_model.return_value = mock_llm
-    mock_langchain_chat = MagicMock()
-    mock_langchain_chat.init_chat_model = mock_init_model
+    runner = LangGraphAgentGraphRunner(graph, {})
+    # Inject pre-compiled state to avoid needing a real LangGraph/LangChain install.
+    runner._compiled = mock_compiled
+    runner._node_keys = {'root-agent'}
+    runner._fn_name_to_config_key = {}
 
     with patch.dict('sys.modules', {
-        'langgraph': MagicMock(),
-        'langgraph.graph': mock_langgraph_graph,
-        'langchain_core': MagicMock(),
         'langchain_core.messages': mock_lc_core_messages,
-        'langchain': MagicMock(),
-        'langchain.chat_models': mock_langchain_chat,
-        'typing_extensions': __import__('typing_extensions'),
     }):
-        runner = LangGraphAgentGraphRunner(graph, {})
         result = await runner.run("find restaurants")
 
-    assert isinstance(result, AgentGraphResult)
-    assert result.output == "langgraph answer"
+    assert isinstance(result, AgentGraphRunnerResult)
+    assert result.content == "langgraph answer"
     assert result.metrics.success is True
-    tracker.track_path.assert_called_once_with([])
-    tracker.track_invocation_success.assert_called_once()
-    tracker.track_duration.assert_called_once()
+    assert result.metrics.path == []
+    assert result.metrics.duration_ms is not None
+
+    # Runner no longer calls graph-level tracker directly.
+    graph_tracker = graph.create_tracker()
+    graph_tracker.track_path.assert_not_called()
+    graph_tracker.track_invocation_success.assert_not_called()
+    graph_tracker.track_duration.assert_not_called()
+
+    # Runner no longer calls node-level trackers directly.
+    node_factory = graph.get_node('root-agent').get_config().create_tracker
+    node_factory.assert_not_called()

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

@@ -11,6 +11,7 @@
 from unittest.mock import AsyncMock, MagicMock, patch
 
 from ldai.agent_graph import AgentGraphDefinition
+from ldai.managed_agent_graph import ManagedAgentGraph
 from ldai.models import AIAgentGraphConfig, AIAgentConfig, Edge, ModelConfig, ProviderConfig
 from ldai.tracker import AIGraphTracker, LDAIConfigTracker
 from ldai.evaluator import Evaluator
@@ -226,16 +227,16 @@ async def test_tracks_node_and_graph_tokens_on_success():
     with patch('ldai_langchain.langgraph_agent_graph_runner.create_langchain_model',
                return_value=_mock_model(fake_response)):
         runner = LangGraphAgentGraphRunner(graph, {})
-        result = await runner.run("What's the weather?")
+        managed = ManagedAgentGraph(runner, graph=graph)
+        result = await managed.run("What's the weather?")
 
     assert result.metrics.success is True
-    assert result.output == 'Sunny.'
+    assert result.content == 'Sunny.'
 
     # Manually simulate what the callback handler would collect and flush
     # (mock models don't fire LangChain callbacks, so we test flush directly)
     mock_ld_client2 = MagicMock()
     graph2 = _make_graph(mock_ld_client2)
-    tracker2 = graph2.create_tracker()
 
     handler = LDMetricsCallbackHandler({'root-agent'}, {})
     node_run_id = uuid4()
@@ -259,7 +260,7 @@ async def test_tracks_node_and_graph_tokens_on_success():
     assert ev2['$ld:ai:generation:success'][0][1] == 1
     assert '$ld:ai:duration:total' in ev2
 
-    # Graph-level events from the real run
+    # Graph-level events from the real run (fired by ManagedAgentGraph)
     ev = _events(mock_ld_client)
     assert ev['$ld:ai:graph:total_tokens'][0][1] == 15
     assert ev['$ld:ai:graph:invocation_success'][0][1] == 1
@@ -277,7 +278,8 @@ async def test_tracks_execution_path():
     with patch('ldai_langchain.langgraph_agent_graph_runner.create_langchain_model',
                return_value=_mock_model(fake_response)):
         runner = LangGraphAgentGraphRunner(graph, {})
-        await runner.run('hello')
+        managed = ManagedAgentGraph(runner, graph=graph)
+        await managed.run('hello')
 
     ev = _events(mock_ld_client)
     path_data = ev['$ld:ai:graph:path'][0][0]
@@ -429,7 +431,8 @@ async def test_tracks_failure_and_latency_on_model_error():
     with patch('ldai_langchain.langgraph_agent_graph_runner.create_langchain_model',
                return_value=error_model):
         runner = LangGraphAgentGraphRunner(graph, {})
-        result = await runner.run('fail')
+        managed = ManagedAgentGraph(runner, graph=graph)
+        result = await managed.run('fail')
 
     assert result.metrics.success is False
 
@@ -461,14 +464,14 @@ def model_factory(node_config, **kwargs):
     with patch('ldai_langchain.langgraph_agent_graph_runner.create_langchain_model',
                side_effect=model_factory):
         runner = LangGraphAgentGraphRunner(graph, {})
-        result = await runner.run('hello')
+        managed = ManagedAgentGraph(runner, graph=graph)
+        result = await managed.run('hello')
 
     assert result.metrics.success is True
 
     # Simulate per-node token events via callback handler (mock models don't fire callbacks)
     mock_ld_client2 = MagicMock()
     graph2 = _make_two_node_graph(mock_ld_client2)
-    tracker2 = graph2.create_tracker()
 
     handler = LDMetricsCallbackHandler({'root-agent', 'child-agent'}, {})
 
@@ -624,7 +627,7 @@ def model_factory(node_config, **kwargs):
         result = await runner.run('hello')
 
     assert result.metrics.success is True
-    assert 'Agent A' in result.output
+    assert 'Agent A' in result.content
     # Agent B's model must never have been invoked — no fan-out
     agent_b_model.ainvoke.assert_not_called()
 
@@ -752,7 +755,7 @@ def model_factory(node_config, **kwargs):
         result = await runner.run('Find info and route to the right agent.')
 
     assert result.metrics.success is True
-    assert 'Agent A' in result.output
+    assert 'Agent A' in result.content
     # Orchestrator must have been called twice: once before tool result, once after
     assert orchestrator_model.ainvoke.call_count == 2
     # Agent B must never have been invoked