fix: collect and await judge eval tasks instead of fire-and-forget

Replace asyncio.create_task fire-and-forget with proper task collection
and awaiting in both OpenAI and LangGraph runners, ensuring judge results
are tracked reliably. Use ContextVar in LangGraph runner to isolate
pending eval task state across concurrent run() calls.

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

@@ -2,6 +2,7 @@
 
 import asyncio
 import time
+from contextvars import ContextVar
 from typing import Annotated, Any, Dict, List, Set, Tuple
 
 from ldai import log
@@ -17,6 +18,9 @@
 )
 from ldai_langchain.langgraph_callback_handler import LDMetricsCallbackHandler
 
+# Per-run eval task accumulator, isolated per concurrent run() call via ContextVar.
+_run_eval_tasks: ContextVar[Dict[str, List[asyncio.Task]]] = ContextVar('_run_eval_tasks')
+
 
 def _make_handoff_tool(child_key: str, description: str) -> Any:
     """
@@ -84,7 +88,6 @@ def __init__(
         self._compiled: Any = None
         self._fn_name_to_config_key: Dict[str, str] = {}
         self._node_keys: Set[str] = set()
-        self._pending_eval_tasks: Dict[str, List[asyncio.Task]] = {}
 
     def _ensure_compiled(self) -> None:
         """Build and cache the compiled graph if not already done."""
@@ -189,7 +192,7 @@ async def invoke(state: WorkflowState) -> dict:
                             response.content if hasattr(response, 'content') else str(response)
                         )
                         task = node_obj.get_config().evaluator.evaluate(input_text, output_text)
-                        self._pending_eval_tasks.setdefault(nk, []).append(task)
+                        _run_eval_tasks.get({}).setdefault(nk, []).append(task)
 
                     return {'messages': [response]}
 
@@ -299,7 +302,8 @@ async def run(self, input: Any) -> AgentGraphResult:
         :param input: The string prompt to send to the agent graph
         :return: AgentGraphResult with the final output and metrics
         """
-        self._pending_eval_tasks = {}
+        pending_eval_tasks: Dict[str, List[asyncio.Task]] = {}
+        token = _run_eval_tasks.set(pending_eval_tasks)
         tracker = self._graph.create_tracker() if self._graph.create_tracker is not None else None
         start_ns = time.perf_counter_ns()
 
@@ -319,7 +323,7 @@ async def run(self, input: Any) -> AgentGraphResult:
             output = extract_last_message_content(messages)
 
             # Flush per-node metrics to LD trackers
-            await handler.flush(self._graph, self._pending_eval_tasks)
+            await handler.flush(self._graph, pending_eval_tasks)
 
             # Graph-level metrics
             if tracker:
@@ -351,3 +355,5 @@ async def run(self, input: Any) -> AgentGraphResult:
                 raw=None,
                 metrics=LDAIMetrics(success=False),
             )
+        finally:
+            _run_eval_tasks.reset(token)

packages/ai-providers/server-ai-openai/src/ldai_openai/openai_agent_graph_runner.py

@@ -1,4 +1,3 @@
-import asyncio
 import re
 import time
 from typing import Any, Dict, List, Optional
@@ -30,6 +29,7 @@ def __init__(self, last_handoff_ns: int, last_node_key: str, input_str: str = ''
         self.last_handoff_ns = last_handoff_ns
         self.last_node_key = last_node_key
         self.input_str = input_str
+        self.pending_eval_tasks: List[tuple] = []
 
 
 class OpenAIAgentGraphRunner(AgentGraphRunner):
@@ -91,6 +91,11 @@ async def run(self, input: Any) -> AgentGraphResult:
             root_agent = self._build_agents(path, state, tracker)
             result = await Runner.run(root_agent, input_str)
             self._flush_final_segment(state, result, input_str)
+            for node_tracker, eval_task in state.pending_eval_tasks:
+                eval_results = await eval_task
+                for r in eval_results:
+                    if r.success:
+                        node_tracker.track_judge_result(r)
             self._track_tool_calls(result)
 
             duration = (time.perf_counter_ns() - start_ns) // 1_000_000
@@ -106,7 +111,7 @@ async def run(self, input: Any) -> AgentGraphResult:
             return AgentGraphResult(
                 output=str(result.final_output),
                 raw=result,
-                metrics=LDAIMetrics(success=True),
+                metrics=LDAIMetrics(success=True, usage=token_usage),
             )
         except Exception as exc:
             if isinstance(exc, ImportError):
@@ -270,19 +275,10 @@ def _handle_handoff(
                 config_tracker.track_duration(int(duration_ms))
             config_tracker.track_success()
 
-            # Fire judge evaluation for the src node (fire-and-forget)
             src_node = self._graph.get_node(src)
             if src_node is not None:
-                evaluator = src_node.get_config().evaluator
-                # Use empty string as output since we don't have the node's final output here
-                eval_task = evaluator.evaluate(input_str, '')
-
-                async def _track(trk, et):
-                    results = await et
-                    for r in results:
-                        if r.success:
-                            trk.track_judge_result(r)
-                asyncio.create_task(_track(config_tracker, eval_task))
+                eval_task = src_node.get_config().evaluator.evaluate(input_str, '')
+                state.pending_eval_tasks.append((config_tracker, eval_task))
 
     def _flush_final_segment(
         self,
@@ -313,19 +309,11 @@ def _flush_final_segment(
         config_tracker.track_duration(int(duration_ms))
         config_tracker.track_success()
 
-        # Fire judge evaluation for the final node (fire-and-forget)
         final_node = self._graph.get_node(state.last_node_key)
         if final_node is not None:
-            evaluator = final_node.get_config().evaluator
             output_str = str(result.final_output) if result is not None else ''
-            eval_task = evaluator.evaluate(input_str, output_str)
-
-            async def _track(trk, et):
-                results = await et
-                for r in results:
-                    if r.success:
-                        trk.track_judge_result(r)
-            asyncio.create_task(_track(config_tracker, eval_task))
+            eval_task = final_node.get_config().evaluator.evaluate(input_str, output_str)
+            state.pending_eval_tasks.append((config_tracker, eval_task))
 
     def _track_tool_calls(self, result: Any) -> None:
         """Track all tool calls from the run result, attributed to the node that called them."""