feat: migrate ManagedAgentGraph to new Runner protocol

jsonbailey · claude · jsonbailey · commit f8fd3a8728b4 · 2026-05-04T09:11:53.000-05:00
ManagedAgentGraph now expects AgentGraphRunnerResult (with GraphMetrics) from
all runners and drives graph-level and per-node LD tracking from result.metrics.
AgentGraphRunner protocol updated to return AgentGraphRunnerResult.

Co-Authored-By: Claude Sonnet 4.6 &lt;noreply@anthropic.com&gt;
diff --git a/packages/sdk/server-ai/src/ldai/managed_agent_graph.py b/packages/sdk/server-ai/src/ldai/managed_agent_graph.py
@@ -2,11 +2,10 @@
 
 from typing import Any, Optional
 
-from ldai.providers import AgentGraphResult, AgentGraphRunner
+from ldai.providers import AgentGraphRunner
 from ldai.providers.types import (
     AgentGraphRunnerResult,
     GraphMetricSummary,
-    LDAIMetrics,
     ManagedGraphResult,
 )
 
@@ -15,19 +14,9 @@ class ManagedAgentGraph:
     """
     LaunchDarkly managed wrapper for AI agent graph execution.
 
-    Holds an AgentGraphRunner and an optional AgentGraphDefinition. Wraps the
-    runner result in a :class:`~ldai.providers.types.ManagedGraphResult` and
-    builds a :class:`~ldai.providers.types.GraphMetricSummary` from the runner's
-    metrics.
-
-    When the runner returns an :class:`~ldai.providers.types.AgentGraphRunnerResult`
-    (new shape), the managed layer drives all graph-level tracking from
-    ``result.metrics``.  When the runner returns the legacy
-    :class:`~ldai.providers.AgentGraphResult`, tracking has already been performed
-    inside the runner; the managed layer simply wraps the result.  This detection
-    branch exists as a deliberate bridge: once PR 11-openai and PR 11-langchain
-    migrate both runners to return ``AgentGraphRunnerResult``, the legacy branch
-    becomes dead code and will be removed in PR 11-langchain's final cleanup commit.
+    Holds an AgentGraphRunner and an AgentGraphDefinition. Delegates execution
+    to the runner, then drives all graph-level and per-node tracking from the
+    returned :class:`~ldai.providers.types.AgentGraphRunnerResult`.
 
     Obtain an instance via ``LDAIClient.create_agent_graph()``.
     """
@@ -41,10 +30,8 @@ def __init__(
         Initialize ManagedAgentGraph.
 
         :param runner: The AgentGraphRunner to delegate execution to
-        :param graph: Optional AgentGraphDefinition used to create the
-            graph-level tracker when the runner returns an
-            :class:`AgentGraphRunnerResult` (new shape).  Not needed for
-            legacy runners that still return :class:`AgentGraphResult`.
+        :param graph: Optional AgentGraphDefinition used to drive graph-level and
+            per-node tracking from the runner result metrics.
         """
         self._runner = runner
         self._graph = graph
@@ -53,50 +40,31 @@ async def run(self, input: Any) -> ManagedGraphResult:
         """
         Run the agent graph with the given input.
 
-        Delegates to the underlying AgentGraphRunner.  The returned type
-        determines which tracking path is taken:
+        Delegates to the underlying AgentGraphRunner, then drives all
+        LaunchDarkly tracking from ``result.metrics``:
 
-        - :class:`AgentGraphRunnerResult` (new shape): the managed layer drives
-          graph-level tracking from ``result.metrics`` via the graph tracker.
-          Per-node tracking from ``result.metrics.node_metrics`` will be wired
-          in a follow-up commit once the runners populate ``node_metrics``.
-        - :class:`AgentGraphResult` (legacy shape): tracking already occurred
-          inside the runner; the managed layer wraps the result without
-          additional tracking.
+        - Graph-level events (path, duration, success/failure, total tokens) via
+          the graph tracker obtained from the graph definition.
+        - Per-node events (tokens, duration, tool calls, success) via per-node
+          trackers for each key present in ``result.metrics.node_metrics``.
 
         :param input: The input prompt or structured input for the graph
         :return: ManagedGraphResult containing the content, metric summary,
-            raw response, and an optional evaluations task (always ``None``
-            for now — per-graph evaluations will be added in a future PR).
+            and raw response.
         """
-        raw_result = await self._runner.run(input)
-
-        if isinstance(raw_result, AgentGraphRunnerResult):
-            # New shape: managed layer drives all tracking.
-            summary = self._build_summary_from_runner_result(raw_result)
-            if self._graph is not None:
-                self._flush_graph_tracking(raw_result, self._graph.create_tracker())
-            return ManagedGraphResult(
-                content=raw_result.content,
-                metrics=summary,
-                raw=raw_result.raw,
-                evaluations=None,
-            )
-
-        # Legacy shape (AgentGraphResult): tracking already happened in the runner.
-        # Build a GraphMetricSummary from the runner result's LDAIMetrics.
-        # path and node_metrics will be populated once graph runners are migrated
-        # to return AgentGraphRunnerResult with GraphMetrics (PR 11-openai/langchain).
-        metrics: LDAIMetrics = raw_result.metrics
-        summary = GraphMetricSummary(
-            success=metrics.success,
-            usage=metrics.usage,
-            duration_ms=getattr(metrics, 'duration_ms', None),
-        )
+        result = await self._runner.run(input)
+
+        summary = self._build_summary_from_runner_result(result)
+
+        if self._graph is not None:
+            graph_tracker = self._graph.create_tracker()
+            self._flush_graph_tracking(result, graph_tracker)
+            self._flush_node_tracking(result)
+
         return ManagedGraphResult(
-            content=raw_result.output,
+            content=result.content,
             metrics=summary,
-            raw=raw_result.raw,
+            raw=result.raw,
             evaluations=None,
         )
 
@@ -117,10 +85,6 @@ def _build_summary_from_runner_result(
     def _flush_graph_tracking(self, result: AgentGraphRunnerResult, tracker: Any) -> None:
         """
         Drive graph-level LaunchDarkly tracking events from runner result metrics.
-
-        Called only when the runner returns the new ``AgentGraphRunnerResult``
-        shape.  Node-level tracking (from ``result.metrics.node_metrics``) will
-        be wired once the runners start populating that field.
         """
         m = result.metrics
         if m.path:
@@ -134,6 +98,36 @@ def _flush_graph_tracking(self, result: AgentGraphRunnerResult, tracker: Any) ->
         if m.usage is not None:
             tracker.track_total_tokens(m.usage)
 
+    def _flush_node_tracking(self, result: AgentGraphRunnerResult) -> None:
+        """
+        Drive per-node LaunchDarkly tracking events from ``result.metrics.node_metrics``.
+
+        For each node key present in ``node_metrics``, obtains the node's
+        config tracker via the graph definition and fires token, duration,
+        tool call, and success/error events.
+        """
+        if self._graph is None:
+            return
+
+        for node_key, node_ldai_metrics in result.metrics.node_metrics.items():
+            node = self._graph.get_node(node_key)
+            if node is None:
+                continue
+            node_tracker = node.get_config().create_tracker()
+            if node_tracker is None:
+                continue
+
+            if node_ldai_metrics.usage is not None:
+                node_tracker.track_tokens(node_ldai_metrics.usage)
+            if node_ldai_metrics.duration_ms is not None:
+                node_tracker.track_duration(node_ldai_metrics.duration_ms)
+            if node_ldai_metrics.tool_calls:
+                node_tracker.track_tool_calls(node_ldai_metrics.tool_calls)
+            if node_ldai_metrics.success:
+                node_tracker.track_success()
+            else:
+                node_tracker.track_error()
+
     def get_agent_graph_runner(self) -> AgentGraphRunner:
         """
         Return the underlying AgentGraphRunner for advanced use.
diff --git a/packages/sdk/server-ai/src/ldai/providers/agent_graph_runner.py b/packages/sdk/server-ai/src/ldai/providers/agent_graph_runner.py
@@ -1,6 +1,6 @@
 from typing import Any, Protocol, runtime_checkable
 
-from ldai.providers.types import AgentGraphResult
+from ldai.providers.types import AgentGraphRunnerResult
 
 
 @runtime_checkable
@@ -18,11 +18,11 @@ class AgentGraphRunner(Protocol):
     the caller just passes input.
     """
 
-    async def run(self, input: Any) -> AgentGraphResult:
+    async def run(self, input: Any) -> AgentGraphRunnerResult:
         """
         Run the agent graph with the given input.
 
         :param input: The input to the agent graph (string prompt or structured input)
-        :return: AgentGraphResult containing the output, raw response, and metrics
+        :return: AgentGraphRunnerResult containing the content, raw response, and GraphMetrics
         """
         ...
diff --git a/packages/sdk/server-ai/tests/test_managed_agent_graph.py b/packages/sdk/server-ai/tests/test_managed_agent_graph.py
@@ -7,27 +7,27 @@
 
 from ldai import LDAIClient, ManagedAgentGraph, ManagedGraphResult
 from ldai.providers.types import AgentGraphRunnerResult, GraphMetrics, LDAIMetrics
-from ldai.providers import AgentGraphResult, AgentGraphRunner, ToolRegistry
+from ldai.providers import AgentGraphRunner, ToolRegistry
 from ldai.tracker import TokenUsage
 
 
 # --- Test doubles ---
 
 class StubAgentGraphRunner(AgentGraphRunner):
-    """Legacy runner that returns AgentGraphResult (old shape)."""
-    def __init__(self, output: str = "stub output"):
-        self._output = output
+    """Runner that returns AgentGraphRunnerResult (new shape)."""
+    def __init__(self, content: str = "stub output"):
+        self._content = content
 
-    async def run(self, input) -> AgentGraphResult:
-        return AgentGraphResult(
-            output=self._output,
+    async def run(self, input) -> AgentGraphRunnerResult:
+        return AgentGraphRunnerResult(
+            content=self._content,
             raw={"input": input},
-            metrics=LDAIMetrics(success=True),
+            metrics=GraphMetrics(success=True),
         )
 
 
-class StubNewShapeRunner(AgentGraphRunner):
-    """New-shape runner that returns AgentGraphRunnerResult."""
+class StubRunnerWithMetrics(AgentGraphRunner):
+    """Runner that returns AgentGraphRunnerResult with full GraphMetrics."""
     def __init__(self, content: str = "new shape output"):
         self._content = content
 
@@ -39,17 +39,28 @@ async def run(self, input) -> AgentGraphRunnerResult:
                 path=["root", "specialist"],
                 duration_ms=42,
                 usage=TokenUsage(total=10, input=5, output=5),
-                node_metrics={},
+                node_metrics={
+                    "root": LDAIMetrics(
+                        success=True,
+                        usage=TokenUsage(total=5, input=3, output=2),
+                        duration_ms=20,
+                    ),
+                    "specialist": LDAIMetrics(
+                        success=True,
+                        usage=TokenUsage(total=5, input=2, output=3),
+                        duration_ms=22,
+                    ),
+                },
             ),
             raw={"input": input},
         )
 
 
-# --- ManagedAgentGraph unit tests (legacy shape) ---
+# --- ManagedAgentGraph unit tests ---
 
 @pytest.mark.asyncio
 async def test_managed_agent_graph_run_delegates_to_runner():
-    """Legacy AgentGraphResult shape: content comes from output field."""
+    """Runner result content is surfaced correctly."""
     runner = StubAgentGraphRunner("hello world")
     managed = ManagedAgentGraph(runner)
     result = await managed.run("test input")
@@ -64,15 +75,14 @@ def test_managed_agent_graph_get_runner():
     assert managed.get_agent_graph_runner() is runner
 
 
-# --- ManagedAgentGraph unit tests (new AgentGraphRunnerResult shape) ---
-
 @pytest.mark.asyncio
-async def test_managed_agent_graph_run_handles_new_shape():
-    """New AgentGraphRunnerResult shape: content and GraphMetrics are surfaced."""
-    runner = StubNewShapeRunner("final answer")
+async def test_managed_agent_graph_run_surfaces_graph_metrics():
+    """GraphMetrics fields are reflected in GraphMetricSummary."""
+    runner = StubRunnerWithMetrics("final answer")
     mock_graph = MagicMock()
     mock_tracker = MagicMock()
     mock_graph.create_tracker = MagicMock(return_value=mock_tracker)
+    mock_graph.get_node = MagicMock(return_value=None)  # no nodes for this test
 
     managed = ManagedAgentGraph(runner, graph=mock_graph)
     result = await managed.run("test input")
@@ -87,12 +97,13 @@ async def test_managed_agent_graph_run_handles_new_shape():
 
 
 @pytest.mark.asyncio
-async def test_managed_agent_graph_new_shape_drives_tracking():
-    """New shape: managed layer calls tracker methods from result.metrics."""
-    runner = StubNewShapeRunner()
+async def test_managed_agent_graph_drives_graph_level_tracking():
+    """Managed layer calls graph tracker methods from result.metrics."""
+    runner = StubRunnerWithMetrics()
     mock_graph = MagicMock()
     mock_tracker = MagicMock()
     mock_graph.create_tracker = MagicMock(return_value=mock_tracker)
+    mock_graph.get_node = MagicMock(return_value=None)
 
     managed = ManagedAgentGraph(runner, graph=mock_graph)
     await managed.run("test input")
@@ -104,16 +115,73 @@ async def test_managed_agent_graph_new_shape_drives_tracking():
 
 
 @pytest.mark.asyncio
-async def test_managed_agent_graph_new_shape_no_graph_skips_tracking():
-    """New shape without graph: no tracking called (graph not available)."""
-    runner = StubNewShapeRunner()
+async def test_managed_agent_graph_drives_per_node_tracking():
+    """Managed layer creates per-node trackers and fires node-level events."""
+    runner = StubRunnerWithMetrics()
+    mock_graph = MagicMock()
+    mock_graph_tracker = MagicMock()
+    mock_graph.create_tracker = MagicMock(return_value=mock_graph_tracker)
+
+    root_tracker = MagicMock()
+    specialist_tracker = MagicMock()
+
+    root_node = MagicMock()
+    root_node.get_config.return_value.create_tracker = MagicMock(return_value=root_tracker)
+    specialist_node = MagicMock()
+    specialist_node.get_config.return_value.create_tracker = MagicMock(return_value=specialist_tracker)
+
+    def get_node(key):
+        return {"root": root_node, "specialist": specialist_node}.get(key)
+
+    mock_graph.get_node = get_node
+
+    managed = ManagedAgentGraph(runner, graph=mock_graph)
+    await managed.run("test input")
+
+    # root node tracking
+    root_tracker.track_tokens.assert_called_once()
+    root_tracker.track_duration.assert_called_once_with(20)
+    root_tracker.track_success.assert_called_once()
+
+    # specialist node tracking
+    specialist_tracker.track_tokens.assert_called_once()
+    specialist_tracker.track_duration.assert_called_once_with(22)
+    specialist_tracker.track_success.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_managed_agent_graph_no_graph_skips_tracking():
+    """Without a graph reference, no tracking is called but run succeeds."""
+    runner = StubRunnerWithMetrics()
     managed = ManagedAgentGraph(runner, graph=None)
-    # Should not raise even without a graph reference
     result = await managed.run("test input")
     assert result.content == "new shape output"
     assert result.metrics.success is True
 
 
+@pytest.mark.asyncio
+async def test_managed_agent_graph_failure_calls_track_invocation_failure():
+    """On a failed run, track_invocation_failure is called instead of success."""
+
+    class FailingRunner(AgentGraphRunner):
+        async def run(self, input) -> AgentGraphRunnerResult:
+            return AgentGraphRunnerResult(
+                content='',
+                raw=None,
+                metrics=GraphMetrics(success=False, duration_ms=5),
+            )
+
+    mock_graph = MagicMock()
+    mock_tracker = MagicMock()
+    mock_graph.create_tracker = MagicMock(return_value=mock_tracker)
+    mock_graph.get_node = MagicMock(return_value=None)
+
+    managed = ManagedAgentGraph(FailingRunner(), graph=mock_graph)
+    result = await managed.run("test input")
+
+    assert result.metrics.success is False
+    mock_tracker.track_invocation_failure.assert_called_once()
+    mock_tracker.track_invocation_success.assert_not_called()
 
 
 # --- LDAIClient.create_agent_graph() integration tests ---
diff --git a/packages/sdk/server-ai/tests/test_runner_abcs.py b/packages/sdk/server-ai/tests/test_runner_abcs.py
