@outputai/llm 0.1.11 → 0.1.12-dev.d521efb.0

package/src/ai_sdk.spec.js

@@ -20,7 +20,9 @@ vi.mock( './ai_model.js', () => ( {
 
 const aiFns = {
   generateText: vi.fn(),
-  streamText: vi.fn()
+  streamText: vi.fn(),
+  tool: vi.fn( def => def ),
+  stepCountIs: vi.fn( n => ( { type: 'stepCount', count: n } ) )
 };
 vi.mock( 'ai', () => ( aiFns ) );
 
@@ -35,6 +37,17 @@ vi.mock( './prompt_loader.js', () => ( {
   loadPrompt: ( ...values ) => loadPromptImpl( ...values )
 } ) );
 
+const loadPromptSkillsImpl = vi.fn();
+const loadColocatedSkillsImpl = vi.fn().mockReturnValue( [] );
+vi.mock( './skill.js', async importOriginal => {
+  const original = await importOriginal();
+  return {
+    ...original,
+    loadPromptSkills: ( ...args ) => loadPromptSkillsImpl( ...args ),
+    loadColocatedSkills: ( ...args ) => loadColocatedSkillsImpl( ...args )
+  };
+} );
+
 const extractSourcesFromStepsImpl = vi.fn().mockReturnValue( [] );
 vi.mock( './source_extraction.js', () => ( {
   extractSourcesFromSteps: ( ...args ) => extractSourcesFromStepsImpl( ...args )
@@ -62,9 +75,13 @@ const cost = 'calculate cost';
 beforeEach( () => {
   emitEventSpy.mockReset();
   loadModelImpl.mockReset().mockReturnValue( 'MODEL' );
-  loadPromptImpl.mockReset().mockReturnValue( basePrompt );
+  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( [] );
 
   const defaultUsage = { inputTokens: 10, outputTokens: 5, totalTokens: 15 };
   aiFns.generateText.mockReset().mockResolvedValue( {
@@ -96,7 +113,7 @@ describe( 'ai_sdk', () => {
     const result = await generateText( { prompt: 'test_prompt@v1' } );
 
     expect( validators.validateGenerateTextArgs ).toHaveBeenCalledWith( { prompt: 'test_prompt@v1' } );
-    expect( loadPromptImpl ).toHaveBeenCalledWith( 'test_prompt@v1', undefined );
+    expect( loadPromptImpl ).toHaveBeenCalledWith( 'test_prompt@v1', undefined, undefined );
     expect( tracingSpies.addEventStart ).toHaveBeenCalledTimes( 1 );
     expect( tracingSpies.addEventEnd ).toHaveBeenCalledTimes( 1 );
     expect( tracingSpies.addEventEnd ).toHaveBeenCalledWith(
@@ -605,6 +622,190 @@ describe( 'ai_sdk', () => {
     );
   } );
 
+  it( 'generateText: loads frontmatter skills from prompt config using promptFileDir', async () => {
+    const frontmatterSkill = { name: 'fm_skill', description: 'FM', instructions: '# FM' };
+    loadPromptImpl.mockReturnValue( {
+      ...basePrompt,
+      promptFileDir: '/some/prompt/dir',
+      config: { ...basePrompt.config, skills: [ './skills/' ] }
+    } );
+    loadPromptSkillsImpl.mockReturnValue( [ frontmatterSkill ] );
+    const { generateText } = await importSut();
+    await generateText( { prompt: 'test_prompt@v1' } );
+
+    expect( loadPromptSkillsImpl ).toHaveBeenCalledWith( [ './skills/' ], '/some/prompt/dir' );
+    const callArgs = aiFns.generateText.mock.calls[0][0];
+    expect( callArgs.tools ).toHaveProperty( 'load_skill' );
+  } );
+
+  it( 'generateText: merges frontmatter skills with caller-provided skills', async () => {
+    const frontmatterSkill = { name: 'fm_skill', description: 'FM', instructions: '# FM' };
+    const callerSkill = { name: 'caller_skill', description: 'Caller', instructions: '# Caller' };
+    loadPromptImpl.mockReturnValue( {
+      ...basePrompt,
+      messages: [ ...basePrompt.messages ],
+      promptFileDir: '/some/prompt/dir',
+      config: { ...basePrompt.config, skills: [ './skills/' ] }
+    } );
+    loadPromptSkillsImpl.mockReturnValue( [ frontmatterSkill ] );
+    const { generateText } = await importSut();
+    await generateText( { prompt: 'test_prompt@v1', skills: [ callerSkill ] } );
+
+    // Skills system message appended, loadPrompt called only once
+    expect( loadPromptImpl ).toHaveBeenCalledTimes( 1 );
+    const callArgs = aiFns.generateText.mock.calls[0][0];
+    const loadSkillResult = callArgs.tools.load_skill.execute( { name: 'caller_skill' } );
+    expect( loadSkillResult ).toBe( '# Caller' );
+  } );
+
+  it( 'generateText: skips frontmatter skill loading when no config.skills', async () => {
+    const { generateText } = await importSut();
+    await generateText( { prompt: 'test_prompt@v1' } );
+
+    expect( loadPromptSkillsImpl ).not.toHaveBeenCalled();
+  } );
+
+  it( 'generateText: skips frontmatter skill loading when no promptFileDir', async () => {
+    loadPromptImpl.mockReturnValue( {
+      ...basePrompt,
+      config: { ...basePrompt.config, skills: [ './skills/' ] }
+      // no promptFileDir
+    } );
+    const { generateText } = await importSut();
+    await generateText( { prompt: 'test_prompt@v1' } );
+
+    expect( loadPromptSkillsImpl ).not.toHaveBeenCalled();
+    const callArgs = aiFns.generateText.mock.calls[0][0];
+    expect( callArgs.tools ).toBeUndefined();
+  } );
+
+  it( 'generateText: appends skills system message when skills present', async () => {
+    const frontmatterSkill = { name: 'fm_skill', description: 'FM skill', instructions: '# FM' };
+    loadPromptImpl.mockReturnValue( {
+      ...basePrompt,
+      messages: [ ...basePrompt.messages ],
+      promptFileDir: '/dir',
+      config: { ...basePrompt.config, skills: [ './skills/' ] }
+    } );
+    loadPromptSkillsImpl.mockReturnValue( [ frontmatterSkill ] );
+    const { generateText } = await importSut();
+    await generateText( { prompt: 'test_prompt@v1', variables: { topic: 'AI' } } );
+
+    // Single loadPrompt call — no two-pass render
+    expect( loadPromptImpl ).toHaveBeenCalledTimes( 1 );
+    // Skills system message inserted into messages
+    const callArgs = aiFns.generateText.mock.calls[0][0];
+    const skillsMsg = callArgs.messages.find( m => m.role === 'system' && m.content.includes( 'fm_skill' ) );
+    expect( skillsMsg ).toBeDefined();
+  } );
+
+  it( 'generateText: appends skills message and load_skill tool when skills provided', async () => {
+    const skills = [
+      { name: 'research', description: 'Research approach', instructions: '# Research\nDo research' }
+    ];
+    loadPromptImpl.mockReturnValue( { ...basePrompt, messages: [ ...basePrompt.messages ] } );
+    const { generateText } = await importSut();
+    await generateText( { prompt: 'test_prompt@v1', skills } );
+
+    expect( loadPromptImpl ).toHaveBeenCalledTimes( 1 );
+    const callArgs = aiFns.generateText.mock.calls[0][0];
+    expect( callArgs.tools ).toHaveProperty( 'load_skill' );
+    const skillsMsg = callArgs.messages.find( m => m.role === 'system' && m.content.includes( 'research' ) );
+    expect( skillsMsg ).toBeDefined();
+  } );
+
+  it( 'generateText: does not inject _system_skills or load_skill when skills is empty', async () => {
+    const { generateText } = await importSut();
+    await generateText( { prompt: 'test_prompt@v1', skills: [] } );
+
+    expect( loadPromptImpl ).toHaveBeenCalledWith( 'test_prompt@v1', undefined, undefined );
+    const callArgs = aiFns.generateText.mock.calls[0][0];
+    expect( callArgs.tools ).toBeUndefined();
+    expect( callArgs.stopWhen ).toBeUndefined();
+  } );
+
+  it( 'generateText: load_skill execute returns instructions for known skill', async () => {
+    const skills = [
+      { name: 'research', description: 'Research', instructions: '# Research\nDetailed steps' }
+    ];
+    const { generateText } = await importSut();
+    await generateText( { prompt: 'test_prompt@v1', skills } );
+
+    const { tools } = aiFns.generateText.mock.calls[0][0];
+    const result = tools.load_skill.execute( { name: 'research' } );
+    expect( result ).toBe( '# Research\nDetailed steps' );
+  } );
+
+  it( 'generateText: load_skill execute returns error for unknown skill', async () => {
+    const skills = [
+      { name: 'research', description: 'Research', instructions: '# Research' }
+    ];
+    const { generateText } = await importSut();
+    await generateText( { prompt: 'test_prompt@v1', skills } );
+
+    const { tools } = aiFns.generateText.mock.calls[0][0];
+    const result = tools.load_skill.execute( { name: 'unknown' } );
+    expect( result ).toMatch( /not found/ );
+    expect( result ).toContain( 'research' );
+  } );
+
+  it( 'generateText: sets stopWhen via maxSteps when skills present', async () => {
+    const skills = [ { name: 'skill', description: 'A skill', instructions: '# Skill' } ];
+    const { generateText } = await importSut();
+    await generateText( { prompt: 'test_prompt@v1', skills, maxSteps: 5 } );
+
+    expect( aiFns.generateText ).toHaveBeenCalledWith(
+      expect.objectContaining( { stopWhen: { type: 'stepCount', count: 5 } } )
+    );
+  } );
+
+  it( 'generateText: defaults maxSteps to 10 when skills present', async () => {
+    const skills = [ { name: 'skill', description: 'A skill', instructions: '# Skill' } ];
+    const { generateText } = await importSut();
+    await generateText( { prompt: 'test_prompt@v1', skills } );
+
+    expect( aiFns.generateText ).toHaveBeenCalledWith(
+      expect.objectContaining( { stopWhen: { type: 'stepCount', count: 10 } } )
+    );
+  } );
+
+  it( 'generateText: merges skill tools with user-provided tools', async () => {
+    const skills = [ { name: 'skill', description: 'A skill', instructions: '# Skill' } ];
+    const userTools = { calculator: { description: 'A calculator' } };
+    const { generateText } = await importSut();
+    await generateText( { prompt: 'test_prompt@v1', skills, tools: userTools } );
+
+    const { tools } = aiFns.generateText.mock.calls[0][0];
+    expect( tools ).toHaveProperty( 'load_skill' );
+    expect( tools ).toHaveProperty( 'calculator' );
+  } );
+
+  it( 'generateText: calls skill function with variables and uses resolved skills', async () => {
+    const resolvedSkill = { name: 'dynamic', description: 'Dynamic skill', instructions: '# Dynamic' };
+    const skillsFn = vi.fn().mockResolvedValue( [ resolvedSkill ] );
+    const vars = { topic: 'AI' };
+    const { generateText } = await importSut();
+    await generateText( { prompt: 'test_prompt@v1', variables: vars, skills: skillsFn } );
+
+    expect( skillsFn ).toHaveBeenCalledWith( vars );
+    expect( loadPromptImpl ).toHaveBeenCalledTimes( 1 );
+    const callArgs = aiFns.generateText.mock.calls[0][0];
+    expect( callArgs.tools ).toHaveProperty( 'load_skill' );
+    const skillsMsg = callArgs.messages.find( m => m.role === 'system' && m.content.includes( 'dynamic' ) );
+    expect( skillsMsg ).toBeDefined();
+  } );
+
+  it( 'generateText: preserves caller stopWhen when skills present', async () => {
+    const skills = [ { name: 'skill', description: 'A skill', instructions: '# Skill' } ];
+    const customStop = { type: 'custom' };
+    const { generateText } = await importSut();
+    await generateText( { prompt: 'test_prompt@v1', skills, stopWhen: customStop } );
+
+    expect( aiFns.generateText ).toHaveBeenCalledWith(
+      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' }

package/src/index.d.ts

@@ -279,7 +279,13 @@ export function generateText<
   args: {
     prompt: string,
     variables?: Record<string, string | number | boolean>,
-    promptDir?: string
+    promptDir?: string,
+    /**
+     * Skill packages to provide to the LLM. Injects `{{ _system_skills }}` and adds the `load_skill` tool.
+     * Can be a static array or a function that receives the resolved variables and returns skills.
+     */
+    skills?: import( './skill.js' ).Skill[] |
+      ( ( variables?: Record<string, string | number | boolean> ) => import( './skill.js' ).Skill[] | Promise<import( './skill.js' ).Skill[]> )
   } & GenerateTextAiSdkOptions<Tools, Output>
 ): Promise<GenerateTextResult<Tools, Output>>;
 
@@ -307,3 +313,90 @@ export function streamText<
     variables?: Record<string, string | number | boolean>
   } & StreamTextAiSdkOptions<Tools, Output>
 ): AIStreamTextResult<Tools, Output>;
+
+export { skill } from './skill.js';
+export type { Skill, SkillsArg } from './skill.js';
+
+/** Pluggable conversation store for multi-turn Agent interactions. */
+export interface ConversationStore {
+  getMessages(): import( 'ai' ).ModelMessage[] | Promise<import( 'ai' ).ModelMessage[]>;
+  addMessages( messages: import( 'ai' ).ModelMessage[] ): void | Promise<void>;
+}
+
+/** Create an in-memory conversation store backed by a closure array. */
+export function createMemoryConversationStore(): ConversationStore;
+
+/**
+ * Agent extends AI SDK's ToolLoopAgent with Output.ai prompt file rendering
+ * and the skill system.
+ *
+ * @example Workflow step — variables per call, stateless
+ * ```ts
+ * const reviewer = new Agent({
+ *   prompt: 'reviewer@v1',
+ *   output: Output.object({ schema: z.object({ summary: z.string() }) }),
+ *   maxSteps: 5
+ * });
+ * const result = await reviewer.generate();
+ * ```
+ *
+ * @example Interactive — fixed setup, conversation history
+ * ```ts
+ * const chatbot = new Agent({
+ *   prompt: 'chatbot@v1',
+ *   conversationStore: createMemoryConversationStore()
+ * });
+ * const r1 = await chatbot.generate({ messages: [{ role: 'user', content: 'Hello' }] });
+ * ```
+ */
+export declare class Agent extends import( 'ai' ).ToolLoopAgent {
+  constructor( params: {
+    /** Prompt file name (e.g. 'my_agent@v1') */
+    prompt: string;
+    /** Override the stack-resolved prompt directory */
+    promptDir?: string;
+    /** Variables to render the prompt template at construction time */
+    variables?: Record<string, unknown>;
+    /** Static skill packages made available to the LLM */
+    skills?: import( './skill.js' ).Skill[];
+    /** AI SDK tools available during the reasoning loop */
+    tools?: ToolSet;
+    /** Maximum tool-loop iterations when stopWhen is not specified (default: 10) */
+    maxSteps?: number;
+    /** Custom stop condition(s) — overrides maxSteps */
+    stopWhen?: import( 'ai' ).StopCondition | import( 'ai' ).StopCondition[];
+    /** Structured output specification */
+    output?: import( 'ai' ).Output<unknown, unknown>;
+    /** Pluggable conversation store — opt-in, stateless by default */
+    conversationStore?: ConversationStore;
+    /** Callback after each step */
+    onStepFinish?: import( 'ai' ).GenerateTextOnStepFinishCallback<ToolSet>;
+    /** Customize each step before execution */
+    prepareStep?: import( 'ai' ).PrepareStepFunction<ToolSet>;
+    /** Generation temperature (overrides prompt file value) */
+    temperature?: number;
+    /** Top-p sampling */
+    topP?: number;
+    /** Top-k sampling */
+    topK?: number;
+    /** Random seed for deterministic output */
+    seed?: number;
+    /** Maximum retry attempts (default: 2) */
+    maxRetries?: number;
+  } );
+
+  /** Run the agent and return when complete. */
+  generate( options?: {
+    messages?: import( 'ai' ).ModelMessage[];
+    abortSignal?: AbortSignal;
+    onStepFinish?: import( 'ai' ).GenerateTextOnStepFinishCallback<ToolSet>;
+  } ): Promise<import( 'ai' ).GenerateTextResult<ToolSet, import( 'ai' ).Output<unknown, unknown>>>;
+
+  /** Stream the agent's response. */
+  stream( options?: {
+    messages?: import( 'ai' ).ModelMessage[];
+    abortSignal?: AbortSignal;
+    onStepFinish?: import( 'ai' ).StreamTextOnStepFinishCallback<ToolSet>;
+    experimental_transform?: import( 'ai' ).StreamTextTransform<ToolSet> | import( 'ai' ).StreamTextTransform<ToolSet>[];
+  } ): Promise<import( 'ai' ).StreamTextResult<ToolSet, import( 'ai' ).Output<unknown, unknown>>>;
+};

package/src/index.js

@@ -1,4 +1,5 @@
 export { generateText, streamText } from './ai_sdk.js';
+export { Agent, createMemoryConversationStore, skill } from './agent.js';
 export { loadPrompt } from './prompt_loader.js';
 export { registerProvider, getRegisteredProviders } from './ai_model.js';
 export { tavilySearch, tavilyExtract, tavilyCrawl, tavilyMap } from '@tavily/ai-sdk';

package/src/load_content.js

@@ -19,25 +19,48 @@ const loadFile = path => {
   }
 };
 
-/**
- * Recursively search for a file by its name and load its content.
- *
- * @param {string} name - Name of the file load its content
- * @param {string} [dir] - The directory to search for the file, defaults to invocation directory
- * @returns {string | null} - File content or null if not found
- */
-export const loadContent = ( name, dir = resolveInvocationDir() ) => {
+const findContent = ( name, dir ) => {
   for ( const entry of scanDir( dir ) ) {
     if ( entry.name === name ) {
-      return loadFile( join( dir, entry.name ) );
+      return { dir, content: loadFile( join( dir, entry.name ) ) };
     }
-
     if ( entry.isDirectory() && !entry.isSymbolicLink() ) {
-      const content = loadContent( name, join( dir, entry.name ) );
-      if ( content ) {
-        return content;
+      const result = findContent( name, join( dir, entry.name ) );
+      if ( result ) {
+        return result;
       }
     }
   }
   return null;
 };
+
+/**
+ * Recursively search for a file by its name and load its content.
+ *
+ * @param {string} name - Name of the file load its content
+ * @param {string} [dir] - The directory to search for the file, defaults to invocation directory
+ * @returns {string | null} - File content or null if not found
+ */
+export const loadContent = ( name, dir = resolveInvocationDir() ) =>
+  findContent( name, dir )?.content ?? null;
+
+/**
+ * Recursively search for a file by name and return the directory containing it.
+ *
+ * @param {string} name - File name to find
+ * @param {string} [dir] - Directory to search, defaults to invocation directory
+ * @returns {string | null} - Directory path containing the file, or null if not found
+ */
+export const findContentDir = ( name, dir = resolveInvocationDir() ) =>
+  findContent( name, dir )?.dir ?? null;
+
+/**
+ * Recursively search for a file by name and return both its content and containing directory.
+ * More efficient than calling loadContent + findContentDir separately (single scan).
+ *
+ * @param {string} name - File name to find
+ * @param {string} [dir] - Directory to search, defaults to invocation directory
+ * @returns {{ content: string, dir: string } | null}
+ */
+export const loadContentWithDir = ( name, dir = resolveInvocationDir() ) =>
+  findContent( name, dir ) ?? null;

package/src/message_utils.js

@@ -0,0 +1,3 @@
+export const ROLE = { SYSTEM: 'system', USER: 'user', ASSISTANT: 'assistant', TOOL: 'tool' };
+export const isRole = role => msg => msg.role === role;
+export const getContent = msg => msg.content;

package/src/prompt_loader.js

@@ -1,6 +1,6 @@
 import { parsePrompt } from './parser.js';
 import { Liquid } from 'liquidjs';
-import { loadContent } from './load_content.js';
+import { loadContentWithDir } from './load_content.js';
 import { validatePrompt } from './prompt_validations.js';
 import { FatalError } from '@outputai/core';
 
@@ -20,26 +20,22 @@ const renderPrompt = ( name, content, values ) => {
  * @param {string} name - Name of the prompt file (without .prompt extension)
  * @param {Record<string, string | number | boolean>} [values] - Variables to interpolate
  * @param {string} [dir] - Directory to search for the prompt file (defaults to stack-resolved invocation dir)
- * @returns {Prompt} Loaded and rendered prompt object
+ * @returns {Prompt} Loaded and rendered prompt object, including promptFileDir
  */
 export const loadPrompt = ( name, values = {}, dir ) => {
-  const promptContent = dir ? loadContent( `${name}.prompt`, dir ) : loadContent( `${name}.prompt` );
-  if ( !promptContent ) {
+  const found = loadContentWithDir( `${name}.prompt`, dir );
+  if ( !found ) {
     throw new FatalError( `Prompt ${name} not found.` );
   }
 
-  const renderedContent = renderPrompt( name, promptContent, values );
+  const renderedContent = renderPrompt( name, found.content, values );
 
   const { config, messages } = parsePrompt( renderedContent );
 
-  const prompt = {
-    name,
-    config,
-    messages
-  };
+  const prompt = { name, config, messages };
 
   validatePrompt( prompt );
 
-  return prompt;
+  return { ...prompt, promptFileDir: found.dir };
 };
 

package/src/prompt_loader.spec.js

@@ -3,14 +3,14 @@ import { loadPrompt } from './prompt_loader.js';
 
 // Mock dependencies that perform I/O or validation
 vi.mock( './load_content.js', () => ( {
-  loadContent: vi.fn()
+  loadContentWithDir: vi.fn()
 } ) );
 
 vi.mock( './prompt_validations.js', () => ( {
   validatePrompt: vi.fn()
 } ) );
 
-import { loadContent } from './load_content.js';
+import { loadContentWithDir } from './load_content.js';
 import { validatePrompt } from './prompt_validations.js';
 
 describe( 'loadPrompt', () => {
@@ -25,7 +25,7 @@ model: claude-3-5-sonnet-20241022
 ---
 <user>Hello {{ name }}!</user>`;
 
-    loadContent.mockReturnValue( promptContent );
+    loadContentWithDir.mockReturnValue( { content: promptContent, dir: '/mock/dir' } );
 
     const result = loadPrompt( 'test', { name: 'World' } );
 
@@ -33,11 +33,11 @@ model: claude-3-5-sonnet-20241022
     expect( result.config ).toEqual( { provider: 'anthropic', model: 'claude-3-5-sonnet-20241022' } );
     expect( result.messages ).toHaveLength( 1 );
     expect( result.messages[0].content ).toBe( 'Hello World!' );
-    expect( validatePrompt ).toHaveBeenCalledWith( result );
+    expect( validatePrompt ).toHaveBeenCalledWith( expect.objectContaining( { name: 'test' } ) );
   } );
 
   it( 'throws error when prompt file not found', () => {
-    loadContent.mockReturnValue( null );
+    loadContentWithDir.mockReturnValue( null );
 
     expect( () => {
       loadPrompt( 'nonexistent' );
@@ -60,7 +60,7 @@ temperature: 0.7
 
 <user>Hello</user>`;
 
-    loadContent.mockReturnValue( promptContent );
+    loadContentWithDir.mockReturnValue( { content: promptContent, dir: '/mock/dir' } );
 
     const result = loadPrompt( 'test', {
       provider_name: 'anthropic',
@@ -80,7 +80,7 @@ model: {{ model }}
 
 <user>Tell me about {{ topic }}</user>`;
 
-    loadContent.mockReturnValue( promptContent );
+    loadContentWithDir.mockReturnValue( { content: promptContent, dir: '/mock/dir' } );
 
     const result = loadPrompt( 'test', {
       provider: 'openai',
@@ -101,7 +101,7 @@ model: claude-3-5-sonnet-20241022
 
 <user>Hello</user>`;
 
-    loadContent.mockReturnValue( promptContent );
+    loadContentWithDir.mockReturnValue( { content: promptContent, dir: '/mock/dir' } );
 
     const result = loadPrompt( 'test', {} );
 
@@ -119,7 +119,7 @@ temperature: 0.7
 
 <user>Hello</user>`;
 
-    loadContent.mockReturnValue( promptContent );
+    loadContentWithDir.mockReturnValue( { content: promptContent, dir: '/mock/dir' } );
 
     const result = loadPrompt( 'test', {
       base_model: 'claude-sonnet',
@@ -139,7 +139,7 @@ temperature: 0.7
 
 <user>Hello</user>`;
 
-    loadContent.mockReturnValue( promptContent );
+    loadContentWithDir.mockReturnValue( { content: promptContent, dir: '/mock/dir' } );
 
     const result = loadPrompt( 'test', {} );
 
@@ -154,7 +154,7 @@ model: claude-3-5-sonnet-20241022
 
 <user>{% if debug %}Debug mode enabled{% else %}Debug mode disabled{% endif %}</user>`;
 
-    loadContent.mockReturnValue( promptContent );
+    loadContentWithDir.mockReturnValue( { content: promptContent, dir: '/mock/dir' } );
 
     const result = loadPrompt( 'test', { debug: true } );
 
@@ -169,7 +169,7 @@ model: claude-3-5-sonnet-20241022
 
 <user>{% if enabled %}Feature enabled{% else %}Feature disabled{% endif %}</user>`;
 
-    loadContent.mockReturnValue( promptContent );
+    loadContentWithDir.mockReturnValue( { content: promptContent, dir: '/mock/dir' } );
 
     const result = loadPrompt( 'test', { enabled: false } );
 

package/src/prompt_loader_validation.spec.js

@@ -2,10 +2,10 @@ import { describe, it, expect, vi, beforeEach } from 'vitest';
 import { loadPrompt } from './prompt_loader.js';
 
 vi.mock( './load_content.js', () => ( {
-  loadContent: vi.fn()
+  loadContentWithDir: vi.fn()
 } ) );
 
-import { loadContent } from './load_content.js';
+import { loadContentWithDir } from './load_content.js';
 
 describe( 'loadPrompt - validation with real schema', () => {
   beforeEach( () => {
@@ -26,7 +26,7 @@ providerOptions:
 
 <user>Hello</user>`;
 
-    loadContent.mockReturnValue( promptContent );
+    loadContentWithDir.mockReturnValue( { content: promptContent, dir: '/mock/dir' } );
 
     const result = loadPrompt( 'test', {} );
 
@@ -44,7 +44,7 @@ max_tokens: 64000
 
 <user>Hello</user>`;
 
-    loadContent.mockReturnValue( promptContent );
+    loadContentWithDir.mockReturnValue( { content: promptContent, dir: '/mock/dir' } );
 
     // Config uses passthrough, so max_tokens is accepted (though ignored by SDK)
     expect( () => {
@@ -64,7 +64,7 @@ options:
 
 <user>Hello</user>`;
 
-    loadContent.mockReturnValue( promptContent );
+    loadContentWithDir.mockReturnValue( { content: promptContent, dir: '/mock/dir' } );
 
     // Config uses passthrough, so 'options' passes through (though not used)
     expect( () => {
@@ -84,7 +84,7 @@ providerOptions:
 
 <user>Hello</user>`;
 
-    loadContent.mockReturnValue( promptContent );
+    loadContentWithDir.mockReturnValue( { content: promptContent, dir: '/mock/dir' } );
 
     // budget_tokens is silently stripped from thinking (unknown field), not rejected
     expect( () => {

package/src/response_utils.js

@@ -0,0 +1,21 @@
+import { extractSourcesFromSteps } from './source_extraction.js';
+import { endTraceWithSuccess } from './trace_utils.js';
+
+export const wrapInOutputResponse = async ( response, { traceId, modelId } ) => {
+  const sourcesFromTools = extractSourcesFromSteps( response.steps );
+  await endTraceWithSuccess( traceId, modelId, response, { sourcesFromTools } );
+
+  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 );
+    }
+  } );
+};