fix: Address code review findings in graph runner and callback handler

- Cache compiled graph in _ensure_compiled() so _build_graph is not called
  on every run() invocation
- Collect node_keys during traversal instead of reaching into _graph._nodes
- Fix unconditional break in make_after_tools_router that made reversed()
  a no-op; replace broken loop with a direct last-message check
- Fix duration tracking key collision: key _node_start_ns by run_id instead
  of node_key so concurrent invocations of the same node don't clobber
  each other's start times
- Warn when a functional-tool node has multiple outgoing edges since only
  the first edge is reachable after the tool loop exits

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

@@ -1,7 +1,7 @@
 """LangGraph agent graph runner for LaunchDarkly AI SDK."""
 
 import time
-from typing import Annotated, Any, Dict, List, Tuple
+from typing import Annotated, Any, Dict, List, Optional, Set, Tuple
 
 from ldai import log
 from ldai.agent_graph import AgentGraphDefinition, AgentGraphNode
@@ -76,13 +76,25 @@ def __init__(self, graph: AgentGraphDefinition, tools: ToolRegistry):
         """
         self._graph = graph
         self._tools = tools
-
-    def _build_graph(self) -> Tuple[Any, Dict[str, str]]:
+        self._compiled: Any = None
+        self._fn_name_to_config_key: Dict[str, str] = {}
+        self._node_keys: Set[str] = set()
+
+    def _ensure_compiled(self) -> None:
+        """Build and cache the compiled graph if not already done."""
+        if self._compiled is None:
+            compiled, fn_name_to_config_key, node_keys = self._build_graph()
+            self._compiled = compiled
+            self._fn_name_to_config_key = fn_name_to_config_key
+            self._node_keys = node_keys
+
+    def _build_graph(self) -> Tuple[Any, Dict[str, str], Set[str]]:
         """
         Build and compile the LangGraph StateGraph from the AgentGraphDefinition.
 
-        :return: Tuple of (compiled_graph, fn_name_to_config_key) where
-            fn_name_to_config_key maps tool function __name__ to LD config key.
+        :return: Tuple of (compiled_graph, fn_name_to_config_key, node_keys) where
+            fn_name_to_config_key maps tool function __name__ to LD config key, and
+            node_keys is the set of all agent node keys in the graph.
         """
         from langchain_core.messages import SystemMessage
         from langgraph.graph import END, START, StateGraph
@@ -99,10 +111,12 @@ class WorkflowState(TypedDict):
         tools_ref = self._tools
         graph_structure: List[str] = []
         fn_name_to_config_key: Dict[str, str] = {}
+        node_keys: Set[str] = set()
 
         def handle_traversal(node: AgentGraphNode, ctx: dict) -> None:
             node_config = node.get_config()
             node_key = node.get_key()
+            node_keys.add(node_key)
             instructions = node_config.instructions if hasattr(node_config, 'instructions') else None
             outgoing_edges = node.get_edges()
 
@@ -190,6 +204,12 @@ async def invoke(state: WorkflowState) -> dict:
                 if not handoff_fns:
                     # No handoff tools: standard loop-back after tool execution.
                     after_loop = outgoing_edges[0].target_config if outgoing_edges else END
+                    if len(outgoing_edges) > 1:
+                        log.warning(
+                            f"Node '{node_key}' has {len(outgoing_edges)} outgoing edges but no handoff "
+                            "tools; only the first edge will be used after the tool loop. "
+                            "Use handoff tools for multi-child routing."
+                        )
                     agent_builder.add_edge(tools_node_key, node_key)
                     agent_builder.add_conditional_edges(
                         node_key,
@@ -212,10 +232,11 @@ async def invoke(state: WorkflowState) -> dict:
 
                     def make_after_tools_router(parent_key: str, ht_names: frozenset):
                         def route(state: WorkflowState) -> str:
-                            for msg in reversed(state['messages']):
-                                if hasattr(msg, 'name') and msg.name:
-                                    return END if msg.name in ht_names else parent_key
-                                break
+                            msgs = state['messages']
+                            if msgs:
+                                last = msgs[-1]
+                                if hasattr(last, 'name') and last.name in ht_names:
+                                    return END
                             return parent_key
                         return route
 
@@ -247,7 +268,7 @@ def route(state: WorkflowState) -> str:
         )
 
         compiled = agent_builder.compile()
-        return compiled, fn_name_to_config_key
+        return compiled, fn_name_to_config_key, node_keys
 
     async def run(self, input: Any) -> AgentGraphResult:
         """
@@ -266,12 +287,10 @@ async def run(self, input: Any) -> AgentGraphResult:
         try:
             from langchain_core.messages import HumanMessage
 
-            compiled, fn_name_to_config_key = self._build_graph()
-
-            node_keys = {node.get_key() for node in self._graph._nodes.values()}
-            handler = LDMetricsCallbackHandler(node_keys, fn_name_to_config_key)
+            self._ensure_compiled()
+            handler = LDMetricsCallbackHandler(self._node_keys, self._fn_name_to_config_key)
 
-            result = await compiled.ainvoke(  # type: ignore[call-overload]
+            result = await self._compiled.ainvoke(  # type: ignore[call-overload]
                 {'messages': [HumanMessage(content=str(input))]},
                 config={'callbacks': [handler], 'recursion_limit': 25},
             )

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

@@ -42,8 +42,8 @@ def __init__(self, node_keys: Set[str], fn_name_to_config_key: Dict[str, str]):
         self._node_tokens: Dict[str, TokenUsage] = {}
         # tool config keys called per node
         self._node_tool_calls: Dict[str, List[str]] = {}
-        # start time (ns) per node — only set while running
-        self._node_start_ns: Dict[str, int] = {}
+        # start time (ns) per active run_id — keyed by run_id to handle re-entrant nodes
+        self._node_start_ns: Dict[UUID, int] = {}
         # accumulated duration (ms) per node
         self._node_duration_ms: Dict[str, int] = {}
         # execution path in order (deduplicated)
@@ -96,7 +96,7 @@ def on_chain_start(
 
         if name in self._node_keys:
             self._run_to_node[run_id] = name
-            self._node_start_ns[name] = time.perf_counter_ns()
+            self._node_start_ns[run_id] = time.perf_counter_ns()
             if name not in self._path_set:
                 self._path.append(name)
                 self._path_set.add(name)
@@ -117,7 +117,7 @@ def on_chain_end(
         node_key = self._run_to_node.get(run_id)
         if node_key is None:
             return
-        start_ns = self._node_start_ns.pop(node_key, None)
+        start_ns = self._node_start_ns.pop(run_id, None)
         if start_ns is not None:
             elapsed_ms = (time.perf_counter_ns() - start_ns) // 1_000_000
             self._node_duration_ms[node_key] = (