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

[viniciusdsmello]

Fixes OPEN-12420, OPEN-12421 and OPEN-12426.

1. Concurrent rows corrupt development-mode traces (OPEN-12420)

The development-mode runner (CLIHandler) processes dataset rows concurrently via Promise.all, but the tracer held trace state in module-level globals. Every row after the first nested into the first row's trace instead of rooting its own.

Measured on a 3-row dataset with one traced step per row and staggered awaits:

row-A | steps: ['Handoffs: step-for-row-A', '  nested: step-for-row-B'] | latency=61
row-B | steps: ['Handoffs: step-for-row-A', '  nested: step-for-row-B'] | latency=None
row-C | steps: ['Handoffs: step-for-row-A', '  nested: step-for-row-B'] | latency=None

All three rows share one identical steps blob (A → B → C nested), only one trace uploads for the whole dataset, and config.json declares latencyColumnName while 2 of 3 rows carry latency: null.

Python is unaffected: OpenlayerModel.run_batch_from_df iterates sequentially, which keeps the same globals safe.

Fix. Trace + step-stack state moved into an AsyncLocalStorage, with every accessor routed through a single ctx() helper. Callers that never establish a context share one mutable defaultContext that behaves exactly like the globals it replaced — so existing integrations and their getCurrentTrace()-after-root assertions are untouched. CLIHandler wraps each row in the new exported runInTraceContext().

enterWith() is not a substitute, which is why the wrapper takes a callback. The synchronous prelude of an async function runs in its caller's context, so enterWith inside dataset.map(async ...) writes into the shared parent and leaks to every sibling:

--- enterWith (auto-root, no wrapper) ---
A: entered-as=A   B: entered-as=A   C: entered-as=A     <- all leaked
--- run() wrapper per row ---
A: entered-as=A   B: entered-as=B   C: entered-as=C     <- isolated

Also fixed alongside it:

• endStep popped the top of the stack rather than the step that ended. Now removes by identity, so a user's own Promise.all over parallel tool calls inside one row cannot evict an unrelated step.
• runFromCLI was fire-and-forget — it returned undefined instead of its promise chain, so nothing could await completion and the process could exit before writeOutput ran.
• endStep resolved the store at invocation time, so a step ended from a different async context (a .then() outside the wrapper, or a framework callback) read a null trace and threw Cannot read properties of null (reading 'steps'). The context is now captured at step creation.

2. inputVariableNames never written to config.json (OPEN-12421)

postProcessTrace returns { traceData, inputVariableNames } as siblings, but CLIHandler unwrapped to .traceData and then read .inputVariableNames off that — always undefined, so it never reached config.json.

3. Cost and tokens read only the root step (OPEN-12426)

postProcessTrace read cost and tokens off the root step alone, cast to ChatCompletionStep. Any agent- or chain-rooted trace carries those on nested LLM steps, so the unchecked cast silently yielded undefined.

Silent in both planes: development mode only declares a column when some row has a numeric value, so an agent-shaped model produced a config.json with no costColumnName or numOfTokenColumnName at all — a test thresholding on either had nothing to evaluate. Monitoring streams the same traceData, so live agent traces published null cost/tokens.

Now totalled across the root and all descendants. A trace already rooted on a chat-completion step reports exactly what it did before, pinned by its own test.

Release note: this changes monitoring output — live agent traces will start reporting real cost/token totals where they previously published nulls.

End-to-end verification

A 3-row batch, each row an agent root wrapping one chat-completion step (tokens: 42, cost: 0.00031).

Before — one shared trace, one upload, latency on row 1 only, no cost/token columns:

config.json: { outputColumnName, inputVariableNames, latencyColumnName }
  Paris                latency=29  cost=None  tokens=None
  William Shakespeare  latency=22  cost=None  tokens=None
  4                    latency=14  cost=None  tokens=None

After — three distinct root step IDs, three uploads, per-row latency, and the cost/token columns present:

config.json: { ..., latencyColumnName, costColumnName, numOfTokenColumnName }
  Paris                latency=16  cost=0.00031  tokens=42
  William Shakespeare  latency=34  cost=0.00031  tokens=42
  4                    latency=13  cost=0.00031  tokens=42

Row inputs (userQuery, groundTruth), output, otherFields, and the nested Agent → OpenAI Chat Completion step tree all serialize correctly.

Verification

• Three new suites — tests/cli-concurrency.test.ts, tests/tracer-context.test.ts, tests/tracer-cost-aggregation.test.ts (9 tests). Each test was watched failing first: the concurrency test failed with row-B nested in row-A nested in row-C; the aggregation tests failed Expected: 42 / Received: undefined; the detached-endStep test failed with the predicted TypeError: Cannot read properties of null (reading 'steps').
• Full suite vs. main on fresh deps: 483 → 492 passing, 68 → 59 failing. Suite-by-suite diff is identical apart from the three new suites flipping to PASS. Remaining failures are pre-existing tests/api-resources/** cases needing a mock server (APIConnectionError), confirmed failing identically on main.
• yarn build passes, including the require("openlayer") / import("openlayer") smoke tests; runInTraceContext and the aggregation are both present in dist.
• tsc --noEmit, eslint, prettier --check clean.

Scope / known gaps

Isolation is automatic for development mode and opt-in for monitoring. An app tracing parallel requests still shares defaultContext unless it wraps each unit of work in runInTraceContext(). Making that automatic is OPEN-12423, deliberately left out — investigation on this branch found it only partly achievable SDK-side:

Integration	SDK-side isolation
tracedTool / tracedAgent	possible — plain async wrapper
traceOpenAI non-streaming	possible
traceOpenAI streaming	not possible — step created inside tracedOutputGenerator, in the consumer's context
tracedQuery (Claude Agent SDK)	not possible — async function*; run() around generator creation is a no-op
langchainCallback	N/A — already isolated via runId-keyed maps

Two further items left alone, noted on OPEN-12420: runFromCLI's .catch(console.error) still swallows write failures, and defaultContext never clears its trace (same as the previous globals — the preserved endStep NOTE depends on it).

openlayer-python has the same root-only cost/token logic (defaulting to 0 rather than omitting), so it under-reports too — worth a parity follow-up.

🤖 Generated with Claude Code

https://claude.ai/code/session_015CvG3PYxCe6fbf8v7iyMRs