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..4c8d73403 100644 --- a/e2e/scenarios/mastra-instrumentation/__snapshots__/mastra-v1260.log-payloads.json +++ b/e2e/scenarios/mastra-instrumentation/__snapshots__/mastra-v1260.log-payloads.json @@ -3,7 +3,6 @@ "metadata": { "scenario": "mastra-instrumentation" }, - "metric_keys": [], "name": "mastra-instrumentation-root", "type": "task" }, @@ -11,7 +10,6 @@ "metadata": { "operation": "generate" }, - "metric_keys": [], "name": "mastra-agent-generate-operation", "type": null }, @@ -22,7 +20,6 @@ "entity_name": "Weather Agent", "entity_type": "agent" }, - "metric_keys": [], "name": "agent run: 'weather-agent'", "output": { "files": [], @@ -60,11 +57,12 @@ "model": "mock-model-id", "provider": "mock-provider" }, - "metric_keys": [ - "completion_tokens", - "prompt_tokens", - "tokens" - ], + "metrics": { + "completion_tokens": 20, + "prompt_tokens": 10, + "time_to_first_token": 0, + "tokens": 30 + }, "name": "llm: 'mock-model-id'", "output": { "files": [], @@ -82,11 +80,11 @@ "entity_name": "Weather Agent", "entity_type": "agent" }, - "metric_keys": [ - "completion_tokens", - "prompt_tokens", - "tokens" - ], + "metrics": { + "completion_tokens": 20, + "prompt_tokens": 10, + "tokens": 30 + }, "name": "step: 0", "output": { "text": "The forecast is sunny.", @@ -100,7 +98,6 @@ "entity_name": "Weather Agent", "entity_type": "agent" }, - "metric_keys": [], "name": "chunk: 'text'", "output": { "text": "The forecast is sunny." @@ -111,7 +108,6 @@ "metadata": { "operation": "stream" }, - "metric_keys": [], "name": "mastra-agent-stream-operation", "type": null }, @@ -122,7 +118,6 @@ "entity_name": "Weather Agent", "entity_type": "agent" }, - "metric_keys": [], "name": "agent run: 'weather-agent'", "output": { "files": [], @@ -160,11 +155,12 @@ "model": "mock-model-id", "provider": "mock-provider" }, - "metric_keys": [ - "completion_tokens", - "prompt_tokens", - "tokens" - ], + "metrics": { + "completion_tokens": 20, + "prompt_tokens": 10, + "time_to_first_token": 0, + "tokens": 30 + }, "name": "llm: 'mock-model-id'", "output": { "files": [], @@ -182,11 +178,11 @@ "entity_name": "Weather Agent", "entity_type": "agent" }, - "metric_keys": [ - "completion_tokens", - "prompt_tokens", - "tokens" - ], + "metrics": { + "completion_tokens": 20, + "prompt_tokens": 10, + "tokens": 30 + }, "name": "step: 0", "output": { "text": "The forecast is sunny.", @@ -200,7 +196,6 @@ "entity_name": "Weather Agent", "entity_type": "agent" }, - "metric_keys": [], "name": "chunk: 'text'", "output": { "text": "The forecast is sunny." @@ -211,7 +206,6 @@ "metadata": { "operation": "tool" }, - "metric_keys": [], "name": "mastra-tool-operation", "type": null }, @@ -219,7 +213,6 @@ "metadata": { "operation": "workflow" }, - "metric_keys": [], "name": "mastra-workflow-operation", "type": null }, @@ -232,7 +225,6 @@ "entity_name": "travel-flow", "entity_type": "workflow_run" }, - "metric_keys": [], "name": "workflow run: 'travel-flow'", "output": { "forecast": "Sunny in Berlin" @@ -248,7 +240,6 @@ "entity_name": "travel-flow", "entity_type": "workflow_step" }, - "metric_keys": [], "name": "workflow step: 'lookup-step'", "output": { "forecast": "Sunny in Berlin" @@ -259,7 +250,6 @@ "metadata": { "scenario": "mastra-instrumentation" }, - "metric_keys": [], "name": "mastra-instrumentation-root", "type": "task" } 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..01bfe9400 100644 --- a/e2e/scenarios/mastra-instrumentation/__snapshots__/mastra-v1260.span-tree.json +++ b/e2e/scenarios/mastra-instrumentation/__snapshots__/mastra-v1260.span-tree.json @@ -42,11 +42,11 @@ "entity_name": "Weather Agent", "entity_type": "agent" }, - "metric_keys": [ - "completion_tokens", - "prompt_tokens", - "tokens" - ] + "metrics": { + "completion_tokens": 20, + "prompt_tokens": 10, + "tokens": 30 + } } ], "input": { @@ -85,11 +85,12 @@ "model": "mock-model-id", "provider": "mock-provider" }, - "metric_keys": [ - "completion_tokens", - "prompt_tokens", - "tokens" - ] + "metrics": { + "completion_tokens": 20, + "prompt_tokens": 10, + "time_to_first_token": 0, + "tokens": 30 + } } ], "input": "What is the weather in Paris?", @@ -138,11 +139,11 @@ "entity_name": "Weather Agent", "entity_type": "agent" }, - "metric_keys": [ - "completion_tokens", - "prompt_tokens", - "tokens" - ] + "metrics": { + "completion_tokens": 20, + "prompt_tokens": 10, + "tokens": 30 + } } ], "input": { @@ -181,11 +182,12 @@ "model": "mock-model-id", "provider": "mock-provider" }, - "metric_keys": [ - "completion_tokens", - "prompt_tokens", - "tokens" - ] + "metrics": { + "completion_tokens": 20, + "prompt_tokens": 10, + "time_to_first_token": 0, + "tokens": 30 + } } ], "input": "Stream the Paris forecast.", 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..cdf5c3c3a 100644 --- a/e2e/scenarios/mastra-instrumentation/__snapshots__/mastra-v1260.span-tree.txt +++ b/e2e/scenarios/mastra-instrumentation/__snapshots__/mastra-v1260.span-tree.txt @@ -55,11 +55,12 @@ span_tree: │ │ "model": "mock-model-id", │ │ "provider": "mock-provider" │ │ } - │ │ metric_keys: [ - │ │ "completion_tokens", - │ │ "prompt_tokens", - │ │ "tokens" - │ │ ] + │ │ metrics: { + │ │ "completion_tokens": 20, + │ │ "prompt_tokens": 10, + │ │ "time_to_first_token": 0, + │ │ "tokens": 30 + │ │ } │ │ └── step: 0 [llm] │ │ output: { │ │ "text": "The forecast is sunny.", @@ -70,11 +71,11 @@ span_tree: │ │ "entity_name": "Weather Agent", │ │ "entity_type": "agent" │ │ } - │ │ metric_keys: [ - │ │ "completion_tokens", - │ │ "prompt_tokens", - │ │ "tokens" - │ │ ] + │ │ metrics: { + │ │ "completion_tokens": 20, + │ │ "prompt_tokens": 10, + │ │ "tokens": 30 + │ │ } │ │ └── chunk: 'text' [llm] │ │ output: { │ │ "text": "The forecast is sunny." @@ -132,11 +133,12 @@ span_tree: │ │ "model": "mock-model-id", │ │ "provider": "mock-provider" │ │ } - │ │ metric_keys: [ - │ │ "completion_tokens", - │ │ "prompt_tokens", - │ │ "tokens" - │ │ ] + │ │ metrics: { + │ │ "completion_tokens": 20, + │ │ "prompt_tokens": 10, + │ │ "time_to_first_token": 0, + │ │ "tokens": 30 + │ │ } │ │ └── step: 0 [llm] │ │ output: { │ │ "text": "The forecast is sunny.", @@ -147,11 +149,11 @@ span_tree: │ │ "entity_name": "Weather Agent", │ │ "entity_type": "agent" │ │ } - │ │ metric_keys: [ - │ │ "completion_tokens", - │ │ "prompt_tokens", - │ │ "tokens" - │ │ ] + │ │ metrics: { + │ │ "completion_tokens": 20, + │ │ "prompt_tokens": 10, + │ │ "tokens": 30 + │ │ } │ │ └── chunk: 'text' [llm] │ │ output: { │ │ "text": "The forecast is sunny." diff --git a/e2e/scenarios/mastra-instrumentation/scenario.test.ts b/e2e/scenarios/mastra-instrumentation/scenario.test.ts index 7fc5ccf5b..17d65c1ed 100644 --- a/e2e/scenarios/mastra-instrumentation/scenario.test.ts +++ b/e2e/scenarios/mastra-instrumentation/scenario.test.ts @@ -89,6 +89,21 @@ function scrubMastraProviderTimestamps(value: Json): Json { return value; } +function summarizeMetrics( + metrics: CapturedLogEvent["metrics"], +): Record | undefined { + if (!metrics) return undefined; + + const entries = Object.entries(metrics) + .filter( + ([key, value]) => key !== "start" && key !== "end" && value !== undefined, + ) + .sort(([left], [right]) => left.localeCompare(right)) + .map(([key, value]) => [key, value as Json] as const); + + return entries.length > 0 ? Object.fromEntries(entries) : undefined; +} + function summarizeMastraPayload(event: CapturedLogEvent): Record { const metadata = event.row.metadata as Record | undefined; const pickedMetadata = Object.fromEntries( @@ -104,13 +119,15 @@ function summarizeMastraPayload(event: CapturedLogEvent): Record { const result: Record = { metadata: Object.keys(pickedMetadata).length > 0 ? (pickedMetadata as Json) : null, - metric_keys: Object.keys(event.metrics ?? {}) - .filter((key) => key !== "start" && key !== "end") - .sort(), name: event.span.name ?? null, type: event.span.type ?? null, }; + const metrics = summarizeMetrics(event.metrics); + if (metrics) { + result.metrics = metrics; + } + if (event.input !== undefined && event.input !== null) { result.input = scrubMastraProviderTimestamps(event.input as Json); } diff --git a/js/src/wrappers/mastra.test.ts b/js/src/wrappers/mastra.test.ts index dcced568b..b83c6b941 100644 --- a/js/src/wrappers/mastra.test.ts +++ b/js/src/wrappers/mastra.test.ts @@ -16,11 +16,11 @@ try { // Best-effort initialization for test environments. } -type ExportedSpan = Parameters< +type MastraExportedSpan = Parameters< BraintrustObservabilityExporter["exportTracingEvent"] >[0]["exportedSpan"]; -const span = (overrides: Partial): ExportedSpan => ({ +const span = (overrides: Partial): MastraExportedSpan => ({ id: "span-1", traceId: "trace-1", name: "agent run", @@ -78,4 +78,67 @@ describe("BraintrustObservabilityExporter", () => { expect(merged[0].span_attributes?.name).toBe("agent run"); expect(merged[0].scores?.quality).toBe(0.9); }); + + // 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 modelSpan = span({ + 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: modelSpan, + }); + await exporter.exportTracingEvent({ + type: "span_ended", + exportedSpan: { ...modelSpan, endTime: 1_000_000_005 }, + }); + await backgroundLogger.flush(); + const events = (await backgroundLogger.drain()) as any[]; + return events.find( + (event) => event.span_attributes?.name === modelSpan.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 28bba3410..a7f0ff607 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; } } @@ -353,8 +381,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) {