Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 14 additions & 0 deletions .changeset/mastra-token-usage-ttft.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,14 @@
---
"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 model span's `completionStartTime` attribute and its start
time, matching the metric Braintrust surfaces for streaming LLM calls. The raw
`completionStartTime` is no longer duplicated into span metadata. Token usage
mapping (prompt/completion/total, cached, cache-creation, and reasoning tokens)
was already at parity with `@mastra/braintrust` and is unchanged.

Ports the TTFT portion of mastra-ai/mastra#11029.
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,7 @@
"metric_keys": [
"completion_tokens",
"prompt_tokens",
"time_to_first_token",
"tokens"
],
"name": "llm: 'mock-model-id'",
Expand Down Expand Up @@ -163,6 +164,7 @@
"metric_keys": [
"completion_tokens",
"prompt_tokens",
"time_to_first_token",
"tokens"
],
"name": "llm: 'mock-model-id'",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,7 @@
"metric_keys": [
"completion_tokens",
"prompt_tokens",
"time_to_first_token",
"tokens"
]
}
Expand Down Expand Up @@ -184,6 +185,7 @@
"metric_keys": [
"completion_tokens",
"prompt_tokens",
"time_to_first_token",
"tokens"
]
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@ span_tree:
│ │ metric_keys: [
│ │ "completion_tokens",
│ │ "prompt_tokens",
│ │ "time_to_first_token",
│ │ "tokens"
│ │ ]
│ │ └── step: 0 [llm]
Expand Down Expand Up @@ -135,6 +136,7 @@ span_tree:
│ │ metric_keys: [
│ │ "completion_tokens",
│ │ "prompt_tokens",
│ │ "time_to_first_token",
│ │ "tokens"
│ │ ]
│ │ └── step: 0 [llm]
Expand Down
100 changes: 100 additions & 0 deletions js/src/wrappers/mastra.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,100 @@
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<string, unknown>,
): Promise<any> {
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 is a metric input, not user-facing metadata.
expect(row?.metadata?.completionStartTime).toBeUndefined();
});

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();
});
});
36 changes: 34 additions & 2 deletions js/src/wrappers/mastra.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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<string, unknown> | 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<string, unknown> {
Expand All @@ -183,6 +207,7 @@ function buildMetadata(exported: MastraExportedSpan): Record<string, unknown> {
if (exported.attributes && isObject(exported.attributes)) {
for (const [key, value] of Object.entries(exported.attributes)) {
if (key === "usage") continue; // surfaced via metrics
if (key === "completionStartTime") continue; // surfaced via metrics (TTFT)
if (value !== undefined) out[key] = value;
}
}
Expand Down Expand Up @@ -347,8 +372,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) {
Expand Down
Loading