@outputai/llm 0.2.1-next.bd54540.0 → 0.2.1-next.e1a91cf.0

package/src/ai_sdk.spec.js

@@ -1,15 +1,24 @@
 import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
 
-const tracingSpies = {
-  addEventStart: vi.fn(),
-  addEventEnd: vi.fn(),
-  addEventError: vi.fn()
-};
-const emitEventSpy = vi.fn();
-vi.mock( '@outputai/core/sdk_activity_integration', () => ( {
-  Tracing: tracingSpies,
-  emitEvent: emitEventSpy
-} ), { virtual: true } );
+const traceMocks = vi.hoisted( () => ( {
+  startTrace: vi.fn( () => 'trace-id' ),
+  endTraceWithError: vi.fn()
+} ) );
+
+const wrapMocks = vi.hoisted( () => ( {
+  wrapTextResponse: vi.fn(),
+  wrapStreamOnFinishResponse: vi.fn()
+} ) );
+
+vi.mock( './utils/trace.js', () => ( {
+  startTrace: ( ...args ) => traceMocks.startTrace( ...args ),
+  endTraceWithError: ( ...args ) => traceMocks.endTraceWithError( ...args )
+} ) );
+
+vi.mock( './utils/response_wrappers.js', () => ( {
+  wrapTextResponse: ( ...args ) => wrapMocks.wrapTextResponse( ...args ),
+  wrapStreamOnFinishResponse: ( ...args ) => wrapMocks.wrapStreamOnFinishResponse( ...args )
+} ) );
 
 const loadModelImpl = vi.fn();
 const loadToolsImpl = vi.fn();
@@ -48,16 +57,6 @@ vi.mock( './skill.js', async importOriginal => {
   };
 } );
 
-const extractSourcesFromStepsImpl = vi.fn().mockReturnValue( [] );
-vi.mock( './source_extraction.js', () => ( {
-  extractSourcesFromSteps: ( ...args ) => extractSourcesFromStepsImpl( ...args )
-} ) );
-
-const calculateLLMCallCostImpl = vi.fn();
-vi.mock( './cost/index.js', () => ( {
-  calculateLLMCallCost: ( ...args ) => calculateLLMCallCostImpl( ...args )
-} ) );
-
 const importSut = async () => import( './ai_sdk.js' );
 
 const basePrompt = {
@@ -70,27 +69,45 @@ const basePrompt = {
   messages: [ { role: 'user', content: 'Hi' } ]
 };
 
-const cost = 'calculate cost';
+/** Mutable payload from `AI.generateText` — identity checks prove `generateText` returns `wrapTextResponse(...)` without substitution. */
+const generateTextAiFixture = {
+  response: {
+    text: 'TEXT',
+    sources: [],
+    usage: { inputTokens: 10, outputTokens: 5, totalTokens: 15 },
+    totalUsage: { inputTokens: 10, outputTokens: 5, totalTokens: 15 },
+    finishReason: 'stop'
+  }
+};
 
 beforeEach( () => {
-  emitEventSpy.mockReset();
   loadModelImpl.mockReset().mockReturnValue( 'MODEL' );
   loadPromptImpl.mockReset().mockReturnValue( { ...basePrompt, messages: [ ...basePrompt.messages ] } );
-  extractSourcesFromStepsImpl.mockReset().mockReturnValue( [] );
-  calculateLLMCallCostImpl.mockReset().mockResolvedValue( cost );
   aiFns.tool.mockReset().mockImplementation( def => def );
   aiFns.stepCountIs.mockReset().mockImplementation( n => ( { type: 'stepCount', count: n } ) );
   loadPromptSkillsImpl.mockReset().mockReturnValue( [] );
   loadColocatedSkillsImpl.mockReset().mockReturnValue( [] );
 
+  traceMocks.startTrace.mockReset().mockReturnValue( 'trace-id' );
+  traceMocks.endTraceWithError.mockReset();
+
+  wrapMocks.wrapTextResponse.mockReset().mockImplementation( async ( { response } ) => response );
+
+  wrapMocks.wrapStreamOnFinishResponse.mockReset().mockImplementation( ( { onFinish } ) => ( {
+    async onFinish( response ) {
+      onFinish?.( response );
+    }
+  } ) );
+
   const defaultUsage = { inputTokens: 10, outputTokens: 5, totalTokens: 15 };
-  aiFns.generateText.mockReset().mockResolvedValue( {
+  generateTextAiFixture.response = {
     text: 'TEXT',
     sources: [],
     usage: defaultUsage,
     totalUsage: defaultUsage,
     finishReason: 'stop'
-  } );
+  };
+  aiFns.generateText.mockReset().mockResolvedValue( generateTextAiFixture.response );
 
   aiFns.streamText.mockReset().mockReturnValue( {
     textStream: 'MOCK_TEXT_STREAM',
@@ -108,27 +125,23 @@ afterEach( async () => {
 } );
 
 describe( 'ai_sdk', () => {
-  it( 'generateText: validates, traces, calls AI and returns text', async () => {
+  it( 'generateText: validates, delegates trace/wrap utils, calls AI and returns wrapTextResponse output', async () => {
     const { generateText } = await importSut();
+    const defaultUsage = { inputTokens: 10, outputTokens: 5, totalTokens: 15 };
     const result = await generateText( { prompt: 'test_prompt@v1' } );
 
     expect( validators.validateGenerateTextArgs ).toHaveBeenCalledWith( { prompt: 'test_prompt@v1' } );
     expect( loadPromptImpl ).toHaveBeenCalledWith( 'test_prompt@v1', undefined, undefined );
-    expect( tracingSpies.addEventStart ).toHaveBeenCalledTimes( 1 );
-    expect( tracingSpies.addEventEnd ).toHaveBeenCalledTimes( 1 );
-    expect( tracingSpies.addEventEnd ).toHaveBeenCalledWith(
-      expect.objectContaining( { details: expect.objectContaining( { cost } ) } )
-    );
-    expect( calculateLLMCallCostImpl ).toHaveBeenCalledWith( {
-      usage: { inputTokens: 10, outputTokens: 5, totalTokens: 15 },
-      modelId: basePrompt.config.model
-    } );
-    const defaultUsage = { inputTokens: 10, outputTokens: 5, totalTokens: 15 };
-    expect( emitEventSpy ).toHaveBeenCalledTimes( 1 );
-    expect( emitEventSpy ).toHaveBeenCalledWith( 'llm:call_cost', {
+    expect( traceMocks.startTrace ).toHaveBeenCalledTimes( 1 );
+    expect( traceMocks.startTrace ).toHaveBeenCalledWith( expect.objectContaining( {
+      name: 'generateText',
+      prompt: 'test_prompt@v1'
+    } ) );
+    expect( wrapMocks.wrapTextResponse ).toHaveBeenCalledTimes( 1 );
+    expect( wrapMocks.wrapTextResponse ).toHaveBeenCalledWith( {
+      traceId: 'trace-id',
       modelId: basePrompt.config.model,
-      cost,
-      usage: defaultUsage
+      response: generateTextAiFixture.response
     } );
 
     expect( loadModelImpl ).toHaveBeenCalledWith( basePrompt );
@@ -138,10 +151,8 @@ describe( 'ai_sdk', () => {
       temperature: 0.3,
       providerOptions: basePrompt.config.providerOptions
     } );
-    expect( result.text ).toBe( 'TEXT' );
-    expect( result.sources ).toEqual( [] );
-    expect( result.usage ).toEqual( { inputTokens: 10, outputTokens: 5, totalTokens: 15 } );
-    expect( result.finishReason ).toBe( 'stop' );
+    expect( result ).toBe( generateTextAiFixture.response );
+    expect( result.totalUsage ).toEqual( defaultUsage );
   } );
 
   it( 'generateText: passes provider-specific options to AI SDK', async () => {
@@ -221,26 +232,16 @@ describe( 'ai_sdk', () => {
     expect( result.response ).toEqual( { id: 'req_123', modelId: 'gpt-4o-2024-05-13' } );
   } );
 
-  it( 'generateText: includes unified result field that matches text', async () => {
-    const { generateText } = await importSut();
-    const response = await generateText( { prompt: 'test_prompt@v1' } );
-
-    expect( response.result ).toBe( 'TEXT' );
-    expect( response.result ).toBe( response.text );
-  } );
-
   it( 'generateText: traces error and rethrows when AI SDK fails', async () => {
     const error = new Error( 'API rate limit exceeded' );
     aiFns.generateText.mockRejectedValueOnce( error );
     const { generateText } = await importSut();
 
     await expect( generateText( { prompt: 'test_prompt@v1' } ) ).rejects.toThrow( 'API rate limit exceeded' );
-    expect( tracingSpies.addEventError ).toHaveBeenCalledWith(
-      expect.objectContaining( { details: error } )
-    );
+    expect( traceMocks.endTraceWithError ).toHaveBeenCalledWith( { traceId: 'trace-id', error } );
   } );
 
-  it( 'generateText: Proxy correctly handles AI SDK response with getter', async () => {
+  it( 'generateText: passes the AI response object through to wrapTextResponse and returns it', async () => {
     const responseWithGetter = {
       _internalText: 'TEXT_FROM_GETTER',
       get text() {
@@ -254,10 +255,15 @@ describe( 'ai_sdk', () => {
     aiFns.generateText.mockResolvedValueOnce( responseWithGetter );
 
     const { generateText } = await importSut();
-    const response = await generateText( { prompt: 'test_prompt@v1' } );
+    const out = await generateText( { prompt: 'test_prompt@v1' } );
 
-    expect( response.text ).toBe( 'TEXT_FROM_GETTER' );
-    expect( response.result ).toBe( 'TEXT_FROM_GETTER' );
+    expect( wrapMocks.wrapTextResponse ).toHaveBeenCalledWith( {
+      traceId: 'trace-id',
+      modelId: basePrompt.config.model,
+      response: responseWithGetter
+    } );
+    expect( out ).toBe( responseWithGetter );
+    expect( out.text ).toBe( 'TEXT_FROM_GETTER' );
   } );
 
   it( 'generateText: passes through AI SDK options like tools and maxRetries', async () => {
@@ -318,13 +324,11 @@ describe( 'ai_sdk', () => {
     );
   } );
 
-  it( 'generateText: .object returns undefined instead of leaking text', async () => {
+  it( 'generateText: does not add structured output fields the AI response lacks', async () => {
     const { generateText } = await importSut();
     const result = await generateText( { prompt: 'test_prompt@v1' } );
 
     expect( result.object ).toBeUndefined();
-    expect( result.text ).toBe( 'TEXT' );
-    expect( result.result ).toBe( 'TEXT' );
   } );
 
   it( 'generateText: passes through unknown future options for forward compatibility', async () => {
@@ -344,13 +348,13 @@ describe( 'ai_sdk', () => {
     );
   } );
 
-  it( 'streamText: validates, traces, calls AI streamText and returns stream result', async () => {
+  it( 'streamText: validates, delegates trace/wrap utils, calls AI streamText and returns stream result', async () => {
     const { streamText } = await importSut();
     const result = streamText( { prompt: 'test_prompt@v1' } );
 
     expect( validators.validateStreamTextArgs ).toHaveBeenCalledWith( { prompt: 'test_prompt@v1' } );
     expect( loadPromptImpl ).toHaveBeenCalledWith( 'test_prompt@v1', undefined );
-    expect( tracingSpies.addEventStart ).toHaveBeenCalledTimes( 1 );
+    expect( traceMocks.startTrace ).toHaveBeenCalledTimes( 1 );
 
     expect( loadModelImpl ).toHaveBeenCalledWith( basePrompt );
     expect( aiFns.streamText ).toHaveBeenCalledWith(
@@ -367,12 +371,18 @@ describe( 'ai_sdk', () => {
     expect( result.fullStream ).toBe( 'MOCK_FULL_STREAM' );
   } );
 
-  it( 'streamText: onFinish callback traces end event and calls user callback', async () => {
+  it( 'streamText: forwards stream onFinish through wrapStreamOnFinishResponse to the user callback', async () => {
     const { streamText } = await importSut();
     const userOnFinish = vi.fn();
 
     streamText( { prompt: 'test_prompt@v1', onFinish: userOnFinish } );
 
+    expect( wrapMocks.wrapStreamOnFinishResponse ).toHaveBeenCalledWith( expect.objectContaining( {
+      traceId: 'trace-id',
+      modelId: basePrompt.config.model,
+      onFinish: userOnFinish
+    } ) );
+
     const callArgs = aiFns.streamText.mock.calls[0][0];
     const usage = { inputTokens: 10, outputTokens: 5, totalTokens: 15 };
     const finishEvent = {
@@ -384,22 +394,7 @@ describe( 'ai_sdk', () => {
     };
     await callArgs.onFinish( finishEvent );
 
-    expect( emitEventSpy ).toHaveBeenCalledTimes( 1 );
-    expect( emitEventSpy ).toHaveBeenCalledWith( 'llm:call_cost', {
-      modelId: basePrompt.config.model,
-      cost,
-      usage
-    } );
-    expect( tracingSpies.addEventEnd ).toHaveBeenCalledWith(
-      expect.objectContaining( {
-        details: {
-          result: 'STREAMED_TEXT',
-          usage,
-          cost,
-          providerMetadata: finishEvent.providerMetadata
-        }
-      } )
-    );
+    expect( userOnFinish ).toHaveBeenCalledTimes( 1 );
     expect( userOnFinish ).toHaveBeenCalledWith( finishEvent );
   } );
 
@@ -413,9 +408,7 @@ describe( 'ai_sdk', () => {
     const error = new Error( 'Stream failed' );
     callArgs.onError( { error } );
 
-    expect( tracingSpies.addEventError ).toHaveBeenCalledWith(
-      expect.objectContaining( { details: error } )
-    );
+    expect( traceMocks.endTraceWithError ).toHaveBeenCalledWith( { traceId: 'trace-id', error } );
     expect( userOnError ).toHaveBeenCalledWith( { error } );
   } );
 
@@ -433,11 +426,6 @@ describe( 'ai_sdk', () => {
       finishReason: 'stop'
     };
     await expect( callArgs.onFinish( finishEvent ) ).resolves.toBeUndefined();
-    expect( emitEventSpy ).toHaveBeenCalledWith( 'llm:call_cost', {
-      modelId: basePrompt.config.model,
-      cost,
-      usage
-    } );
     expect( () => callArgs.onError( { error: new Error( 'fail' ) } ) ).not.toThrow();
   } );
 
@@ -499,15 +487,11 @@ describe( 'ai_sdk', () => {
 
     streamText( { prompt: 'test_prompt@v1', variables: vars } );
 
-    expect( tracingSpies.addEventStart ).toHaveBeenCalledWith( {
-      kind: 'llm',
+    expect( traceMocks.startTrace ).toHaveBeenCalledWith( {
       name: 'streamText',
-      id: expect.stringContaining( 'streamText-' ),
-      details: {
-        prompt: 'test_prompt@v1',
-        variables: vars,
-        loadedPrompt: basePrompt
-      }
+      prompt: 'test_prompt@v1',
+      variables: vars,
+      loadedPrompt: basePrompt
     } );
   } );
 
@@ -519,9 +503,7 @@ describe( 'ai_sdk', () => {
     const { streamText } = await importSut();
 
     expect( () => streamText( { prompt: 'test_prompt@v1' } ) ).toThrow( syncError );
-    expect( tracingSpies.addEventError ).toHaveBeenCalledWith(
-      expect.objectContaining( { details: syncError } )
-    );
+    expect( traceMocks.endTraceWithError ).toHaveBeenCalledWith( { traceId: 'trace-id', error: syncError } );
   } );
 
   it( 'streamText: passes variables to prompt loader', async () => {
@@ -533,95 +515,6 @@ describe( 'ai_sdk', () => {
     expect( loadPromptImpl ).toHaveBeenCalledWith( 'test_prompt@v1', vars );
   } );
 
-  it( 'generateText: merges tool-extracted sources into response.sources', async () => {
-    const extracted = [
-      { type: 'source', sourceType: 'url', id: 'abc123', url: 'https://tool.com/1', title: 'Tool 1' },
-      { type: 'source', sourceType: 'url', id: 'def456', url: 'https://tool.com/2', title: 'Tool 2' }
-    ];
-    extractSourcesFromStepsImpl.mockReturnValue( extracted );
-
-    const usageTools = { inputTokens: 10, outputTokens: 5, totalTokens: 15 };
-    aiFns.generateText.mockResolvedValueOnce( {
-      text: 'answer',
-      sources: [],
-      steps: [ { toolResults: [] } ],
-      usage: usageTools,
-      totalUsage: usageTools,
-      finishReason: 'stop'
-    } );
-
-    const { generateText } = await importSut();
-    const result = await generateText( { prompt: 'test_prompt@v1' } );
-
-    expect( result.sources ).toEqual( extracted );
-  } );
-
-  it( 'generateText: deduplicates extracted sources against native sources', async () => {
-    const nativeSources = [
-      { type: 'source', sourceType: 'url', id: 'native1', url: 'https://shared.com', title: 'Native' }
-    ];
-    const extracted = [
-      { type: 'source', sourceType: 'url', id: 'ext1', url: 'https://shared.com', title: 'Extracted' },
-      { type: 'source', sourceType: 'url', id: 'ext2', url: 'https://unique.com', title: 'Unique' }
-    ];
-    extractSourcesFromStepsImpl.mockReturnValue( extracted );
-
-    const usageDedup = { inputTokens: 10, outputTokens: 5, totalTokens: 15 };
-    aiFns.generateText.mockResolvedValueOnce( {
-      text: 'answer',
-      sources: nativeSources,
-      steps: [ { toolResults: [] } ],
-      usage: usageDedup,
-      totalUsage: usageDedup,
-      finishReason: 'stop'
-    } );
-
-    const { generateText } = await importSut();
-    const result = await generateText( { prompt: 'test_prompt@v1' } );
-
-    expect( result.sources ).toHaveLength( 2 );
-    expect( result.sources[0].url ).toBe( 'https://shared.com' );
-    expect( result.sources[0].title ).toBe( 'Native' );
-    expect( result.sources[1].url ).toBe( 'https://unique.com' );
-  } );
-
-  it( 'generateText: returns native sources unchanged when no tool sources extracted', async () => {
-    const nativeSources = [
-      { type: 'source', sourceType: 'url', id: 'n1', url: 'https://native.com', title: 'Native' }
-    ];
-    extractSourcesFromStepsImpl.mockReturnValue( [] );
-
-    const usageNative = { inputTokens: 10, outputTokens: 5, totalTokens: 15 };
-    aiFns.generateText.mockResolvedValueOnce( {
-      text: 'answer',
-      sources: nativeSources,
-      usage: usageNative,
-      totalUsage: usageNative,
-      finishReason: 'stop'
-    } );
-
-    const { generateText } = await importSut();
-    const result = await generateText( { prompt: 'test_prompt@v1' } );
-
-    expect( result.sources ).toEqual( nativeSources );
-  } );
-
-  it( 'generateText: includes costs from cost module in trace details', async () => {
-    const customCost = { total: 0.02, components: { input: { value: 0.01 }, output: { value: 0.01 } } };
-    calculateLLMCallCostImpl.mockResolvedValueOnce( customCost );
-
-    const { generateText } = await importSut();
-    await generateText( { prompt: 'test_prompt@v1' } );
-
-    expect( calculateLLMCallCostImpl ).toHaveBeenCalledWith( {
-      usage: { inputTokens: 10, outputTokens: 5, totalTokens: 15 },
-      modelId: basePrompt.config.model
-    } );
-    expect( tracingSpies.addEventEnd ).toHaveBeenCalledWith(
-      expect.objectContaining( { details: expect.objectContaining( { cost: customCost } ) } )
-    );
-  } );
-
   it( 'generateText: loads frontmatter skills from prompt config using promptFileDir', async () => {
     const frontmatterSkill = { name: 'fm_skill', description: 'FM', instructions: '# FM' };
     loadPromptImpl.mockReturnValue( {
@@ -805,30 +698,4 @@ describe( 'ai_sdk', () => {
       expect.objectContaining( { stopWhen: customStop } )
     );
   } );
-
-  it( 'generateText: includes sourcesFromTools in trace details', async () => {
-    const extracted = [
-      { type: 'source', sourceType: 'url', id: 'abc', url: 'https://t.com', title: 'T' }
-    ];
-    extractSourcesFromStepsImpl.mockReturnValue( extracted );
-
-    const usageSources = { inputTokens: 10, outputTokens: 5, totalTokens: 15 };
-    aiFns.generateText.mockResolvedValueOnce( {
-      text: 'TEXT',
-      sources: [],
-      steps: [],
-      usage: usageSources,
-      totalUsage: usageSources,
-      finishReason: 'stop'
-    } );
-
-    const { generateText } = await importSut();
-    await generateText( { prompt: 'test_prompt@v1' } );
-
-    expect( tracingSpies.addEventEnd ).toHaveBeenCalledWith(
-      expect.objectContaining( {
-        details: expect.objectContaining( { sourcesFromTools: extracted } )
-      } )
-    );
-  } );
 } );

package/src/cost/index.js

@@ -4,41 +4,10 @@ import Decimal from 'decimal.js';
 const M = 1_000_000;
 const calcCost = ( tokens, ppm ) => Decimal( tokens ?? 0 ).div( M ).mul( ppm ).toNumber();
 
-/**
- * Calculates the input cost based on the input value
- */
-const calculateInput = ( { tokens, cost } ) =>
-  !Number.isFinite( cost.input ) ? { value: null, message: 'Missing input cost' } : { value: calcCost( tokens, cost.input ) };
-
-/**
- * Calculates the input cost based on the cache_read
- */
-const calculateCachedInput = ( { tokens, cost } ) =>
-  !Number.isFinite( cost.cache_read ) ? { value: null, message: 'Missing cache input cost' } : { value: calcCost( tokens, cost.cache_read ) };
-
-/**
- * Calculates the output cost based on the output value
- */
-const calculateOutput = ( { tokens, cost } ) =>
-  !Number.isFinite( cost.output ) ? { value: null, message: 'Missing output' } : { value: calcCost( tokens, cost.output ) };
-
-/**
- * Calculates the reasoning cost based on the reasoning token's
- * If there isn't reasoning costs, this means this providers doesn't differentiate reasoning vs output,
- * so don't calculate it as the price is included in output
- */
-const calculateReasoning = ( { tokens, cost } ) =>
-  Number.isFinite( cost.reasoning ) ? { value: calcCost( tokens, cost.reasoning ) } : undefined;
-
-/**
- * Calculates the total cost based on the components
- */
-const calculateTotal = components => Object.values( components ).reduce( ( v, e ) => v.plus( e?.value ? e.value : 0 ), Decimal( 0 ) ).toNumber();
-
 /**
  * Calculates the cost of an llm call based on the model and usage.
  * @param {object} args
- * @param {string} args.modelId - Name of the mode, provider prefix is optional
+ * @param {string} args.modelId - Name of the model, provider prefix is optional
  * @param {object} args.usage - Usage, as returned from AI SDK
  * @returns {object} The cost with total value and components
  */
@@ -58,14 +27,14 @@ export const calculateLLMCallCost = async ( { modelId, usage } ) => {
 
     const nonCachedTokens = inputTokens - ( cachedInputTokens ?? 0 );
 
-    const components = {
-      input: calculateInput( { tokens: nonCachedTokens, cost } ),
-      cachedInput: calculateCachedInput( { tokens: cachedInputTokens, cost } ),
-      output: calculateOutput( { tokens: outputTokens, cost } ),
-      reasoning: calculateReasoning( { tokens: reasoningTokens ?? 0, cost } )
-    };
-
-    return { total: calculateTotal( components ), components };
+    const components = [
+      Number.isFinite( cost.input ) ? { name: 'input_tokens', value: calcCost( nonCachedTokens, cost.input ) } : false,
+      Number.isFinite( cost.cache_read ) ? { name: 'input_cached_tokens', value: calcCost( cachedInputTokens, cost.cache_read ) } : false,
+      Number.isFinite( cost.output ) ? { name: 'output_tokens', value: calcCost( outputTokens, cost.output ) } : false,
+      /* When there aren't reasoning costs, the providers doesn't differentiate reasoning vs output, so the price is included in the output */
+      Number.isFinite( cost.reasoning ) ? { name: 'reasoning_tokens', value: calcCost( reasoningTokens, cost.reasoning ) } : false
+    ].filter( v => !!v );
+    return { total: components.reduce( ( v, e ) => v.plus( e.value ), Decimal( 0 ) ).toNumber(), components };
   } catch ( error ) {
     console.error( 'Error calculating LLM call costs', error );
     return { total: null, message: `Error calculating LLM call costs: ${error.constructor.name} - ${error.message}` };

package/src/cost/index.spec.js

@@ -47,10 +47,11 @@ describe( 'calculateLLMCallCost', () => {
     } );
 
     expect( result.total ).toBe( 7 );
-    expect( result.components.input ).toEqual( { value: 2 } );
-    expect( result.components.cachedInput ).toEqual( { value: 0 } );
-    expect( result.components.output ).toEqual( { value: 5 } );
-    expect( result.components.reasoning ).toBeUndefined();
+    expect( result.components ).toEqual( [
+      { name: 'input_tokens', value: 2 },
+      { name: 'input_cached_tokens', value: 0 },
+      { name: 'output_tokens', value: 5 }
+    ] );
   } );
 
   it( 'splits input into non-cached and cached at respective rates', async () => {
@@ -62,13 +63,15 @@ describe( 'calculateLLMCallCost', () => {
       usage: { inputTokens: 1_000_000, cachedInputTokens: 500_000, outputTokens: 100_000 }
     } );
 
-    expect( result.components.input ).toEqual( { value: 2 } );
-    expect( result.components.cachedInput ).toEqual( { value: 0.5 } );
-    expect( result.components.output ).toEqual( { value: 1 } );
+    expect( result.components ).toEqual( [
+      { name: 'input_tokens', value: 2 },
+      { name: 'input_cached_tokens', value: 0.5 },
+      { name: 'output_tokens', value: 1 }
+    ] );
     expect( result.total ).toBeCloseTo( 3.5 );
   } );
 
-  it( 'sets cachedInput to null when model has no cache_read', async () => {
+  it( 'omits cached component when model has no cache_read (non-cached rate applies to full input minus cached)', async () => {
     mockFetchModelsPricing.mockResolvedValue( new Map( [ [ 'no-cache', { input: 2, output: 10 } ] ] ) );
 
     const result = await calculateLLMCallCost( {
@@ -76,12 +79,14 @@ describe( 'calculateLLMCallCost', () => {
       usage: { inputTokens: 1_000_000, cachedInputTokens: 200_000, outputTokens: 0 }
     } );
 
-    expect( result.components.input ).toEqual( { value: 1.6 } );
-    expect( result.components.cachedInput ).toEqual( { value: null, message: 'Missing cache input cost' } );
+    expect( result.components ).toEqual( [
+      { name: 'input_tokens', value: 1.6 },
+      { name: 'output_tokens', value: 0 }
+    ] );
     expect( result.total ).toBe( 1.6 );
   } );
 
-  it( 'sets input to null and message when pricing has no input', async () => {
+  it( 'omits input component when pricing has no input rate', async () => {
     mockFetchModelsPricing.mockResolvedValue( new Map( [ [ 'out-only', { output: 10 } ] ] ) );
 
     const result = await calculateLLMCallCost( {
@@ -90,11 +95,12 @@ describe( 'calculateLLMCallCost', () => {
     } );
 
     expect( result.total ).toBe( 0.0005 );
-    expect( result.components.input ).toEqual( { value: null, message: 'Missing input cost' } );
-    expect( result.components.output ).toEqual( { value: 0.0005 } );
+    expect( result.components ).toEqual( [
+      { name: 'output_tokens', value: 0.0005 }
+    ] );
   } );
 
-  it( 'sets output to null and message when pricing has no output', async () => {
+  it( 'omits output component when pricing has no output rate', async () => {
     mockFetchModelsPricing.mockResolvedValue( new Map( [ [ 'in-only', { input: 1 } ] ] ) );
 
     const result = await calculateLLMCallCost( {
@@ -103,8 +109,9 @@ describe( 'calculateLLMCallCost', () => {
     } );
 
     expect( result.total ).toBe( 0.0001 );
-    expect( result.components.input ).toEqual( { value: 0.0001 } );
-    expect( result.components.output ).toEqual( { value: null, message: 'Missing output' } );
+    expect( result.components ).toEqual( [
+      { name: 'input_tokens', value: 0.0001 }
+    ] );
   } );
 
   it( 'uses reasoning cost when present', async () => {
@@ -119,7 +126,11 @@ describe( 'calculateLLMCallCost', () => {
     } );
 
     expect( result.total ).toBeCloseTo( 0.0033 );
-    expect( result.components.reasoning ).toEqual( { value: 0.003 } );
+    expect( result.components ).toEqual( [
+      { name: 'input_tokens', value: 0.0001 },
+      { name: 'output_tokens', value: 0.0002 },
+      { name: 'reasoning_tokens', value: 0.003 }
+    ] );
   } );
 
   it( 'omits reasoning component when reasoning cost missing (included in output)', async () => {
@@ -131,10 +142,13 @@ describe( 'calculateLLMCallCost', () => {
     } );
 
     expect( result.total ).toBeCloseTo( 0.0003 );
-    expect( result.components.reasoning ).toBeUndefined();
+    expect( result.components ).toEqual( [
+      { name: 'input_tokens', value: 0.0001 },
+      { name: 'output_tokens', value: 0.0002 }
+    ] );
   } );
 
-  it( 'Calculate reasoning component when reasoningTokens is zero', async () => {
+  it( 'includes reasoning component with zero when reasoningTokens is zero', async () => {
     mockFetchModelsPricing.mockResolvedValue( new Map( [ [
       'full',
       { input: 2, output: 8, reasoning: 60 }
@@ -145,7 +159,11 @@ describe( 'calculateLLMCallCost', () => {
       usage: { inputTokens: 100, outputTokens: 50, reasoningTokens: 0 }
     } );
 
-    expect( result.components.reasoning ).toEqual( { value: 0 } );
+    expect( result.components ).toEqual( [
+      { name: 'input_tokens', value: 0.0002 },
+      { name: 'output_tokens', value: 0.0004 },
+      { name: 'reasoning_tokens', value: 0 }
+    ] );
     expect( result.total ).toBeCloseTo( 0.0006 );
   } );
 
@@ -158,5 +176,9 @@ describe( 'calculateLLMCallCost', () => {
     } );
 
     expect( result.total ).toBe( 0 );
+    expect( result.components ).toEqual( [
+      { name: 'input_tokens', value: 0 },
+      { name: 'output_tokens', value: 0 }
+    ] );
   } );
 } );