@mastra/observability 1.18.4 → 1.19.0-alpha.1

package/dist/index.js

@@ -1041,6 +1041,7 @@ var MetricsContextImpl = class {
 				labels: filteredLabels,
 				correlationContext: this.correlationContext,
 				costContext,
+				...options?.usageId ? { usageId: options.usageId } : {},
 				metadata: this.metadata
 			}
 		};
@@ -1742,13 +1743,13 @@ function emitUsageMetrics(attrs, usage, metrics) {
 		metricCosts = /* @__PURE__ */ new Map();
 	}
 	const labels = attrs.usageIncomplete ? { usageIncomplete: "true" } : void 0;
+	const usageId = generateSignalId();
 	const emit = (name, value) => {
 		const costContext = metricCosts.get(name);
-		if (!costContext) {
-			metrics.emit(name, value, labels);
-			return;
-		}
-		metrics.emit(name, value, labels, { costContext });
+		metrics.emit(name, value, labels, costContext ? {
+			costContext,
+			usageId
+		} : { usageId });
 	};
 	for (const sample of getTokenMetricSamples(usage)) emit(sample.name, sample.value);
 }
@@ -3179,6 +3180,7 @@ var BaseSpan = class BaseSpan {
 	exportSpan(includeInternalSpans) {
 		const hideInput = this.traceState?.hideInput ?? false;
 		const hideOutput = this.traceState?.hideOutput ?? false;
+		const nestedUnderParent = this.traceState?.nestedUnderParent ?? false;
 		return {
 			id: this.id,
 			traceId: this.traceId,
@@ -3200,7 +3202,8 @@ var BaseSpan = class BaseSpan {
 			...this.isInternal ? { isInternal: true } : {},
 			parentSpanId: this.getParentSpanId(includeInternalSpans),
 			externalParentSpanId: this.getExportedExternalParentSpanId(includeInternalSpans),
-			...this.isRootSpan && this.tags?.length ? { tags: this.tags } : {}
+			...nestedUnderParent ? { nestedUnderParent: true } : {},
+			...this.isRootSpan && !nestedUnderParent && this.tags?.length ? { tags: this.tags } : {}
 		};
 	}
 	get externalTraceId() {
@@ -3213,7 +3216,7 @@ var BaseSpan = class BaseSpan {
 	async executeInContext(fn) {
 		const bridge = this.observabilityInstance.getBridge();
 		if (bridge?.executeInContext) {
-			const bridgeContextSpan = this.isInternal ? this.getParentSpan(false) : this;
+			const bridgeContextSpan = this.isInternal || this.isExcluded ? this.getParentSpan(false) : this;
 			return bridge.executeInContext(bridgeContextSpan?.id ?? this.id, fn);
 		}
 		return fn();
@@ -3225,7 +3228,7 @@ var BaseSpan = class BaseSpan {
 	executeInContextSync(fn) {
 		const bridge = this.observabilityInstance.getBridge();
 		if (bridge?.executeInContextSync) {
-			const bridgeContextSpan = this.isInternal ? this.getParentSpan(false) : this;
+			const bridgeContextSpan = this.isInternal || this.isExcluded ? this.getParentSpan(false) : this;
 			return bridge.executeInContextSync(bridgeContextSpan?.id ?? this.id, fn);
 		}
 		return fn();
@@ -3270,7 +3273,7 @@ var DefaultSpan = class extends BaseSpan {
 			return;
 		}
 		const bridge = observabilityInstance.getBridge();
-		if (bridge && !this.isInternal) {
+		if (bridge && !this.isInternal && !this.isExcluded) {
 			const bridgeIds = bridge.createSpan(options);
 			if (bridgeIds) {
 				this.id = bridgeIds.spanId;
@@ -4747,6 +4750,11 @@ var BaseObservabilityInstance = class extends MastraBase {
 			tags,
 			requestContext
 		});
+		if (!options.parent && traceState?.nestedUnderParent && (!tracingOptions?.traceId || span.traceId !== tracingOptions.traceId || !tracingOptions.parentSpanId || !(span instanceof BaseSpan) || span.externalParentSpanId !== tracingOptions.parentSpanId)) {
+			const { nestedUnderParent: _ignored, ...ownTraceState } = traceState;
+			span.traceState = ownTraceState;
+			this.logger.debug("[Observability] Ignoring tracingOptions.nestUnderParent: the run did not join the requested trace under the requested parent");
+		}
 		if (rest.type === SpanType.MODEL_GENERATION && this.config.excludeSpanTypes?.includes(SpanType.MODEL_GENERATION)) {
 			const attrs = rest.attributes;
 			const model = resolveModelId(attrs?.responseModel, attrs?.model);
@@ -4782,6 +4790,7 @@ var BaseObservabilityInstance = class extends MastraBase {
 			traceId: cached.traceId,
 			spanId: cached.id,
 			parentSpanId: cached.parentSpanId,
+			externalParentSpanId: cached.externalParentSpanId,
 			startTime: cached.startTime instanceof Date ? cached.startTime : new Date(cached.startTime),
 			input: cached.input,
 			attributes: cached.attributes,
@@ -4789,7 +4798,11 @@ var BaseObservabilityInstance = class extends MastraBase {
 			entityType: cached.entityType,
 			entityId: cached.entityId,
 			entityName: cached.entityName,
-			tracingPolicy: cached.isInternal ? { internal: InternalSpans.ALL } : void 0
+			tracingPolicy: cached.isInternal ? { internal: InternalSpans.ALL } : void 0,
+			traceState: cached.nestedUnderParent ? {
+				requestContextKeys: [],
+				nestedUnderParent: true
+			} : void 0
 		});
 		this.wireSpanLifecycle(span);
 		return span;
@@ -4993,11 +5006,13 @@ var BaseObservabilityInstance = class extends MastraBase {
 		const allKeys = [...configuredKeys, ...additionalKeys];
 		const hideInput = tracingOptions?.hideInput;
 		const hideOutput = tracingOptions?.hideOutput;
-		if (allKeys.length === 0 && !hideInput && !hideOutput) return;
+		const nestedUnderParent = tracingOptions?.nestUnderParent === true;
+		if (allKeys.length === 0 && !hideInput && !hideOutput && !nestedUnderParent) return;
 		return {
 			requestContextKeys: allKeys,
 			...hideInput !== void 0 && { hideInput },
-			...hideOutput !== void 0 && { hideOutput }
+			...hideOutput !== void 0 && { hideOutput },
+			...nestedUnderParent && { nestedUnderParent }
 		};
 	}
 	/**
@@ -5869,6 +5884,12 @@ var TraceData = class {
 		this.#waitingForParent.delete(args.spanId);
 	}
 	/**
+	* Get the span IDs that have events waiting on them.
+	*/
+	getWaitingForSpanIds() {
+		return [...this.#waitingForParent.keys()];
+	}
+	/**
 	* Get total count of events in all waiting queues.
 	*/
 	waitingQueueSize() {
@@ -6052,6 +6073,8 @@ var TrackingExporter = class extends BaseExporter {
 	#hardCapEnforcementInProgress = false;
 	/** Map of traceId to scheduled cleanup timeout */
 	#pendingCleanups = /* @__PURE__ */ new Map();
+	/** Waiting-queue replays started by setImmediate that haven't finished yet */
+	#activeReplays = /* @__PURE__ */ new Set();
 	/** Subclass configuration with resolved values */
 	config;
 	/** Maximum attempts to process a queued event before dropping */
@@ -6079,12 +6102,13 @@ var TrackingExporter = class extends BaseExporter {
 	*/
 	#scheduleProcessWaitingForRoot(traceId) {
 		setImmediate(() => {
-			this.#processWaitingForRoot(traceId).catch((error) => {
+			if (this.#shutdownStarted) return;
+			this.#trackReplay(this.#processWaitingForRoot(traceId).then(() => {}, (error) => {
 				this.logger.error(`${this.name}: Error processing waiting-for-root queue`, {
 					error,
 					traceId
 				});
-			});
+			}));
 		});
 	}
 	/**
@@ -6093,24 +6117,59 @@ var TrackingExporter = class extends BaseExporter {
 	*/
 	#scheduleProcessWaitingFor(traceId, spanId) {
 		setImmediate(() => {
-			this.#processWaitingFor(traceId, spanId).catch((error) => {
+			if (this.#shutdownStarted) return;
+			this.#trackReplay(this.#processWaitingFor(traceId, spanId).then(() => {}, (error) => {
 				this.logger.error(`${this.name}: Error processing waiting queue`, {
 					error,
 					traceId,
 					spanId
 				});
-			});
+			}));
 		});
 	}
 	/**
+	* Track a running replay so flush() and shutdown() can wait for it.
+	*/
+	#trackReplay(replay) {
+		this.#activeReplays.add(replay);
+		replay.finally(() => this.#activeReplays.delete(replay));
+	}
+	/**
+	* Process every queued event whose dependency has been built, without waiting for the
+	* setImmediate-scheduled replays (which never fire under fake timers). Waits for replays
+	* that are already running, then repeats until a pass makes no progress, so cascades
+	* (root -> parent -> child -> end) are fully drained.
+	*/
+	async #drainWaitingQueues() {
+		while (true) {
+			while (this.#activeReplays.size > 0) await Promise.all(this.#activeReplays);
+			let progressed = false;
+			for (const [traceId, traceData] of this.#traceMap) try {
+				if (traceData.isRootProcessed() && await this.#processWaitingForRoot(traceId) > 0) progressed = true;
+				for (const spanId of traceData.getWaitingForSpanIds()) {
+					if (!traceData.hasSpan({ spanId })) continue;
+					if (await this.#processWaitingFor(traceId, spanId) > 0) progressed = true;
+				}
+			} catch (error) {
+				this.logger.error(`${this.name}: Error draining waiting queues`, {
+					error,
+					traceId
+				});
+			}
+			if (!progressed && this.#activeReplays.size === 0) return;
+		}
+	}
+	/**
 	* Process all events waiting for root span.
+	* Takes the queue before processing, so a concurrent call can't process the same events.
+	* @returns The number of events that left the queue (processed, moved or dropped)
 	*/
 	async #processWaitingForRoot(traceId) {
-		if (this.#shutdownStarted) return;
 		const traceData = this.#traceMap.get(traceId);
-		if (!traceData) return;
+		if (!traceData) return 0;
 		const queue = traceData.getEventsWaitingForRoot();
-		if (queue.length === 0) return;
+		if (queue.length === 0) return 0;
+		traceData.clearWaitingForRoot();
 		this.logger.debug(`${this.name}: Processing ${queue.length} events waiting for root`, { traceId });
 		const toKeep = [];
 		const now = Date.now();
@@ -6146,23 +6205,25 @@ var TrackingExporter = class extends BaseExporter {
 				else toKeep.push(queuedEvent);
 			}
 		}
-		traceData.clearWaitingForRoot();
 		for (const event of toKeep) traceData.addToWaitingQueue({
 			event: event.event,
 			waitingFor: "root",
 			attempts: event.attempts,
 			queuedAt: event.queuedAt
 		});
+		return queue.length - toKeep.length;
 	}
 	/**
 	* Process events waiting for a specific parent span.
+	* Takes the queue before processing, so a concurrent call can't process the same events.
+	* @returns The number of events that left the queue (processed, moved or dropped)
 	*/
 	async #processWaitingFor(traceId, spanId) {
-		if (this.#shutdownStarted) return;
 		const traceData = this.#traceMap.get(traceId);
-		if (!traceData) return;
+		if (!traceData) return 0;
 		const queue = traceData.getEventsWaitingFor({ spanId });
-		if (queue.length === 0) return;
+		if (queue.length === 0) return 0;
+		traceData.clearWaitingFor({ spanId });
 		this.logger.debug(`${this.name}: Processing ${queue.length} events waiting for span`, {
 			traceId,
 			spanId
@@ -6192,13 +6253,13 @@ var TrackingExporter = class extends BaseExporter {
 			queuedEvent.attempts++;
 			if (!await this.#tryProcessQueuedEvent(queuedEvent, traceData)) toKeep.push(queuedEvent);
 		}
-		traceData.clearWaitingFor({ spanId });
 		for (const event of toKeep) traceData.addToWaitingQueue({
 			event: event.event,
 			waitingFor: spanId,
 			attempts: event.attempts,
 			queuedAt: event.queuedAt
 		});
+		return queue.length - toKeep.length;
 	}
 	/**
 	* Try to process a queued event.
@@ -6658,11 +6719,15 @@ var TrackingExporter = class extends BaseExporter {
 	* This is useful in serverless environments where you need to ensure spans
 	* are exported before the runtime instance is terminated.
 	*
+	* Queued events whose dependencies have been built are processed first, so
+	* a shutdown() right after flush() only aborts spans that are still open.
+	*
 	* Subclasses should override _flush() to implement vendor-specific flush logic.
 	*/
 	async flush() {
 		if (this.isDisabled) return;
 		this.logger.debug(`${this.name}: Flushing`);
+		if (!this.#shutdownStarted) await this.#drainWaitingQueues();
 		await this._flush();
 	}
 	/**
@@ -6682,6 +6747,7 @@ var TrackingExporter = class extends BaseExporter {
 	async shutdown() {
 		if (this.isDisabled) return;
 		this.#shutdownStarted = true;
+		await this.#drainWaitingQueues();
 		await this._preShutdown();
 		for (const [traceId, timeout] of this.#pendingCleanups) {
 			clearTimeout(timeout);