@insightfactory.ai/insightfactory-databricks-langgraph-tracer 1.0.0-dev.0

package/dist/databricks-tracer.js

@@ -0,0 +1,841 @@
+/**
+ * Custom LangChain `BaseTracer` that emits the shared schema to Databricks
+ * MLflow via `@mlflow/core`.
+ *
+ * This is the TypeScript counterpart of the Python `DatabricksLangGraphTracer`
+ * and emits the **identical** shared trace/span schema
+ * (`schema/trace-schema.json`):
+ *
+ * - **Root run create** → trace **tags** (`source`, `langgraph.run_id` /
+ *   `graph_id` / `env` / `api_revision`, conditional `langgraph.thread_id`) and
+ *   trace **metadata** (`mlflow.trace.session` / `mlflow.trace.user` when a
+ *   thread / user is present).
+ * - **Child run create** → `langgraph.node` / `langgraph.step` on the live node
+ *   span (from LangGraph's run metadata).
+ * - **LLM run update (end)** → resolved `mlflow.llm.model` / `mlflow.llm.provider`,
+ *   the reserved `mlflow.chat.tokenUsage` (incl. the cache-read / cache-creation
+ *   slots), the non-reserved token-count attrs, the reserved 3-key
+ *   `mlflow.llm.cost` plus the non-reserved cache cost line-items, and the
+ *   OpenInference-style `llm.*` dashboard keys (`llm.model_name` /
+ *   `llm.model_provider` / `llm.usage.prompt_tokens_cost` /
+ *   `llm.usage.completion_tokens_cost` — the keys the Databricks experiment
+ *   Overview cost charts aggregate) — accumulating a per-root cost aggregate
+ *   (totals **and** a per-model cost/token rollup).
+ * - **Root run update (end), before `span.end()`** → the reserved 3-key
+ *   `mlflow.trace.cost` metadata rollup, the per-model `cost.by_model` tag
+ *   (JSON), and the `cost.unknown_model` tag.
+ *
+ * Uses `startSpan` with an explicit `parent` ({@link LiveSpan}) so nested runs
+ * keep hierarchy without relying on the OTEL active context (callbacks are not
+ * executed inside `withSpan`). {@link BaseTracer} invokes `persistRun` only for
+ * root runs; every run (root included) is closed from `onRunUpdate`.
+ */
+import { DEFAULT_SPAN_NAME, SpanStatusCode, SpanType, startSpan, } from "@mlflow/core";
+import { BaseTracer } from "@langchain/core/tracers/base";
+import { calculateLLMCost, getPricingProvider, qualifyDeploymentName, spanCostPayload, } from "./cost.js";
+import { SpanAttrKey, TraceMetadataKey, TraceTagKey } from "./generated-keys.js";
+import { writeTraceFields } from "./trace-metadata.js";
+/** Maps LangChain `run_type` strings to MLflow {@link SpanType} (default CHAIN). */
+export function mapLangChainRunTypeToMlflowSpanType(runType) {
+    switch (runType) {
+        case "llm":
+            return SpanType.LLM;
+        case "tool":
+            return SpanType.TOOL;
+        case "chain":
+            return SpanType.CHAIN;
+        case "retriever":
+            return SpanType.RETRIEVER;
+        default:
+            return SpanType.CHAIN;
+    }
+}
+/** Best-effort JSON-serializable copy for span inputs/outputs (handles cycles). */
+export function serializeTracePayload(value) {
+    if (value === undefined)
+        return undefined;
+    try {
+        return JSON.parse(JSON.stringify(value));
+    }
+    catch {
+        return { _unserializable: String(value) };
+    }
+}
+/**
+ * True when a URL-typed content field holds an inline `data:` URI. Only `image_url`
+ * is ambiguous: it may carry an inline `data:` URI (megabytes that bloat the trace
+ * — externalize it) *or* a remote `http(s)://` URL (a small reference that does not
+ * bloat the trace — keep it verbatim). The base64-typed fields (`file_data` /
+ * `source.data` / `base64` / `data`) are inline by contract; remote images in those
+ * representations use a *separate* key (Anthropic `source.url`, standard-block
+ * top-level `url`) the size probe never matches, so they are already left untouched.
+ */
+function isInlinePayload(value) {
+    return value.slice(0, 5).toLowerCase() === "data:";
+}
+/**
+ * The multimodal content part's *inline* binary payload *string* (`undefined` if
+ * none) — the data-URI / base64 text that would otherwise be stored in the span
+ * input. Returning the value (not just its length) lets callers detect a resolver
+ * that echoes the exact payload back under a neutral key. Covers the encodings seen
+ * across providers and LangChain representations:
+ *
+ * - OpenAI / LangChain `image_url` — nested `image_url.url` (chat completions) or a
+ *   bare `image_url` string (Responses API `input_image`), counted only when it is
+ *   an inline `data:` URI; a remote `http(s)://` URL returns `undefined` and is
+ *   preserved verbatim (see {@link isInlinePayload});
+ * - OpenAI `file` — nested `file.file_data` data URI (chat completions) or a
+ *   top-level `file_data` (Responses API `input_file`);
+ * - Anthropic `image` / `document` — `source.data` base64;
+ * - LangChain v1 standard content blocks — top-level `base64` (or `data`).
+ *
+ * `text` parts, remote image URLs, and any part without such a payload return
+ * `undefined` and are never externalized.
+ */
+function contentPartPayloadValue(part) {
+    const imageUrl = part.image_url;
+    if (imageUrl && typeof imageUrl === "object") {
+        const url = imageUrl.url;
+        if (typeof url === "string" && isInlinePayload(url))
+            return url;
+    }
+    if (typeof imageUrl === "string" && isInlinePayload(imageUrl))
+        return imageUrl;
+    const file = part.file;
+    if (file && typeof file === "object") {
+        const fileData = file.file_data;
+        if (typeof fileData === "string")
+            return fileData;
+    }
+    const source = part.source;
+    if (source && typeof source === "object") {
+        const data = source.data;
+        if (typeof data === "string")
+            return data;
+    }
+    // Top-level payload keys: OpenAI Responses API `input_file.file_data`,
+    // LangChain v1 standard content blocks `base64`, and the generic `data`.
+    for (const key of ["file_data", "base64", "data"]) {
+        const value = part[key];
+        if (typeof value === "string")
+            return value;
+    }
+    return undefined;
+}
+/**
+ * Encoded length of a part's inline payload (`0` if none) — the value reported as
+ * `bytes` in the placeholder. See {@link contentPartPayloadValue}.
+ */
+function contentPartPayloadSize(part) {
+    return contentPartPayloadValue(part)?.length ?? 0;
+}
+/**
+ * True if `obj` (a resolver result) carries an inline binary payload anywhere. A
+ * resolver is meant to return a *reference*, not bytes. A shallow
+ * {@link contentPartPayloadSize} check on the result catches `return part` and
+ * `{ ...part }`, but a result that tucks the payload elsewhere hides it from a
+ * top-level probe. Walking the whole result catches it regardless of where the
+ * resolver put it, via two string-leaf checks:
+ *
+ * - a bare `data:` URI leaf (`{ ref: part.image_url.url }`); and
+ * - a leaf equal to `originalPayload` — the exact inline payload of the part being
+ *   externalized — which catches a resolver echoing the *raw base64* of an Anthropic
+ *   `source.data` / standard-block `base64` part under a neutral key
+ *   (`{ ref: part.source.data }`), where the string has no `data:` prefix to flag.
+ *   (A large *fabricated* string unrelated to the part is out of scope — ref size is
+ *   the consumer's responsibility, see the README.)
+ *
+ * Cycle-safe: `path` tracks the containers on the current branch, so a
+ * self-referential result (`result.self = result`) is detected as a back-edge and
+ * treated as a leak (it cannot be verified payload-free, and a cyclic object is not
+ * serializable into the trace anyway) instead of recursing to a stack overflow. The
+ * path is per-branch, so a legitimate shared (diamond) reference off the current
+ * branch is not mistaken for a cycle.
+ */
+function containsInlinePayload(obj, originalPayload, path = new Set()) {
+    if (typeof obj === "string") {
+        return isInlinePayload(obj) || (originalPayload !== undefined && obj === originalPayload);
+    }
+    if (!obj || typeof obj !== "object")
+        return false;
+    if (path.has(obj))
+        return true; // cycle -> cannot verify -> reject (placeholder)
+    if (!Array.isArray(obj) && contentPartPayloadSize(obj) > 0) {
+        return true;
+    }
+    const nextPath = new Set(path).add(obj);
+    const children = Array.isArray(obj) ? obj : Object.values(obj);
+    return children.some((child) => containsInlinePayload(child, originalPayload, nextPath));
+}
+/**
+ * Externalize one content part if it carries an inline binary payload (else
+ * recurse). Returns `part` unchanged (same reference) when nothing changed, so
+ * callers can detect "no change" by identity and skip copying.
+ *
+ * The resolver result is itself checked (recursively): if it carries an inline
+ * payload anywhere — a resolver that returns the part, spreads it (`{ ...part, ref }`),
+ * or nests it (`{ ref, original: part }`) would reintroduce the very bytes it was
+ * meant to strip — it is rejected and the placeholder is used, so the "no inline
+ * bytes are stored" invariant holds regardless of resolver behaviour.
+ *
+ * The text cap (`maxStringChars`) is intentionally *not* applied to the resolver
+ * result or the `_omitted` placeholder: a resolver's reference is the consumer's to
+ * keep small (see the README), and the placeholder is already tiny. The cap reaches
+ * non-payload parts (e.g. a huge `text` part) via the recursive {@link externalizeWalk}.
+ */
+function externalizeContentPart(part, index, resolver, metadata, maxStringChars) {
+    if (!part || typeof part !== "object" || Array.isArray(part)) {
+        return externalizeWalk(part, resolver, metadata, maxStringChars);
+    }
+    const p = part;
+    const payload = contentPartPayloadValue(p);
+    if (payload === undefined)
+        return externalizeWalk(part, resolver, metadata, maxStringChars);
+    const size = payload.length;
+    if (resolver) {
+        let replacement;
+        try {
+            replacement = resolver(p, { index, bytes: size, metadata });
+        }
+        catch (err) {
+            // Resilient by design: a faulty resolver must not break tracing.
+            console.warn("[DatabricksLangGraphTracer] contentRefResolver failed:", err);
+            replacement = undefined;
+        }
+        if (replacement && typeof replacement === "object" && !Array.isArray(replacement)) {
+            let leaks;
+            try {
+                // Pass the part's own payload so a resolver echoing the raw base64 back
+                // (under any key) is caught, not just data: URIs.
+                leaks = containsInlinePayload(replacement, payload);
+            }
+            catch (err) {
+                console.warn("[DatabricksLangGraphTracer] contentRefResolver result validation failed:", err);
+                leaks = true; // cannot verify -> reject
+            }
+            if (!leaks)
+                return replacement;
+            console.warn("[DatabricksLangGraphTracer] contentRefResolver result still carries an " +
+                "inline payload (or could not be verified); using placeholder instead");
+        }
+    }
+    return typeof p.type === "string"
+        ? { type: p.type, _omitted: true, bytes: size }
+        : { _omitted: true, bytes: size };
+}
+/** Externalize a message `content` array (copy-on-write; same array when unchanged). */
+function externalizeContentList(parts, resolver, metadata, maxStringChars) {
+    let result = parts;
+    for (let i = 0; i < parts.length; i++) {
+        const newPart = externalizeContentPart(parts[i], i, resolver, metadata, maxStringChars);
+        if (newPart !== parts[i]) {
+            if (result === parts)
+                result = [...parts];
+            result[i] = newPart;
+        }
+    }
+    return result;
+}
+/**
+ * Plain-text span-content cap. Large *text* leaves — not multimodal
+ * payloads — also bloat traces past the SQL inline read limit (e.g. a classification
+ * vocabulary / aggregated result set threaded through every fan-out span's inputs and
+ * outputs). When `maxStringChars` is set, any string leaf longer than it is replaced
+ * with a compact `{ _truncated: true, chars: N, bytes: M, preview: "…" }` placeholder
+ * so the recorded trace stays readable.
+ *
+ * Opt-in / default off: generic text truncation costs debuggability, so the consumer
+ * chooses the threshold (constructor option or `DATABRICKS_TRACING_MAX_STRING_CHARS`).
+ * It caps individual string *values*, never keys or our short structural fields, and
+ * is a per-leaf mitigation — not a hard per-trace byte budget (enough sub-cap leaves
+ * can still sum past the limit), so pair it with consumer-side payload reduction for
+ * the heaviest spans. Distinct from the removed multimodal `maxContentPartBytes`: that
+ * gated inline binary payloads; this caps arbitrary text.
+ */
+const TRUNCATION_PREVIEW_CHARS = 256;
+const MAX_STRING_CHARS_ENV = "DATABRICKS_TRACING_MAX_STRING_CHARS";
+/**
+ * Resolve the text cap: explicit option wins, else `DATABRICKS_TRACING_MAX_STRING_CHARS`.
+ * Returns `undefined` (cap disabled) when neither is set, or when the value is not a
+ * positive integer. A non-positive / unparseable env value is ignored (warned), not
+ * thrown, so a stray override never breaks tracing startup.
+ */
+export function resolveMaxStringChars(explicit) {
+    if (explicit !== undefined)
+        return explicit > 0 ? explicit : undefined;
+    const raw = process.env[MAX_STRING_CHARS_ENV]?.trim();
+    if (!raw)
+        return undefined;
+    // Mirror Python's int(): an optional sign and digits with single underscores
+    // allowed as separators (no leading / trailing / doubled underscore). Validate
+    // before parsing — Number.parseInt accepts a prefix ("50_000" -> 50, "100abc" ->
+    // 100), which would silently truncate traces far more aggressively than intended.
+    if (!/^[+-]?\d(?:_?\d)*$/.test(raw)) {
+        console.warn(`[DatabricksLangGraphTracer] ${MAX_STRING_CHARS_ENV}=${raw} is not an integer; ignoring`);
+        return undefined;
+    }
+    const value = Number.parseInt(raw.replace(/_/g, ""), 10);
+    return value > 0 ? value : undefined;
+}
+/**
+ * Replace an over-cap string leaf with a compact placeholder; else return it as-is
+ * (same reference, so the copy-on-write walk treats it as "no change"). The
+ * placeholder records the original length (UTF-16 code units) and UTF-8 byte length
+ * plus a short head preview (itself bounded by the cap, so it never exceeds the
+ * allowance), keeping the value identifiable in the trace without its full size.
+ */
+function truncateLongString(value, maxStringChars) {
+    if (maxStringChars === undefined || value.length <= maxStringChars)
+        return value;
+    return {
+        _truncated: true,
+        chars: value.length,
+        bytes: Buffer.byteLength(value, "utf8"),
+        preview: value.slice(0, Math.min(TRUNCATION_PREVIEW_CHARS, maxStringChars)),
+    };
+}
+/**
+ * Externalize multimodal content parts' inline payloads and cap text.
+ * Public entry point. Takes a **shallow copy** of `metadata` once (when payloads are
+ * externalized with a resolver) so a resolver mutating `ctx.metadata` cannot corrupt
+ * the live run metadata, which the caller passes by reference and which is shared with
+ * sibling / child spans (resilient by design — a faulty resolver must not affect
+ * tracing). Then walks via {@link externalizeWalk}.
+ *
+ * `externalizePayloads` controls the multimodal-part stripping:
+ * - **inputs** (`true`) — the #21 behaviour: inline image / PDF / file payloads in
+ *   message `content` arrays are replaced with references (inputs are LLM messages,
+ *   where the `content` / `data` / `base64` shapes are genuinely multimodal).
+ * - **outputs** (`false`) — only the text cap is applied. Span outputs are arbitrary
+ *   chain / tool return values; the multimodal probe matches a bare `data` / `base64`
+ *   string in any `content` list, so running it would strip legitimate non-binary
+ *   output to an `_omitted` placeholder. A genuinely huge inline payload in an output
+ *   is still bounded by the text cap when it is set.
+ */
+export function externalizeContentRefs(obj, resolver, metadata, maxStringChars, externalizePayloads = true) {
+    const meta = externalizePayloads && resolver ? { ...metadata } : metadata;
+    return externalizeWalk(obj, resolver, meta, maxStringChars, externalizePayloads);
+}
+/**
+ * Recursively externalize multimodal inline payloads and cap over-long text. When
+ * `externalizePayloads` is set, walks message-content arrays — the `content` key
+ * whose value is an array, as the serialized message nests it under `kwargs` — and
+ * replaces every part carrying an inline binary payload with a reference (see
+ * {@link externalizeContentPart}). At every string leaf, when `maxStringChars` is
+ * set, an over-long value is replaced with a compact `{ _truncated: … }` placeholder
+ * (see {@link truncateLongString}); keys are never touched (only values).
+ *
+ * Copy-on-write: returns the *same* object when nothing beneath it changed, and
+ * allocates new containers only along the path to a replaced part / truncated leaf.
+ * So inputs / outputs with no multimodal parts and no over-long text (every chain /
+ * tool span, text-only chat turns) do no copying, and the input is never mutated —
+ * the live model message is untouched.
+ */
+function externalizeWalk(obj, resolver, metadata, maxStringChars, externalizePayloads = true) {
+    if (typeof obj === "string")
+        return truncateLongString(obj, maxStringChars);
+    if (Array.isArray(obj)) {
+        let result = obj;
+        for (let i = 0; i < obj.length; i++) {
+            const newItem = externalizeWalk(obj[i], resolver, metadata, maxStringChars, externalizePayloads);
+            if (newItem !== obj[i]) {
+                if (result === obj)
+                    result = [...obj];
+                result[i] = newItem;
+            }
+        }
+        return result;
+    }
+    if (obj && typeof obj === "object") {
+        const src = obj;
+        let result = src;
+        for (const [key, value] of Object.entries(src)) {
+            const newValue = externalizePayloads && key === "content" && Array.isArray(value)
+                ? externalizeContentList(value, resolver, metadata, maxStringChars)
+                : externalizeWalk(value, resolver, metadata, maxStringChars, externalizePayloads);
+            if (newValue !== value) {
+                if (result === src)
+                    result = { ...src };
+                result[key] = newValue;
+            }
+        }
+        return result;
+    }
+    return obj;
+}
+function toFiniteNumber(v) {
+    return typeof v === "number" && Number.isFinite(v) ? v : undefined;
+}
+function runMetadata(run) {
+    const meta = run.extra?.metadata;
+    return meta && typeof meta === "object" ? meta : {};
+}
+function getMessageKwargs(run) {
+    const outputs = run.outputs;
+    const generations = outputs?.generations;
+    if (!Array.isArray(generations))
+        return undefined;
+    const firstGroup = generations[0];
+    if (!Array.isArray(firstGroup))
+        return undefined;
+    const firstGen = firstGroup[0];
+    if (!firstGen || typeof firstGen !== "object")
+        return undefined;
+    const message = firstGen.message;
+    if (!message || typeof message !== "object")
+        return undefined;
+    const kwargs = message.kwargs;
+    if (kwargs && typeof kwargs === "object")
+        return kwargs;
+    return message;
+}
+/**
+ * Extract normalized token usage from an LLM run end payload, extended for the
+ * Anthropic `cache_creation` category to match the Python `_extract_usage`.
+ *
+ * Prefers LangChain normalized `usage_metadata`
+ * (`input_token_details.cache_read` / `.cache_creation`,
+ * `output_token_details.reasoning`); falls back to raw `response_metadata.usage`
+ * (`input_tokens_details.cached_tokens` / `cache_read_input_tokens`,
+ * `cache_creation_input_tokens`, `output_tokens_details.reasoning_tokens`).
+ * Missing cache/reasoning sub-fields default to `0`. Returns `undefined` when
+ * neither source is present.
+ */
+export function extractUsage(run) {
+    const kwargs = getMessageKwargs(run);
+    const normalize = (input, output, total, cached, cacheCreation, reasoning) => {
+        const inputN = toFiniteNumber(input);
+        const outputN = toFiniteNumber(output);
+        if (inputN === undefined || outputN === undefined)
+            return undefined;
+        const totalN = toFiniteNumber(total) ?? inputN + outputN;
+        return {
+            input_tokens: inputN,
+            output_tokens: outputN,
+            total_tokens: totalN,
+            cached_tokens: toFiniteNumber(cached) ?? 0,
+            cache_creation_tokens: toFiniteNumber(cacheCreation) ?? 0,
+            reasoning_tokens: toFiniteNumber(reasoning) ?? 0,
+        };
+    };
+    if (kwargs) {
+        const usageMetadata = kwargs.usage_metadata;
+        if (usageMetadata && typeof usageMetadata === "object") {
+            const inputDetails = usageMetadata.input_token_details;
+            const outputDetails = usageMetadata.output_token_details;
+            const result = normalize(usageMetadata.input_tokens, usageMetadata.output_tokens, usageMetadata.total_tokens, inputDetails?.cache_read, inputDetails?.cache_creation, outputDetails?.reasoning);
+            if (result)
+                return result;
+        }
+        const responseMeta = kwargs.response_metadata;
+        const usage = responseMeta?.usage;
+        if (usage && typeof usage === "object") {
+            const inputDetails = usage.input_tokens_details;
+            const outputDetails = usage.output_tokens_details;
+            // Azure-raw nests cached under input_tokens_details; Anthropic-raw exposes
+            // cache_read_input_tokens / cache_creation_input_tokens at the usage root.
+            const cached = inputDetails && "cached_tokens" in inputDetails
+                ? inputDetails.cached_tokens
+                : usage.cache_read_input_tokens;
+            const result = normalize(usage.input_tokens, usage.output_tokens, usage.total_tokens, cached, usage.cache_creation_input_tokens, outputDetails?.reasoning_tokens);
+            if (result)
+                return result;
+        }
+    }
+    return undefined;
+}
+/**
+ * Derive the canonical LiteLLM catalog key for an LLM run, aligned with the
+ * Python `_extract_model`.
+ *
+ * Sources, in order: invocation params `model` (probed both top-level on
+ * `run.extra.invocation_params` and nested under `run.extra.metadata`), then
+ * `response_metadata.model_name` — each qualified via `qualify` (default
+ * {@link qualifyDeploymentName}; consumers with non-canonical deployment
+ * names supply a resolver via the tracer's `pricingKeyResolver` option).
+ */
+export function extractModel(run, qualify = qualifyDeploymentName) {
+    const extra = run.extra;
+    let invocationParams = extra?.invocation_params;
+    if (!invocationParams || typeof invocationParams !== "object") {
+        const meta = extra?.metadata;
+        invocationParams = meta?.invocation_params;
+    }
+    if (invocationParams && typeof invocationParams === "object") {
+        const model = invocationParams.model;
+        if (typeof model === "string" && model.length > 0) {
+            return qualify(model);
+        }
+    }
+    const kwargs = getMessageKwargs(run);
+    const responseMeta = kwargs?.response_metadata;
+    const modelName = responseMeta?.model_name;
+    if (typeof modelName === "string" && modelName.length > 0) {
+        return qualify(modelName);
+    }
+    return undefined;
+}
+function createModelRollup() {
+    return {
+        input_cost: 0,
+        output_cost: 0,
+        total_cost: 0,
+        input_tokens: 0,
+        output_tokens: 0,
+        cache_read_input_tokens: 0,
+        cache_creation_input_tokens: 0,
+    };
+}
+function createTraceAggregate(rootTraceId) {
+    return {
+        rootTraceId,
+        inputCost: 0,
+        outputCost: 0,
+        totalCost: 0,
+        unknownModels: new Set(),
+        hasAnyCost: false,
+        byModel: new Map(),
+    };
+}
+export class DatabricksLangGraphTracer extends BaseTracer {
+    name = "databricks_langgraph";
+    source;
+    qualifyModel;
+    contentRefResolver;
+    maxStringChars;
+    spanMap = new Map();
+    traceAggregates = new Map();
+    rootByRunId = new Map();
+    constructor(options = {}) {
+        super();
+        this.source = options.source ?? "langgraph";
+        const resolver = options.pricingKeyResolver;
+        this.qualifyModel = resolver
+            ? (name) => resolver(name) ?? qualifyDeploymentName(name)
+            : qualifyDeploymentName;
+        this.contentRefResolver = options.contentRefResolver;
+        this.maxStringChars = resolveMaxStringChars(options.maxStringChars);
+    }
+    async onRunCreate(run) {
+        try {
+            const spanType = mapLangChainRunTypeToMlflowSpanType(run.run_type);
+            const parentSpan = run.parent_run_id !== undefined ? this.spanMap.get(run.parent_run_id) : undefined;
+            // Capture structured chat-model message inputs (the TS BaseTracer already
+            // records `{ messages }` rather than flattened prompts), then externalize the
+            // inline payload of every multimodal image / PDF / file part to a reference and
+            // cap over-long text so it does not bloat the trace (issues #16/#21/#23). A
+            // no-op for text-only inputs under the cap; runs on the serialized copy, so the
+            // live model message is untouched.
+            let inputs = serializeTracePayload(run.inputs);
+            if (inputs !== undefined) {
+                inputs = externalizeContentRefs(inputs, this.contentRefResolver, runMetadata(run), this.maxStringChars);
+            }
+            const span = startSpan({
+                name: run.name?.trim() ? run.name : DEFAULT_SPAN_NAME,
+                spanType,
+                inputs: inputs,
+                parent: parentSpan,
+            });
+            this.spanMap.set(run.id, span);
+            if (run.parent_run_id === undefined) {
+                // ROOT
+                this.traceAggregates.set(run.id, createTraceAggregate(span.traceId));
+                this.rootByRunId.set(run.id, run.id); // self-map
+                this.writeRootTraceFields(run, span);
+            }
+            else {
+                // child — propagate the root id for cost aggregation
+                const rootId = this.rootByRunId.get(run.parent_run_id);
+                if (rootId)
+                    this.rootByRunId.set(run.id, rootId);
+                this.writeChildLangGraphContext(run, span);
+            }
+        }
+        catch (err) {
+            console.warn("[DatabricksLangGraphTracer] onRunCreate failed:", err);
+        }
+    }
+    async onRunUpdate(run) {
+        if (run.end_time === undefined)
+            return;
+        const span = this.spanMap.get(run.id);
+        if (!span)
+            return;
+        // Enrichment, status, and the root rollup are each best-effort: a failure in
+        // one must degrade only its own fields, never prevent span finalization
+        // below (resilient by design, per-run). If span.end() were skipped,
+        // the finally would drop the only reference and the span would never export.
+        if (run.run_type === "llm") {
+            try {
+                this.enrichLLMSpan(run, span);
+            }
+            catch (err) {
+                console.warn("[DatabricksLangGraphTracer] LLM span enrichment failed:", err);
+            }
+        }
+        if (run.error) {
+            try {
+                span.setStatus(SpanStatusCode.ERROR, run.error);
+            }
+            catch (err) {
+                console.warn("[DatabricksLangGraphTracer] setStatus failed:", err);
+            }
+        }
+        // Flush the root rollup BEFORE span.end(): the MLflow span processor pops the
+        // trace from InMemoryTraceManager synchronously inside onEnd, after which
+        // writeTraceFields silently no-ops.
+        if (run.parent_run_id === undefined) {
+            try {
+                const aggregate = this.traceAggregates.get(run.id);
+                if (aggregate)
+                    this.writeRootTraceRollup(aggregate);
+            }
+            catch (err) {
+                console.warn("[DatabricksLangGraphTracer] root trace rollup failed:", err);
+            }
+        }
+        try {
+            // Cap over-long text on a copy of the outputs — the #23 spans that breach the
+            // inline read limit (final_output / aggregate_results) are output-side. Unlike
+            // inputs, the multimodal payload stripping is NOT run here
+            // (externalizePayloads=false): outputs are arbitrary chain / tool JSON, so the
+            // data/base64 probe would mis-strip legitimate non-binary output (see
+            // externalizeContentRefs).
+            let outputs = serializeTracePayload(run.outputs ?? {});
+            try {
+                outputs = externalizeContentRefs(outputs, this.contentRefResolver, runMetadata(run), this.maxStringChars, false);
+            }
+            catch (err) {
+                // Resilient by design: an externalization failure must never skip
+                // span.end(). Fall back to the serialized outputs so the span still exports.
+                console.warn("[DatabricksLangGraphTracer] output externalization failed; recording raw:", err);
+                outputs = serializeTracePayload(run.outputs ?? {});
+            }
+            span.end({ outputs: outputs });
+        }
+        catch (err) {
+            console.warn("[DatabricksLangGraphTracer] span.end failed:", err);
+        }
+        finally {
+            this.spanMap.delete(run.id);
+            // Root aggregate is keyed by root run id; a no-op for child/LLM runs.
+            this.traceAggregates.delete(run.id);
+            // Unconditionally clear this run's root mapping (self-map for a root,
+            // own->root lookup for a child). Never delete the *value*: sibling children
+            // of a still-running root rely on the root's self-mapping staying put.
+            this.rootByRunId.delete(run.id);
+        }
+    }
+    /**
+     * Write the root trace tags + session/user metadata (Python
+     * `_write_root_trace_fields`). Always emits `source` + the `langgraph.run_id` /
+     * `graph_id` / `env` / `api_revision` tags. A thread id additionally produces
+     * the `langgraph.thread_id` tag and the reserved `mlflow.trace.session`
+     * metadata; a user identity produces the reserved `mlflow.trace.user` metadata.
+     */
+    writeRootTraceFields(run, span) {
+        const meta = runMetadata(run);
+        const tags = {
+            [TraceTagKey.SOURCE]: this.source,
+            [TraceTagKey.LANGGRAPH_RUN_ID]: run.id,
+            [TraceTagKey.LANGGRAPH_GRAPH_ID]: typeof meta.graph_id === "string" ? meta.graph_id : "agent",
+            [TraceTagKey.LANGGRAPH_ENV]: process.env.AGENT_ENV ?? "unknown",
+            [TraceTagKey.LANGGRAPH_API_REVISION]: process.env.AGENT_API_REVISION ?? "unknown",
+        };
+        const metadata = {};
+        const threadId = typeof meta.thread_id === "string" ? meta.thread_id : undefined;
+        if (threadId) {
+            tags[TraceTagKey.LANGGRAPH_THREAD_ID] = threadId;
+            metadata[TraceMetadataKey.TRACE_SESSION] = threadId;
+        }
+        // User identity: LangGraph's auth middleware populates `langgraph_auth_user_id`;
+        // callers may instead pass `user_id` / `user`.
+        const user = pickString(meta, "langgraph_auth_user_id") ??
+            pickString(meta, "user_id") ??
+            pickString(meta, "user");
+        if (user)
+            metadata[TraceMetadataKey.TRACE_USER] = user;
+        writeTraceFields(span.traceId, {
+            tags,
+            metadata: Object.keys(metadata).length > 0 ? metadata : undefined,
+        });
+    }
+    /**
+     * Write `langgraph.node` / `langgraph.step` onto a child span. LangGraph stamps
+     * these onto every per-node run's metadata (including the node's nested LLM
+     * run), mirroring the Python `_write_child_langgraph_context`.
+     */
+    writeChildLangGraphContext(run, span) {
+        const meta = runMetadata(run);
+        const node = meta.langgraph_node;
+        if (typeof node === "string" && node.length > 0) {
+            span.setAttribute(SpanAttrKey.LANGGRAPH_NODE, node);
+        }
+        const step = meta.langgraph_step;
+        if (typeof step === "number" && Number.isFinite(step)) {
+            span.setAttribute(SpanAttrKey.LANGGRAPH_STEP, step);
+        }
+    }
+    /**
+     * Write model / provider / token-usage / cost attrs on a live LLM span (Python
+     * `_enrich_llm_span`). Token usage is written whenever usage is extractable;
+     * model / provider whenever the model resolves. Cost needs both: the reserved
+     * 3-key `mlflow.llm.cost` plus the non-reserved cache-read / cache-creation
+     * cost line-items (Option B) are written and folded into the root aggregate; a
+     * resolved-but-unpriceable model is recorded for the `cost.unknown_model` tag.
+     * Whenever the model resolves with usage, the span is also folded into the
+     * per-model `cost.by_model` rollup (tokens always; cost when priceable).
+     *
+     * The OpenInference-style `llm.*` duplicates (`llm.model_name` /
+     * `llm.model_provider` / `llm.usage.prompt_tokens_cost` /
+     * `llm.usage.completion_tokens_cost`) mirror the MLflow-convention attributes
+     * because the Databricks experiment Overview dashboard aggregates cost from
+     * those keys via `variant_get` on the UC span table — without them, Cost
+     * Breakdown / Cost Over Time render $0.00. The cost values must
+     * stay numbers: `variant_get(..., 'DOUBLE')` returns NULL for strings.
+     */
+    enrichLLMSpan(run, span) {
+        const usage = extractUsage(run);
+        const model = extractModel(run, this.qualifyModel);
+        if (usage)
+            this.writeTokenUsage(span, usage);
+        if (model) {
+            const provider = getPricingProvider(model) ?? "azure";
+            span.setAttribute(SpanAttrKey.LLM_MODEL, model);
+            span.setAttribute(SpanAttrKey.LLM_PROVIDER, provider);
+            span.setAttribute(SpanAttrKey.LLM_MODEL_NAME, model);
+            span.setAttribute(SpanAttrKey.LLM_MODEL_PROVIDER, provider);
+        }
+        if (!usage || !model)
+            return;
+        const rootRunId = this.rootByRunId.get(run.id);
+        const aggregate = rootRunId !== undefined ? this.traceAggregates.get(rootRunId) : undefined;
+        const cost = calculateLLMCost(model, usage);
+        if (cost) {
+            span.setAttribute(SpanAttrKey.LLM_COST, spanCostPayload(cost));
+            // Non-reserved cost line-items (Option B): the cache-read / cache-creation
+            // split the reserved 3-key cost can't carry.
+            span.setAttribute(SpanAttrKey.INPUT_TOKENS_CACHED_COST, cost.cached_input_cost);
+            span.setAttribute(SpanAttrKey.INPUT_TOKENS_CACHE_CREATION_COST, cost.cache_creation_cost);
+            // Dashboard cost keys (numbers, summed by the Overview charts). These
+            // mirror the reserved mlflow.llm.cost input/output split, so
+            // prompt_tokens_cost is cache-inclusive (folds in the cache-read /
+            // cache-creation cost) — not a pure prompt-token figure. The cache split
+            // lives on the gen_ai.usage.* slots above.
+            span.setAttribute(SpanAttrKey.LLM_PROMPT_TOKENS_COST, cost.input_cost);
+            span.setAttribute(SpanAttrKey.LLM_COMPLETION_TOKENS_COST, cost.output_cost);
+        }
+        if (aggregate) {
+            // Fold this span into the per-model rollup (the cost.by_model tag). Tokens
+            // always contribute — even for a resolved-but-unpriceable model — so the
+            // per-model token split stays complete; cost contributes only when the
+            // model priced (otherwise it's recorded for the cost.unknown_model tag,
+            // matching the trace-level rollup).
+            let bucket = aggregate.byModel.get(model);
+            if (!bucket) {
+                bucket = createModelRollup();
+                aggregate.byModel.set(model, bucket);
+            }
+            bucket.input_tokens += usage.input_tokens || 0;
+            bucket.output_tokens += usage.output_tokens || 0;
+            bucket.cache_read_input_tokens += usage.cached_tokens || 0;
+            bucket.cache_creation_input_tokens += usage.cache_creation_tokens || 0;
+            if (cost) {
+                bucket.input_cost += cost.input_cost;
+                bucket.output_cost += cost.output_cost;
+                bucket.total_cost += cost.total_cost;
+                aggregate.inputCost += cost.input_cost;
+                aggregate.outputCost += cost.output_cost;
+                aggregate.totalCost += cost.total_cost;
+                aggregate.hasAnyCost = true;
+            }
+            else {
+                aggregate.unknownModels.add(model);
+            }
+        }
+    }
+    /**
+     * Write token-usage span attributes. The reserved `mlflow.chat.tokenUsage`
+     * attribute carries MLflow's `TokenUsageKey` shape extended with the two cache
+     * slots — `{ input_tokens, output_tokens, total_tokens, cache_read_input_tokens,
+     * cache_creation_input_tokens }` — so they survive Databricks ingestion and
+     * render in the UI usage breakdown. The non-reserved `gen_ai.usage.*`
+     * token-count attrs mirror the Python `_write_token_usage`.
+     */
+    writeTokenUsage(span, usage) {
+        const inputTokens = usage.input_tokens || 0;
+        const outputTokens = usage.output_tokens || 0;
+        const totalTokens = usage.total_tokens || inputTokens + outputTokens;
+        const cached = usage.cached_tokens || 0;
+        const cacheCreation = usage.cache_creation_tokens || 0;
+        const reasoning = usage.reasoning_tokens || 0;
+        span.setAttribute(SpanAttrKey.TOKEN_USAGE, {
+            input_tokens: inputTokens,
+            output_tokens: outputTokens,
+            total_tokens: totalTokens,
+            cache_read_input_tokens: cached,
+            cache_creation_input_tokens: cacheCreation,
+        });
+        span.setAttribute(SpanAttrKey.INPUT_TOKENS_CACHED, cached);
+        span.setAttribute(SpanAttrKey.INPUT_TOKENS_CACHE_CREATION, cacheCreation);
+        span.setAttribute(SpanAttrKey.OUTPUT_TOKENS_REASONING, reasoning);
+    }
+    /**
+     * Write the trace-level cost rollups + `cost.unknown_model` tag (Python
+     * `_write_root_trace_rollup`). `mlflow.trace.cost` is the reserved MLflow
+     * `CostKey` shape — exactly `{ input_cost, output_cost, total_cost }`;
+     * Databricks strips any extra key on ingestion. `total_cost` already accounts
+     * for the cached discount and cache-creation premium; the per-line-item split
+     * lives on the spans.
+     *
+     * `cost.by_model` is the non-reserved per-model breakdown: emitted
+     * as a JSON string under a **trace tag** (like `cost.unknown_model`) rather
+     * than a custom metadata key — tags are proven to survive Databricks UC
+     * ingestion and stay queryable from the `*_trace_unified` view. Each bucket
+     * carries cost and tokens; the per-model totals reconcile with
+     * `mlflow.trace.cost` / `mlflow.chat.tokenUsage` within rounding. Models are
+     * emitted in sorted order, mirroring the Python tracer (the JSON parses to the
+     * same structure; the serialized bytes differ — Python `json.dumps` inserts
+     * `", "` / `": "` separators and renders integral zero floats as `0.0`).
+     *
+     * Must run before the root `span.end()` (see {@link onRunUpdate}).
+     */
+    writeRootTraceRollup(aggregate) {
+        if (aggregate.hasAnyCost) {
+            writeTraceFields(aggregate.rootTraceId, {
+                metadata: {
+                    [TraceMetadataKey.TRACE_COST]: JSON.stringify({
+                        input_cost: aggregate.inputCost,
+                        output_cost: aggregate.outputCost,
+                        total_cost: aggregate.totalCost,
+                    }),
+                },
+            });
+        }
+        if (aggregate.byModel.size > 0) {
+            const byModel = {};
+            for (const model of [...aggregate.byModel.keys()].sort()) {
+                byModel[model] = aggregate.byModel.get(model);
+            }
+            writeTraceFields(aggregate.rootTraceId, {
+                tags: { [TraceTagKey.COST_BY_MODEL]: JSON.stringify(byModel) },
+            });
+        }
+        if (aggregate.unknownModels.size > 0) {
+            writeTraceFields(aggregate.rootTraceId, {
+                tags: {
+                    [TraceTagKey.COST_UNKNOWN_MODEL]: [...aggregate.unknownModels].sort().join(","),
+                },
+            });
+        }
+    }
+    async persistRun(_run) {
+        // No-op: span lifecycle is handled in onRunCreate / onRunUpdate. Root runs
+        // also receive onRunUpdate (Python `_persist_run` parallel).
+    }
+}
+function pickString(meta, key) {
+    const v = meta[key];
+    return typeof v === "string" && v.length > 0 ? v : undefined;
+}
+//# sourceMappingURL=databricks-tracer.js.map