Skip to content

Commit 890b24c

Browse files
tulga-bytesTylerBarnesMastra Code (openai/gpt-5.5)
authored
feat: propagate cache metrics through harness token usage (#14746)
## Summary Propagate `cachedInputTokens` and `cacheCreationInputTokens` through the Mastra harness token usage pipeline so downstream consumers (e.g., auggie-v2 print mode) can include cache hit/miss metrics in session output. ## Problem The AI SDK provider already extracts cache read tokens from the backend (`cache_read_input_tokens` in TOKEN_USAGE nodes), and `normalizeUsage()` makes `cachedInputTokens` available on the `LanguageModelUsage` type. However, the harness `processStream()` step-finish handler only extracted `promptTokens` and `completionTokens`, completely ignoring cache fields. This meant cache metrics were stripped before reaching any consumer. ## Changes ### `types.ts` - Added `cachedInputTokens?: number` and `cacheCreationInputTokens?: number` to `TokenUsage` interface - Updated `defaultDisplayState()` to initialize cache fields to 0 ### `harness.ts` - Added `emptyTokenUsage()` helper for consistent zero-initialization with all 5 fields - Updated `processStream()` step-finish handler to extract `cachedInputTokens` from usage - Updated all 7 reset sites to use `emptyTokenUsage()` - Updated `loadThreadMetadata()` to restore cache fields from saved thread metadata - Updated `getTokenUsage()` return type from inline 3-field type to `TokenUsage` - Updated `usage_update` event emission and display state handlers to propagate cache fields ### `token-usage.test.ts` - Added test: extracts `cachedInputTokens` from usage - Added test: accumulates `cachedInputTokens` across multiple steps - Added test: defaults cache fields to 0 when not present <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Token usage now includes cached-input and cache-creation-input token metrics; these metrics are shown in the UI, included in persisted thread metadata, emitted in lifecycle and streaming updates, and returned by the public token-usage API. * **Tests** * Added tests ensuring cached token fields accumulate correctly across streaming steps and default to zero when absent. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Tyler Barnes <tylerdbarnes@gmail.com> Co-authored-by: Mastra Code (openai/gpt-5.5) <noreply@mastra.ai>
1 parent 7329046 commit 890b24c

5 files changed

Lines changed: 34 additions & 10 deletions

File tree

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
---
2+
'@mastra/core': patch
3+
---
4+
5+
Propagate cache metrics (`cachedInputTokens`, `cacheCreationInputTokens`) through harness token usage. The step-finish handler now extracts `cachedInputTokens` from AI SDK usage and propagates it through `usage_update` events, `getTokenUsage()`, display state, and thread metadata persistence.

packages/core/src/harness/display-state.test.ts

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ import { InMemoryStore } from '../storage/mock';
55
import { ChunkFrom } from '../stream/types';
66
import { Harness } from './harness';
77
import type { HarnessEvent } from './types';
8-
import { defaultDisplayState } from './types';
8+
import { createEmptyTokenUsage, defaultDisplayState } from './types';
99

1010
function createHarness(storage?: InMemoryStore) {
1111
const agent = new Agent({
@@ -31,7 +31,7 @@ describe('defaultDisplayState', () => {
3131
const ds = defaultDisplayState();
3232
expect(ds.isRunning).toBe(false);
3333
expect(ds.currentMessage).toBeNull();
34-
expect(ds.tokenUsage).toEqual({ promptTokens: 0, completionTokens: 0, totalTokens: 0 });
34+
expect(ds.tokenUsage).toEqual(createEmptyTokenUsage());
3535
expect(ds.activeTools).toBeInstanceOf(Map);
3636
expect(ds.activeTools.size).toBe(0);
3737
expect(ds.toolInputBuffers).toBeInstanceOf(Map);
@@ -71,7 +71,7 @@ describe('Harness.getDisplayState()', () => {
7171
const ds = harness.getDisplayState();
7272
expect(ds.isRunning).toBe(false);
7373
expect(ds.currentMessage).toBeNull();
74-
expect(ds.tokenUsage).toEqual({ promptTokens: 0, completionTokens: 0, totalTokens: 0 });
74+
expect(ds.tokenUsage).toEqual(createEmptyTokenUsage());
7575
expect(ds.activeTools.size).toBe(0);
7676
expect(ds.pendingApproval).toBeNull();
7777
expect(ds.pendingQuestion).toBeNull();
@@ -1321,7 +1321,7 @@ describe('resetThreadDisplayState', () => {
13211321
expect(harness.getDisplayState().tokenUsage.totalTokens).toBe(150);
13221322

13231323
emit(harness, { type: 'thread_created', thread: { id: 'new', title: 'New' } } as any);
1324-
expect(harness.getDisplayState().tokenUsage).toEqual({ promptTokens: 0, completionTokens: 0, totalTokens: 0 });
1324+
expect(harness.getDisplayState().tokenUsage).toEqual(createEmptyTokenUsage());
13251325
});
13261326

13271327
it('preserves isRunning across thread_created', () => {

packages/core/src/harness/harness.ts

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ import {
3333
taskWriteTool,
3434
} from './tools';
3535
import type { TaskItemSnapshot } from './tools';
36-
import { defaultDisplayState, defaultOMProgressState } from './types';
36+
import { createEmptyTokenUsage, defaultDisplayState, defaultOMProgressState } from './types';
3737
import type {
3838
AvailableModel,
3939
HeartbeatHandler,
@@ -57,10 +57,6 @@ import type {
5757
ToolCategory,
5858
} from './types';
5959

60-
function createEmptyTokenUsage(): TokenUsage {
61-
return { promptTokens: 0, completionTokens: 0, totalTokens: 0 };
62-
}
63-
6460
function getUsageNumber(usage: Record<string, unknown>, key: string): number | undefined {
6561
const value = usage[key];
6662
if (typeof value === 'number' && Number.isFinite(value)) {
@@ -1057,6 +1053,8 @@ export class Harness<TState = {}> {
10571053
promptTokens: savedUsage.promptTokens ?? 0,
10581054
completionTokens: savedUsage.completionTokens ?? 0,
10591055
totalTokens: savedUsage.totalTokens ?? 0,
1056+
cachedInputTokens: savedUsage.cachedInputTokens ?? 0,
1057+
cacheCreationInputTokens: savedUsage.cacheCreationInputTokens ?? 0,
10601058
};
10611059
} else {
10621060
this.tokenUsage = createEmptyTokenUsage();

packages/core/src/harness/token-usage.test.ts

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -248,4 +248,14 @@ describe('step-finish token usage extraction', () => {
248248
raw: { step: 2 },
249249
});
250250
});
251+
252+
it('defaults cache usage fields to 0 when not present in usage', async () => {
253+
const usage = { inputTokens: 100, outputTokens: 50, totalTokens: 150 };
254+
255+
await (harness as any).processStream({ fullStream: mockStream(usage) });
256+
257+
const tokenUsage = harness.getTokenUsage();
258+
expect(tokenUsage.cachedInputTokens).toBe(0);
259+
expect(tokenUsage.cacheCreationInputTokens).toBe(0);
260+
});
251261
});

packages/core/src/harness/types.ts

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -438,6 +438,17 @@ export interface TokenUsage {
438438
raw?: unknown;
439439
}
440440

441+
/** Creates a zero-initialized TokenUsage object. */
442+
export function createEmptyTokenUsage(): TokenUsage {
443+
return {
444+
promptTokens: 0,
445+
completionTokens: 0,
446+
totalTokens: 0,
447+
cachedInputTokens: 0,
448+
cacheCreationInputTokens: 0,
449+
};
450+
}
451+
441452
// =============================================================================
442453
// Observational Memory Progress
443454
// =============================================================================
@@ -646,7 +657,7 @@ export function defaultDisplayState(): HarnessDisplayState {
646657
return {
647658
isRunning: false,
648659
currentMessage: null,
649-
tokenUsage: { promptTokens: 0, completionTokens: 0, totalTokens: 0 },
660+
tokenUsage: createEmptyTokenUsage(),
650661
activeTools: new Map(),
651662
toolInputBuffers: new Map(),
652663
pendingApproval: null,

0 commit comments

Comments
 (0)