@mastra/datadog 1.4.15 → 1.4.16-alpha.1

package/dist/index.js

@@ -59,6 +59,20 @@ function formatUsageMetrics(usage) {
 * Utility functions for Datadog LLM Observability Exporter
 */
 /**
+* Datadog span links for Mastra span links, such as the caller of a served MCP
+* request. Each link is resolved from a W3C `traceparent`, so only full-length
+* W3C IDs are kept, and links dd-trace can't resolve are skipped.
+*/
+function toDatadogLinks(links) {
+	const ddLinks = [];
+	for (const link of links ?? []) {
+		if (!/^[0-9a-f]{32}$/.test(link.traceId) || !/^[0-9a-f]{16}$/.test(link.spanId)) continue;
+		const context = tracer.extract("text_map", { traceparent: `00-${link.traceId}-${link.spanId}-01` });
+		if (context) ddLinks.push({ context });
+	}
+	return ddLinks;
+}
+/**
 * Maps Mastra SpanTypes to Datadog LLMObs span kinds for the legacy hierarchy
 * (no `model-inference-span` feature). MODEL_STEP is the actual API call here
 * because MODEL_INFERENCE doesn't exist.
@@ -309,10 +323,18 @@ var DatadogBridge = class extends BaseExporter {
 	name = "datadog-bridge";
 	config;
 	ddSpanMap = /* @__PURE__ */ new Map();
+	/** Keys of the links each dd span was started with, so links set later are added once at finish. */
+	startLinkKeys = /* @__PURE__ */ new WeakMap();
 	traceContext = /* @__PURE__ */ new Map();
 	openSpanCounts = /* @__PURE__ */ new Map();
+	flushGraceMs;
+	/** True when spans or evaluations were handed to dd-trace since the last flush() */
+	pendingSend = false;
+	/** Grace wait in progress, shared by flush() calls that overlap it */
+	graceWait = null;
 	constructor(config = {}) {
 		super(config);
+		this.flushGraceMs = config.flushGraceMs ?? 2e3;
 		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";
@@ -372,10 +394,13 @@ var DatadogBridge = class extends BaseExporter {
 				if (apmParentDdSpan) parentSource = "active-scope";
 			}
 			if (!llmobsParentDdSpan && externalParentId) llmobsParentDdSpan = apmParentDdSpan;
+			const links = toDatadogLinks(options.links);
 			const ddSpan = tracer.startSpan(options.name, {
 				...apmParentDdSpan ? { childOf: apmParentDdSpan } : {},
-				...options.startTime ? { startTime: toDate(options.startTime).getTime() } : {}
+				...options.startTime ? { startTime: toDate(options.startTime).getTime() } : {},
+				...links.length > 0 ? { links } : {}
 			});
+			if (options.links?.length) this.startLinkKeys.set(ddSpan, new Set(options.links.map((link) => `${link.traceId}:${link.spanId}`)));
 			const ddContext = ddSpan.context?.();
 			const spanId = ddContext?.toSpanId?.(true) ?? generateSpanId();
 			const traceId = ddContext?.toTraceId?.(true) ?? (externalParentId ? options.parent?.traceId ?? generateTraceId() : generateTraceId());
@@ -416,6 +441,7 @@ var DatadogBridge = class extends BaseExporter {
 			return;
 		}
 		try {
+			this.pendingSend = true;
 			tracer.llmobs.submitEvaluation({
 				traceId: score.traceId,
 				spanId: toDatadogSpanId(score.spanId)
@@ -538,7 +564,20 @@ var DatadogBridge = class extends BaseExporter {
 			});
 		}
 		try {
-			if (typeof ddSpan.finish === "function") ddSpan.finish(endTime.getTime());
+			const startKeys = this.startLinkKeys.get(ddSpan);
+			const lateLinks = toDatadogLinks(span.links?.filter((link) => !startKeys?.has(`${link.traceId}:${link.spanId}`)));
+			if (lateLinks.length > 0 && typeof ddSpan.addLinks === "function") ddSpan.addLinks(lateLinks);
+		} catch (error) {
+			this.logger.error("[DatadogBridge] Failed to add span links", {
+				error,
+				spanId: span.id
+			});
+		}
+		try {
+			if (typeof ddSpan.finish === "function") {
+				this.pendingSend = true;
+				ddSpan.finish(endTime.getTime());
+			}
 		} catch (error) {
 			this.logger.error("[DatadogBridge] Failed to finish dd span", {
 				error,
@@ -643,6 +682,7 @@ var DatadogBridge = class extends BaseExporter {
 		} catch (e) {
 			this.logger.error("Error flushing llmobs", { error: e });
 		}
+		await this.waitForBackgroundSend();
 		try {
 			await flushApmExporter();
 			this.logger.debug("Datadog APM exporter flushed");
@@ -650,11 +690,32 @@ var DatadogBridge = class extends BaseExporter {
 			this.logger.error("Error flushing Datadog APM exporter", { error: e });
 		}
 	}
+	/**
+	* `llmobs.flush()` returns before the background writer's HTTP request
+	* completes and dd-trace gives nothing to await, so wait a bounded time
+	* instead. Starts one wait per batch of new data, so the repeated flushes
+	* during shutdown don't stack. A flush that overlaps a running wait joins
+	* it, so no caller resolves before the data it produced had its wait.
+	*/
+	async waitForBackgroundSend() {
+		if (!(this.flushGraceMs > 0) || !Number.isFinite(this.flushGraceMs)) return;
+		if (this.pendingSend) {
+			this.pendingSend = false;
+			const wait = new Promise((resolve) => setTimeout(resolve, this.flushGraceMs)).finally(() => {
+				if (this.graceWait === wait) this.graceWait = null;
+			});
+			this.graceWait = wait;
+		}
+		await this.graceWait;
+	}
 	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();
+				if (typeof ddSpan.finish === "function") {
+					this.pendingSend = true;
+					ddSpan.finish();
+				}
 			} catch {}
 		}
 		this.ddSpanMap.clear();
@@ -719,8 +780,14 @@ var DatadogExporter = class extends BaseExporter {
 	config;
 	traceContext = /* @__PURE__ */ new Map();
 	traceState = /* @__PURE__ */ new Map();
+	flushGraceMs;
+	/** True when spans or evaluations were handed to dd-trace since the last flush() */
+	pendingSend = false;
+	/** Grace wait in progress, shared by flush() calls that overlap it */
+	graceWait = null;
 	constructor(config = {}) {
 		super(config);
+		this.flushGraceMs = config.flushGraceMs ?? 2e3;
 		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";
@@ -898,6 +965,7 @@ var DatadogExporter = class extends BaseExporter {
 			return;
 		}
 		try {
+			this.pendingSend = true;
 			tracer.llmobs.submitEvaluation({
 				traceId: exported.traceId,
 				spanId: exported.spanId
@@ -932,6 +1000,25 @@ var DatadogExporter = class extends BaseExporter {
 		} catch (e) {
 			this.logger.error("Error flushing llmobs", { error: e });
 		}
+		await this.waitForBackgroundSend();
+	}
+	/**
+	* `llmobs.flush()` returns before the background writer's HTTP request
+	* completes and dd-trace gives nothing to await, so wait a bounded time
+	* instead. Starts one wait per batch of new data, so the repeated flushes
+	* during shutdown don't stack. A flush that overlaps a running wait joins
+	* it, so no caller resolves before the data it produced had its wait.
+	*/
+	async waitForBackgroundSend() {
+		if (!(this.flushGraceMs > 0) || !Number.isFinite(this.flushGraceMs)) return;
+		if (this.pendingSend) {
+			this.pendingSend = false;
+			const wait = new Promise((resolve) => setTimeout(resolve, this.flushGraceMs)).finally(() => {
+				if (this.graceWait === wait) this.graceWait = null;
+			});
+			this.graceWait = wait;
+		}
+		await this.graceWait;
 	}
 	/**
 	* Gracefully shuts down the exporter.
@@ -1108,6 +1195,7 @@ var DatadogExporter = class extends BaseExporter {
 			provider: ownAttrs?.provider ?? inheritedModelAttrs?.provider
 		};
 		const startTime = toDate(span.startTime);
+		const links = toDatadogLinks(span.links);
 		const endTime = span.endTime ? toDate(span.endTime) : span.isEvent ? startTime : /* @__PURE__ */ new Date();
 		return {
 			traceOptions: {
@@ -1117,7 +1205,8 @@ var DatadogExporter = class extends BaseExporter {
 				userId: traceCtx.userId,
 				startTime,
 				...kind === "llm" && attrs?.model ? { modelName: attrs.model } : {},
-				...kind === "llm" && attrs?.provider ? { modelProvider: attrs.provider } : {}
+				...kind === "llm" && attrs?.provider ? { modelProvider: attrs.provider } : {},
+				...links.length > 0 ? { links } : {}
 			},
 			endTimeMs: endTime.getTime()
 		};
@@ -1134,6 +1223,7 @@ var DatadogExporter = class extends BaseExporter {
 			model: span.attributes?.model,
 			provider: span.attributes?.provider
 		} : inheritedModelAttrs;
+		this.pendingSend = true;
 		tracer.llmobs.trace(traceOptions, (ddSpan) => {
 			const annotations = this.buildAnnotations(span);
 			if (Object.keys(annotations).length > 0) tracer.llmobs.annotate(ddSpan, annotations);