diff --git a/.changeset/mastra-token-usage-ttft.md b/.changeset/mastra-token-usage-ttft.md new file mode 100644 index 000000000..f8a7f210f --- /dev/null +++ b/.changeset/mastra-token-usage-ttft.md @@ -0,0 +1,12 @@ +--- +"braintrust": patch +--- + +fix(mastra): Record `time_to_first_token` for streaming Mastra model spans + +The Mastra observability exporter now derives `time_to_first_token` (in +seconds) from a streaming model span's `completionStartTime` attribute and its +start time, matching the metric Braintrust surfaces for other streaming LLM +calls. + +Ports the TTFT portion of mastra-ai/mastra#11029. diff --git a/e2e/scenarios/mastra-instrumentation/__snapshots__/mastra-v1260.log-payloads.json b/e2e/scenarios/mastra-instrumentation/__snapshots__/mastra-v1260.log-payloads.json index 75d3077e7..279193455 100644 --- a/e2e/scenarios/mastra-instrumentation/__snapshots__/mastra-v1260.log-payloads.json +++ b/e2e/scenarios/mastra-instrumentation/__snapshots__/mastra-v1260.log-payloads.json @@ -63,6 +63,7 @@ "metric_keys": [ "completion_tokens", "prompt_tokens", + "time_to_first_token", "tokens" ], "name": "llm: 'mock-model-id'", @@ -163,6 +164,7 @@ "metric_keys": [ "completion_tokens", "prompt_tokens", + "time_to_first_token", "tokens" ], "name": "llm: 'mock-model-id'", diff --git a/e2e/scenarios/mastra-instrumentation/__snapshots__/mastra-v1260.span-tree.json b/e2e/scenarios/mastra-instrumentation/__snapshots__/mastra-v1260.span-tree.json index e7ee35483..c1520ef3b 100644 --- a/e2e/scenarios/mastra-instrumentation/__snapshots__/mastra-v1260.span-tree.json +++ b/e2e/scenarios/mastra-instrumentation/__snapshots__/mastra-v1260.span-tree.json @@ -88,6 +88,7 @@ "metric_keys": [ "completion_tokens", "prompt_tokens", + "time_to_first_token", "tokens" ] } @@ -184,6 +185,7 @@ "metric_keys": [ "completion_tokens", "prompt_tokens", + "time_to_first_token", "tokens" ] } diff --git a/e2e/scenarios/mastra-instrumentation/__snapshots__/mastra-v1260.span-tree.txt b/e2e/scenarios/mastra-instrumentation/__snapshots__/mastra-v1260.span-tree.txt index e94ee3248..180de6568 100644 --- a/e2e/scenarios/mastra-instrumentation/__snapshots__/mastra-v1260.span-tree.txt +++ b/e2e/scenarios/mastra-instrumentation/__snapshots__/mastra-v1260.span-tree.txt @@ -58,6 +58,7 @@ span_tree: │ │ metric_keys: [ │ │ "completion_tokens", │ │ "prompt_tokens", + │ │ "time_to_first_token", │ │ "tokens" │ │ ] │ │ └── step: 0 [llm] @@ -135,6 +136,7 @@ span_tree: │ │ metric_keys: [ │ │ "completion_tokens", │ │ "prompt_tokens", + │ │ "time_to_first_token", │ │ "tokens" │ │ ] │ │ └── step: 0 [llm] diff --git a/js/src/wrappers/mastra.test.ts b/js/src/wrappers/mastra.test.ts new file mode 100644 index 000000000..0423baa99 --- /dev/null +++ b/js/src/wrappers/mastra.test.ts @@ -0,0 +1,101 @@ +import { + afterEach, + beforeAll, + beforeEach, + describe, + expect, + test, +} from "vitest"; +import { configureNode } from "../node/config"; +import { _exportsForTestingOnly, initLogger } from "../logger"; +import { BraintrustObservabilityExporter } from "./mastra"; + +try { + configureNode(); +} catch { + // Best-effort initialization for test environments. +} + +type MastraExportedSpan = Parameters< + BraintrustObservabilityExporter["exportTracingEvent"] +>[0]["exportedSpan"]; + +describe("BraintrustObservabilityExporter model metrics", () => { + let backgroundLogger: ReturnType< + typeof _exportsForTestingOnly.useTestBackgroundLogger + >; + + beforeAll(async () => { + await _exportsForTestingOnly.simulateLoginForTests(); + }); + + beforeEach(() => { + backgroundLogger = _exportsForTestingOnly.useTestBackgroundLogger(); + initLogger({ projectId: "test-project-id", projectName: "mastra.test.ts" }); + }); + + afterEach(() => { + _exportsForTestingOnly.clearTestBackgroundLogger(); + }); + + // Run a model span's full lifecycle through the exporter and return the logged + // row. Both events are required: onEnd is a no-op without a prior onStart. + async function logModelSpan( + attributes: Record, + ): Promise { + const span: MastraExportedSpan = { + id: "span-1", + traceId: "trace-1", + name: "llm: 'mock-model'", + type: "model_generation", + startTime: 1_000_000_000, + attributes, + }; + const exporter = new BraintrustObservabilityExporter(); + await exporter.exportTracingEvent({ + type: "span_started", + exportedSpan: span, + }); + await exporter.exportTracingEvent({ + type: "span_ended", + exportedSpan: { ...span, endTime: 1_000_000_005 }, + }); + await backgroundLogger.flush(); + const events = (await backgroundLogger.drain()) as any[]; + return events.find((event) => event.span_attributes?.name === span.name); + } + + test("maps token usage and derives time_to_first_token in seconds", async () => { + // Span starts at t=1_000_000_000ms; first token at +750ms → 0.75s TTFT. + const row = await logModelSpan({ + usage: { + inputTokens: 100, + outputTokens: 50, + inputDetails: { cacheRead: 40, cacheWrite: 10 }, + outputDetails: { reasoning: 20 }, + }, + completionStartTime: new Date(1_000_000_750), + }); + + expect(row?.metrics).toMatchObject({ + prompt_tokens: 100, + completion_tokens: 50, + tokens: 150, + prompt_cached_tokens: 40, + prompt_cache_creation_tokens: 10, + completion_reasoning_tokens: 20, + }); + expect(row?.metrics?.time_to_first_token).toBeCloseTo(0.75, 5); + // completionStartTime feeds the TTFT metric, but the raw value stays in + // metadata for backward compatibility (earlier releases surfaced it there). + expect(row?.metadata?.completionStartTime).toBeDefined(); + }); + + test("omits time_to_first_token for non-streaming spans", async () => { + const row = await logModelSpan({ + usage: { inputTokens: 10, outputTokens: 5 }, + }); + expect(row?.metrics?.tokens).toBe(15); + expect(row?.metrics?.time_to_first_token).toBeUndefined(); + }); +}); diff --git a/js/src/wrappers/mastra.ts b/js/src/wrappers/mastra.ts index 9946a9643..0771f465b 100644 --- a/js/src/wrappers/mastra.ts +++ b/js/src/wrappers/mastra.ts @@ -170,6 +170,30 @@ function modelMetrics( return Object.keys(out).length > 0 ? out : undefined; } +/** + * Compute time-to-first-token for streaming model spans. Mastra records + * `completionStartTime` (the wall-clock time the first token/chunk arrived) in + * a `MODEL_GENERATION` / `MODEL_INFERENCE` span's attributes; Braintrust + * expects `time_to_first_token` as the elapsed **seconds** between the span + * start and that first token. Returns undefined for non-streaming spans (no + * `completionStartTime`) or when either timestamp is unusable. + */ +function timeToFirstTokenSeconds( + attributes: Record | undefined, + spanStartSeconds: number | undefined, +): number | undefined { + if (!isObject(attributes)) return undefined; + if (spanStartSeconds === undefined) return undefined; + const raw = attributes.completionStartTime; + const completionStart = + raw instanceof Date || typeof raw === "string" || typeof raw === "number" + ? epochSeconds(raw) + : undefined; + if (completionStart === undefined) return undefined; + const ttft = completionStart - spanStartSeconds; + return Number.isFinite(ttft) && ttft >= 0 ? ttft : undefined; +} + /** Build the metadata payload Braintrust shows on the span, merging * Mastra's own `metadata`, `attributes` (sans usage), and entity fields. */ function buildMetadata(exported: MastraExportedSpan): Record { @@ -183,6 +207,10 @@ function buildMetadata(exported: MastraExportedSpan): Record { if (exported.attributes && isObject(exported.attributes)) { for (const [key, value] of Object.entries(exported.attributes)) { if (key === "usage") continue; // surfaced via metrics + // `completionStartTime` is also surfaced as the `time_to_first_token` + // metric, but we keep the raw value in metadata too: earlier released + // versions exposed it there, so dropping it would be a backward- + // incompatible removal of a consumer-visible field. if (value !== undefined) out[key] = value; } } @@ -347,8 +375,15 @@ export class BraintrustObservabilityExporter implements MastraObservabilityExpor } const metrics = modelMetrics(exported.attributes); - if (metrics) { - event.metrics = metrics; + const ttft = timeToFirstTokenSeconds( + exported.attributes, + epochSeconds(exported.startTime), + ); + if (metrics || ttft !== undefined) { + event.metrics = { + ...(metrics ?? {}), + ...(ttft !== undefined ? { time_to_first_token: ttft } : {}), + }; } if (Object.keys(event).length > 0) {