fix(tracing): isolate trace context per async execution context (OPEN-12420, OPEN-12421, OPEN-12426)

src/lib/core/cli.ts

@@ -13,7 +13,7 @@
 import { program } from 'commander';
 import * as fs from 'fs';
 import * as path from 'path';
-import { getCurrentTrace, postProcessTrace } from '../tracing/tracer';
+import { getCurrentTrace, postProcessTrace, runInTraceContext } from '../tracing/tracer';
 
 // Define shared interfaces and utilities here
 export interface RunReturn {
@@ -141,40 +141,48 @@ class CLIHandler {
     const { data: dataset } = loadDataset(datasetPath);
 
     // Process each item in the dataset dynamically
-    Promise.all<Output>(
-      dataset.map(async (item: any) => {
-        try {
-          const result = await this.run(item);
-          // Merge the original item fields with the result
-          const traceData = getCurrentTrace() ?? undefined;
-          const postProcessedTrace =
-            typeof traceData === 'undefined' || traceData === null ?
-              undefined
-            : postProcessTrace(traceData)?.traceData;
-
-          const output: Output = {
-            ...item,
-            ...result.otherFields,
-            output: result.output,
-            steps: traceData?.toJSON(),
-            latency: postProcessedTrace?.latency,
-            cost: postProcessedTrace?.cost,
-            tokens: postProcessedTrace?.tokens,
-            metadata: {
-              ...(postProcessedTrace?.metadata ?? {}),
-              inputVariableNames: postProcessedTrace?.inputVariableNames,
-            },
-          };
-
-          return output;
-        } catch (error) {
-          console.error('Error processing dataset: ', error);
-          return {
-            ...item,
-            error: error instanceof Error ? error.message : String(error),
-          };
-        }
-      }),
+    return Promise.all<Output>(
+      // Each row runs in its own trace context so concurrent rows root their
+      // own traces rather than nesting into whichever row started first.
+      dataset.map((item: any) =>
+        runInTraceContext(async () => {
+          try {
+            const result = await this.run(item);
+            // Merge the original item fields with the result
+            const traceData = getCurrentTrace() ?? undefined;
+            // postProcessTrace returns { traceData, inputVariableNames } as
+            // siblings, so keep the whole result — the names are not on
+            // traceData itself.
+            const processed =
+              typeof traceData === 'undefined' || traceData === null ?
+                undefined
+              : postProcessTrace(traceData);
+            const postProcessedTrace = processed?.traceData;
+
+            const output: Output = {
+              ...item,
+              ...result.otherFields,
+              output: result.output,
+              steps: traceData?.toJSON(),
+              latency: postProcessedTrace?.latency,
+              cost: postProcessedTrace?.cost,
+              tokens: postProcessedTrace?.tokens,
+              metadata: {
+                ...(postProcessedTrace?.metadata ?? {}),
+                inputVariableNames: processed?.inputVariableNames,
+              },
+            };
+
+            return output;
+          } catch (error) {
+            console.error('Error processing dataset: ', error);
+            return {
+              ...item,
+              error: error instanceof Error ? error.message : String(error),
+            };
+          }
+        }),
+      ),
     )
       .then((results) => {
         /*

src/lib/tracing/index.d.ts

@@ -9,6 +9,7 @@ export {
   startAgentStep,
   addHandoffStepToTrace,
   configure,
+  runInTraceContext,
   replayBufferedTraces,
   getBufferStatus,
   clearOfflineBuffer,

src/lib/tracing/tracer.ts

@@ -15,8 +15,42 @@ import {
 import Openlayer, { type ClientOptions } from '../../index';
 import type { DataStreamParams } from '../../resources/inference-pipelines/data';
 import { OfflineBuffer } from './offlineBuffer';
+import { AsyncLocalStorage } from 'node:async_hooks';
 
-let currentTrace: Trace | null = null;
+/**
+ * Trace context: the in-flight trace plus the stack of open steps.
+ *
+ * This lives in an AsyncLocalStorage so that concurrently traced units of work
+ * (dataset rows in development mode, parallel requests in monitoring) each root
+ * their own trace instead of nesting into whichever one happened to start first.
+ *
+ * Callers that never establish a context share `defaultContext`, which behaves
+ * exactly like the module-level globals this replaced.
+ */
+interface TraceContext {
+  trace: Trace | null;
+  stepStack: Step[];
+}
+
+const traceStorage = new AsyncLocalStorage<TraceContext>();
+
+const defaultContext: TraceContext = { trace: null, stepStack: [] };
+
+function ctx(): TraceContext {
+  return traceStorage.getStore() ?? defaultContext;
+}
+
+/**
+ * Runs `fn` in an isolated trace context, so steps created inside it can never
+ * attach to a trace started outside it (or by a sibling running concurrently).
+ *
+ * Note this must wrap the work with a callback boundary — `enterWith()` is not
+ * a substitute, because the synchronous prelude of an async function runs in
+ * its *caller's* context and would leak the store to sibling calls.
+ */
+export function runInTraceContext<T>(fn: () => T): T {
+  return traceStorage.run({ trace: null, stepStack: [] }, fn);
+}
 
 // Lazy-initialized Openlayer client to ensure environment variables are loaded
 let client: Openlayer | null = null;
@@ -118,16 +152,9 @@ function getOpenlayerClient(): Openlayer | null {
 }
 
 export function getCurrentTrace(): Trace | null {
-  return currentTrace;
-}
-
-function setCurrentTrace(trace: Trace | null) {
-  currentTrace = trace;
+  return ctx().trace;
 }
 
-// Function to create a new step
-const stepStack: Step[] = [];
-
 function createStep(
   name: string,
   stepType: StepType = StepType.USER_CALL,
@@ -144,22 +171,27 @@ function createStep(
   const inferencePipelineId =
     openlayerInferencePipelineId || configuredPipelineId || process.env['OPENLAYER_INFERENCE_PIPELINE_ID'];
 
+  // Bind the context at creation time. endStep may be invoked from a different
+  // async context than the one that opened the step — frameworks call
+  // completion callbacks from wherever they happen to be — so resolving the
+  // store again inside endStep could reach the wrong context, or none at all.
+  const stepCtx = ctx();
+
   const parentStep = getCurrentStep();
   const isRootStep = parentStep === null;
 
   if (isRootStep) {
     console.debug('Starting a new trace...');
     console.debug(`Adding step ${name} as the root step`);
-    const currentTrace = new Trace();
-    setCurrentTrace(currentTrace);
-    currentTrace.addStep(newStep);
+    const newTrace = new Trace();
+    stepCtx.trace = newTrace;
+    newTrace.addStep(newStep);
   } else {
     console.debug(`Adding step ${name} as a nested step to ${parentStep!.name}`);
-    currentTrace = getCurrentTrace()!;
     parentStep!.addNestedStep(newStep);
   }
 
-  stepStack.push(newStep);
+  stepCtx.stepStack.push(newStep);
 
   const endStep = () => {
     // Calculate latency for this step before removing from stack
@@ -173,12 +205,19 @@ function createStep(
       }
     }
 
-    stepStack.pop(); // Remove the current step from the stack
+    // Remove *this* step, not whatever is on top: within a single context a
+    // user's own Promise.all over parallel tool calls can end steps out of
+    // order, and popping blindly would evict an unrelated step.
+    const { stepStack } = stepCtx;
+    const stepIndex = stepStack.lastIndexOf(newStep);
+    if (stepIndex !== -1) {
+      stepStack.splice(stepIndex, 1);
+    }
     console.debug(`Ending step ${newStep.name}`);
 
     if (isRootStep) {
       console.debug('Ending the trace...');
-      const traceData = getCurrentTrace();
+      const traceData = stepCtx.trace;
 
       // NOTE: currentTrace is intentionally NOT reset here — integrations and
       // tests inspect the completed trace via getCurrentTrace() after the root
@@ -193,6 +232,7 @@ function createStep(
 }
 
 export function getCurrentStep(): Step | null | undefined {
+  const { stepStack } = ctx();
   const currentStep = stepStack.length > 0 ? stepStack[stepStack.length - 1] : null;
   return currentStep;
 }
@@ -686,6 +726,17 @@ export function addGuardrailStepToTrace(params: {
   return { step, endStep };
 }
 
+/**
+ * Totals a numeric ChatCompletionStep field across a step and all of its
+ * descendants. Only chat-completion steps carry `cost` / `tokens`; every other
+ * step type contributes 0.
+ */
+function sumStepField(step: Step, field: 'cost' | 'tokens'): number {
+  const own = (step as ChatCompletionStep)[field];
+  const ownValue = typeof own === 'number' ? own : 0;
+  return step.steps.reduce((total, nested) => total + sumStepField(nested, field), ownValue);
+}
+
 export function postProcessTrace(traceObj: Trace): { traceData: any; inputVariableNames: string[] } {
   const rootStep = traceObj.steps[0];
 
@@ -700,8 +751,11 @@ export function postProcessTrace(traceObj: Trace): { traceData: any; inputVariab
     output: rootStep!.output,
     groundTruth: rootStep!.groundTruth,
     latency: rootStep!.latency,
-    cost: (rootStep as ChatCompletionStep)!.cost,
-    tokens: (rootStep as ChatCompletionStep)!.tokens,
+    // Totals for the whole trace, not just the root: an agent- or chain-rooted
+    // trace carries its cost and tokens on nested LLM steps, and reading only
+    // the root silently reported nothing for them.
+    cost: sumStepField(rootStep!, 'cost'),
+    tokens: sumStepField(rootStep!, 'tokens'),
     steps: processed_steps,
     metadata: rootStep!.metadata,
   };

tests/cli-concurrency.test.ts

@@ -0,0 +1,130 @@
+import * as fs from 'fs';
+import * as os from 'os';
+import * as path from 'path';
+
+/**
+ * These tests exercise the development-mode runner (`CLIHandler`) over a
+ * multi-row dataset. Regression coverage for OPEN-12420.
+ */
+
+const makeTempDir = (): string => fs.mkdtempSync(path.join(os.tmpdir(), 'ol-cli-conc-'));
+
+const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms));
+
+/**
+ * `commander`'s `program` is a module singleton, so each run needs a fresh
+ * module graph. cli.ts and tracer.ts must be imported after the same reset so
+ * they share one tracer instance.
+ */
+const freshModules = async () => {
+  jest.resetModules();
+  const cli = await import('../src/lib/core/cli');
+  const tracer = await import('../src/lib/tracing/tracer');
+  return { CLIHandler: cli.default, tracer };
+};
+
+const writeDataset = (dir: string, rows: unknown[]): string => {
+  const datasetPath = path.join(dir, 'dataset.json');
+  fs.writeFileSync(datasetPath, JSON.stringify(rows), 'utf8');
+  return datasetPath;
+};
+
+describe('CLIHandler dataset run', () => {
+  const originalArgv = process.argv;
+  const originalDisablePublish = process.env['OPENLAYER_DISABLE_PUBLISH'];
+
+  beforeEach(() => {
+    process.env['OPENLAYER_DISABLE_PUBLISH'] = 'true';
+  });
+
+  afterEach(() => {
+    process.argv = originalArgv;
+    if (originalDisablePublish === undefined) {
+      delete process.env['OPENLAYER_DISABLE_PUBLISH'];
+    } else {
+      process.env['OPENLAYER_DISABLE_PUBLISH'] = originalDisablePublish;
+    }
+  });
+
+  it('resolves only after the output files are written', async () => {
+    const dir = makeTempDir();
+    const datasetPath = writeDataset(dir, [{ userQuery: 'row-A' }]);
+    const outputDir = path.join(dir, 'out');
+
+    const { CLIHandler } = await freshModules();
+    const handler = new CLIHandler(async ({ userQuery }: { userQuery: string }) => {
+      await sleep(10);
+      return { output: `out-${userQuery}`, otherFields: {} };
+    });
+
+    process.argv = ['node', 'probe', '--dataset-path', datasetPath, '--output-dir', outputDir];
+
+    await handler.runFromCLI();
+
+    expect(fs.existsSync(path.join(outputDir, 'dataset.json'))).toBe(true);
+    expect(fs.existsSync(path.join(outputDir, 'config.json'))).toBe(true);
+  });
+
+  it('gives each row its own root trace instead of nesting rows into each other', async () => {
+    const dir = makeTempDir();
+    const datasetPath = writeDataset(dir, [
+      { userQuery: 'row-A' },
+      { userQuery: 'row-B' },
+      { userQuery: 'row-C' },
+    ]);
+    const outputDir = path.join(dir, 'out');
+
+    const { CLIHandler, tracer } = await freshModules();
+
+    // Staggered delays make the rows genuinely interleave: row-A is still
+    // in-flight when row-B and row-C start their steps.
+    const delays: Record<string, number> = { 'row-A': 60, 'row-B': 30, 'row-C': 10 };
+
+    const handler = new CLIHandler(async ({ userQuery }: { userQuery: string }) => {
+      const { endStep } = tracer.addChainStepToTrace({
+        name: `step-for-${userQuery}`,
+        inputs: { userQuery },
+      });
+      await sleep(delays[userQuery]!);
+      endStep();
+      return { output: `out-${userQuery}`, otherFields: {} };
+    });
+
+    process.argv = ['node', 'probe', '--dataset-path', datasetPath, '--output-dir', outputDir];
+
+    await handler.runFromCLI();
+
+    const rows = JSON.parse(fs.readFileSync(path.join(outputDir, 'dataset.json'), 'utf8'));
+    expect(rows).toHaveLength(3);
+
+    for (const row of rows) {
+      // Each row owns exactly one root step, and that step is its own.
+      expect(row.steps).toHaveLength(1);
+      expect(row.steps[0].name).toBe(`Handoffs: step-for-${row.userQuery}`);
+      // No other row leaked in as a nested step.
+      expect(row.steps[0].steps ?? []).toHaveLength(0);
+      // Per-row latency is recorded for every row, not just the first.
+      expect(typeof row.latency).toBe('number');
+    }
+  });
+
+  it('writes inputVariableNames into the generated config', async () => {
+    const dir = makeTempDir();
+    const datasetPath = writeDataset(dir, [{ userQuery: 'row-A' }]);
+    const outputDir = path.join(dir, 'out');
+
+    const { CLIHandler, tracer } = await freshModules();
+    const handler = new CLIHandler(async ({ userQuery }: { userQuery: string }) => {
+      const { endStep } = tracer.addChainStepToTrace({ name: 'step', inputs: { userQuery } });
+      endStep();
+      return { output: `out-${userQuery}`, otherFields: {} };
+    });
+
+    process.argv = ['node', 'probe', '--dataset-path', datasetPath, '--output-dir', outputDir];
+
+    await handler.runFromCLI();
+
+    const config = JSON.parse(fs.readFileSync(path.join(outputDir, 'config.json'), 'utf8'));
+    expect(config.inputVariableNames).toEqual(['userQuery']);
+  });
+});