@outputai/llm 0.6.1-dev.aab2335.0 → 0.6.1-next.0d08ff5.0

package/src/utils/response_wrappers.spec.js

@@ -1,18 +1,22 @@
 import { describe, it, expect, vi, beforeEach } from 'vitest';
+import { generateImage } from 'ai';
+import { MockImageModelV3 } from 'ai/test';
+import textResponseFixture from './__fixtures__/text_response.json' with { type: 'json' };
+import streamResponseFixture from './__fixtures__/stream_response.json' with { type: 'json' };
+import imageResponseFixture from './__fixtures__/image_response.json' with { type: 'json' };
 
 const mocks = vi.hoisted( () => ( {
+  combineSources: vi.fn(),
   extractSourcesFromSteps: vi.fn(),
   calculateLLMCallCost: vi.fn(),
-  endTraceWithSuccess: vi.fn()
+  endTraceWithSuccess: vi.fn(),
+  calculateBase64FileSize: vi.fn()
 } ) );
 
-vi.mock( './source_extraction.js', async importOriginal => {
-  const actual = await importOriginal();
-  return {
-    ...actual,
-    extractSourcesFromSteps: mocks.extractSourcesFromSteps
-  };
-} );
+vi.mock( './source_extraction.js', () => ( {
+  combineSources: mocks.combineSources,
+  extractSourcesFromSteps: mocks.extractSourcesFromSteps
+} ) );
 
 vi.mock( '../cost/index.js', () => ( {
   calculateLLMCallCost: mocks.calculateLLMCallCost
@@ -22,7 +26,29 @@ vi.mock( './trace.js', () => ( {
   endTraceWithSuccess: mocks.endTraceWithSuccess
 } ) );
 
-import { wrapTextResponse, wrapStreamOnFinishResponse } from './response_wrappers.js';
+vi.mock( './image.js', () => ( {
+  calculateBase64FileSize: mocks.calculateBase64FileSize
+} ) );
+
+import { wrapTextResponse, wrapStreamOnFinishResponse, wrapImageResponse } from './response_wrappers.js';
+
+const clone = value => structuredClone( value );
+
+const makeAiSdkImageResponse = () => generateImage( {
+  model: new MockImageModelV3( {
+    doGenerate: async () => ( {
+      images: imageResponseFixture.images.map( image => Buffer.from( image.base64Data, 'base64' ) ),
+      warnings: imageResponseFixture.warnings,
+      response: {
+        ...imageResponseFixture.responses[0],
+        timestamp: new Date( imageResponseFixture.responses[0].timestamp )
+      },
+      providerMetadata: imageResponseFixture.providerMetadata,
+      usage: imageResponseFixture.usage
+    } )
+  } ),
+  prompt: 'fixture prompt'
+} );
 
 describe( 'wrapTextResponse', () => {
   const traceId = 'trace-1';
@@ -33,53 +59,47 @@ describe( 'wrapTextResponse', () => {
     vi.clearAllMocks();
     mocks.extractSourcesFromSteps.mockReturnValue( [] );
     mocks.calculateLLMCallCost.mockResolvedValue( mockCost );
+    mocks.combineSources.mockReturnValue( [] );
   } );
 
-  it( 'proxies result to text and attaches cost', async () => {
-    const response = {
-      text: 'hello',
-      totalUsage: { inputTokens: 10, outputTokens: 5 },
-      steps: [],
-      sources: []
-    };
+  it( 'uses a text response fixture to calculate cost, end trace, and attach cost', async () => {
+    const response = clone( textResponseFixture );
 
     const wrapped = await wrapTextResponse( { traceId, modelId, response } );
 
-    expect( wrapped.result ).toBe( 'hello' );
-    expect( wrapped.text ).toBe( 'hello' );
+    expect( wrapped.result ).toBe( response.text );
     expect( wrapped.cost ).toEqual( mockCost );
     expect( mocks.calculateLLMCallCost ).toHaveBeenCalledWith( {
       usage: response.totalUsage,
       modelId
     } );
+    expect( mocks.extractSourcesFromSteps ).toHaveBeenCalledWith( response.steps );
     expect( mocks.endTraceWithSuccess ).toHaveBeenCalledWith( {
       traceId,
-      modelId,
-      response,
+      usage: response.totalUsage,
       cost: mockCost,
+      result: response.text,
+      providerMetadata: response.providerMetadata,
       sourcesFromTools: []
     } );
   } );
 
-  it( 'leaves sources unchanged when no tool-extracted sources (same as raw response)', async () => {
+  it( 'leaves sources unchanged when no tool-extracted sources are found', async () => {
+    const response = clone( streamResponseFixture );
     const nativeSources = [
       { type: 'source', sourceType: 'url', id: 'n1', url: 'https://native.test', title: 'Native' }
     ];
+    response.sources = nativeSources;
     mocks.extractSourcesFromSteps.mockReturnValue( [] );
 
-    const response = {
-      text: 'x',
-      totalUsage: {},
-      steps: [],
-      sources: nativeSources
-    };
-
     const wrapped = await wrapTextResponse( { traceId, modelId, response } );
 
     expect( wrapped.sources ).toBe( nativeSources );
+    expect( mocks.combineSources ).not.toHaveBeenCalled();
   } );
 
-  it( 'merges tool-extracted sources with response sources by url when tool sources exist', async () => {
+  it( 'delegates source merging when tool-extracted sources exist', async () => {
+    const response = clone( streamResponseFixture );
     const toolSource = {
       type: 'source',
       sourceType: 'url',
@@ -94,40 +114,18 @@ describe( 'wrapTextResponse', () => {
       url: 'https://example.com/b',
       title: 'B'
     };
+    const mergedSources = [ toolSource, responseSource ];
+    response.sources = [ responseSource ];
     mocks.extractSourcesFromSteps.mockReturnValue( [ toolSource ] );
-
-    const response = {
-      text: 'x',
-      totalUsage: {},
-      steps: [ {} ],
-      sources: [ responseSource ]
-    };
-
-    const wrapped = await wrapTextResponse( { traceId, modelId, response } );
-
-    expect( wrapped.sources ).toEqual( expect.arrayContaining( [ toolSource, responseSource ] ) );
-    expect( wrapped.sources ).toHaveLength( 2 );
-  } );
-
-  it( 'when tool sources overlap urls, later entry in merge wins', async () => {
-    const url = 'https://example.com/same';
-    mocks.extractSourcesFromSteps.mockReturnValue( [
-      { type: 'source', sourceType: 'url', id: '1', url, title: 'from-tool' }
-    ] );
-
-    const response = {
-      text: 'x',
-      totalUsage: {},
-      steps: [],
-      sources: [
-        { type: 'source', sourceType: 'url', id: '2', url, title: 'from-response' }
-      ]
-    };
+    mocks.combineSources.mockReturnValue( mergedSources );
 
     const wrapped = await wrapTextResponse( { traceId, modelId, response } );
 
-    expect( wrapped.sources ).toHaveLength( 1 );
-    expect( wrapped.sources[0].title ).toBe( 'from-response' );
+    expect( wrapped.sources ).toBe( mergedSources );
+    expect( mocks.combineSources ).toHaveBeenCalledWith( {
+      sourcesFromTools: [ toolSource ],
+      sourcesFromResponse: [ responseSource ]
+    } );
   } );
 } );
 
@@ -139,15 +137,12 @@ describe( 'wrapStreamOnFinishResponse', () => {
   beforeEach( () => {
     vi.clearAllMocks();
     mocks.calculateLLMCallCost.mockResolvedValue( mockCost );
+    mocks.extractSourcesFromSteps.mockReturnValue( [] );
   } );
 
-  it( 'onFinish finishes trace, proxies result and cost for user callback, and forwards other props', async () => {
+  it( 'uses the stream response fixture to finish trace and call the user callback with a proxied response', async () => {
     const userOnFinish = vi.fn();
-    const response = {
-      text: 'done',
-      totalUsage: { inputTokens: 1 },
-      finishReason: 'stop'
-    };
+    const response = clone( streamResponseFixture );
 
     const callbacks = wrapStreamOnFinishResponse( {
       traceId,
@@ -159,14 +154,79 @@ describe( 'wrapStreamOnFinishResponse', () => {
 
     expect( mocks.endTraceWithSuccess ).toHaveBeenCalledWith( {
       traceId,
-      modelId,
-      response,
-      cost: mockCost
+      usage: response.totalUsage,
+      cost: mockCost,
+      result: response.text,
+      providerMetadata: response.providerMetadata,
+      sourcesFromTools: []
     } );
     expect( userOnFinish ).toHaveBeenCalledTimes( 1 );
     const proxied = userOnFinish.mock.calls[0][0];
-    expect( proxied.result ).toBe( 'done' );
+    expect( proxied.result ).toBe( response.text );
     expect( proxied.cost ).toEqual( mockCost );
-    expect( proxied.finishReason ).toBe( 'stop' );
+    expect( proxied.finishReason ).toBe( response.finishReason );
+    expect( mocks.extractSourcesFromSteps ).toHaveBeenCalledWith( response.steps );
+  } );
+
+  it( 'finishes trace even when no user onFinish callback is provided', async () => {
+    const response = clone( streamResponseFixture );
+
+    const callbacks = wrapStreamOnFinishResponse( {
+      traceId,
+      modelId
+    } );
+
+    await callbacks.onFinish( response );
+
+    expect( mocks.endTraceWithSuccess ).toHaveBeenCalledWith( {
+      traceId,
+      usage: response.totalUsage,
+      cost: mockCost,
+      result: response.text,
+      providerMetadata: response.providerMetadata,
+      sourcesFromTools: []
+    } );
+    expect( mocks.calculateLLMCallCost ).toHaveBeenCalledWith( {
+      usage: response.totalUsage,
+      modelId
+    } );
+  } );
+} );
+
+describe( 'wrapImageResponse', () => {
+  const traceId = 'image-trace';
+  const modelId = 'image-model';
+  const mockCost = { total: 0.003, components: [] };
+
+  beforeEach( () => {
+    vi.clearAllMocks();
+    mocks.calculateLLMCallCost.mockResolvedValue( mockCost );
+    mocks.calculateBase64FileSize.mockReturnValue( 1234 );
+  } );
+
+  it( 'uses an image response fixture to trace image metadata and attach cost', async () => {
+    const response = await makeAiSdkImageResponse();
+
+    const wrapped = await wrapImageResponse( { traceId, modelId, response } );
+
+    expect( wrapped.result ).toBe( response.images[0] );
+    expect( wrapped.cost ).toEqual( mockCost );
+    expect( mocks.calculateLLMCallCost ).toHaveBeenCalledWith( {
+      usage: response.usage,
+      modelId
+    } );
+    expect( mocks.calculateBase64FileSize ).toHaveBeenCalledWith( response.images[0].base64 );
+    expect( mocks.endTraceWithSuccess ).toHaveBeenCalledWith( {
+      traceId,
+      usage: response.usage,
+      cost: mockCost,
+      result: [
+        {
+          size: 1234,
+          mediaType: response.images[0].mediaType
+        }
+      ],
+      providerMetadata: response.providerMetadata
+    } );
   } );
 } );

package/src/utils/source_extraction.js

@@ -3,9 +3,16 @@ import { createHash } from 'node:crypto';
 /**
  * Checks whether a tool result looks like a search response (has a `results` array with `url` strings).
  */
-const isSearchResult = result => !!result?.results?.[0]?.url;
+const isSearchResult = v => Array.isArray( v?.results ) && v.results.some( r => typeof r?.url === 'string' );
 
-const toSource = ( { url, title } ) => ( {
+/**
+ * Builds the final source shape
+ * @param {object} args
+ * @param {string} args.url
+ * @param {string} args.title
+ * @returns {object} final object
+ */
+const buildSource = ( { url, title } ) => ( {
   type: 'source',
   sourceType: 'url',
   id: createHash( 'sha256' ).update( url ).digest( 'hex' ).slice( 0, 16 ),
@@ -19,34 +26,17 @@ const toSource = ( { url, title } ) => ( {
  * 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.
  *
- * Best-effort: returns empty array on any error rather than throwing.
- *
  * @param {Array} steps - AI SDK response steps (response.steps)
  * @returns {Array<{ type: string, sourceType: string, id: string, url: string, title: string }>}
  */
-export function extractSourcesFromSteps( steps ) {
-  try {
-    if ( !Array.isArray( steps ) || steps.length === 0 ) {
-      return [];
-    }
-
-    const seen = new Set();
-    return steps
-      .flatMap( step => Array.isArray( step.toolResults ) ? step.toolResults : [] )
-      .flatMap( toolResult => isSearchResult( toolResult.output ) ? toolResult.output.results : [] )
-      .filter( item => {
-        if ( !item.url || seen.has( item.url ) ) {
-          return false;
-        }
-        seen.add( item.url );
-        return true;
-      } )
-      .map( toSource );
-  } catch ( error ) {
-    console.warn( '[output-llm] source extraction failed, returning empty sources', error );
-    return [];
-  }
-}
+export const extractSourcesFromSteps = steps =>
+  ( Array.isArray( steps ) ? steps : [] )
+    .flatMap( step => Array.isArray( step?.toolResults ) ? step.toolResults : [] )
+    .flatMap( toolResult => isSearchResult( toolResult?.output ) ? toolResult.output.results : [] )
+    .filter( item => typeof item?.url === 'string' && item.url.length > 0 ) // Ignore non string or empty string urls
+    .reduce( ( map, v ) => map.has( v.url ) ? map : map.set( v.url, v ), new Map() ) // deduplicate, by keeping first entry
+    .values().toArray()
+    .map( v => buildSource( v ) );
 
 /**
  * Merge sources used tools and sources from AI SDK response into a single list

package/src/utils/trace.js

@@ -10,11 +10,10 @@ export const endTraceWithError = ( { traceId, error } ) => {
   Tracing.addEventError( { id: traceId, details: error } );
 };
 
-export const endTraceWithSuccess = ( { traceId, modelId, response, cost, ...extra } ) => {
-  const { totalUsage: usage, text: result, providerMetadata } = response;
+export const endTraceWithSuccess = ( { traceId, result, cost, ...extra } ) => {
   if ( cost ) {
     Tracing.addEventAttribute( { eventId: traceId, attribute: cost } );
     emitEvent( 'cost:llm:request', cost );
   }
-  Tracing.addEventEnd( { id: traceId, details: { result, usage, providerMetadata, ...extra } } );
+  Tracing.addEventEnd( { id: traceId, details: { result, ...extra } } );
 };

package/src/utils/trace.spec.js

@@ -51,19 +51,16 @@ describe( 'trace utils', () => {
   } );
 
   describe( 'endTraceWithSuccess', () => {
-    it( 'adds cost attribute, emits cost attribute, and ends the trace with response fields and extra details', () => {
+    it( 'adds cost attribute, emits cost event, and ends the trace with normalized details', () => {
       const cost = { type: 'llm:usage', modelId: 'my-model', total: 0.01, usage: [] };
       const usage = { inputTokens: 2, outputTokens: 3 };
-      const response = {
-        text: 'hello',
-        totalUsage: usage,
-        providerMetadata: { provider: 'x' }
-      };
 
       endTraceWithSuccess( {
         traceId: 'trace-a',
         modelId: 'my-model',
-        response,
+        result: 'hello',
+        usage,
+        providerMetadata: { provider: 'x' },
         cost,
         sourcesFromTools: [ { url: 'https://u.test', title: '' } ]
       } );
@@ -77,6 +74,7 @@ describe( 'trace utils', () => {
         id: 'trace-a',
         details: {
           result: 'hello',
+          modelId: 'my-model',
           usage,
           providerMetadata: { provider: 'x' },
           sourcesFromTools: [ { url: 'https://u.test', title: '' } ]
@@ -86,16 +84,13 @@ describe( 'trace utils', () => {
 
     it( 'does not emit or add an attribute when cost is missing', () => {
       const usage = { inputTokens: 2, outputTokens: 3 };
-      const response = {
-        text: 'hello',
-        totalUsage: usage,
-        providerMetadata: { provider: 'x' }
-      };
 
       endTraceWithSuccess( {
         traceId: 'trace-no-cost',
         modelId: 'my-model',
-        response
+        result: 'hello',
+        usage,
+        providerMetadata: { provider: 'x' }
       } );
 
       expect( tracing.addEventAttribute ).not.toHaveBeenCalled();
@@ -104,6 +99,7 @@ describe( 'trace utils', () => {
         id: 'trace-no-cost',
         details: {
           result: 'hello',
+          modelId: 'my-model',
           usage,
           providerMetadata: { provider: 'x' }
         }

package/src/validations.js

@@ -1,8 +1,56 @@
+import { Buffer } from 'node:buffer';
 import { ValidationError, z } from '@outputai/core';
 
+const skillArgSchema = z.object( {
+  name: z.string().min( 1 ),
+  description: z.string().optional(),
+  instructions: z.string().min( 1 )
+} ).strict();
+
 const generateTextArgsSchema = z.object( {
-  prompt: z.string(),
-  variables: z.any().optional()
+  prompt: z.string().min( 1 ),
+  variables: z.any().optional(),
+  promptDir: z.string().min( 1 ).optional(),
+  skills: z.union( [ z.array( skillArgSchema ), z.function() ] ).optional(),
+  maxSteps: z.number().int().positive().optional()
+} );
+
+const base64StringSchema = z.string()
+  .min( 1 )
+  .regex(
+    /^(?:[A-Za-z0-9+/]{4})*(?:[A-Za-z0-9+/]{2}(?:==)?|[A-Za-z0-9+/]{3}=?)?$/,
+    'Image strings must be raw base64 data.'
+  );
+
+const imageDataSchema = z.union( [
+  z.instanceof( Buffer ),
+  z.instanceof( Uint8Array ),
+  z.instanceof( ArrayBuffer ),
+  base64StringSchema
+] );
+
+const imageInputSchema = z.union( [
+  imageDataSchema,
+  z.object( {
+    data: imageDataSchema,
+    mediaType: z.string().min( 1 ).optional()
+  } ).strict()
+] );
+
+const generateImageArgsSchema = z.object( {
+  prompt: z.string().min( 1 ),
+  variables: z.any().optional(),
+  promptDir: z.string().min( 1 ).optional(),
+  images: z.array( imageInputSchema ).min( 1 ).optional(),
+  mask: imageInputSchema.optional()
+} ).superRefine( ( args, ctx ) => {
+  if ( args.mask && !args.images ) {
+    ctx.addIssue( {
+      code: 'custom',
+      path: [ 'mask' ],
+      message: 'mask requires images.'
+    } );
+  }
 } );
 
 function validateSchema( schema, input, errorPrefix ) {
@@ -19,3 +67,7 @@ export function validateGenerateTextArgs( args ) {
 export function validateStreamTextArgs( args ) {
   validateSchema( generateTextArgsSchema, args, 'Invalid streamText() arguments' );
 }
+
+export function validateGenerateImageArgs( args ) {
+  validateSchema( generateImageArgsSchema, args, 'Invalid generateImage() arguments' );
+}

package/src/validations.spec.js

@@ -0,0 +1,166 @@
+import { describe, it, expect } from 'vitest';
+import { ValidationError } from '@outputai/core';
+import { validateGenerateTextArgs, validateStreamTextArgs, validateGenerateImageArgs } from './validations.js';
+
+describe( 'validateGenerateTextArgs', () => {
+  it( 'accepts a prompt with optional variables', () => {
+    expect( () => validateGenerateTextArgs( {
+      prompt: 'summary@v1',
+      variables: { topic: 'testing' },
+      promptDir: '/prompts',
+      skills: [ { name: 'style', description: 'Style', instructions: '# Style' } ],
+      maxSteps: 4
+    } ) ).not.toThrow();
+  } );
+
+  it( 'accepts a dynamic skills function', () => {
+    expect( () => validateGenerateTextArgs( {
+      prompt: 'summary@v1',
+      skills: () => [ { name: 'style', description: 'Style', instructions: '# Style' } ]
+    } ) ).not.toThrow();
+  } );
+
+  it( 'throws ValidationError with generateText prefix for invalid args', () => {
+    expect( () => validateGenerateTextArgs( {
+      variables: { topic: 'testing' }
+    } ) ).toThrow( ValidationError );
+
+    expect( () => validateGenerateTextArgs( {
+      prompt: 123
+    } ) ).toThrow( /Invalid generateText\(\) arguments/ );
+  } );
+
+  it( 'throws ValidationError for invalid promptDir, skills, or maxSteps', () => {
+    expect( () => validateGenerateTextArgs( {
+      prompt: 'summary@v1',
+      promptDir: ''
+    } ) ).toThrow( ValidationError );
+
+    expect( () => validateGenerateTextArgs( {
+      prompt: 'summary@v1',
+      skills: [ { name: '', description: 'Style', instructions: '# Style' } ]
+    } ) ).toThrow( ValidationError );
+
+    expect( () => validateGenerateTextArgs( {
+      prompt: 'summary@v1',
+      maxSteps: 0
+    } ) ).toThrow( ValidationError );
+  } );
+} );
+
+describe( 'validateStreamTextArgs', () => {
+  it( 'accepts a prompt with optional variables', () => {
+    expect( () => validateStreamTextArgs( {
+      prompt: 'summary@v1',
+      variables: [ 'arrays are accepted by the current schema' ],
+      promptDir: '/prompts',
+      skills: [ { name: 'style', description: 'Style', instructions: '# Style' } ],
+      maxSteps: 4
+    } ) ).not.toThrow();
+  } );
+
+  it( 'accepts a dynamic skills function', () => {
+    expect( () => validateStreamTextArgs( {
+      prompt: 'summary@v1',
+      skills: () => [ { name: 'style', description: 'Style', instructions: '# Style' } ]
+    } ) ).not.toThrow();
+  } );
+
+  it( 'throws ValidationError with streamText prefix for invalid args', () => {
+    expect( () => validateStreamTextArgs( {} ) ).toThrow( ValidationError );
+    expect( () => validateStreamTextArgs( { prompt: null } ) ).toThrow( /Invalid streamText\(\) arguments/ );
+  } );
+
+  it( 'throws ValidationError for invalid promptDir, skills, or maxSteps', () => {
+    expect( () => validateStreamTextArgs( {
+      prompt: 'summary@v1',
+      promptDir: ''
+    } ) ).toThrow( ValidationError );
+
+    expect( () => validateStreamTextArgs( {
+      prompt: 'summary@v1',
+      skills: [ { name: 'style', description: 'Style', instructions: '' } ]
+    } ) ).toThrow( ValidationError );
+
+    expect( () => validateStreamTextArgs( {
+      prompt: 'summary@v1',
+      maxSteps: 1.5
+    } ) ).toThrow( ValidationError );
+  } );
+} );
+
+describe( 'validateGenerateImageArgs', () => {
+  it( 'accepts text-to-image args without images or mask', () => {
+    expect( () => validateGenerateImageArgs( {
+      prompt: 'image@v1',
+      variables: { topic: 'race cars' },
+      promptDir: '/prompts'
+    } ) ).not.toThrow();
+  } );
+
+  it( 'accepts all supported image input shapes', () => {
+    const buffer = Buffer.from( 'image-bytes' );
+    const uint8Array = new Uint8Array( [ 1, 2, 3 ] );
+    const arrayBuffer = new ArrayBuffer( 3 );
+    const paddedBase64 = 'aW1hZ2UtYnl0ZXM=';
+    const unpaddedBase64 = 'aW1hZ2U';
+
+    expect( () => validateGenerateImageArgs( {
+      prompt: 'image@v1',
+      images: [
+        buffer,
+        uint8Array,
+        arrayBuffer,
+        paddedBase64,
+        unpaddedBase64,
+        { data: buffer, mediaType: 'image/png' },
+        { data: uint8Array },
+        { data: arrayBuffer, mediaType: 'image/jpeg' },
+        { data: paddedBase64, mediaType: 'image/webp' }
+      ],
+      mask: { data: Buffer.from( 'mask-bytes' ), mediaType: 'image/png' }
+    } ) ).not.toThrow();
+  } );
+
+  it( 'throws ValidationError for invalid image args', () => {
+    expect( () => validateGenerateImageArgs( {
+      prompt: ''
+    } ) ).toThrow( /Invalid generateImage\(\) arguments/ );
+
+    expect( () => validateGenerateImageArgs( {
+      prompt: 'image@v1',
+      images: []
+    } ) ).toThrow( ValidationError );
+
+    expect( () => validateGenerateImageArgs( {
+      prompt: 'image@v1',
+      images: [ { data: null } ]
+    } ) ).toThrow( ValidationError );
+
+    expect( () => validateGenerateImageArgs( {
+      prompt: 'image@v1',
+      images: [ { data: 'aW1hZ2U=', mediaType: '' } ]
+    } ) ).toThrow( ValidationError );
+  } );
+
+  it( 'rejects image strings that are not raw base64 data', () => {
+    for ( const image of [
+      'https://example.com/image.png',
+      'data:image/png;base64,aW1hZ2U=',
+      'not base64',
+      'abcde'
+    ] ) {
+      expect( () => validateGenerateImageArgs( {
+        prompt: 'image@v1',
+        images: [ image ]
+      } ) ).toThrow( /Image strings must be raw base64 data/ );
+    }
+  } );
+
+  it( 'requires images when mask is provided', () => {
+    expect( () => validateGenerateImageArgs( {
+      prompt: 'image@v1',
+      mask: Buffer.from( 'mask-bytes' )
+    } ) ).toThrow( /mask requires images/ );
+  } );
+} );

package/src/parser.js

@@ -1,28 +0,0 @@
-import matter from 'gray-matter';
-import { FatalError } from '@outputai/core';
-
-export function parsePrompt( raw ) {
-  const { data: config, content } = matter( raw );
-
-  if ( !content || content.trim() === '' ) {
-    throw new FatalError( 'Prompt file has no content after frontmatter' );
-  }
-
-  const infoExtractor = /<(system|user|assistant|tool)>([\s\S]*?)<\/\1>/gm;
-  const messages = [ ...content.matchAll( infoExtractor ) ].map(
-    ( [ _, role, text ] ) => ( { role, content: text.trim() } )
-  );
-
-  if ( messages.length === 0 ) {
-    const contentPreview = content.substring( 0, 200 );
-    const ellipsis = content.length > 200 ? '...' : '';
-
-    throw new FatalError(
-      `No valid message blocks found in prompt file.
-Expected format: <system>...</system>, <user>...</user>, etc.
-Content preview: ${contentPreview}${ellipsis}`
-    );
-  }
-
-  return { config, messages };
-}