@mastra/datadog 1.0.22 → 1.1.0-alpha.1

package/dist/index.js

@@ -1,9 +1,9 @@
 import { SpanType } from '@mastra/core/observability';
 import { omitKeys } from '@mastra/core/utils';
-import { BaseExporter } from '@mastra/observability';
-import tracer2 from 'dd-trace';
+import { BaseExporter, getExternalParentId } from '@mastra/observability';
+import tracer3 from 'dd-trace';
 
-// src/tracing.ts
+// src/bridge.ts
 
 // src/metrics.ts
 function formatUsageMetrics(usage) {
@@ -50,16 +50,16 @@ function ensureTracer(config) {
   if (config.apiKey) {
     process.env.DD_API_KEY = config.apiKey;
   }
-  const alreadyStarted = tracer2._tracer?.started;
+  const alreadyStarted = tracer3._tracer?.started;
   if (!alreadyStarted) {
-    tracer2.init({
+    tracer3.init({
       service: config.service || config.mlApp,
       env: config.env || process.env.DD_ENV,
       // Disable automatic integrations by default to avoid surprise instrumentation
       plugins: config.integrationsEnabled ?? false
     });
   }
-  tracer2.llmobs.enable({
+  tracer3.llmobs.enable({
     mlApp: config.mlApp,
     agentlessEnabled: config.agentless
   });
@@ -136,7 +136,440 @@ function formatOutput(output, spanType) {
   return safeStringify(output);
 }
 
-// src/tracing.ts
+// src/bridge.ts
+function flushApmExporter() {
+  const exporterFlush = tracer3?._tracer?._exporter?.flush;
+  if (typeof exporterFlush !== "function") {
+    return Promise.resolve();
+  }
+  return new Promise((resolve, reject) => {
+    let settled = false;
+    const done = (error) => {
+      if (settled) return;
+      settled = true;
+      if (error) {
+        reject(error);
+        return;
+      }
+      resolve();
+    };
+    try {
+      exporterFlush.call(tracer3._tracer._exporter, done);
+    } catch (error) {
+      reject(error);
+    }
+  });
+}
+var DatadogBridge = class extends BaseExporter {
+  name = "datadog-bridge";
+  config;
+  ddSpanMap = /* @__PURE__ */ new Map();
+  traceContext = /* @__PURE__ */ new Map();
+  openSpanCounts = /* @__PURE__ */ new Map();
+  constructor(config = {}) {
+    super(config);
+    const mlApp = config.mlApp ?? process.env.DD_LLMOBS_ML_APP;
+    const apiKey = config.apiKey ?? process.env.DD_API_KEY;
+    const site = config.site ?? process.env.DD_SITE ?? "datadoghq.com";
+    const env = config.env ?? process.env.DD_ENV;
+    const envAgentless = process.env.DD_LLMOBS_AGENTLESS_ENABLED?.toLowerCase();
+    const agentless = config.agentless ?? (envAgentless === "true" || envAgentless === "1" ? true : false);
+    if (!mlApp) {
+      this.setDisabled(`Missing required mlApp. Set DD_LLMOBS_ML_APP environment variable or pass mlApp in config.`);
+      this.config = config;
+      return;
+    }
+    if (agentless && !apiKey) {
+      this.setDisabled(
+        `Missing required apiKey for agentless mode. Set DD_API_KEY environment variable or pass apiKey in config.`
+      );
+      this.config = config;
+      return;
+    }
+    this.config = { ...config, mlApp, site, apiKey, agentless, env };
+    ensureTracer({
+      mlApp,
+      site,
+      apiKey,
+      agentless,
+      service: config.service,
+      env,
+      integrationsEnabled: config.integrationsEnabled ?? true
+    });
+    this.logger.info("Datadog bridge initialized", { mlApp, site, agentless });
+  }
+  /**
+   * Create a dd-trace span eagerly when a Mastra span is constructed.
+   */
+  createSpan(options) {
+    if (this.isDisabled) return void 0;
+    try {
+      let apmParentDdSpan = void 0;
+      let llmobsParentDdSpan = void 0;
+      let parentSource = "none";
+      const externalParentId = getExternalParentId(options);
+      if (externalParentId) {
+        apmParentDdSpan = this.ddSpanMap.get(externalParentId);
+        llmobsParentDdSpan = apmParentDdSpan;
+        if (apmParentDdSpan) {
+          parentSource = "external-parent";
+        }
+      }
+      if (!apmParentDdSpan) {
+        apmParentDdSpan = tracer3.scope().active() ?? void 0;
+        if (apmParentDdSpan) {
+          parentSource = "active-scope";
+        }
+      }
+      if (!llmobsParentDdSpan && externalParentId) {
+        llmobsParentDdSpan = apmParentDdSpan;
+      }
+      const ddSpan = tracer3.startSpan(options.name, {
+        ...apmParentDdSpan ? { childOf: apmParentDdSpan } : {},
+        ...options.startTime ? { startTime: toDate(options.startTime).getTime() } : {}
+      });
+      const ddContext = ddSpan.context?.();
+      const spanId = ddContext?.toSpanId?.(true) ?? generateSpanId();
+      const traceId = ddContext?.toTraceId?.(true) ?? (externalParentId ? options.parent?.traceId ?? generateTraceId() : generateTraceId());
+      const parentContext = apmParentDdSpan?.context?.();
+      const parentSpanId = parentContext?.toSpanId?.(true) ?? externalParentId;
+      this.captureTraceContext(traceId, options);
+      this.openSpanCounts.set(traceId, (this.openSpanCounts.get(traceId) ?? 0) + 1);
+      this.ddSpanMap.set(spanId, ddSpan);
+      this.registerLlmObsSpan(ddSpan, traceId, options, llmobsParentDdSpan);
+      this.logger.debug(
+        `[DatadogBridge.createSpan] Created APM span [spanId=${spanId}] [traceId=${traceId}] [parentSpanId=${parentSpanId}] [type=${options.type}] [mapSize=${this.ddSpanMap.size}] [parentSource=${parentSource}] [externalParentId=${externalParentId ?? "none"}]`
+      );
+      return { spanId, traceId, parentSpanId };
+    } catch (error) {
+      this.logger.error("[DatadogBridge] Failed to create span:", error);
+      return void 0;
+    }
+  }
+  /**
+   * Execute an async function within the dd-trace context of a Mastra span.
+   */
+  executeInContext(spanId, fn) {
+    return this.executeWithSpanContext(spanId, fn);
+  }
+  /**
+   * Execute a synchronous function within the dd-trace context of a Mastra span.
+   */
+  executeInContextSync(spanId, fn) {
+    return this.executeWithSpanContext(spanId, fn);
+  }
+  executeWithSpanContext(spanId, fn) {
+    const ddSpan = this.ddSpanMap.get(spanId);
+    if (ddSpan) {
+      return tracer3.scope().activate(ddSpan, () => {
+        const llmobs = tracer3.llmobs;
+        if (typeof llmobs?._activate === "function") {
+          return llmobs._activate(ddSpan, void 0, fn);
+        }
+        return fn();
+      });
+    }
+    return fn();
+  }
+  /**
+   * Handle tracing events from the observability bus.
+   */
+  async _exportTracingEvent(event) {
+    if (this.isDisabled) return;
+    try {
+      const span = event.exportedSpan;
+      if (span.isEvent) {
+        if (event.type === "span_started" || event.type === "span_ended") {
+          this.annotateAndFinishSpan(span);
+        }
+        return;
+      }
+      switch (event.type) {
+        case "span_started":
+        case "span_updated":
+          return;
+        case "span_ended":
+          this.annotateAndFinishSpan(span);
+          return;
+      }
+    } catch (error) {
+      this.logger.error("Datadog bridge error", {
+        error,
+        eventType: event.type,
+        spanId: event.exportedSpan?.id,
+        spanName: event.exportedSpan?.name
+      });
+    }
+  }
+  registerLlmObsSpan(ddSpan, traceId, options, parentDdSpan) {
+    const tagger = this.getLlmObsTagger();
+    if (!tagger) return;
+    try {
+      const kind = kindFor(options.type);
+      const ownAttrs = options.attributes;
+      const inheritedModelAttrs = options.parent?.type === SpanType.MODEL_GENERATION ? options.parent.attributes : void 0;
+      const traceCtx = this.resolveTraceContext(traceId, options);
+      tagger.registerLLMObsSpan(ddSpan, {
+        parent: parentDdSpan,
+        kind,
+        name: options.name,
+        userId: traceCtx.userId,
+        sessionId: traceCtx.sessionId,
+        ...kind === "llm" && (ownAttrs?.model ?? inheritedModelAttrs?.model) ? { modelName: ownAttrs?.model ?? inheritedModelAttrs?.model } : {},
+        ...kind === "llm" && (ownAttrs?.provider ?? inheritedModelAttrs?.provider) ? { modelProvider: ownAttrs?.provider ?? inheritedModelAttrs?.provider } : {}
+      });
+    } catch (error) {
+      this.logger.warn("[DatadogBridge] Failed to register LLMObs span", {
+        error,
+        spanName: options.name,
+        spanType: options.type
+      });
+    }
+  }
+  getLlmObsTagger() {
+    const tagger = tracer3.llmobs?._tagger;
+    if (tagger && typeof tagger.registerLLMObsSpan === "function") {
+      return tagger;
+    }
+    return void 0;
+  }
+  /**
+   * Annotate the eagerly-created dd span and finish it using the final Mastra
+   * span state.
+   */
+  annotateAndFinishSpan(span) {
+    const ddSpan = this.ddSpanMap.get(span.id);
+    if (!ddSpan) {
+      this.logger.warn("[DatadogBridge] No dd span found when finalizing Mastra span", {
+        spanId: span.id,
+        spanName: span.name,
+        spanType: span.type,
+        mapSize: this.ddSpanMap.size
+      });
+      return;
+    }
+    const endTime = span.endTime ? toDate(span.endTime) : span.isEvent ? toDate(span.startTime) : /* @__PURE__ */ new Date();
+    const annotations = this.buildAnnotations(span);
+    try {
+      if (Object.keys(annotations).length > 0 && tracer3.llmobs?.annotate) {
+        tracer3.llmobs.annotate(ddSpan, annotations);
+      }
+    } catch (error) {
+      this.logger.error("[DatadogBridge] Failed to annotate span before finish", {
+        error,
+        spanId: span.id,
+        spanName: span.name
+      });
+    }
+    try {
+      if (span.errorInfo) {
+        this.setErrorTags(ddSpan, span.errorInfo);
+      }
+    } catch (error) {
+      this.logger.error("[DatadogBridge] Failed to set error tags on dd span", {
+        error,
+        spanId: span.id,
+        spanName: span.name
+      });
+    }
+    try {
+      if (typeof ddSpan.finish === "function") {
+        ddSpan.finish(endTime.getTime());
+      }
+    } catch (error) {
+      this.logger.error("[DatadogBridge] Failed to finish dd span", {
+        error,
+        spanId: span.id,
+        spanName: span.name
+      });
+    } finally {
+      this.ddSpanMap.delete(span.id);
+      this.releaseTraceContext(span.traceId);
+    }
+  }
+  captureTraceContext(traceId, options) {
+    const existing = this.traceContext.get(traceId);
+    const next = {
+      userId: firstString(options.metadata?.userId, options.parent?.metadata?.userId, existing?.userId),
+      sessionId: firstString(options.metadata?.sessionId, options.parent?.metadata?.sessionId, existing?.sessionId)
+    };
+    if (next.userId || next.sessionId) {
+      this.traceContext.set(traceId, next);
+    }
+  }
+  resolveTraceContext(traceId, options) {
+    const stored = this.traceContext.get(traceId);
+    return {
+      userId: firstString(options.metadata?.userId, options.parent?.metadata?.userId, stored?.userId),
+      sessionId: firstString(options.metadata?.sessionId, options.parent?.metadata?.sessionId, stored?.sessionId)
+    };
+  }
+  releaseTraceContext(traceId) {
+    const nextCount = (this.openSpanCounts.get(traceId) ?? 1) - 1;
+    if (nextCount > 0) {
+      this.openSpanCounts.set(traceId, nextCount);
+      return;
+    }
+    this.openSpanCounts.delete(traceId);
+    this.traceContext.delete(traceId);
+  }
+  buildAnnotations(span) {
+    const annotations = {};
+    if (span.input !== void 0) {
+      annotations.inputData = formatInput(span.input, span.type);
+    }
+    if (span.output !== void 0) {
+      annotations.outputData = formatOutput(span.output, span.type);
+    }
+    if (span.type === SpanType.MODEL_STEP) {
+      const usage = span.attributes?.usage;
+      const metrics = formatUsageMetrics(usage);
+      if (metrics) {
+        annotations.metrics = metrics;
+      }
+    }
+    const knownFields = ["usage", "model", "provider", "parameters"];
+    const otherAttributes = omitKeys(span.attributes ?? {}, knownFields);
+    const contextKeySet = new Set(this.config.requestContextKeys ?? []);
+    const flatContextTags = {};
+    const remainingMetadata = {};
+    for (const [key, value] of Object.entries(span.metadata ?? {})) {
+      if (contextKeySet.has(key)) {
+        flatContextTags[key] = value;
+      } else {
+        remainingMetadata[key] = value;
+      }
+    }
+    const remainingAttributes = {};
+    for (const [key, value] of Object.entries(otherAttributes)) {
+      if (contextKeySet.has(key)) {
+        if (!(key in flatContextTags)) {
+          flatContextTags[key] = value;
+        }
+      } else {
+        remainingAttributes[key] = value;
+      }
+    }
+    const combinedMetadata = {
+      ...remainingMetadata,
+      ...remainingAttributes
+    };
+    if (span.errorInfo) {
+      combinedMetadata["error.message"] = span.errorInfo.message;
+    }
+    if (Object.keys(combinedMetadata).length > 0) {
+      annotations.metadata = combinedMetadata;
+    }
+    const tags = {
+      ...flatContextTags
+    };
+    if (span.tags?.length) {
+      for (const tag of span.tags) {
+        const colonIndex = tag.indexOf(":");
+        if (colonIndex > 0) {
+          tags[tag.substring(0, colonIndex)] = tag.substring(colonIndex + 1);
+        } else {
+          tags[tag] = true;
+        }
+      }
+    }
+    if (span.errorInfo) {
+      tags.error = true;
+      if (span.errorInfo.id) {
+        tags["error.id"] = span.errorInfo.id;
+      }
+      if (span.errorInfo.domain) {
+        tags["error.domain"] = span.errorInfo.domain;
+      }
+      if (span.errorInfo.category) {
+        tags["error.category"] = span.errorInfo.category;
+      }
+    }
+    if (Object.keys(tags).length > 0) {
+      annotations.tags = tags;
+    }
+    return annotations;
+  }
+  setErrorTags(ddSpan, errorInfo) {
+    ddSpan.setTag("error", true);
+    ddSpan.setTag("error.message", errorInfo.message);
+    ddSpan.setTag("error.type", errorInfo.name ?? errorInfo.category ?? "Error");
+    if (errorInfo.stack) {
+      ddSpan.setTag("error.stack", errorInfo.stack);
+    }
+  }
+  async flush() {
+    if (this.isDisabled) return;
+    if (tracer3.llmobs?.flush) {
+      try {
+        await tracer3.llmobs.flush();
+        this.logger.debug("Datadog llmobs flushed");
+      } catch (e) {
+        this.logger.error("Error flushing llmobs", { error: e });
+      }
+    }
+    try {
+      await flushApmExporter();
+      this.logger.debug("Datadog APM exporter flushed");
+    } catch (e) {
+      this.logger.error("Error flushing Datadog APM exporter", { error: e });
+    }
+  }
+  async shutdown() {
+    for (const [spanId, ddSpan] of this.ddSpanMap) {
+      this.logger.warn(`[DatadogBridge] Force-finishing APM span that was not properly closed [id=${spanId}]`);
+      try {
+        if (typeof ddSpan.finish === "function") {
+          ddSpan.finish();
+        }
+      } catch {
+      }
+    }
+    this.ddSpanMap.clear();
+    this.openSpanCounts.clear();
+    this.traceContext.clear();
+    await this.flush();
+    if (tracer3.llmobs?.disable) {
+      try {
+        tracer3.llmobs.disable();
+      } catch (e) {
+        this.logger.error("Error disabling llmobs", { error: e });
+      }
+    }
+    await super.shutdown();
+  }
+};
+function firstString(...values) {
+  for (const value of values) {
+    if (typeof value === "string") {
+      return value;
+    }
+  }
+  return void 0;
+}
+function fillRandomBytes(bytes) {
+  try {
+    const webCrypto = globalThis.crypto;
+    if (webCrypto?.getRandomValues) {
+      webCrypto.getRandomValues.call(webCrypto, bytes);
+      return;
+    }
+  } catch {
+  }
+  for (let i = 0; i < bytes.length; i++) {
+    bytes[i] = Math.floor(Math.random() * 256);
+  }
+}
+function generateSpanId() {
+  const bytes = new Uint8Array(8);
+  fillRandomBytes(bytes);
+  return Array.from(bytes, (byte) => byte.toString(16).padStart(2, "0")).join("");
+}
+function generateTraceId() {
+  const bytes = new Uint8Array(16);
+  fillRandomBytes(bytes);
+  return Array.from(bytes, (byte) => byte.toString(16).padStart(2, "0")).join("");
+}
 var MAX_TRACE_LIFETIME_MS = 30 * 60 * 1e3;
 var REGULAR_CLEANUP_INTERVAL_MS = 1 * 60 * 1e3;
 var DatadogExporter = class extends BaseExporter {
@@ -180,7 +613,7 @@ var DatadogExporter = class extends BaseExporter {
    * Main entry point for tracing events from Mastra.
    */
   async _exportTracingEvent(event) {
-    if (this.isDisabled || !tracer2.llmobs) return;
+    if (this.isDisabled || !tracer3.llmobs) return;
     try {
       const span = event.exportedSpan;
       if (span.isEvent) {
@@ -330,21 +763,14 @@ var DatadogExporter = class extends BaseExporter {
    * are exported before the runtime instance is terminated.
    */
   async flush() {
-    if (this.isDisabled || !tracer2.llmobs) return;
-    if (tracer2.llmobs?.flush) {
+    if (this.isDisabled || !tracer3.llmobs) return;
+    if (tracer3.llmobs?.flush) {
       try {
-        await tracer2.llmobs.flush();
+        await tracer3.llmobs.flush();
         this.logger.debug("Datadog llmobs flushed");
       } catch (e) {
         this.logger.error("Error flushing llmobs", { error: e });
       }
-    } else if (tracer2.flush) {
-      try {
-        await tracer2.flush();
-        this.logger.debug("Datadog tracer flushed");
-      } catch (e) {
-        this.logger.error("Error flushing tracer", { error: e });
-      }
     }
   }
   /**
@@ -368,9 +794,9 @@ var DatadogExporter = class extends BaseExporter {
     }
     this.traceState.clear();
     await this.flush();
-    if (tracer2.llmobs?.disable) {
+    if (tracer3.llmobs?.disable) {
       try {
-        tracer2.llmobs.disable();
+        tracer3.llmobs.disable();
       } catch (e) {
         this.logger.error("Error disabling llmobs", { error: e });
       }
@@ -559,15 +985,15 @@ var DatadogExporter = class extends BaseExporter {
       model: span.attributes?.model,
       provider: span.attributes?.provider
     } : inheritedModelAttrs;
-    tracer2.llmobs.trace(traceOptions, (ddSpan) => {
+    tracer3.llmobs.trace(traceOptions, (ddSpan) => {
       const annotations = this.buildAnnotations(span);
       if (Object.keys(annotations).length > 0) {
-        tracer2.llmobs.annotate(ddSpan, annotations);
+        tracer3.llmobs.annotate(ddSpan, annotations);
       }
       if (span.errorInfo) {
         this.setErrorTags(ddSpan, span.errorInfo);
       }
-      const exported = tracer2.llmobs.exportSpan ? tracer2.llmobs.exportSpan(ddSpan) : void 0;
+      const exported = tracer3.llmobs.exportSpan ? tracer3.llmobs.exportSpan(ddSpan) : void 0;
       state.contexts.set(span.id, { ddSpan, exported });
       for (const child of node.children) {
         this.emitSpanTree(child, state, childInheritedModelAttrs);
@@ -583,28 +1009,28 @@ var DatadogExporter = class extends BaseExporter {
    */
   emitSingleSpan(span, state, parent) {
     const { traceOptions, endTimeMs } = this.buildSpanOptions(span);
-    const runTrace = () => tracer2.llmobs.trace(traceOptions, (ddSpan) => {
+    const runTrace = () => tracer3.llmobs.trace(traceOptions, (ddSpan) => {
       const annotations = this.buildAnnotations(span);
       if (Object.keys(annotations).length > 0) {
-        tracer2.llmobs.annotate(ddSpan, annotations);
+        tracer3.llmobs.annotate(ddSpan, annotations);
       }
       if (span.errorInfo) {
         this.setErrorTags(ddSpan, span.errorInfo);
       }
-      const exported = tracer2.llmobs.exportSpan ? tracer2.llmobs.exportSpan(ddSpan) : void 0;
+      const exported = tracer3.llmobs.exportSpan ? tracer3.llmobs.exportSpan(ddSpan) : void 0;
       state.contexts.set(span.id, { ddSpan, exported });
       if (typeof ddSpan.finish === "function") {
         ddSpan.finish(endTimeMs);
       }
     });
     if (parent) {
-      tracer2.scope().activate(parent, runTrace);
+      tracer3.scope().activate(parent, runTrace);
     } else {
       runTrace();
     }
   }
 };
 
-export { DatadogExporter };
+export { DatadogBridge, DatadogExporter };
 //# sourceMappingURL=index.js.map
 //# sourceMappingURL=index.js.map