@outputai/llm 0.1.12-next.ecd4dda.0 → 0.1.13-next.04243eb.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@outputai/llm",
-  "version": "0.1.12-next.ecd4dda.0",
+  "version": "0.1.13-next.04243eb.0",
   "description": "Framework abstraction to interact with LLM models",
   "type": "module",
   "main": "src/index.js",
@@ -9,20 +9,20 @@
     "./src"
   ],
   "dependencies": {
-    "@ai-sdk/amazon-bedrock": "4.0.83",
-    "@ai-sdk/anthropic": "3.0.64",
-    "@ai-sdk/azure": "3.0.49",
-    "@ai-sdk/google-vertex": "4.0.95",
-    "@ai-sdk/openai": "3.0.48",
-    "@ai-sdk/perplexity": "3.0.26",
+    "@ai-sdk/amazon-bedrock": "4.0.89",
+    "@ai-sdk/anthropic": "3.0.66",
+    "@ai-sdk/azure": "3.0.51",
+    "@ai-sdk/google-vertex": "4.0.102",
+    "@ai-sdk/openai": "3.0.50",
+    "@ai-sdk/perplexity": "3.0.27",
     "@exalabs/ai-sdk": "2.0.1",
     "@perplexity-ai/ai-sdk": "0.1.2",
     "@tavily/ai-sdk": "0.4.1",
-    "ai": "6.0.141",
+    "ai": "6.0.146",
     "decimal.js": "10.6.0",
     "gray-matter": "4.0.3",
     "liquidjs": "10.25.2",
-    "@outputai/core": "0.1.12-next.ecd4dda.0"
+    "@outputai/core": "0.1.13-next.04243eb.0"
   },
   "license": "Apache-2.0",
   "publishConfig": {

package/src/agent.js

@@ -0,0 +1,100 @@
+import { ValidationError } from '@outputai/core';
+import { resolveInvocationDir } from '@outputai/core/sdk_utils';
+import { ToolLoopAgent as AIToolLoopAgent, stepCountIs } from 'ai';
+import { hydratePromptTemplate, loadAiSdkOptionsFromPrompt } from './ai_sdk.js';
+import { startTrace, endTraceWithError, traceStreamCallbacks } from './trace_utils.js';
+import { wrapInOutputResponse } from './response_utils.js';
+import { ROLE, isRole, getContent } from './message_utils.js';
+
+export { skill } from './skill.js';
+
+export const createMemoryConversationStore = () => {
+  const messages = [];
+  return {
+    getMessages: () => messages,
+    addMessages: newMessages => messages.push( ...newMessages )
+  };
+};
+
+export class Agent extends AIToolLoopAgent {
+  _prompt;
+  _modelId;
+  _initialMessages;
+  _store;
+
+  constructor( {
+    prompt, promptDir, variables = {}, skills = [], tools = {},
+    stopWhen, maxSteps = 10, conversationStore, ...rest
+  } ) {
+    if ( !prompt ) {
+      throw new ValidationError( 'Agent requires a prompt' );
+    }
+
+    // Must be captured synchronously — Temporal async activity execution
+    // breaks the call stack, so resolveInvocationDir() fails if called lazily.
+    const resolvedPromptDir = promptDir ?? resolveInvocationDir();
+
+    const { loadedPrompt, tools: mergedTools } =
+      hydratePromptTemplate( prompt, variables, resolvedPromptDir, skills, tools );
+
+    const { messages: allMessages, ...constructorOptions } = loadAiSdkOptionsFromPrompt( loadedPrompt );
+
+    // Extract system messages as `instructions` for the ToolLoopAgent constructor
+    // and keep user messages for generate() calls — avoids provider errors
+    // with multiple system messages during multi-step tool loops
+    const systemContent = allMessages.filter( isRole( ROLE.SYSTEM ) ).map( getContent ).join( '\n\n' );
+
+    super( {
+      ...constructorOptions,
+      ...( systemContent ? { instructions: systemContent } : {} ),
+      ...( Object.keys( mergedTools ).length > 0 ? { tools: mergedTools } : {} ),
+      stopWhen: stopWhen ?? stepCountIs( maxSteps ),
+      ...rest
+    } );
+
+    this._prompt = prompt;
+    this._modelId = loadedPrompt.config.model;
+    this._initialMessages = allMessages.filter( isRole( ROLE.USER ) );
+    this._store = conversationStore ?? null;
+  }
+
+  async _preSendHook( userMessages ) {
+    const priorMessages = this._store ? await this._store.getMessages() : [];
+    return [ ...this._initialMessages, ...priorMessages, ...userMessages ];
+  }
+
+  async _postSendHook( userMessages, result ) {
+    if ( this._store ) {
+      await this._store.addMessages( [ ...userMessages, ...( result.response?.messages ?? [] ) ] );
+    }
+  }
+
+  async generate( { messages: userMessages = [], ...callOptions } = {} ) {
+    const traceId = startTrace( 'Agent.generate', { prompt: this._prompt } );
+    try {
+      const messages = await this._preSendHook( userMessages );
+      const result = await super.generate( { messages, ...callOptions } );
+      const wrapped = await wrapInOutputResponse( result, { traceId, modelId: this._modelId } );
+      await this._postSendHook( userMessages, wrapped );
+      return wrapped;
+    } catch ( error ) {
+      endTraceWithError( traceId, error );
+      throw error;
+    }
+  }
+
+  async stream( { messages: userMessages = [], onFinish, onError, ...callOptions } = {} ) {
+    const traceId = startTrace( 'Agent.stream', { prompt: this._prompt } );
+    try {
+      const messages = await this._preSendHook( userMessages );
+      return super.stream( {
+        messages,
+        ...callOptions,
+        ...traceStreamCallbacks( traceId, this._modelId, { onFinish, onError } )
+      } );
+    } catch ( error ) {
+      endTraceWithError( traceId, error );
+      throw error;
+    }
+  }
+}

package/src/agent.spec.js

@@ -0,0 +1,304 @@
+import { describe, it, expect, vi, beforeEach } from 'vitest';
+import { mkdtempSync } from 'node:fs';
+import { join } from 'node:path';
+import { tmpdir } from 'node:os';
+
+// ─── Mocks ────────────────────────────────────────────────────────────────────
+
+const state = vi.hoisted( () => ( { promptDir: '' } ) );
+
+vi.mock( '@outputai/core/sdk_utils', () => ( {
+  resolveInvocationDir: () => state.promptDir
+} ) );
+
+const superGenerateImpl = vi.fn();
+const superStreamImpl = vi.fn();
+const superConstructorSpy = vi.fn();
+
+vi.mock( 'ai', () => {
+  class MockToolLoopAgent {
+    constructor( options ) {
+      superConstructorSpy( options );
+    }
+
+    async generate( ...args ) {
+      return superGenerateImpl( ...args );
+    }
+
+    stream( ...args ) {
+      return superStreamImpl( ...args );
+    }
+  }
+  return {
+    ToolLoopAgent: MockToolLoopAgent,
+    stepCountIs: vi.fn( n => ( { _stepCount: n } ) ),
+    tool: vi.fn( def => def )
+  };
+} );
+
+const hydratePromptTemplateImpl = vi.fn();
+const loadAiSdkOptionsImpl = vi.fn();
+vi.mock( './ai_sdk.js', () => ( {
+  hydratePromptTemplate: ( ...args ) => hydratePromptTemplateImpl( ...args ),
+  loadAiSdkOptionsFromPrompt: ( ...args ) => loadAiSdkOptionsImpl( ...args )
+} ) );
+
+const startTraceImpl = vi.fn( () => 'trace-id' );
+const endTraceWithErrorImpl = vi.fn();
+const traceStreamCallbacksImpl = vi.fn( () => ( {} ) );
+vi.mock( './trace_utils.js', () => ( {
+  startTrace: ( ...args ) => startTraceImpl( ...args ),
+  endTraceWithError: ( ...args ) => endTraceWithErrorImpl( ...args ),
+  traceStreamCallbacks: ( ...args ) => traceStreamCallbacksImpl( ...args )
+} ) );
+
+const wrapInOutputResponseImpl = vi.fn( response => response );
+vi.mock( './response_utils.js', () => ( {
+  wrapInOutputResponse: ( ...args ) => wrapInOutputResponseImpl( ...args )
+} ) );
+
+vi.mock( './skill.js', () => ( {
+  skill: vi.fn( ( { name, description, instructions } ) => ( { name, description: description ?? name, instructions } ) ),
+  buildLoadSkillTool: vi.fn( skills => ( { _loadSkillTool: true, skills } ) )
+} ) );
+
+// ─── Defaults ─────────────────────────────────────────────────────────────────
+
+const defaultMessages = [ { role: 'user', content: 'test message' } ];
+const defaultPromptMeta = {
+  config: { model: 'claude-sonnet-4-6' },
+  messages: defaultMessages,
+  promptFileDir: '/mock/dir'
+};
+
+const importSut = () => import( './agent.js' );
+
+beforeEach( () => {
+  state.promptDir = mkdtempSync( join( tmpdir(), 'agent-test-' ) );
+  vi.clearAllMocks();
+
+  hydratePromptTemplateImpl.mockReturnValue( {
+    loadedPrompt: defaultPromptMeta,
+    allVariables: {},
+    tools: {}
+  } );
+  loadAiSdkOptionsImpl.mockReturnValue( {
+    model: { _modelId: 'claude-sonnet-4-6' },
+    messages: defaultMessages
+  } );
+  superGenerateImpl.mockResolvedValue( { text: 'response', response: { messages: [] } } );
+  superStreamImpl.mockReturnValue( { textStream: 'stream' } );
+} );
+
+// ─── Tests ────────────────────────────────────────────────────────────────────
+
+describe( 'skill()', () => {
+  it( 'creates a skill object with name, description, instructions', async () => {
+    const { skill } = await importSut();
+    const s = skill( { name: 'my_skill', description: 'Does stuff', instructions: '# Do stuff\nStep 1' } );
+    expect( s ).toEqual( { name: 'my_skill', description: 'Does stuff', instructions: '# Do stuff\nStep 1' } );
+  } );
+} );
+
+describe( 'Agent — construction', () => {
+  it( 'throws ValidationError when prompt is missing', async () => {
+    const { Agent } = await importSut();
+    expect( () => new Agent( {} ) ).toThrow( /requires a prompt/ );
+  } );
+
+  it( 'constructs successfully with a prompt', async () => {
+    const { Agent } = await importSut();
+    expect( () => new Agent( { prompt: 'test@v1' } ) ).not.toThrow();
+  } );
+
+  it( 'calls AIToolLoopAgent constructor once at construction time', async () => {
+    const { Agent } = await importSut();
+    new Agent( { prompt: 'test@v1' } );
+    expect( superConstructorSpy ).toHaveBeenCalledTimes( 1 );
+  } );
+
+  it( 'passes model and stopWhen to AIToolLoopAgent constructor', async () => {
+    const { Agent } = await importSut();
+    new Agent( { prompt: 'test@v1' } );
+    expect( superConstructorSpy ).toHaveBeenCalledWith( expect.objectContaining( {
+      model: expect.objectContaining( { _modelId: 'claude-sonnet-4-6' } ),
+      stopWhen: expect.objectContaining( { _stepCount: 10 } )
+    } ) );
+  } );
+
+  it( 'uses resolveInvocationDir when promptDir not provided', async () => {
+    const { Agent } = await importSut();
+    new Agent( { prompt: 'test@v1' } );
+    expect( hydratePromptTemplateImpl ).toHaveBeenCalledWith( 'test@v1', {}, state.promptDir, [], {} );
+  } );
+
+  it( 'uses explicitly provided promptDir', async () => {
+    const explicitDir = mkdtempSync( join( tmpdir(), 'explicit-' ) );
+    const { Agent } = await importSut();
+    new Agent( { prompt: 'test@v1', promptDir: explicitDir } );
+    expect( hydratePromptTemplateImpl ).toHaveBeenCalledWith( 'test@v1', {}, explicitDir, [], {} );
+  } );
+
+  it( 'passes construction-time variables to hydratePromptTemplate', async () => {
+    const { Agent } = await importSut();
+    new Agent( { prompt: 'test@v1', variables: { persona: 'writer' } } );
+    expect( hydratePromptTemplateImpl ).toHaveBeenCalledWith(
+      'test@v1', { persona: 'writer' }, state.promptDir, [], {}
+    );
+  } );
+} );
+
+describe( 'Agent.generate() — messages', () => {
+  it( 'passes initialMessages from construction', async () => {
+    const { Agent } = await importSut();
+    const agent = new Agent( { prompt: 'test@v1' } );
+    await agent.generate();
+    expect( superGenerateImpl ).toHaveBeenCalledWith( expect.objectContaining( {
+      messages: defaultMessages
+    } ) );
+  } );
+
+  it( 'appends extra messages after initial messages', async () => {
+    const { Agent } = await importSut();
+    const agent = new Agent( { prompt: 'test@v1' } );
+    const extraMsg = { role: 'assistant', content: 'prior turn' };
+    await agent.generate( { messages: [ extraMsg ] } );
+    expect( superGenerateImpl ).toHaveBeenCalledWith( {
+      messages: [ ...defaultMessages, extraMsg ]
+    } );
+  } );
+} );
+
+describe( 'Agent.generate() — reuse', () => {
+  it( 'does not call AIToolLoopAgent constructor again on subsequent generate() calls', async () => {
+    const { Agent } = await importSut();
+    const agent = new Agent( { prompt: 'test@v1' } );
+    expect( superConstructorSpy ).toHaveBeenCalledTimes( 1 );
+
+    await agent.generate();
+    await agent.generate();
+    await agent.generate();
+
+    expect( superConstructorSpy ).toHaveBeenCalledTimes( 1 );
+  } );
+} );
+
+describe( 'Agent.generate() — conversation store', () => {
+  it( 'does not use store when none provided (stateless)', async () => {
+    const { Agent } = await importSut();
+    const agent = new Agent( { prompt: 'test@v1' } );
+    await agent.generate();
+    // No error, no store interaction — just prompt messages
+    expect( superGenerateImpl ).toHaveBeenCalledWith( {
+      messages: defaultMessages
+    } );
+  } );
+
+  it( 'loads prior messages from store before calling super.generate', async () => {
+    const priorMessages = [ { role: 'user', content: 'hi' }, { role: 'assistant', content: 'hello' } ];
+    const store = {
+      getMessages: vi.fn( () => priorMessages ),
+      addMessages: vi.fn()
+    };
+
+    const { Agent } = await importSut();
+    const agent = new Agent( { prompt: 'test@v1', conversationStore: store } );
+    await agent.generate( { messages: [ { role: 'user', content: 'new msg' } ] } );
+
+    expect( store.getMessages ).toHaveBeenCalled();
+    expect( superGenerateImpl ).toHaveBeenCalledWith( {
+      messages: [ ...defaultMessages, ...priorMessages, { role: 'user', content: 'new msg' } ]
+    } );
+  } );
+
+  it( 'appends user messages and response messages to store after generate()', async () => {
+    const responseMessages = [ { role: 'assistant', content: 'reply' } ];
+    superGenerateImpl.mockResolvedValue( { text: 'reply', response: { messages: responseMessages } } );
+    const store = {
+      getMessages: vi.fn( () => [] ),
+      addMessages: vi.fn()
+    };
+
+    const { Agent } = await importSut();
+    const agent = new Agent( { prompt: 'test@v1', conversationStore: store } );
+    await agent.generate( { messages: [ { role: 'user', content: 'ask' } ] } );
+
+    expect( store.addMessages ).toHaveBeenCalledWith( [
+      { role: 'user', content: 'ask' },
+      { role: 'assistant', content: 'reply' }
+    ] );
+  } );
+
+  it( 'supports async store methods', async () => {
+    const store = {
+      getMessages: vi.fn( async () => [] ),
+      addMessages: vi.fn( async () => {} )
+    };
+
+    const { Agent } = await importSut();
+    const agent = new Agent( { prompt: 'test@v1', conversationStore: store } );
+    await agent.generate();
+
+    expect( store.getMessages ).toHaveBeenCalled();
+    expect( store.addMessages ).toHaveBeenCalled();
+  } );
+} );
+
+describe( 'createMemoryConversationStore()', () => {
+  it( 'starts with empty messages', async () => {
+    const { createMemoryConversationStore } = await importSut();
+    const store = createMemoryConversationStore();
+    expect( store.getMessages() ).toEqual( [] );
+  } );
+
+  it( 'accumulates messages across addMessages calls', async () => {
+    const { createMemoryConversationStore } = await importSut();
+    const store = createMemoryConversationStore();
+    store.addMessages( [ { role: 'user', content: 'hi' } ] );
+    store.addMessages( [ { role: 'assistant', content: 'hello' } ] );
+    expect( store.getMessages() ).toEqual( [
+      { role: 'user', content: 'hi' },
+      { role: 'assistant', content: 'hello' }
+    ] );
+  } );
+} );
+
+describe( 'Agent.stream()', () => {
+  it( 'uses pre-rendered messages when no variables provided', async () => {
+    const { Agent } = await importSut();
+    const agent = new Agent( { prompt: 'test@v1' } );
+    await agent.stream();
+    expect( superStreamImpl ).toHaveBeenCalledWith( {
+      messages: defaultMessages
+    } );
+  } );
+
+  it( 'loads prior messages from store', async () => {
+    const priorMessages = [ { role: 'user', content: 'old' } ];
+    const store = {
+      getMessages: vi.fn( () => priorMessages ),
+      addMessages: vi.fn()
+    };
+
+    const { Agent } = await importSut();
+    const agent = new Agent( { prompt: 'test@v1', conversationStore: store } );
+    await agent.stream( { messages: [ { role: 'user', content: 'new' } ] } );
+
+    expect( superStreamImpl ).toHaveBeenCalledWith( {
+      messages: [ ...defaultMessages, ...priorMessages, { role: 'user', content: 'new' } ]
+    } );
+  } );
+
+  it( 'does not auto-append to store', async () => {
+    const store = {
+      getMessages: vi.fn( () => [] ),
+      addMessages: vi.fn()
+    };
+
+    const { Agent } = await importSut();
+    const agent = new Agent( { prompt: 'test@v1', conversationStore: store } );
+    await agent.stream();
+
+    expect( store.addMessages ).not.toHaveBeenCalled();
+  } );
+} );

package/src/ai_sdk.js

@@ -1,19 +1,13 @@
-import { Tracing, emitEvent } from '@outputai/core/sdk_activity_integration';
 import { loadModel, loadTools } from './ai_model.js';
 import * as AI from 'ai';
+import { stepCountIs } from 'ai';
 import { validateGenerateTextArgs, validateStreamTextArgs } from './validations.js';
 import { loadPrompt } from './prompt_loader.js';
-import { extractSourcesFromSteps } from './source_extraction.js';
-import { calculateLLMCallCost } from './cost/index.js';
-
-// Starts the LLM trace, with the start event
-const startTrace = ( name, details ) => {
-  const traceId = `${name}-${Date.now()}`;
-  Tracing.addEventStart( { kind: 'llm', name, id: traceId, details } );
-  return traceId;
-};
+import { buildSystemSkillsVar, buildLoadSkillTool, loadPromptSkills, loadColocatedSkills } from './skill.js';
+import { startTrace, endTraceWithError, traceStreamCallbacks } from './trace_utils.js';
+import { wrapInOutputResponse } from './response_utils.js';
 
-const loadAiSdkOptionsFromPrompt = prompt => {
+export const loadAiSdkOptionsFromPrompt = prompt => {
   const options = {
     model: loadModel( prompt ),
     messages: prompt.messages,
@@ -36,87 +30,66 @@ const loadAiSdkOptionsFromPrompt = prompt => {
   return options;
 };
 
-/**
- * Use an LLM model to generate text.
- *
- * Accepts additional AI SDK options (tools, maxRetries, seed, etc.) that are passed through
- * to the underlying provider. Options from the prompt file can be overridden at call time.
- *
- * @param {object} args - Generation arguments
- * @param {string} args.prompt - Prompt file name
- * @param {Record<string, string | number>} [args.variables] - Variables to interpolate
- * @param {object} [args.tools] - AI SDK tools the model can call
- * @param {'auto'|'none'|'required'|object} [args.toolChoice] - Tool selection strategy
- * @param {number} [args.maxRetries] - Max retry attempts (default: 2)
- * @param {number} [args.seed] - Seed for deterministic output
- * @param {AbortSignal} [args.abortSignal] - Signal to abort the request
- * @throws {ValidationError} If the prompt config is invalid (e.g., snake_case fields)
- * @throws {FatalError} If the prompt file is not found or template rendering fails
- * @returns {Promise<GenerateTextResult>} AI SDK response with text, toolCalls, and metadata
- */
-export async function generateText( { prompt, variables, promptDir, ...extraAiSdkOptions } ) {
-  validateGenerateTextArgs( { prompt, variables } );
-  const loadedPrompt = promptDir ? loadPrompt( prompt, variables, promptDir ) : loadPrompt( prompt, variables );
-  const traceId = startTrace( 'generateText', { prompt, variables, loadedPrompt } );
+export const hydratePromptTemplate = ( prompt, variables, promptDir, callerSkills, callerTools = {} ) => {
+  const meta = loadPrompt( prompt, variables, promptDir );
+
+  // Resolve skills: explicit frontmatter paths > colocated auto-discovery
+  const hasExplicitSkills = meta.config.skills && meta.promptFileDir;
+  const frontmatterSkills = hasExplicitSkills ?
+    loadPromptSkills( meta.config.skills, meta.promptFileDir ) :
+    [];
+  const autoSkills = !hasExplicitSkills && meta.promptFileDir ?
+    loadColocatedSkills( meta.promptFileDir ) :
+    [];
+  const resolvedSkills = [ ...frontmatterSkills, ...autoSkills, ...callerSkills ];
+
+  const tools = resolvedSkills.length > 0 ?
+    { load_skill: buildLoadSkillTool( resolvedSkills ), ...callerTools } :
+    callerTools;
+
+  const skillsMessage = resolvedSkills.length > 0 ?
+    { role: 'system', content: buildSystemSkillsVar( resolvedSkills ) } :
+    null;
+
+  if ( skillsMessage ) {
+    // Merge into existing system message to avoid provider errors with multiple system messages
+    const systemMsg = meta.messages.find( m => m.role === 'system' );
+    if ( systemMsg ) {
+      systemMsg.content = `${systemMsg.content}\n\n${skillsMessage.content}`;
+    } else {
+      meta.messages.unshift( skillsMessage );
+    }
+  }
+
+  return { loadedPrompt: meta, allVariables: variables, tools };
+};
+
+export async function generateText( { prompt, variables, promptDir, skills = [], maxSteps = 10, ...extraAiSdkOptions } ) {
+  const callerSkills = typeof skills === 'function' ? await skills( variables ) : skills;
+  const { loadedPrompt, allVariables, tools } =
+    hydratePromptTemplate( prompt, variables, promptDir, callerSkills, extraAiSdkOptions.tools );
+  const hasTools = Object.keys( tools ).length > 0;
+
+  validateGenerateTextArgs( { prompt, variables: allVariables } );
+
+  const traceId = startTrace( 'generateText', { prompt, variables: allVariables, loadedPrompt } );
   const { model: modelId } = loadedPrompt.config;
 
   try {
     const response = await AI.generateText( {
       ...loadAiSdkOptionsFromPrompt( loadedPrompt ),
-      ...extraAiSdkOptions
-    } );
-    const { text: result, totalUsage: usage, providerMetadata } = response;
-    const sourcesFromTools = extractSourcesFromSteps( response.steps );
-    const cost = await calculateLLMCallCost( { usage, modelId } );
-
-    emitEvent( 'llm:call_cost', { modelId, cost, usage } );
-    Tracing.addEventEnd( { id: traceId, details: { result, usage, cost, providerMetadata, sourcesFromTools } } );
-
-    // Creates a proxy over the response from AI SDK to add
-    // - result: a shortcut for the actual SDK response (eg. .text);
-    // - sources: a way to retrieve the computed sources;
-    // It uses proxies instead of spreading the response object because AI SDK uses getters
-    // (and they don't deconstruct properly in JS).
-    return new Proxy( response, {
-      get( target, prop, receiver ) {
-        if ( prop === 'result' ) {
-          return target.text;
-        }
-        if ( prop === 'sources' && sourcesFromTools.length > 0 ) {
-          const responseSources = Array.isArray( target[prop] ) ? target[prop] : [];
-          const byUrl = new Map( [ ...sourcesFromTools, ...responseSources ].map( s => [ s.url, s ] ) );
-          return [ ...byUrl.values() ];
-        }
-        return Reflect.get( target, prop, receiver );
-      }
+      ...extraAiSdkOptions,
+      ...( hasTools ? { tools } : {} ),
+      ...( hasTools && !extraAiSdkOptions.stopWhen ? { stopWhen: stepCountIs( maxSteps ) } : {} )
     } );
+    return wrapInOutputResponse( response, { traceId, modelId } );
   } catch ( error ) {
-    Tracing.addEventError( { id: traceId, details: error } );
+    endTraceWithError( traceId, error );
     throw error;
   }
 }
 
-/**
- * Use an LLM model to stream text generation.
- *
- * Accepts additional AI SDK options (tools, onChunk, onFinish, onError, etc.) that are passed
- * through to the underlying provider. Options from the prompt file can be overridden at call time.
- *
- * @param {object} args - Generation arguments
- * @param {string} args.prompt - Prompt file name
- * @param {Record<string, string | number>} [args.variables] - Variables to interpolate
- * @param {object} [args.tools] - AI SDK tools the model can call
- * @param {'auto'|'none'|'required'|object} [args.toolChoice] - Tool selection strategy
- * @param {number} [args.maxRetries] - Max retry attempts (default: 2)
- * @param {AbortSignal} [args.abortSignal] - Signal to abort the request
- * @param {Function} [args.onChunk] - Callback for each stream chunk
- * @param {Function} [args.onFinish] - Callback when stream finishes (called after internal tracing)
- * @param {Function} [args.onError] - Callback when stream errors (called after internal tracing)
- * @throws {ValidationError} If required arguments are missing or prompt file has invalid config
- * @throws {FatalError} If the prompt file is not found or template rendering fails
- * @returns {AIStreamTextResult} AI SDK stream result with textStream, fullStream, and metadata promises (synchronous)
- */
-export function streamText( { prompt, variables, onFinish: userOnFinish, onError: userOnError, ...restOptions } ) {
+export function streamText( { prompt, variables, onFinish, onError, ...restOptions } ) {
   validateStreamTextArgs( { prompt, variables } );
   const loadedPrompt = loadPrompt( prompt, variables );
   const traceId = startTrace( 'streamText', { prompt, variables, loadedPrompt } );
@@ -126,20 +99,10 @@ export function streamText( { prompt, variables, onFinish: userOnFinish, onError
     return AI.streamText( {
       ...loadAiSdkOptionsFromPrompt( loadedPrompt ),
       ...restOptions,
-      async onFinish( response ) {
-        const { text: result, totalUsage: usage, providerMetadata } = response;
-        const cost = await calculateLLMCallCost( { usage, modelId } );
-        emitEvent( 'llm:call_cost', { modelId, cost, usage } );
-        Tracing.addEventEnd( { id: traceId, details: { result, usage, cost, providerMetadata } } );
-        userOnFinish?.( response );
-      },
-      onError( event ) {
-        Tracing.addEventError( { id: traceId, details: event.error } );
-        userOnError?.( event );
-      }
+      ...traceStreamCallbacks( traceId, modelId, { onFinish, onError } )
     } );
   } catch ( error ) {
-    Tracing.addEventError( { id: traceId, details: error } );
+    endTraceWithError( traceId, error );
     throw error;
   }
 }