@mastra/braintrust 1.0.0-beta.9 → 1.0.0

package/dist/index.js

@@ -1,10 +1,140 @@
 import { SpanType } from '@mastra/core/observability';
 import { omitKeys } from '@mastra/core/utils';
-import { BaseExporter } from '@mastra/observability';
+import { TrackingExporter } from '@mastra/observability';
 import { initLogger, currentSpan } from 'braintrust';
 
 // src/tracing.ts
 
+// src/formatter.ts
+function removeNullish(obj) {
+  return Object.fromEntries(Object.entries(obj).filter(([_, v]) => v != null));
+}
+function convertContentPart(part) {
+  if (!part || typeof part !== "object") {
+    return null;
+  }
+  switch (part.type) {
+    case "text":
+      return part.text || null;
+    case "image":
+      return "[image]";
+    case "file": {
+      const filePart = part;
+      if (filePart.filename || filePart.name) {
+        return `[file: ${filePart.filename || filePart.name}]`;
+      }
+      return "[file]";
+    }
+    case "reasoning": {
+      const reasoningPart = part;
+      if (typeof reasoningPart.text === "string" && reasoningPart.text.length > 0) {
+        return `[reasoning: ${reasoningPart.text.substring(0, 100)}${reasoningPart.text.length > 100 ? "..." : ""}]`;
+      }
+      return "[reasoning]";
+    }
+    case "tool-call":
+      return null;
+    case "tool-result":
+      return null;
+    default: {
+      const unknownPart = part;
+      if (typeof unknownPart.text === "string") {
+        return unknownPart.text;
+      }
+      if (typeof unknownPart.content === "string") {
+        return unknownPart.content;
+      }
+      return `[${unknownPart.type || "unknown"}]`;
+    }
+  }
+}
+function serializeToolResult(resultData) {
+  if (typeof resultData === "string") {
+    return resultData;
+  }
+  if (resultData && typeof resultData === "object" && "value" in resultData) {
+    const valueData = resultData.value;
+    return typeof valueData === "string" ? valueData : JSON.stringify(valueData);
+  }
+  if (resultData === void 0 || resultData === null) {
+    return "";
+  }
+  try {
+    return JSON.stringify(resultData);
+  } catch {
+    return "[unserializable result]";
+  }
+}
+function convertAISDKMessage(message) {
+  if (!message || typeof message !== "object") {
+    return message;
+  }
+  const { role, content, ...rest } = message;
+  if (typeof content === "string") {
+    return message;
+  }
+  if (Array.isArray(content)) {
+    if (content.length === 0) {
+      return { role, content: "", ...rest };
+    }
+    if (role === "user" || role === "system") {
+      const contentParts = content.map((part) => convertContentPart(part)).filter(Boolean);
+      return {
+        role,
+        content: contentParts.length > 0 ? contentParts.join("\n") : "",
+        ...rest
+      };
+    }
+    if (role === "assistant") {
+      const contentParts = content.filter((part) => part?.type !== "tool-call").map((part) => convertContentPart(part)).filter(Boolean);
+      const toolCallParts = content.filter((part) => part?.type === "tool-call");
+      const result = {
+        role,
+        content: contentParts.length > 0 ? contentParts.join("\n") : "",
+        ...rest
+      };
+      if (toolCallParts.length > 0) {
+        result.tool_calls = toolCallParts.map((tc) => {
+          const toolCall = tc;
+          const toolCallId = toolCall.toolCallId;
+          const toolName = toolCall.toolName;
+          const args = toolCall.args ?? toolCall.input;
+          let argsString;
+          if (typeof args === "string") {
+            argsString = args;
+          } else if (args !== void 0 && args !== null) {
+            argsString = JSON.stringify(args);
+          } else {
+            argsString = "{}";
+          }
+          return {
+            id: toolCallId,
+            type: "function",
+            function: {
+              name: toolName,
+              arguments: argsString
+            }
+          };
+        });
+      }
+      return result;
+    }
+    if (role === "tool") {
+      const toolResult = content.find((part) => part?.type === "tool-result");
+      if (toolResult) {
+        const resultData = toolResult.output ?? toolResult.result;
+        const resultContent = serializeToolResult(resultData);
+        return {
+          role: "tool",
+          content: resultContent,
+          tool_call_id: toolResult.toolCallId
+        };
+      }
+    }
+  }
+  return message;
+}
+
 // src/metrics.ts
 function formatUsageMetrics(usage) {
   const metrics = {};
@@ -29,8 +159,61 @@ function formatUsageMetrics(usage) {
   return metrics;
 }
 
+// src/thread-reconstruction.ts
+function reconstructThreadOutput(threadData, originalOutput) {
+  const messages = [];
+  const sortedSteps = [...threadData].sort((a, b) => a.stepIndex - b.stepIndex);
+  for (const step of sortedSteps) {
+    const sortedToolCalls = step.toolCalls ? [...step.toolCalls].sort((a, b) => {
+      if (!a.startTime || !b.startTime) return 0;
+      return a.startTime.getTime() - b.startTime.getTime();
+    }) : [];
+    if (sortedToolCalls.length > 0) {
+      messages.push({
+        role: "assistant",
+        content: step.text || "",
+        tool_calls: sortedToolCalls.map((tc) => {
+          const cleanArgs = tc.args && typeof tc.args === "object" && !Array.isArray(tc.args) ? removeNullish(tc.args) : tc.args;
+          return {
+            id: tc.toolCallId,
+            type: "function",
+            function: {
+              name: tc.toolName,
+              arguments: typeof cleanArgs === "string" ? cleanArgs : JSON.stringify(cleanArgs)
+            }
+          };
+        })
+      });
+      for (const tc of sortedToolCalls) {
+        if (tc.result !== void 0) {
+          messages.push({
+            role: "tool",
+            content: typeof tc.result === "string" ? tc.result : JSON.stringify(tc.result),
+            tool_call_id: tc.toolCallId
+          });
+        }
+      }
+    } else if (step.text) {
+      messages.push({
+        role: "assistant",
+        content: step.text
+      });
+    }
+  }
+  if (messages.length > 0) {
+    const lastMessage = messages[messages.length - 1];
+    const originalText = originalOutput?.text;
+    if (originalText && lastMessage.role === "tool") {
+      messages.push({
+        role: "assistant",
+        content: originalText
+      });
+    }
+  }
+  return messages;
+}
+
 // src/tracing.ts
-var MASTRA_TRACE_ID_METADATA_KEY = "mastra-trace-id";
 var DEFAULT_SPAN_TYPE = "task";
 var SPAN_TYPE_EXCEPTIONS = {
   [SpanType.MODEL_GENERATION]: "llm",
@@ -42,246 +225,250 @@ var SPAN_TYPE_EXCEPTIONS = {
 function mapSpanType(spanType) {
   return SPAN_TYPE_EXCEPTIONS[spanType] ?? DEFAULT_SPAN_TYPE;
 }
-var BraintrustExporter = class extends BaseExporter {
+var BraintrustExporter = class extends TrackingExporter {
   name = "braintrust";
-  traceMap = /* @__PURE__ */ new Map();
-  config;
   // Flags and logger for context-aware mode
-  useProvidedLogger;
-  providedLogger;
-  constructor(config) {
-    super(config);
-    if (config.braintrustLogger) {
-      this.useProvidedLogger = true;
-      this.providedLogger = config.braintrustLogger;
-      this.config = config;
+  #useProvidedLogger;
+  #providedLogger;
+  #localLogger;
+  constructor(config = {}) {
+    const resolvedApiKey = config.apiKey ?? process.env.BRAINTRUST_API_KEY;
+    const resolvedEndpoint = config.endpoint ?? process.env.BRAINTRUST_ENDPOINT;
+    super({
+      ...config,
+      apiKey: resolvedApiKey,
+      endpoint: resolvedEndpoint
+    });
+    this.#useProvidedLogger = !!config.braintrustLogger;
+    if (this.#useProvidedLogger) {
+      this.#providedLogger = config.braintrustLogger;
     } else {
-      if (!config.apiKey) {
-        this.setDisabled(`Missing required credentials (apiKey: ${!!config.apiKey})`);
-        this.config = null;
-        this.useProvidedLogger = false;
+      if (!this.config.apiKey) {
+        this.setDisabled(
+          `Missing required API key. Set BRAINTRUST_API_KEY environment variable or pass apiKey in config.`
+        );
         return;
       }
-      this.useProvidedLogger = false;
-      this.config = config;
+      this.#localLogger = void 0;
     }
   }
-  async _exportTracingEvent(event) {
-    if (event.exportedSpan.isEvent) {
-      await this.handleEventSpan(event.exportedSpan);
-      return;
+  async getLocalLogger() {
+    if (this.#localLogger) {
+      return this.#localLogger;
     }
-    switch (event.type) {
-      case "span_started":
-        await this.handleSpanStarted(event.exportedSpan);
-        break;
-      case "span_updated":
-        await this.handleSpanUpdateOrEnd(event.exportedSpan, false);
-        break;
-      case "span_ended":
-        await this.handleSpanUpdateOrEnd(event.exportedSpan, true);
-        break;
+    try {
+      const logger = await initLogger({
+        projectName: this.config.projectName ?? "mastra-tracing",
+        apiKey: this.config.apiKey,
+        appUrl: this.config.endpoint,
+        ...this.config.tuningParameters
+      });
+      this.#localLogger = logger;
+      return logger;
+    } catch (err) {
+      this.logger.error("Braintrust exporter: Failed to initialize logger", { error: err });
+      this.setDisabled("Failed to initialize Braintrust logger");
     }
   }
-  async handleSpanStarted(span) {
-    if (span.isRootSpan) {
-      if (this.useProvidedLogger) {
-        await this.initLoggerOrUseContext(span);
-      } else {
-        await this.initLoggerPerTrace(span);
-      }
-    }
-    const method = "handleSpanStarted";
-    const spanData = this.getSpanData({ span, method });
-    if (!spanData) {
-      return;
-    }
-    if (!span.isEvent) {
-      spanData.activeIds.add(span.id);
-    }
-    const braintrustParent = this.getBraintrustParent({ spanData, span, method });
-    if (!braintrustParent) {
-      return;
-    }
+  startSpan(args) {
+    const { parent, span } = args;
     const payload = this.buildSpanPayload(span);
-    const braintrustSpan = braintrustParent.startSpan({
+    const braintrustSpan = parent.startSpan({
       spanId: span.id,
       name: span.name,
       type: mapSpanType(span.type),
-      ...payload
-    });
-    braintrustSpan.log({
-      metadata: {
-        [MASTRA_TRACE_ID_METADATA_KEY]: span.traceId
-      },
-      ...span.isRootSpan && span.tags?.length ? { tags: span.tags } : {}
+      startTime: span.startTime.getTime() / 1e3,
+      event: {
+        id: span.id,
+        // Use Mastra span ID as Braintrust row ID for logFeedback() compatibility
+        ...payload
+      }
     });
-    spanData.spans.set(span.id, braintrustSpan);
+    const isModelGeneration = span.type === SpanType.MODEL_GENERATION;
+    return {
+      span: braintrustSpan,
+      spanType: span.type,
+      threadData: isModelGeneration ? [] : void 0,
+      pendingToolResults: isModelGeneration ? /* @__PURE__ */ new Map() : void 0
+    };
   }
-  async handleSpanUpdateOrEnd(span, isEnd) {
-    const method = isEnd ? "handleSpanEnd" : "handleSpanUpdate";
-    const spanData = this.getSpanData({ span, method });
-    if (!spanData) {
-      return;
-    }
-    const braintrustSpan = spanData.spans.get(span.id);
-    if (!braintrustSpan) {
-      this.logger.warn("Braintrust exporter: No Braintrust span found for span update/end", {
-        traceId: span.traceId,
-        spanId: span.id,
-        spanName: span.name,
-        spanType: span.type,
-        isRootSpan: span.isRootSpan,
-        parentSpanId: span.parentSpanId,
-        method
-      });
-      return;
-    }
-    braintrustSpan.log(this.buildSpanPayload(span));
-    if (isEnd) {
-      if (span.endTime) {
-        braintrustSpan.end({ endTime: span.endTime.getTime() / 1e3 });
+  async _buildRoot(_args) {
+    if (this.#useProvidedLogger) {
+      const externalSpan = currentSpan();
+      if (externalSpan && externalSpan.id) {
+        return externalSpan;
       } else {
-        braintrustSpan.end();
-      }
-      if (!span.isEvent) {
-        spanData.activeIds.delete(span.id);
-      }
-      if (spanData.activeIds.size === 0 && !spanData.isExternal) {
-        this.traceMap.delete(span.traceId);
+        return this.#providedLogger;
       }
+    } else {
+      return this.getLocalLogger();
     }
   }
-  async handleEventSpan(span) {
+  async _buildSpan(args) {
+    const { span, traceData } = args;
     if (span.isRootSpan) {
-      this.logger.debug("Braintrust exporter: Creating logger for event", {
-        traceId: span.traceId,
-        spanId: span.id,
-        spanName: span.name,
-        method: "handleEventSpan"
-      });
-      if (this.useProvidedLogger) {
-        await this.initLoggerOrUseContext(span);
-      } else {
-        await this.initLoggerPerTrace(span);
+      const root = traceData.getRoot();
+      if (root) {
+        return this.startSpan({ parent: root, span });
+      }
+    } else {
+      const parent = traceData.getParent(args);
+      if (parent) {
+        const parentSpan = "span" in parent ? parent.span : parent;
+        return this.startSpan({ parent: parentSpan, span });
       }
     }
-    const method = "handleEventSpan";
-    const spanData = this.getSpanData({ span, method });
+  }
+  async _buildEvent(args) {
+    const spanData = await this._buildSpan(args);
     if (!spanData) {
       return;
     }
-    const braintrustParent = this.getBraintrustParent({ spanData, span, method });
-    if (!braintrustParent) {
+    spanData.span.end({ endTime: args.span.startTime.getTime() / 1e3 });
+    return spanData.span;
+  }
+  async _updateSpan(args) {
+    const { span, traceData } = args;
+    const spanData = traceData.getSpan({ spanId: span.id });
+    if (!spanData) {
       return;
     }
-    const payload = this.buildSpanPayload(span);
-    const braintrustSpan = braintrustParent.startSpan({
-      spanId: span.id,
-      name: span.name,
-      type: mapSpanType(span.type),
-      startTime: span.startTime.getTime() / 1e3,
-      ...payload
-    });
-    braintrustSpan.end({ endTime: span.startTime.getTime() / 1e3 });
+    spanData.span.log(this.buildSpanPayload(span, false));
   }
-  initTraceMap(params) {
-    const { traceId, isExternal, logger } = params;
-    if (this.traceMap.has(traceId)) {
-      this.logger.debug("Braintrust exporter: Reusing existing trace from local map", { traceId });
+  async _finishSpan(args) {
+    const { span, traceData } = args;
+    const spanData = traceData.getSpan({ spanId: span.id });
+    if (!spanData) {
       return;
     }
-    this.traceMap.set(traceId, {
-      logger,
-      spans: /* @__PURE__ */ new Map(),
-      activeIds: /* @__PURE__ */ new Set(),
-      isExternal
+    if (span.type === SpanType.MODEL_STEP) {
+      this.accumulateModelStepData(span, traceData);
+    } else if (span.type === SpanType.TOOL_CALL) {
+      this.accumulateToolCallResult(span, traceData);
+    }
+    const payload = span.type === SpanType.MODEL_GENERATION ? this.buildModelGenerationPayload(span, spanData) : this.buildSpanPayload(span, false);
+    spanData.span.log(payload);
+    if (span.endTime) {
+      spanData.span.end({ endTime: span.endTime.getTime() / 1e3 });
+    } else {
+      spanData.span.end();
+    }
+  }
+  async _abortSpan(args) {
+    const { span: spanData, reason } = args;
+    spanData.span.log({
+      error: reason.message,
+      metadata: { errorDetails: reason }
     });
+    spanData.span.end();
   }
+  // ==============================================================================
+  // Thread view reconstruction helpers
+  // ==============================================================================
   /**
-   * Creates a new logger per trace using config credentials
+   * Walk up the tree to find the MODEL_GENERATION ancestor span.
+   * Returns the BraintrustSpanData if found, undefined otherwise.
    */
-  async initLoggerPerTrace(span) {
-    if (this.traceMap.has(span.traceId)) {
-      this.logger.debug("Braintrust exporter: Reusing existing trace from local map", { traceId: span.traceId });
-      return;
-    }
-    try {
-      const loggerInstance = await initLogger({
-        projectName: this.config.projectName ?? "mastra-tracing",
-        apiKey: this.config.apiKey,
-        appUrl: this.config.endpoint,
-        ...this.config.tuningParameters
-      });
-      this.initTraceMap({ logger: loggerInstance, isExternal: false, traceId: span.traceId });
-    } catch (err) {
-      this.logger.error("Braintrust exporter: Failed to initialize logger", { error: err, traceId: span.traceId });
-      this.setDisabled("Failed to initialize Braintrust logger");
+  findModelGenerationAncestor(spanId, traceData) {
+    let currentId = spanId;
+    while (currentId) {
+      const parentId = traceData.getParentId({ spanId: currentId });
+      if (!parentId) return void 0;
+      const parentSpanData = traceData.getSpan({ spanId: parentId });
+      if (parentSpanData?.spanType === SpanType.MODEL_GENERATION) {
+        return parentSpanData;
+      }
+      currentId = parentId;
     }
+    return void 0;
   }
   /**
-   * Uses provided logger and detects external Braintrust spans.
-   * If a Braintrust span is detected (from logger.traced() or Eval()), attaches to it.
-   * Otherwise, uses the provided logger instance.
+   * Accumulate MODEL_STEP data to the parent MODEL_GENERATION's threadData.
+   * Called when a MODEL_STEP span ends.
    */
-  async initLoggerOrUseContext(span) {
-    if (this.traceMap.has(span.traceId)) {
-      this.logger.debug("Braintrust exporter: Reusing existing trace from local map", { traceId: span.traceId });
+  accumulateModelStepData(span, traceData) {
+    const modelGenSpanData = this.findModelGenerationAncestor(span.id, traceData);
+    if (!modelGenSpanData?.threadData) {
       return;
     }
-    const braintrustSpan = currentSpan();
-    if (braintrustSpan && braintrustSpan.id) {
-      this.initTraceMap({ logger: braintrustSpan, isExternal: true, traceId: span.traceId });
-    } else {
-      this.initTraceMap({ logger: this.providedLogger, isExternal: false, traceId: span.traceId });
-    }
+    const output = span.output;
+    const attributes = span.attributes;
+    const stepData = {
+      stepSpanId: span.id,
+      stepIndex: attributes?.stepIndex ?? 0,
+      text: output?.text,
+      toolCalls: output?.toolCalls?.map((tc) => ({
+        toolCallId: tc.toolCallId,
+        toolName: tc.toolName,
+        args: tc.args
+      }))
+    };
+    modelGenSpanData.threadData.push(stepData);
   }
-  getSpanData(options) {
-    const { span, method } = options;
-    if (this.traceMap.has(span.traceId)) {
-      return this.traceMap.get(span.traceId);
+  /**
+   * Store TOOL_CALL result in parent MODEL_GENERATION's pendingToolResults.
+   * Called when a TOOL_CALL span ends.
+   * Results are merged into threadData when MODEL_GENERATION ends.
+   */
+  accumulateToolCallResult(span, traceData) {
+    const modelGenSpanData = this.findModelGenerationAncestor(span.id, traceData);
+    if (!modelGenSpanData?.pendingToolResults) {
+      return;
     }
-    this.logger.warn("Braintrust exporter: No span data found for span", {
-      traceId: span.traceId,
-      spanId: span.id,
-      spanName: span.name,
-      spanType: span.type,
-      isRootSpan: span.isRootSpan,
-      parentSpanId: span.parentSpanId,
-      method
+    const input = span.input;
+    const toolCallId = input?.toolCallId;
+    if (!toolCallId) {
+      return;
+    }
+    modelGenSpanData.pendingToolResults.set(toolCallId, {
+      result: span.output,
+      startTime: span.startTime
     });
   }
-  getBraintrustParent(options) {
-    const { spanData, span, method } = options;
-    const parentId = span.parentSpanId;
-    if (!parentId) {
-      return spanData.logger;
+  /**
+   * Build the payload for MODEL_GENERATION span, reconstructing output from threadData if available.
+   */
+  buildModelGenerationPayload(span, spanData) {
+    const basePayload = this.buildSpanPayload(span, false);
+    const threadData = spanData.threadData;
+    if (!threadData || threadData.length === 0) {
+      return basePayload;
     }
-    if (spanData.spans.has(parentId)) {
-      return spanData.spans.get(parentId);
+    if (spanData.pendingToolResults && spanData.pendingToolResults.size > 0) {
+      for (const step of threadData) {
+        if (step.toolCalls) {
+          for (const toolCall of step.toolCalls) {
+            const pendingResult = spanData.pendingToolResults.get(toolCall.toolCallId);
+            if (pendingResult) {
+              toolCall.result = pendingResult.result;
+              toolCall.startTime = pendingResult.startTime;
+            }
+          }
+        }
+      }
     }
-    if (parentId && !spanData.spans.has(parentId)) {
-      return spanData.logger;
+    const hasToolCalls = threadData.some((step) => step.toolCalls && step.toolCalls.length > 0);
+    if (!hasToolCalls) {
+      return basePayload;
     }
-    this.logger.warn("Braintrust exporter: No parent data found for span", {
-      traceId: span.traceId,
-      spanId: span.id,
-      spanName: span.name,
-      spanType: span.type,
-      isRootSpan: span.isRootSpan,
-      parentSpanId: span.parentSpanId,
-      method
-    });
+    const reconstructedOutput = reconstructThreadOutput(threadData, span.output);
+    return {
+      ...basePayload,
+      output: reconstructedOutput
+    };
   }
   /**
    * Transforms MODEL_GENERATION input to Braintrust Thread view format.
+   * Converts AI SDK messages (v4/v5) to OpenAI Chat Completion format, which Braintrust requires
+   * for proper rendering of threads (fixes #11023).
    */
   transformInput(input, spanType) {
     if (spanType === SpanType.MODEL_GENERATION) {
-      if (input && Array.isArray(input.messages)) {
-        return input.messages;
-      } else if (input && typeof input === "object" && "content" in input) {
-        return [{ role: input.role, content: input.content }];
+      if (Array.isArray(input)) {
+        return input.map((msg) => convertAISDKMessage(msg));
+      }
+      if (input && typeof input === "object" && "messages" in input && Array.isArray(input.messages)) {
+        return input.messages.map((msg) => convertAISDKMessage(msg));
       }
     }
     return input;
@@ -291,12 +478,15 @@ var BraintrustExporter = class extends BaseExporter {
    */
   transformOutput(output, spanType) {
     if (spanType === SpanType.MODEL_GENERATION) {
+      if (!output || typeof output !== "object") {
+        return output;
+      }
       const { text, ...rest } = output;
-      return { role: "assistant", content: text, ...rest };
+      return { role: "assistant", content: text, ...removeNullish(rest) };
     }
     return output;
   }
-  buildSpanPayload(span) {
+  buildSpanPayload(span, isCreate = true) {
     const payload = {};
     if (span.input !== void 0) {
       payload.input = this.transformInput(span.input, span.type);
@@ -304,11 +494,17 @@ var BraintrustExporter = class extends BaseExporter {
     if (span.output !== void 0) {
       payload.output = this.transformOutput(span.output, span.type);
     }
+    if (isCreate && span.isRootSpan && span.tags?.length) {
+      payload.tags = span.tags;
+    }
     payload.metrics = {};
     payload.metadata = {
-      spanType: span.type,
-      ...span.metadata
+      ...span.metadata,
+      spanType: span.type
     };
+    if (isCreate) {
+      payload.metadata["mastra-trace-id"] = span.traceId;
+    }
     const attributes = span.attributes ?? {};
     if (span.type === SpanType.MODEL_GENERATION) {
       const modelAttr = attributes;
@@ -343,20 +539,9 @@ var BraintrustExporter = class extends BaseExporter {
     if (Object.keys(payload.metrics).length === 0) {
       delete payload.metrics;
     }
+    payload.metadata = removeNullish(payload.metadata);
     return payload;
   }
-  async shutdown() {
-    if (!this.config) {
-      return;
-    }
-    for (const [_traceId, spanData] of this.traceMap) {
-      for (const [_spanId, span] of spanData.spans) {
-        span.end();
-      }
-    }
-    this.traceMap.clear();
-    await super.shutdown();
-  }
 };
 
 export { BraintrustExporter };