@outputai/llm 0.13.0 → 0.13.1-next.318b2e6.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@outputai/llm",
-  "version": "0.13.0",
+  "version": "0.13.1-next.318b2e6.0",
   "description": "Framework abstraction to interact with LLM models",
   "type": "module",
   "main": "src/index.js",
@@ -22,7 +22,7 @@
     "gray-matter": "4.0.3",
     "liquidjs": "10.27.2",
     "undici": "8.9.0",
-    "@outputai/core": "0.13.0"
+    "@outputai/core": "0.13.1-next.318b2e6.0"
   },
   "devDependencies": {
     "@ai-sdk/amazon-bedrock": "5.0.57",

package/src/agent.js

@@ -1,6 +1,6 @@
 import { ToolLoopAgent as AIToolLoopAgent } from 'ai';
 import { loadAiSdkTextOptions } from './ai_sdk_options.js';
-import { wrapGeneration, wrapStream } from './utils/wrap.js';
+import { wrapTextGeneration, wrapStream } from './utils/wrap.js';
 import { Role } from './consts.js';
 import { drainStream } from './utils/stream.js';
 import { loadPrompt } from './prompt/loader.js';
@@ -48,15 +48,16 @@ export class Agent {
     const { messages, abortSignal, toolChoice } = Validator.parseAgentGenerateArgs( args );
     const combinedMessages = await this.#combineWithPreviousMessages( messages );
 
-    return wrapGeneration( {
+    return wrapTextGeneration( {
       name: 'Agent.generate',
       prompt: this.#prompt,
-      fn: async () => {
+      fn: async ( { onStepEndHook } ) => {
         const response = await this.#agent.generate( {
           messages: combinedMessages,
           allowSystemInMessages: true,
           ...( abortSignal && { abortSignal } ),
-          ...( toolChoice && { toolChoice } )
+          ...( toolChoice && { toolChoice } ),
+          onStepEnd: onStepEndHook
         } );
         if ( response.finishReason !== 'error' ) {
           await this.#storeMessages( messages.concat( response.responseMessages ?? [] ) );
@@ -73,10 +74,10 @@ export class Agent {
     const { messages, abortSignal, toolChoice, onChunk } = Validator.parseAgentGenerateWithStreamingArgs( args );
     const combinedMessages = await this.#combineWithPreviousMessages( messages );
 
-    return wrapGeneration( {
+    return wrapTextGeneration( {
       name: 'Agent.generateWithStreaming',
       prompt: this.#prompt,
-      fn: async () => {
+      fn: async ( { onStepEndHook } ) => {
         const state = { response: null };
         const stream = await this.#agent.stream( {
           messages: combinedMessages,
@@ -84,6 +85,7 @@ export class Agent {
           ...( onChunk && { onChunk } ),
           ...( abortSignal && { abortSignal } ),
           ...( toolChoice && { toolChoice } ),
+          onStepEnd: onStepEndHook,
           onEnd: res => {
             state.response = res;
           },
@@ -93,7 +95,7 @@ export class Agent {
         await drainStream( stream, abortSignal );
 
         if ( !state.response ) {
-          throw new Error( 'Agent streaming generation completed without a response.' );
+          throw new Error( 'Streaming completed without a response.' );
         }
 
         state.response.output = await stream.output;
@@ -114,20 +116,21 @@ export class Agent {
       name: 'Agent.stream',
       prompt: this.#prompt,
       abortSignal,
-      fn: ( { onEndHook, onErrorHook } ) => this.#agent.stream( {
+      fn: ( { onEndHook, onErrorHook, onStepEndHook, onAbortHook } ) => this.#agent.stream( {
         messages: combinedMessages,
         allowSystemInMessages: true,
         ...( onChunk && { onChunk } ),
         ...( abortSignal && { abortSignal } ),
         ...( toolChoice && { toolChoice } ),
+        onStepEnd: onStepEndHook,
+        onAbort: onAbortHook,
         onEnd: response =>
           onEndHook( response, async parsedResponse => {
             if ( response.finishReason !== 'error' ) {
               await this.#storeMessages( messages.concat( response.responseMessages ?? [] ) )
-                .catch( error => Logger.error( 'Agent.stream message store persistence failed', {
-                  namespace: 'LLM',
-                  error: error instanceof Error ? error.message : String( error )
-                } ) );
+                .catch( error =>
+                  Logger.error( 'Message store persistence failed', { namespace: 'LLM', error: error?.message || String( error ) } )
+                );
             }
             await onEnd?.( parsedResponse );
           } ),

package/src/generate.js

@@ -1,6 +1,6 @@
 import * as AI from 'ai';
 import { loadPrompt } from './prompt/loader.js';
-import { wrapGeneration, wrapStream } from './utils/wrap.js';
+import { wrapImageGeneration, wrapStream, wrapTextGeneration } from './utils/wrap.js';
 import { loadAiSdkTextOptions, loadAiSdkImageOptions } from './ai_sdk_options.js';
 import { drainStream } from './utils/stream.js';
 import { loadSkills } from './utils/skills.js';
@@ -11,10 +11,13 @@ export const generateText = async args => {
   const prompt = promptObject ?? loadPrompt( promptFile, variables, promptDir );
   const skills = loadSkills( prompt );
 
-  return wrapGeneration( {
+  return wrapTextGeneration( {
     name: 'generateText',
     prompt,
-    fn: () => AI.generateText( loadAiSdkTextOptions( { prompt, skills, ...aiOptions } ) )
+    fn: ( { onStepEndHook } ) => AI.generateText( {
+      ...loadAiSdkTextOptions( { prompt, skills, ...aiOptions } ),
+      onStepEnd: onStepEndHook
+    } )
   } );
 };
 
@@ -27,9 +30,11 @@ export const streamText = args => {
     name: 'streamText',
     prompt,
     abortSignal: aiOptions.abortSignal,
-    fn: ( { onEndHook, onErrorHook } ) => AI.streamText( {
+    fn: ( { onEndHook, onErrorHook, onStepEndHook, onAbortHook } ) => AI.streamText( {
       ...loadAiSdkTextOptions( { prompt, skills, ...aiOptions } ),
       ...( onChunk && { onChunk } ),
+      onStepEnd: onStepEndHook,
+      onAbort: onAbortHook,
       onEnd: response => onEndHook( response, onEnd ),
       onError: event => onErrorHook( event, error => onError?.( { ...event, error } ) )
     } )
@@ -44,14 +49,15 @@ export const generateTextWithStreaming = async args => {
   const prompt = promptObject ?? loadPrompt( promptFile, variables, promptDir );
   const skills = loadSkills( prompt );
 
-  return wrapGeneration( {
+  return wrapTextGeneration( {
     name: 'generateTextWithStreaming',
     prompt,
-    fn: async () => {
+    fn: async ( { onStepEndHook } ) => {
       const state = { response: null };
       const stream = AI.streamText( {
         ...loadAiSdkTextOptions( { prompt, skills, ...aiOptions } ),
         ...( onChunk && { onChunk } ),
+        onStepEnd: onStepEndHook,
         onEnd: res => {
           state.response = res;
         },
@@ -61,7 +67,7 @@ export const generateTextWithStreaming = async args => {
       await drainStream( stream, aiOptions.abortSignal );
 
       if ( !state.response ) {
-        throw new Error( 'Streaming generation completed without a response.' );
+        throw new Error( 'Streaming completed without a response.' );
       }
 
       state.response.output = await stream.output;
@@ -74,7 +80,7 @@ export const generateImage = async args => {
   const { promptFile, promptObject, promptDir, variables, ...aiOptions } = Validator.parseGenerateImageArgs( args );
   const prompt = promptObject ?? loadPrompt( promptFile, variables, promptDir );
 
-  return wrapGeneration( {
+  return wrapImageGeneration( {
     name: 'generateImage',
     prompt,
     fn: () => AI.generateImage( loadAiSdkImageOptions( { prompt, ...aiOptions } ) )

package/src/utils/error_handler.js

@@ -25,6 +25,10 @@ const nonRetryableAiSdkErrorTypes = [
  * @returns {object} A new Error
  */
 export const mapAiError = error => {
+  if ( !( error instanceof Error ) ) {
+    return error;
+  }
+
   if ( error instanceof FatalError ) {
     return error;
   }

package/src/utils/image.js

@@ -1,10 +1,37 @@
+import { Logger } from '@outputai/core';
+
 /**
  * Get the approximate file size from a base64 string.
  * @param {string} b64data
  * @returns {number} Size in bytes
  */
 export const calculateBase64FileSize = b64data => {
+  if ( typeof b64data !== 'string' ) {
+    return null;
+  }
   const baseSize = b64data.length * ( 3 / 4 );
   const paddingSize = [ b64data.at( -2 ), b64data.at( -1 ) ].filter( v => v === '=' ).length;
   return baseSize - paddingSize;
 };
+
+/**
+ * Return a serialized version of the images from the AI SDK response.
+ * It contains only the file size and the media-type
+ *
+ * @param {object} response
+ * @returns {Array<{size: number, mediaType: string}>}
+ */
+export const serializeImagesFromResponse = response => {
+  try {
+    if ( !Array.isArray( response?.images ) ) {
+      return [];
+    }
+    return response.images
+      .filter( image => image !== null && typeof image === 'object' )
+      .map( ( { base64, mediaType } ) => ( { size: calculateBase64FileSize( base64 ), mediaType } ) );
+
+  } catch ( error ) {
+    Logger.error( 'Image serialization failed', { namespace: 'LLM', error: error?.message || String( error ) } );
+    return [];
+  }
+};

package/src/utils/metering.js

@@ -0,0 +1,125 @@
+import { parseLLMUsage } from './usage.js';
+import { calculateCosts } from './cost.js';
+import { Tracing, Event } from '@outputai/core/sdk/runtime';
+import { Logger } from '@outputai/core';
+import { convertCostToLegacy } from './legacy_cost_attribute.js';
+
+/**
+ * Collects the usage of a single LLM call and bills it once.
+ *
+ * Steps are recorded as the SDK lifecycle reports them, and the response is handed to `bill()` by
+ * the paths that have one, so a call that throws after the model ran is still billed for what it
+ * spent. `bill()` parses the usage, costs it, attaches both as trace attributes and emits the
+ * metering events; it never throws.
+ */
+export class Metering {
+  #recordedSteps = [];
+  #billed = false;
+  #prompt;
+  #traceId;
+  #attributes = {
+    usage: null,
+    cost: null,
+    legacy: null
+  };
+
+  async #createAttributes( { usage, steps } ) {
+    this.#attributes.usage = parseLLMUsage( { prompt: this.#prompt, usage, steps } );
+    if ( this.#attributes.usage ) {
+      this.#attributes.cost = await calculateCosts( this.#attributes.usage );
+      if ( this.#attributes.cost ) {
+        // @TEMP Preserve the deprecated event and trace attribute for legacy consumers.
+        this.#attributes.legacy = convertCostToLegacy( this.#attributes.cost );
+      }
+    }
+  }
+
+  #attachAttributes() {
+    if ( this.#attributes.usage ) {
+      Tracing.addEventAttribute( { eventId: this.#traceId, attribute: this.#attributes.usage } );
+    }
+    if ( this.#attributes.cost ) {
+      Tracing.addEventAttribute( { eventId: this.#traceId, attribute: this.#attributes.cost } );
+    }
+    if ( this.#attributes.legacy ) {
+      Tracing.addEventAttribute( { eventId: this.#traceId, attribute: this.#attributes.legacy } );
+    }
+  }
+
+  #emitEvents() {
+    if ( this.#attributes.usage ) {
+      Event.emit( 'llm:generation:metering', structuredClone( { cost: this.#attributes.cost, usage: this.#attributes.usage } ) );
+    }
+    if ( this.#attributes.legacy ) {
+      Event.emit( 'cost:llm:request', structuredClone( this.#attributes.legacy ) );
+    }
+  }
+
+  /**
+   * @param {object} args
+   * @param {string} args.traceId - Trace event the usage and cost attributes are attached to
+   * @param {object} args.prompt - Loaded prompt (`config.provider` / `config.model` used for cost)
+   */
+  constructor( { traceId, prompt } ) {
+    this.#traceId = traceId;
+    this.#prompt = prompt;
+  }
+
+  /**
+   * Records a completed step, used as the usage source when the call ends without a response.
+   *
+   * @param {object} step - AI SDK step
+   */
+  recordStep = step => {
+    if ( !this.#billed ) {
+      this.#recordedSteps.push( step );
+    }
+  };
+
+  /**
+   * Bills the recorded usage: parses it, costs it, attaches the trace attributes and emits the
+   * metering events. Later calls and later records are ignored, so it is safe to call from every
+   * path that can end the call. Read the result from `attributes`.
+   *
+   * An empty collector is not a final answer: the SDK reports a stream error before the steps that
+   * preceded it, so billing stays open until there is usage or a step to report.
+   *
+   * @param {object} [response] - AI SDK response, whose aggregate usage and steps take precedence
+   *   over the recorded steps; omitted by the paths that end without one
+   * @returns {Promise<void>}
+   */
+  bill = async ( response = null ) => {
+    if ( this.#billed ) {
+      return;
+    }
+
+    try {
+      const responseSteps = response?.steps;
+      const steps = Array.isArray( responseSteps ) && responseSteps.length > 0 ? responseSteps : this.#recordedSteps;
+      const usage = response?.usage;
+
+      if ( !usage && steps.length === 0 ) {
+        return;
+      }
+
+      this.#billed = true;
+
+      await this.#createAttributes( { usage, steps } );
+      this.#attachAttributes();
+      this.#emitEvents();
+
+    } catch ( error ) {
+      Logger.error( 'Metering failed', { namespace: 'LLM', error: error?.message || String( error ) } );
+    }
+  };
+
+  /**
+   * Attributes produced by `bill()`; each one is null until billed, and stays null when there is
+   * nothing to report.
+   *
+   * @returns {{ usage: object|null, cost: object|null, legacy: object|null }}
+   */
+  get attributes() {
+    return this.#attributes;
+  }
+}

package/src/utils/sources.js

@@ -1,4 +1,5 @@
 import { createHash } from 'node:crypto';
+import { Logger } from '@outputai/core';
 
 /** Builds the final source shape */
 const buildSource = ( { url, title } ) => {
@@ -15,34 +16,58 @@ const buildSource = ( { url, title } ) => {
 /** Return value it is array, otherwise return [] */
 const asArray = v => Array.isArray( v ) ? v : [];
 
+const isNonBlankUrl = url => typeof url === 'string' && url.trim().length > 0;
+
 /**
- * Extracts source URLs from search tool results embedded in AI SDK step data.
+ * Extracts source from search tool results embedded in AI SDK step data.
  *
  * Detects any tool result containing a `results[]` array whose items have a `url` string field.
  * This covers perplexitySearch, tavilySearch, exaSearch, and any future tool with the same shape.
+ * Deduplicate key is "url"
  *
  * @param {Array} steps - AI SDK response steps (response.steps)
- * @returns {Array<{ type: string, sourceType: string, id: string, url: string, title: string }>}
+ * @returns {Array<{ key: string, source: object }>}
  */
 const extractSourcesFromSteps = steps =>
   asArray( steps )
     .flatMap( step => asArray( step?.toolResults ) )
     .flatMap( toolResult => asArray( toolResult?.output?.results ) )
-    .filter( item => typeof item?.url === 'string' && !!item.url.trim() )
-    .map( v => buildSource( v ) );
+    .filter( item => isNonBlankUrl( item?.url ) )
+    .map( item => {
+      const source = buildSource( item );
+      return { key: source.url, source };
+    } );
 
 /**
- * Extract sources from tools usage and final response, deduplicate, and return.
+ * Extract only the valid sources from the response (has url or id).
+ * Deduplicate key is `url` when set, otherwise `id` (document sources have `id` and no `url`).
  *
- * Dedup key is `url` when set, otherwise `id` (document sources have `id` and no `url`).
+ * @param {Array} sources
+ * @returns {Array<{ key: string, source: object }>}
+ */
+const extractValidSourcesFromResponse = sources =>
+  asArray( sources )
+    .filter( item => item !== null && typeof item === 'object' && ( isNonBlankUrl( item.url ) || typeof item.id === 'string' ) )
+    .map( item => ( { key: isNonBlankUrl( item.url ) ? item.url.trim() : item.id, source: item } ) );
+
+/** Deduplicate wrapped sources by their key */
+const deduplicateSources = sources => new Map( sources.map( item => [ item.key, item.source ] ) ).values().toArray();
+
+/**
+ * Extract sources from tools usage and final response, deduplicate, and return.
  * Response sources are preferred over tools when the key matches.
  *
  * @param {object} response AI SDK response
  * @returns {object[]} Merged sources
  */
 export const extractSources = response => {
-  const { steps, sources: sourcesFromResponse } = response;
-  const sourcesFromTools = extractSourcesFromSteps( steps );
-  const allSources = sourcesFromTools.concat( asArray( sourcesFromResponse ) );
-  return new Map( allSources.map( s => [ s.url ?? s.id, s ] ) ).values().toArray();
+  try {
+    const { steps, sources: sourcesFromResponse } = response ?? {};
+    const sourcesFromTools = extractSourcesFromSteps( asArray( steps ) );
+    const validSourcesFromResponse = extractValidSourcesFromResponse( sourcesFromResponse );
+    return deduplicateSources( sourcesFromTools.concat( validSourcesFromResponse ) );
+  } catch ( error ) {
+    Logger.error( 'Sources extraction failed', { namespace: 'LLM', error: error?.message || String( error ) } );
+    return [];
+  }
 };

package/src/utils/stream.js

@@ -1,22 +1,75 @@
+import { setTimeout as delay } from 'node:timers/promises';
+
+/** How long a failed stream may go quiet before the drain gives up on its remaining parts */
+const FAILURE_DRAIN_TIMEOUT_MS = 250;
+
+const TIMED_OUT = Symbol( 'drain-timed-out' );
+
+/** Extracts the error an abort or error part should surface, or null for every other part */
+const extractError = ( part, abortSignal ) => {
+  if ( part?.type === 'abort' ) {
+    const reason = abortSignal?.reason;
+    return reason instanceof Error ?
+      reason :
+      new Error( part.reason ?? 'Streaming aborted.', { cause: reason } );
+  }
+
+  if ( part?.type === 'error' ) {
+    return part.error instanceof Error ?
+      part.error :
+      new Error( part.error ? String( part.error ) : 'Streaming failed.', { cause: part.error } );
+  }
+
+  if ( abortSignal?.aborted ) {
+    return abortSignal.reason instanceof Error ?
+      abortSignal.reason :
+      new Error( 'Streaming aborted.', { cause: abortSignal.reason } );
+  }
+
+  return null;
+};
+
+/** Reads the next part, resolving with `TIMED_OUT` when the stream goes quiet for longer than `timeoutMs` */
+const nextWithin = async ( iterator, timeoutMs ) =>
+  Promise.race( [ iterator.next(), delay( timeoutMs, TIMED_OUT, { ref: false } ) ] );
+
 /**
- * Consumes a streaming result until it completes. Throws on abort or error parts.
+ * Consumes a streaming result until it completes, then throws the first abort or error part it saw.
  * Callers must not throw from the AI SDK `onError` callback; that errors the stream before these parts are delivered.
  *
+ * The draining continues past a failure on purpose: the AI SDK fires its lifecycle callbacks as parts flow through the
+ * stream, so `onStepEnd` and `onEnd` - and with them the usage of everything already spent - only arrive if something
+ * keeps reading. A stalled provider can leave the stream open forever, so once a failure is captured every read waits
+ * at most `FAILURE_DRAIN_TIMEOUT_MS` for the next part, and the drain gives up and throws anyway when none arrives.
+ * A healthy stream ends immediately after a failure, so the timeout is not part of the normal path; a successful
+ * stream is never timed out.
+ *
  * @param {object} stream - AI SDK stream result with `stream`
  * @param {AbortSignal} [abortSignal] - Used to recover the original abort reason
  */
 export const drainStream = async ( stream, abortSignal ) => {
-  for await ( const part of stream.stream ) {
-    if ( part.type === 'abort' ) {
-      const reason = abortSignal?.reason;
-      throw reason instanceof Error ?
-        reason :
-        new Error( part.reason ?? 'Streaming generation aborted.', { cause: reason } );
+  const iterator = stream.stream[Symbol.asyncIterator]();
+  const state = { error: null };
+
+  while ( true ) {
+    const hasDeadline = state.error || abortSignal?.aborted;
+    const result = await ( hasDeadline ? nextWithin( iterator, FAILURE_DRAIN_TIMEOUT_MS ) : iterator.next() );
+
+    if ( result.done ) {
+      break;
+    }
+
+    if ( result?.value || abortSignal?.aborted ) {
+      state.error ??= extractError( result.value, abortSignal );
     }
-    if ( part.type === 'error' ) {
-      throw part.error instanceof Error ?
-        part.error :
-        new Error( part.error ? String( part.error ) : 'Streaming generation failed.', { cause: part.error } );
+
+    if ( result === TIMED_OUT ) {
+      iterator.return?.()?.catch( () => {} );
+      break;
     }
   }
+
+  if ( state.error ) {
+    throw state.error;
+  }
 };

package/src/utils/usage.js

@@ -1,6 +1,7 @@
 import Decimal from 'decimal.js';
 import { Tracing } from '@outputai/core/sdk/runtime';
 import { parseGroundingUsage } from './grounding.js';
+import { extractUsageFromSteps } from './usage_tools.js';
 
 const exists = v => Number.isSafeInteger( v ) && v >= 0;
 
@@ -58,19 +59,37 @@ export class LLMGenerationUsage extends Tracing.Attribute.BaseAttribute {
 }
 
 /**
- * Converts raw AI SDK usage to LLMGenerationUsage data
+ * Returns usage from response or reconstruct it from steps.
+ *
+ * The token check is load-bearing: aborted and truncated calls still report an aggregate, but with
+ * every token count undefined, and the steps are then the only record of what was already spent.
+ */
+const resolveUsage = ( { usage, steps } ) => {
+  if ( usage && ( exists( usage.inputTokens ) || exists( usage.outputTokens ) ) ) {
+    return usage;
+  }
+  return Array.isArray( steps ) && steps.length > 0 ? extractUsageFromSteps( steps ) : null;
+};
+
+/**
+ * Converts raw AI SDK usage to LLMGenerationUsage data.
+ *
+ * Reads `response.usage`, falling back to the sum of `response.steps` usage when the aggregate is
+ * missing tokens (truncated or failed streams). Grounding is aggregated from `response.steps`, so a
+ * partial `{ steps }` response is accepted when only the steps are known.
  *
  * @param {object} args
  * @param {object} args.prompt - Output prompt with model configuration
- * @param {object} args.prompt.config - Prompt model configuration
- * @param {string} args.prompt.config.provider - Id of the provider
- * @param {string} args.prompt.config.model - Id of the model
- * @param {object} args.usage - AI SDK usage with aggregate token counts and optional input/output token details
- * @param {object[]} [args.steps] - AI SDK steps, each with its own `providerMetadata`, used for per-request charges
+ * @param {object} [args.usage] - AI SDK response usage
+ * @param {object} [args.steps] - AI SDK response steps
  *
  * @returns {LLMGenerationUsage | null} LLM generation usage with input, output, total and detailed breakdown
  */
-export const parseLLMUsage = ( { prompt, usage, steps } ) => {
+export const parseLLMUsage = ( { prompt, usage: usageArg = null, steps = [] } ) => {
+  const usage = resolveUsage( { usage: usageArg, steps } );
+  if ( !usage ) {
+    return null;
+  }
   const { provider: providerId, model: modelId } = prompt.config;
   const { inputTokens, inputTokenDetails, outputTokens, outputTokenDetails } = usage;
   const { noCacheTokens, cacheReadTokens, cacheWriteTokens } = inputTokenDetails ?? {};
@@ -118,13 +137,15 @@ export const parseLLMUsage = ( { prompt, usage, steps } ) => {
   // Grounding is billed per step, so aggregate across every step rather than only the final one.
   // Per-query families (Gemini 3) return queries.length per step; per-prompt families (Gemini 2.x)
   // return 1 per grounded step. Summing yields total queries and grounded-step count respectively.
-  const grounding = ( steps ?? [] )
-    .filter( step => step?.providerMetadata )
-    .map( step => parseGroundingUsage( modelId, step.providerMetadata ) )
-    .filter( Boolean );
-  if ( grounding.length > 0 ) {
-    const amount = grounding.reduce( ( sum, g ) => sum + g.amount, 0 );
-    items.push( new LLMGenerationUsageItem( LLMGenerationUsageItem.Group.TOOLS, grounding[0].label, amount ) );
+  if ( Array.isArray( steps ) ) {
+    const grounding = steps
+      .filter( step => step?.providerMetadata )
+      .map( step => parseGroundingUsage( modelId, step.providerMetadata ) )
+      .filter( Boolean );
+    if ( grounding.length > 0 ) {
+      const amount = grounding.reduce( ( sum, g ) => sum + g.amount, 0 );
+      items.push( new LLMGenerationUsageItem( LLMGenerationUsageItem.Group.TOOLS, grounding[0].label, amount ) );
+    }
   }
 
   if ( items.length === 0 ) {

package/src/utils/usage_tools.js

@@ -0,0 +1,23 @@
+const sum = ( a, b ) => {
+  if ( Number.isFinite( a ) ) {
+    return a + ( Number.isFinite( b ) ? b : 0 );
+  } else {
+    return Number.isFinite( b ) ? b : undefined;
+  }
+};
+
+const isPlainObject = value => value !== null && typeof value === 'object';
+
+/** Deep numeric merge of two AI SDK usage objects; `raw` is provider-shaped and never summed */
+const mergeUsage = ( total, usage ) => Object.entries( usage )
+  .filter( ( [ key ] ) => key !== 'raw' )
+  .reduce( ( merged, [ key, value ] ) => ( {
+    ...merged,
+    [key]: isPlainObject( value ) ? mergeUsage( merged[key] ?? {}, value ) : sum( merged[key], value )
+  } ), total );
+
+/** Sums per-step AI SDK usage */
+export const extractUsageFromSteps = steps => steps
+  .map( step => step?.usage )
+  .filter( Boolean )
+  .reduce( ( total, usage ) => mergeUsage( total, usage ), {} );

package/src/utils/wrap.js

@@ -1,12 +1,10 @@
 import { extractSources } from './sources.js';
-import { parseLLMUsage } from './usage.js';
-import { calculateCosts } from './cost.js';
-import { calculateBase64FileSize } from './image.js';
-import { Tracing, Event } from '@outputai/core/sdk/runtime';
+import { serializeImagesFromResponse } from './image.js';
+import { Tracing } from '@outputai/core/sdk/runtime';
 import { mapAiError } from './error_handler.js';
 import { isPromise } from 'node:util/types';
-import { Logger } from '@outputai/core';
-import { convertCostToLegacy } from './legacy_cost_attribute.js';
+import { FatalError, Logger } from '@outputai/core';
+import { Metering } from './metering.js';
 import { randomBytes } from 'node:crypto';
 
 /** Creates a proxy of the AI SDK response and attach virtual getters to it based on an object map */
@@ -23,44 +21,93 @@ const startTrace = ( { name, prompt } ) => {
   return traceId;
 };
 
-/** Map AI Error, add an error trace entry and return the new error */
-const handleError = ( { traceId, error: originalError } ) => {
+/** Handle AI SDK errors: map them, add an error trace entry and return the new error */
+const handleAiSdkError = ( { traceId, error: originalError } ) => {
   const error = mapAiError( originalError );
   Tracing.addEventError( { id: traceId, details: error } );
   return error;
 };
 
-/** Normalize raw AI SDK usage, calculate cost, attach trace attributes, and emit metering events */
-const handleMetering = async ( { traceId, usage: sdkUsage, prompt, steps } ) => {
-  const usageAttribute = parseLLMUsage( { usage: sdkUsage, prompt, steps } );
-  if ( !usageAttribute ) {
-    return null;
+/**
+ * Handle a failure of our own response handling: wrap it, add an error trace entry and return it.
+ * Fatal on purpose - such a failure is deterministic, so a retry pays for the model call again and
+ * fails the same way.
+ */
+const handleResponseError = ( { traceId, error: cause } ) => {
+  const error = new FatalError( 'AI SDK response handling failed.', { cause } );
+  Tracing.addEventError( { id: traceId, details: error } );
+  return error;
+};
+
+/** Invokes an async function, wrapped in a try/catch, return { result, error } */
+const inlineTryAsync = async fn => {
+  try {
+    return { result: await fn(), error: null };
+  } catch ( error ) {
+    return { result: null, error };
   }
+};
 
-  const costAttribute = await calculateCosts( usageAttribute );
-  if ( costAttribute ) {
-    Tracing.addEventAttribute( { eventId: traceId, attribute: costAttribute } );
-    // @TEMP Preserve the deprecated event and trace attribute for legacy consumers.
-    const legacyPayload = convertCostToLegacy( costAttribute );
-    if ( legacyPayload ) {
-      Event.emit( 'cost:llm:request', structuredClone( legacyPayload ) );
-      Tracing.addEventAttribute( { eventId: traceId, attribute: legacyPayload } );
-    }
+/**
+ * Runs a completing AI SDK text call (`generateText`, `generateTextWithStreaming`,
+ * `Agent.generate`, `Agent.generateWithStreaming`): starts the LLM trace, runs `fn`, ends the trace
+ * and returns the response proxied with `result` (`text`), `cost` and merged `sources`.
+ *
+ * `fn` receives `onStepEndHook` to wire into the SDK `onStepEnd`, so the usage of every completed
+ * step is collected as the call progresses and still lands when `fn` throws after the model ran.
+ * The response `fn` returns is billed directly, and its aggregate usage takes precedence over the
+ * collected steps. Trace output keeps raw `response.usage`; normalized usage and cost are trace
+ * attributes.
+ *
+ * Reading the billed response is guarded: a throw from `fn` is mapped as an SDK error, while a
+ * failure of our own handling becomes a `FatalError` recorded on the trace. The proxy is built
+ * before the trace ends, so a call that fails there never leaves an `end` entry next to an `error`.
+ *
+ * @param {object} args
+ * @param {string} args.name - Trace event name
+ * @param {object} args.prompt - Loaded prompt (`config.provider` / `config.model` used for cost)
+ * @param {( wiring: { onStepEndHook: Function } ) => Promise<object>} args.fn - AI SDK call; wire
+ *   `onStepEndHook` into the call `onStepEnd` and return the SDK response
+ * @returns {Promise<object>} Proxied SDK response
+ */
+export const wrapTextGeneration = async ( { name, prompt, fn } ) => {
+  const traceId = startTrace( { name, prompt } );
+
+  const metering = new Metering( { traceId, prompt } );
+
+  const onStepEndHook = metering.recordStep;
+
+  const { result: response, error } = await inlineTryAsync( () => fn( { onStepEndHook } ) );
+
+  await metering.bill( response );
+
+  if ( error ) {
+    throw handleAiSdkError( { traceId, error } );
   }
 
-  Tracing.addEventAttribute( { eventId: traceId, attribute: usageAttribute } );
-  Event.emit( 'llm:generation:metering', structuredClone( { cost: costAttribute, usage: usageAttribute } ) );
-  return costAttribute;
+  try {
+    const { text: result, finalStep, usage } = response;
+    const providerMetadata = finalStep?.providerMetadata;
+    const sources = extractSources( response );
+
+    // Create proxy first, as this could theoretically fail
+    const responseProxy = createResponseProxy( { response, properties: { cost: metering.attributes.cost, sources, result } } );
+    Tracing.addEventEnd( { id: traceId, details: { result, usage, providerMetadata, sources } } );
+
+    return responseProxy;
+  } catch ( error ) {
+    throw handleResponseError( { traceId, error } );
+  }
 };
 
 /**
- * Runs a completing AI SDK call (`generateText`, `generateImage`, `generateTextWithStreaming`,
- * `Agent.generate`, `Agent.generateWithStreaming`): start the LLM trace, run `fn`, attach cost
- * (and sources on text), end the trace, and return a proxied response.
+ * Runs `generateImage`: starts the LLM trace, meters the returned usage, ends the trace and returns
+ * the response proxied with `result` (the first `image`) and `cost`.
  *
- * Text responses get `result` (`text`), `cost`, and merged `sources`. Image responses get
- * `result` (`image`) and `cost`. Trace output keeps raw `response.usage`; normalized usage and
- * cost are trace attributes.
+ * Image calls expose no step lifecycle, so usage is read from the response. With no
+ * step loop and no output parsing, a throw from `fn` happens before any usage exists, so it is
+ * mapped and recorded without metering. Reading the billed response is guarded the same way
+ * `wrapTextGeneration` guards it: a failure there is a `FatalError` recorded after billing.
  *
  * @param {object} args
  * @param {string} args.name - Trace event name
@@ -68,131 +115,162 @@ const handleMetering = async ( { traceId, usage: sdkUsage, prompt, steps } ) =>
  * @param {() => Promise<object>} args.fn - AI SDK call; must return the SDK response
  * @returns {Promise<object>} Proxied SDK response
  */
-export const wrapGeneration = async ( { name, prompt, fn } ) => {
+export const wrapImageGeneration = async ( { name, prompt, fn } ) => {
   const traceId = startTrace( { name, prompt } );
 
-  try {
-    const response = await fn();
-    const { usage } = response;
-    const cost = await handleMetering( { traceId, usage, prompt, steps: response.steps } );
+  const metering = new Metering( { traceId, prompt } );
 
-    if ( Array.isArray( response.images ) ) {
-      const { image, images, providerMetadata } = response;
-      const mappedImages = images.map( ( { mediaType, base64 } ) => ( { size: calculateBase64FileSize( base64 ), mediaType } ) );
+  const { result: response, error } = await inlineTryAsync( fn );
 
-      Tracing.addEventEnd( { id: traceId, details: { result: mappedImages, usage, providerMetadata } } );
-      return createResponseProxy( { response, properties: { cost, result: image } } );
-    } else {
-      const { text: result, finalStep } = response;
-      const { providerMetadata } = finalStep;
-      const sources = extractSources( response );
+  if ( error ) {
+    throw handleAiSdkError( { traceId, error } );
+  }
 
-      Tracing.addEventEnd( { id: traceId, details: { result, usage, providerMetadata, sources } } );
-      return createResponseProxy( { response, properties: { cost, sources, result } } );
-    }
+  await metering.bill( response );
+
+  try {
+    const { usage, image: result, providerMetadata } = response;
+    const serializeImages = serializeImagesFromResponse( response );
 
+    // Create proxy first, as this could theoretically fail
+    const responseProxy = createResponseProxy( { response, properties: { cost: metering.attributes.cost, result } } );
+
+    Tracing.addEventEnd( { id: traceId, details: { result: serializeImages, usage, providerMetadata } } );
+    return responseProxy;
   } catch ( error ) {
-    throw handleError( { traceId, error } );
+    throw handleResponseError( { traceId, error } );
+  }
+};
+
+/** Awaits a consumer callback, logging and swallowing its failures: the hooks are fire and forget */
+const invokeCallback = async ( cb, args, cbName ) => {
+  if ( typeof cb === 'function' ) {
+    try {
+      await cb( ...args );
+    } catch ( e ) {
+      Logger.error( `Stream ${cbName}() callback failed`, { namespace: 'LLM', error: e?.message || String( e ) } );
+    }
   }
 };
 
 /**
  * Starts an LLM trace around a live AI SDK stream (`streamText`, `Agent.stream`).
  *
- * `fn` receives `onEndHook(response, callback)` and `onErrorHook(event, callback)`. Call
- * `onEndHook` from the SDK `onEnd`: it wraps the response, ends the trace, then invokes
- * `callback` without rethrowing. Call `onErrorHook` from SDK `onError`: it maps the error,
- * records it, and invokes `callback` without rethrowing. Callback failures are logged.
+ * `fn` receives `onEndHook(response, callback)` and `onErrorHook(event, callback)` to wire into the
+ * SDK `onEnd` / `onError`, plus `onStepEndHook` and `onAbortHook` for `onStepEnd` / `onAbort`. Both
+ * terminal hooks bill, end or record the trace, then invoke `callback`; a `callback` that is not a
+ * function is skipped, and callback failures are logged, never rethrown.
+ *
+ * Neither hook throws, because the SDK invokes them through `notify`, which swallows callback
+ * failures and discards their return: a failure while reading the response would vanish, so it is
+ * logged and recorded as a trace error here instead. An `onError`
+ * event carrying no `Error` is normalized to `Streaming failed.` with the original value as `cause`,
+ * so the trace always closes and the consumer callback always runs.
+ *
+ * Metering is fed by `onStepEndHook`, which collects step usage so the tokens survive the exits
+ * that report no response: `NoOutputGeneratedError` skips `onEnd`, and an abort fires neither
+ * terminal hook. `onAbortHook` bills that case, and the SDK awaits it before enqueueing the abort
+ * part, so billing from there always lands before the consumer observes the cancellation.
  *
- * A throw or rejected Promise from `fn` (stream creation / Agent setup) is mapped and recorded
- * on the LLM trace. A non-Promise return (the `streamText` stream) is returned immediately.
- * When `abortSignal` aborts, its reason is recorded as an LLM trace error. The listener is removed
- * when the stream ends, reports an error, or fails during setup.
+ * A throw or rejected Promise from `fn` (stream creation / Agent setup) is mapped and recorded on
+ * the trace; a non-Promise return (the `streamText` stream) is returned as is, with a rejection of
+ * its `output` promise recorded without being swallowed from the consumer. An abort on
+ * `abortSignal` records its reason as a trace error, and the listener is removed when the stream
+ * ends, reports an error, or fails during setup. An abort the SDK raises on its own timeout leaves
+ * `abortSignal` untouched, so `onAbort` records that one instead.
  *
  * @param {object} args
  * @param {string} args.name - Trace event name
  * @param {object} args.prompt - Loaded prompt (`config.provider` / `config.model` used for cost)
  * @param {AbortSignal} [args.abortSignal] - Optional signal used to trace stream cancellation
- * @param {(hooks: { onEndHook: Function, onErrorHook: Function }) => object | Promise<object>} args.fn -
- *   Return the SDK stream, or a Promise of that stream (`Agent.stream`); wire the hooks into SDK
- *   `onEnd` / `onError`
+ * @param {( wiring: { onEndHook: Function, onErrorHook: Function, onStepEndHook: Function,
+ *   onAbortHook: Function } ) => object | Promise<object>} args.fn - Wire the hooks into the call
+ *   and return the SDK stream, or a Promise of that stream (`Agent.stream`)
  * @returns {object | Promise<object>} Value returned by `fn`; a rejected Promise is remapped
  */
 export const wrapStream = ( { name, prompt, abortSignal, fn } ) => {
   const traceId = startTrace( { name, prompt } );
-
-  const onAbortHook = reason =>
-    handleError( {
-      traceId,
-      error: reason instanceof Error ? reason : new Error( 'Streaming aborted.', { cause: reason } )
-    } );
+  const metering = new Metering( { traceId, prompt } );
 
   const handleAbort = () => {
-    if ( !abortSignal ) {
-      return () => {};
-    }
+    const error = abortSignal?.reason instanceof Error ? abortSignal.reason : new Error( 'Streaming aborted.', { cause: abortSignal?.reason } );
+    handleAiSdkError( { traceId, error } );
+  };
 
+  if ( abortSignal ) {
     if ( abortSignal.aborted ) {
-      onAbortHook( abortSignal.reason );
-      return () => {};
+      handleAbort();
+    } else {
+      abortSignal.addEventListener( 'abort', handleAbort, { once: true } );
     }
+  }
+  const removeAbortListener = () => abortSignal?.removeEventListener( 'abort', handleAbort );
 
-    const listener = () => onAbortHook( abortSignal.reason );
-    abortSignal.addEventListener( 'abort', listener, { once: true } );
+  const onStepEndHook = metering.recordStep;
 
-    return () => abortSignal.removeEventListener( 'abort', listener );
-  };
+  const onAbortHook = async () => {
+    removeAbortListener();
 
-  const removeAbortListener = handleAbort();
+    await metering.bill();
+    // abortion was caused by timeout and not signal
+    if ( !abortSignal?.aborted ) {
+      handleAiSdkError( { traceId, error: new Error( 'Streaming timed out.' ) } );
+    }
+  };
 
   const onEndHook = async ( response, callback ) => {
-    const state = { proxyResponse: null };
     removeAbortListener();
+
+    await metering.bill( response );
+
     try {
-      const { text: result, finalStep, usage, steps } = response;
-      const { providerMetadata } = finalStep;
-      const cost = await handleMetering( { traceId, usage, prompt, steps } );
+      const { text: result, finalStep, usage } = response;
+      const providerMetadata = finalStep?.providerMetadata;
       const sources = extractSources( response );
+
+      const proxyResponse = createResponseProxy( { response, properties: { cost: metering.attributes.cost, sources, result } } );
       Tracing.addEventEnd( { id: traceId, details: { result, usage, providerMetadata, sources } } );
-      state.proxyResponse = createResponseProxy( { response, properties: { cost, sources, result } } );
-    } catch ( error ) {
-      Logger.error( 'Stream onEnd() handler failed', { namespace: 'LLM', error: error?.message ?? String( error ) } );
-      throw handleError( { traceId, error } );
-    }
 
-    try {
-      await callback?.( state.proxyResponse );
-    } catch ( callbackError ) {
-      // ignore these as this callback is fire and forget
-      Logger.error( 'Stream onEnd() callback failed', {
-        namespace: 'LLM',
-        error: callbackError instanceof Error ? callbackError.message : String( callbackError )
-      } );
+      await invokeCallback( callback, [ proxyResponse ], 'onEnd' );
+    } catch ( error ) {
+      Logger.error( 'AI SDK response handling failed', { namespace: 'LLM', error: error?.message || String( error ) } );
+      handleResponseError( { traceId, error } );
     }
   };
 
   const onErrorHook = async ( event, callback ) => {
-    const error = handleError( { traceId, error: event.error } );
     removeAbortListener();
+    await metering.bill();
+
+    const error = event?.error instanceof Error ? event.error : new Error( 'Streaming failed.', { cause: event?.error } );
+    const mappedError = handleAiSdkError( { traceId, error } );
+    await invokeCallback( callback, [ mappedError ], 'onError' );
+  };
+
+  const handleStreamError = error => {
+    removeAbortListener();
+    return handleAiSdkError( { traceId, error } );
+  };
+
+  const addStreamOutputErrorHandler = stream => {
+    stream.output?.catch( error => handleAiSdkError( { traceId, error } ) );
+    return stream;
+  };
+
+  const createStream = () => {
     try {
-      await callback?.( error );
-    } catch ( callbackError ) {
-      // ignore these as this callback is fire and forget
-      Logger.error( 'Stream onError() callback failed', {
-        namespace: 'LLM',
-        error: callbackError instanceof Error ? callbackError.message : String( callbackError )
-      } );
+      return fn( { onEndHook, onErrorHook, onStepEndHook, onAbortHook } );
+    } catch ( error ) {
+      throw handleStreamError( error );
     }
   };
 
-  try {
-    const stream = fn( { onEndHook, onErrorHook } );
-    return isPromise( stream ) ? stream.catch( error => {
-      removeAbortListener();
-      throw handleError( { traceId, error } );
-    } ) : stream;
-  } catch ( error ) {
-    removeAbortListener();
-    throw handleError( { traceId, error } );
-  }
+  const stream = createStream();
+
+  /** Streams can be a promise or not, for both cases add the handler to catch .output errors, and if it is a promise add a .catch handler to observe errors there as well */
+  return isPromise( stream ) ?
+    stream.then( addStreamOutputErrorHandler ).catch( e => {
+      throw handleStreamError( e );
+    } ) :
+    addStreamOutputErrorHandler( stream );
 };