fix: Graph tracking refactor — ManagedAgentGraph drives tracking for new runner shape (#154)

Author: jsonbailey

packages/sdk/server-ai/poetry.lock

@@ -799,7 +799,7 @@ async def create_agent_graph(
         if not runner:
             return None
 
-        return ManagedAgentGraph(runner)
+        return ManagedAgentGraph(graph, runner)
 
     def agents(
         self,

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

@@ -1,61 +1,101 @@
 """ManagedAgentGraph — LaunchDarkly managed wrapper for agent graph execution."""
 
-from typing import Any
+from typing import Any, Dict
 
-from ldai.providers import AgentGraphResult, AgentGraphRunner
-from ldai.providers.types import GraphMetricSummary, ManagedGraphResult
+from ldai.agent_graph import AgentGraphDefinition
+from ldai.providers import AgentGraphRunner
+from ldai.providers.types import (
+    AgentGraphRunnerResult,
+    LDAIMetrics,
+    ManagedGraphResult,
+)
+from ldai.tracker import LDAIMetricSummary
 
 
 class ManagedAgentGraph:
     """
     LaunchDarkly managed wrapper for AI agent graph execution.
 
-    Holds an AgentGraphRunner. Wraps the runner result in a
-    :class:`~ldai.providers.types.ManagedGraphResult` and builds a
-    :class:`~ldai.providers.types.GraphMetricSummary` from the runner's metrics.
+    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()``.
     """
 
     def __init__(
         self,
+        graph: AgentGraphDefinition,
         runner: AgentGraphRunner,
     ):
         """
         Initialize ManagedAgentGraph.
 
+        :param graph: The AgentGraphDefinition used to drive graph-level and
+            per-node tracking from the runner result metrics.
         :param runner: The AgentGraphRunner to delegate execution to
         """
+        self._graph = graph
         self._runner = runner
 
     async def run(self, input: Any) -> ManagedGraphResult:
         """
         Run the agent graph with the given input.
 
+        Delegates to the underlying AgentGraphRunner, then drives all
+        LaunchDarkly tracking from ``result.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 (currently always ``None`` for graphs —
-            per-graph evaluations will be added in a future PR).
+        :return: ManagedGraphResult containing the content, metric summary,
+            and raw response.
         """
-        result: AgentGraphResult = await self._runner.run(input)
-
-        # 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).
-        metrics = result.metrics
-        summary = GraphMetricSummary(
-            success=metrics.success,
-            usage=metrics.usage,
-            duration_ms=getattr(metrics, 'duration_ms', None),
+        graph_tracker = self._graph.create_tracker()
+        result = await graph_tracker.track_graph_metrics_of_async(
+            lambda r: r.metrics,
+            lambda: self._runner.run(input),
         )
 
+        summary = graph_tracker.get_summary()
+        summary.node_metrics = self._track_node_metrics(result.metrics.node_metrics)
+
         return ManagedGraphResult(
-            content=result.output,
+            content=result.content,
             metrics=summary,
             raw=result.raw,
             evaluations=None,
         )
 
+    def _track_node_metrics(
+        self, node_metrics: Dict[str, LDAIMetrics]
+    ) -> Dict[str, LDAIMetricSummary]:
+        """
+        Drive per-node LaunchDarkly tracking events and collect node metric summaries.
+
+        For each node key present in ``node_metrics``, obtains the node's
+        config tracker via the graph definition, fires tracking events, and
+        returns a map of node key to the tracker's metric summary.
+        """
+        node_summaries: Dict[str, LDAIMetricSummary] = {}
+        for node_key, node_ldai_metrics in node_metrics.items():
+            node = self._graph.get_node(node_key)
+            if node is None:
+                continue
+            node_tracker = node.get_config().create_tracker()
+
+            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()
+
+            node_summaries[node_key] = node_tracker.get_summary()
+        return node_summaries
+
     def get_agent_graph_runner(self) -> AgentGraphRunner:
         """
         Return the underlying AgentGraphRunner for advanced use.

packages/sdk/server-ai/src/ldai/managed_agent_graph.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
         """
         ...

packages/sdk/server-ai/src/ldai/providers/agent_graph_runner.py

@@ -110,8 +110,8 @@ class GraphMetrics:
 class GraphMetricSummary:
     """Contains a summary of metrics for an agent graph run."""
 
-    success: bool
-    """Whether the graph run succeeded."""
+    success: Optional[bool] = None
+    """Whether the graph run succeeded. Absent if invocation status has not been tracked."""
 
     path: List[str] = field(default_factory=list)
     """Ordered list of node keys visited during the run."""
@@ -122,8 +122,8 @@ class GraphMetricSummary:
     usage: Optional[TokenUsage] = None
     """Optional aggregate token usage information across all nodes in the graph run."""
 
-    node_metrics: Dict[str, LDAIMetrics] = field(default_factory=dict)
-    """Per-node metrics keyed by node key."""
+    node_metrics: Dict[str, LDAIMetricSummary] = field(default_factory=dict)
+    """Per-node metric summaries keyed by node key."""
 
     resumption_token: Optional[str] = None
     """Optional resumption token from the graph tracker for cross-process resumption."""