@outputai/llm 0.6.1-next.5d7e612.0 → 0.6.1-next.fc6a93e.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@outputai/llm",
-  "version": "0.6.1-next.5d7e612.0",
+  "version": "0.6.1-next.fc6a93e.0",
   "description": "Framework abstraction to interact with LLM models",
   "type": "module",
   "main": "src/index.js",
@@ -23,7 +23,7 @@
     "gray-matter": "4.0.3",
     "liquidjs": "10.25.7",
     "undici": "8.1.0",
-    "@outputai/core": "0.6.1-next.5d7e612.0"
+    "@outputai/core": "0.6.1-next.fc6a93e.0"
   },
   "license": "Apache-2.0",
   "publishConfig": {

package/src/ai_sdk.js

@@ -2,12 +2,13 @@ import { types as utilTypes } from 'node:util';
 import * as AI from 'ai';
 import { stepCountIs } from 'ai';
 import { ValidationError } from '@outputai/core';
-import { validateGenerateTextArgs, validateStreamTextArgs } from './validations.js';
+import { validateGenerateTextArgs, validateStreamTextArgs, validateGenerateImageArgs } from './validations.js';
 import { loadPrompt } from './prompt/loader.js';
 import { startTrace, endTraceWithError } from './utils/trace.js';
 import { wrapTextResponse, wrapStreamOnFinishResponse, wrapImageResponse } from './utils/response_wrappers.js';
 import { loadAiSdkTextOptions, loadAiSdkImageOptions } from './ai_sdk_options.js';
 import { prepareTextPrompt } from './prompt/prepare_text.js';
+import { mapAiError } from './utils/error_handler.js';
 
 export async function generateText( { prompt, variables, promptDir, skills = [], maxSteps = 10, ...aiSdkArgs } ) {
   validateGenerateTextArgs( { prompt, variables, promptDir, skills, maxSteps } );
@@ -27,13 +28,14 @@ export async function generateText( { prompt, variables, promptDir, skills = [],
       ...( tools && !aiSdkArgs.stopWhen ? { stopWhen: stepCountIs( maxSteps ) } : {} )
     } );
     return wrapTextResponse( { traceId, modelId, response } );
-  } catch ( error ) {
+  } catch ( originalError ) {
+    const error = mapAiError( originalError );
     endTraceWithError( { traceId, error } );
     throw error;
   }
 }
 
-export function streamText( { prompt, variables, promptDir, skills = [], maxSteps = 10, onFinish, onError, ...aiSdkArgs } ) {
+export function streamText( { prompt, variables, promptDir, skills = [], maxSteps = 10, onFinish, onError: _onError, ...aiSdkArgs } ) {
   validateStreamTextArgs( { prompt, variables, promptDir, skills, maxSteps } );
 
   const parsedSkills = typeof skills === 'function' ? skills( variables ) : skills;
@@ -54,30 +56,34 @@ export function streamText( { prompt, variables, promptDir, skills = [], maxStep
       ...( tools && !aiSdkArgs.stopWhen ? { stopWhen: stepCountIs( maxSteps ) } : {} ),
       ...wrapStreamOnFinishResponse( { traceId, modelId, onFinish } ),
       onError( event ) {
-        endTraceWithError( { traceId, error: event.error } );
-        onError?.( event );
+        const error = mapAiError( event.error );
+        endTraceWithError( { traceId, error } );
+        _onError?.( { ...event, error } );
       }
     } );
-  } catch ( error ) {
+  } catch ( originalError ) {
+    const error = mapAiError( originalError );
     endTraceWithError( { traceId, error } );
     throw error;
   }
 }
 
-export async function generateImage( { prompt, variables, promptDir, ...aiSdkArgs } ) {
-  const loadedPrompt = loadPrompt( prompt, variables, promptDir );
+export async function generateImage( { prompt, variables, promptDir, images, mask, ...aiSdkArgs } ) {
+  validateGenerateImageArgs( { prompt, variables, promptDir, images, mask } );
 
+  const loadedPrompt = loadPrompt( prompt, variables, promptDir );
   const traceId = startTrace( { name: 'generateImage', prompt, variables, loadedPrompt } );
   const { model: modelId } = loadedPrompt.config;
 
   try {
     const response = await AI.generateImage( {
-      ...loadAiSdkImageOptions( loadedPrompt ),
+      ...loadAiSdkImageOptions( { prompt: loadedPrompt, images, mask } ),
       maxRetries: 0,
       ...aiSdkArgs
     } );
     return wrapImageResponse( { traceId, modelId, response } );
-  } catch ( error ) {
+  } catch ( originalError ) {
+    const error = mapAiError( originalError );
     endTraceWithError( { traceId, error } );
     throw error;
   }

package/src/ai_sdk.spec.js

@@ -9,7 +9,8 @@ const aiFns = vi.hoisted( () => ( {
 
 const validators = vi.hoisted( () => ( {
   validateGenerateTextArgs: vi.fn(),
-  validateStreamTextArgs: vi.fn()
+  validateStreamTextArgs: vi.fn(),
+  validateGenerateImageArgs: vi.fn()
 } ) );
 
 const promptMocks = vi.hoisted( () => ( {
@@ -33,6 +34,10 @@ const wrapMocks = vi.hoisted( () => ( {
   wrapImageResponse: vi.fn()
 } ) );
 
+const errorMocks = vi.hoisted( () => ( {
+  mapAiError: vi.fn( error => error )
+} ) );
+
 vi.mock( 'ai', () => aiFns );
 
 vi.mock( './validations.js', () => validators );
@@ -61,6 +66,10 @@ vi.mock( './utils/response_wrappers.js', () => ( {
   wrapImageResponse: ( ...args ) => wrapMocks.wrapImageResponse( ...args )
 } ) );
 
+vi.mock( './utils/error_handler.js', () => ( {
+  mapAiError: ( ...args ) => errorMocks.mapAiError( ...args )
+} ) );
+
 const importSut = async () => import( './ai_sdk.js' );
 
 const loadedPrompt = {
@@ -86,15 +95,29 @@ const streamResult = {
   fullStream: 'FULL_STREAM'
 };
 
+const imageOptions = {
+  model: 'IMAGE_MODEL',
+  prompt: {
+    text: 'Generate an image'
+  },
+  providerOptions: { openai: { quality: 'high' } }
+};
+
+const imageResponse = {
+  images: [ { mediaType: 'image/png', base64: 'aW1hZ2U=' } ],
+  usage: { inputTokens: 1, outputTokens: 2 }
+};
+
 describe( 'ai_sdk', () => {
   beforeEach( () => {
     aiFns.generateText.mockReset().mockResolvedValue( textResponse );
     aiFns.streamText.mockReset().mockReturnValue( streamResult );
-    aiFns.generateImage.mockReset();
+    aiFns.generateImage.mockReset().mockResolvedValue( imageResponse );
     aiFns.stepCountIs.mockReset().mockImplementation( count => ( { type: 'step-count', count } ) );
 
     validators.validateGenerateTextArgs.mockReset();
     validators.validateStreamTextArgs.mockReset();
+    validators.validateGenerateImageArgs.mockReset();
 
     promptMocks.loadPrompt.mockReset().mockReturnValue( loadedPrompt );
     promptMocks.prepareTextPrompt.mockReset().mockReturnValue( {
@@ -103,7 +126,7 @@ describe( 'ai_sdk', () => {
     } );
 
     optionMocks.loadAiSdkTextOptions.mockReset().mockReturnValue( textOptions );
-    optionMocks.loadAiSdkImageOptions.mockReset();
+    optionMocks.loadAiSdkImageOptions.mockReset().mockReturnValue( imageOptions );
 
     traceMocks.startTrace.mockReset().mockReturnValue( 'trace-id' );
     traceMocks.endTraceWithError.mockReset();
@@ -112,7 +135,9 @@ describe( 'ai_sdk', () => {
     wrapMocks.wrapStreamOnFinishResponse.mockReset().mockReturnValue( {
       onFinish: vi.fn()
     } );
-    wrapMocks.wrapImageResponse.mockReset();
+    wrapMocks.wrapImageResponse.mockReset().mockResolvedValue( { wrapped: imageResponse } );
+
+    errorMocks.mapAiError.mockReset().mockImplementation( error => error );
   } );
 
   afterEach( async () => {
@@ -234,13 +259,16 @@ describe( 'ai_sdk', () => {
 
     it( 'traces and rethrows AI SDK errors', async () => {
       const error = new Error( 'Provider failed' );
+      const mappedError = new Error( 'Mapped provider failed' );
       aiFns.generateText.mockRejectedValueOnce( error );
+      errorMocks.mapAiError.mockReturnValueOnce( mappedError );
       const { generateText } = await importSut();
 
-      await expect( generateText( { prompt: 'test@v1' } ) ).rejects.toThrow( error );
+      await expect( generateText( { prompt: 'test@v1' } ) ).rejects.toThrow( mappedError );
+      expect( errorMocks.mapAiError ).toHaveBeenCalledWith( error );
       expect( traceMocks.endTraceWithError ).toHaveBeenCalledWith( {
         traceId: 'trace-id',
-        error
+        error: mappedError
       } );
     } );
   } );
@@ -371,16 +399,19 @@ describe( 'ai_sdk', () => {
       const { streamText } = await importSut();
       const onError = vi.fn();
       const error = new Error( 'Stream failed' );
+      const mappedError = new Error( 'Mapped stream failed' );
+      errorMocks.mapAiError.mockReturnValueOnce( mappedError );
 
       streamText( { prompt: 'test@v1', onError } );
       const callOptions = aiFns.streamText.mock.calls[0][0];
       callOptions.onError( { error } );
 
+      expect( errorMocks.mapAiError ).toHaveBeenCalledWith( error );
       expect( traceMocks.endTraceWithError ).toHaveBeenCalledWith( {
         traceId: 'trace-id',
-        error
+        error: mappedError
       } );
-      expect( onError ).toHaveBeenCalledWith( { error } );
+      expect( onError ).toHaveBeenCalledWith( { error: mappedError } );
     } );
 
     it( 'does not pass the raw onFinish or onError callbacks to AI SDK', async () => {
@@ -410,15 +441,116 @@ describe( 'ai_sdk', () => {
 
     it( 'traces and rethrows synchronous AI SDK errors', async () => {
       const error = new Error( 'Invalid model' );
+      const mappedError = new Error( 'Mapped invalid model' );
       aiFns.streamText.mockImplementationOnce( () => {
         throw error;
       } );
+      errorMocks.mapAiError.mockReturnValueOnce( mappedError );
       const { streamText } = await importSut();
 
-      expect( () => streamText( { prompt: 'test@v1' } ) ).toThrow( error );
+      expect( () => streamText( { prompt: 'test@v1' } ) ).toThrow( mappedError );
+      expect( errorMocks.mapAiError ).toHaveBeenCalledWith( error );
+      expect( traceMocks.endTraceWithError ).toHaveBeenCalledWith( {
+        traceId: 'trace-id',
+        error: mappedError
+      } );
+    } );
+  } );
+
+  describe( 'generateImage', () => {
+    it( 'validates, loads prompt, traces, calls AI SDK, and wraps the response', async () => {
+      const { generateImage } = await importSut();
+      const variables = { scene: 'race cars' };
+      const images = [ Buffer.from( 'image-bytes' ) ];
+      const mask = Buffer.from( 'mask-bytes' );
+
+      const result = await generateImage( {
+        prompt: 'image@v1',
+        variables,
+        promptDir: '/prompts',
+        images,
+        mask,
+        n: 2,
+        providerOptions: { openai: { background: 'transparent' } }
+      } );
+
+      expect( validators.validateGenerateImageArgs ).toHaveBeenCalledWith( {
+        prompt: 'image@v1',
+        variables,
+        promptDir: '/prompts',
+        images,
+        mask
+      } );
+      expect( promptMocks.loadPrompt ).toHaveBeenCalledWith( 'image@v1', variables, '/prompts' );
+      expect( traceMocks.startTrace ).toHaveBeenCalledWith( {
+        name: 'generateImage',
+        prompt: 'image@v1',
+        variables,
+        loadedPrompt
+      } );
+      expect( optionMocks.loadAiSdkImageOptions ).toHaveBeenCalledWith( {
+        prompt: loadedPrompt,
+        images,
+        mask
+      } );
+      expect( aiFns.generateImage ).toHaveBeenCalledWith( {
+        ...imageOptions,
+        maxRetries: 0,
+        n: 2,
+        providerOptions: { openai: { background: 'transparent' } }
+      } );
+      expect( wrapMocks.wrapImageResponse ).toHaveBeenCalledWith( {
+        traceId: 'trace-id',
+        modelId: 'test-model',
+        response: imageResponse
+      } );
+      expect( result ).toEqual( { wrapped: imageResponse } );
+    } );
+
+    it( 'supports text-to-image calls without images or mask', async () => {
+      const { generateImage } = await importSut();
+
+      await generateImage( { prompt: 'image@v1' } );
+
+      expect( validators.validateGenerateImageArgs ).toHaveBeenCalledWith( {
+        prompt: 'image@v1',
+        variables: undefined,
+        promptDir: undefined,
+        images: undefined,
+        mask: undefined
+      } );
+      expect( optionMocks.loadAiSdkImageOptions ).toHaveBeenCalledWith( {
+        prompt: loadedPrompt,
+        images: undefined,
+        mask: undefined
+      } );
+    } );
+
+    it( 'propagates validation errors before loading or tracing', async () => {
+      const validationError = new Error( 'Invalid image args' );
+      validators.validateGenerateImageArgs.mockImplementationOnce( () => {
+        throw validationError;
+      } );
+      const { generateImage } = await importSut();
+
+      await expect( generateImage( { prompt: '' } ) ).rejects.toThrow( validationError );
+      expect( promptMocks.loadPrompt ).not.toHaveBeenCalled();
+      expect( traceMocks.startTrace ).not.toHaveBeenCalled();
+      expect( aiFns.generateImage ).not.toHaveBeenCalled();
+    } );
+
+    it( 'traces and rethrows AI SDK errors', async () => {
+      const error = new Error( 'Image provider failed' );
+      const mappedError = new Error( 'Mapped image provider failed' );
+      aiFns.generateImage.mockRejectedValueOnce( error );
+      errorMocks.mapAiError.mockReturnValueOnce( mappedError );
+      const { generateImage } = await importSut();
+
+      await expect( generateImage( { prompt: 'image@v1' } ) ).rejects.toThrow( mappedError );
+      expect( errorMocks.mapAiError ).toHaveBeenCalledWith( error );
       expect( traceMocks.endTraceWithError ).toHaveBeenCalledWith( {
         traceId: 'trace-id',
-        error
+        error: mappedError
       } );
     } );
   } );

package/src/ai_sdk_options.js

@@ -39,13 +39,17 @@ export const loadAiSdkTextOptions = prompt => {
  * @param {object} prompt - Loaded prompt object
  * @returns {object} Options for AI SDK image calls
  */
-export const loadAiSdkImageOptions = prompt => {
+export const loadAiSdkImageOptions = ( { prompt, images, mask } ) => {
   if ( !prompt.instructions ) {
     throw new FatalError( `Prompt "${prompt.name}" has no instructions. Image prompts must use plain instructions.` );
   }
   const options = {
     model: loadImageModel( prompt ),
-    prompt: prompt.instructions,
+    prompt: ( images || mask ) ? {
+      text: prompt.instructions,
+      ...( images && { images } ),
+      ...( mask && { mask } )
+    } : prompt.instructions,
     providerOptions: prompt.config.providerOptions
   };
   for ( const key of [ 'n', 'maxImagesPerCall', 'size', 'aspectRatio', 'seed' ] ) {

package/src/ai_sdk_options.spec.js

@@ -100,6 +100,8 @@ describe( 'ai_sdk_options', () => {
   } );
 
   it( 'maps loaded prompts to AI SDK image options', async () => {
+    const images = [ Buffer.from( 'image-bytes' ) ];
+    const mask = Buffer.from( 'mask-bytes' );
     const prompt = makeImagePrompt( {
       n: 2,
       maxImagesPerCall: 1,
@@ -112,14 +114,18 @@ describe( 'ai_sdk_options', () => {
     } );
 
     const { loadAiSdkImageOptions } = await importSut();
-    const result = loadAiSdkImageOptions( prompt );
+    const result = loadAiSdkImageOptions( { prompt, images, mask } );
 
     expect( loadImageModelImpl ).toHaveBeenCalledWith( prompt );
     expect( loadModelImpl ).not.toHaveBeenCalled();
     expect( loadToolsImpl ).not.toHaveBeenCalled();
     expect( result ).toEqual( {
       model: 'IMAGE_MODEL',
-      prompt: 'Generate a cinematic image of a NASCAR race at sunset.',
+      prompt: {
+        text: 'Generate a cinematic image of a NASCAR race at sunset.',
+        images,
+        mask
+      },
       providerOptions: prompt.config.providerOptions,
       n: 2,
       maxImagesPerCall: 1,
@@ -135,7 +141,7 @@ describe( 'ai_sdk_options', () => {
     const prompt = makeImagePrompt( { seed: 0 } );
 
     const { loadAiSdkImageOptions } = await importSut();
-    const result = loadAiSdkImageOptions( prompt );
+    const result = loadAiSdkImageOptions( { prompt } );
 
     expect( result ).toEqual( {
       model: 'IMAGE_MODEL',
@@ -150,7 +156,7 @@ describe( 'ai_sdk_options', () => {
 
     const { loadAiSdkImageOptions } = await importSut();
 
-    expect( () => loadAiSdkImageOptions( prompt ) ).toThrow(
+    expect( () => loadAiSdkImageOptions( { prompt } ) ).toThrow(
       'Prompt "test@v1" has no instructions.'
     );
     expect( loadImageModelImpl ).not.toHaveBeenCalled();

package/src/index.d.ts

@@ -200,6 +200,9 @@ export type StreamTextAiSdkOptions<
  * `model` and `prompt` are omitted because Output supplies them from the prompt file.
  */
 export type GenerateImageAiSdkOptions = Omit<Parameters<typeof aiGenerateImage>[0], 'model' | 'prompt'>;
+type GenerateImagePrompt = Parameters<typeof aiGenerateImage>[0]['prompt'];
+type GenerateImagePromptWithImages = Exclude<GenerateImagePrompt, string>;
+type GenerateImageInput = GenerateImagePromptWithImages['images'][number];
 
 /** Agent {@link Agent.stream} options: same as AI SDK plus wrapped `onFinish` (adds `cost`). */
 export type OutputAgentStreamParameters = Omit<AgentStreamParameters<never, ToolSet>, 'onFinish'> & {
@@ -275,6 +278,10 @@ export type GenerateImageParameters = {
   variables?: Record<string, string | number | boolean>;
   /** Override the stack-resolved prompt directory */
   promptDir?: string;
+  /** Runtime image inputs for image-to-image generation */
+  images?: GenerateImageInput[];
+  /** Optional mask for image editing */
+  mask?: GenerateImagePromptWithImages['mask'];
 } & GenerateImageAiSdkOptions;
 
 /** A source extracted from search tool results during multi-step LLM execution. */

package/src/utils/error_handler.js

@@ -0,0 +1,104 @@
+import {
+  APICallError,
+  InvalidArgumentError,
+  InvalidDataContentError,
+  InvalidPromptError,
+  LoadAPIKeyError,
+  LoadSettingError,
+  NoImageGeneratedError,
+  NoObjectGeneratedError,
+  NoSuchModelError,
+  NoSuchProviderError,
+  UnsupportedFunctionalityError
+} from 'ai';
+import { FatalError } from '@outputai/core';
+
+/**
+ * Recursively search an error cause chain until finds an error which is instance of given prototype.
+ *
+ * @param {object} error - Error instance.
+ * @param {Function|string} _class - Target constructor or constructor name.
+ * @param {number} depth - Current depth, search up to 10 causes deep.
+ * @returns {object|null} - Error or null if not found.
+ */
+export const findInstanceInCauseChain = ( error, _class, depth = 0 ) => {
+  if ( !error || typeof error !== 'object' ) {
+    return null;
+  }
+  if ( typeof _class === 'string' && error.constructor.name === _class ) {
+    return error;
+  }
+  if ( typeof _class === 'function' && error instanceof _class ) {
+    return error;
+  }
+  if ( depth >= 10 ) {
+    return null;
+  }
+  return error.cause ? findInstanceInCauseChain( error.cause, _class, depth + 1 ) : null;
+};
+
+const toFatalError = ( error, extraMessage = '' ) => new FatalError(
+  `AI-SDK fatal error${extraMessage ? ` (${extraMessage})` : ''}: ${error.message}`,
+  { cause: error }
+);
+
+export const mapAiError = error => {
+  if ( error instanceof FatalError ) {
+    return error;
+  }
+
+  // NoObjectGeneratedError can be thrown when the response doesn't match the schema
+  // This adds a wrapper to that error serializing the first zod validation in the message, to make it easier to debug.
+  if ( NoObjectGeneratedError.isInstance( error ) && error.message.includes( 'No object generated: response did not match schema.' ) ) {
+    const zodError = findInstanceInCauseChain( error, 'ZodError' );
+    if ( zodError && zodError.issues?.length > 0 ) {
+      const { path, message } = zodError.issues[0];
+      const wrapper = new Error( `${error.message} First issue is "${message}" at path [${path.join( ', ' )}].`, { cause: error } );
+      wrapper.name = 'NoObjectGeneratedError';
+      return wrapper;
+    }
+    return error;
+  }
+
+  if ( APICallError.isInstance( error ) && !error.isRetryable ) {
+    // Non-retryable API failures are already classified by AI SDK as permanent provider failures.
+    return toFatalError( error, error.statusCode ? `HTTP ${error.statusCode}` : '' );
+  }
+  if ( InvalidArgumentError.isInstance( error ) ) {
+    // Invalid call settings are deterministic caller bugs, so retrying the same activity cannot fix them.
+    return toFatalError( error );
+  }
+  if ( InvalidDataContentError.isInstance( error ) ) {
+    // Invalid media content has the wrong local shape/encoding and will fail again with the same input.
+    return toFatalError( error );
+  }
+  if ( InvalidPromptError.isInstance( error ) ) {
+    // Invalid prompt structure is a deterministic request-construction error.
+    return toFatalError( error );
+  }
+  if ( LoadAPIKeyError.isInstance( error ) ) {
+    // Missing or invalid API key configuration will not change during an activity retry.
+    return toFatalError( error );
+  }
+  if ( LoadSettingError.isInstance( error ) ) {
+    // Missing or invalid provider settings are deployment/configuration problems.
+    return toFatalError( error );
+  }
+  if ( NoImageGeneratedError.isInstance( error ) ) {
+    // Image generation completed provider calls but collected zero images; repeating identical input is not useful.
+    return toFatalError( error );
+  }
+  if ( NoSuchProviderError.isInstance( error ) ) {
+    // A missing provider id is a deterministic provider registry/configuration error.
+    return toFatalError( error );
+  }
+  if ( NoSuchModelError.isInstance( error ) ) {
+    // A missing model id is a deterministic provider/model configuration error.
+    return toFatalError( error );
+  }
+  if ( UnsupportedFunctionalityError.isInstance( error ) ) {
+    // The selected model/output mode does not support the requested feature.
+    return toFatalError( error );
+  }
+  return error;
+};

package/src/utils/error_handler.spec.js

@@ -0,0 +1,285 @@
+import { describe, expect, it } from 'vitest';
+import {
+  APICallError,
+  InvalidArgumentError,
+  InvalidDataContentError,
+  InvalidMessageRoleError,
+  InvalidPromptError,
+  InvalidToolApprovalError,
+  InvalidToolInputError,
+  LoadAPIKeyError,
+  LoadSettingError,
+  MessageConversionError,
+  NoImageGeneratedError,
+  NoOutputGeneratedError,
+  NoObjectGeneratedError,
+  NoSuchModelError,
+  NoSuchProviderError,
+  ToolCallNotFoundForApprovalError,
+  ToolCallRepairError,
+  UnsupportedFunctionalityError
+} from 'ai';
+import { FatalError } from '@outputai/core';
+import { findInstanceInCauseChain, mapAiError } from './error_handler.js';
+
+const makeApiCallError = ( input = {} ) => new APICallError( {
+  message: 'Provider rejected the request',
+  url: 'https://provider.test/v1/generate',
+  requestBodyValues: {},
+  responseHeaders: {},
+  responseBody: '{"error":"bad request"}',
+  ...input
+} );
+
+const fatalAiSdkErrors = [
+  [
+    'InvalidArgumentError',
+    () => new InvalidArgumentError( {
+      parameter: 'temperature',
+      value: 'hot',
+      message: 'temperature must be a number'
+    } )
+  ],
+  [
+    'InvalidDataContentError',
+    () => new InvalidDataContentError( { content: { bad: true } } )
+  ],
+  [
+    'InvalidPromptError',
+    () => new InvalidPromptError( { prompt: {}, message: 'prompt or messages must be defined' } )
+  ],
+  [
+    'LoadAPIKeyError',
+    () => new LoadAPIKeyError( { message: 'Missing API key' } )
+  ],
+  [
+    'LoadSettingError',
+    () => new LoadSettingError( { message: 'Missing setting' } )
+  ],
+  [
+    'NoImageGeneratedError',
+    () => new NoImageGeneratedError( { responses: [] } )
+  ],
+  [
+    'NoSuchModelError',
+    () => new NoSuchModelError( { modelId: 'missing-model', modelType: 'languageModel' } )
+  ],
+  [
+    'NoSuchProviderError',
+    () => new NoSuchProviderError( {
+      modelId: 'missing-provider:model',
+      modelType: 'languageModel',
+      providerId: 'missing-provider',
+      availableProviders: [ 'openai' ]
+    } )
+  ],
+  [
+    'UnsupportedFunctionalityError',
+    () => new UnsupportedFunctionalityError( { functionality: 'image masks' } )
+  ]
+];
+
+const preservedAiSdkErrors = [
+  [
+    'InvalidMessageRoleError',
+    () => new InvalidMessageRoleError( { role: 'critic' } )
+  ],
+  [
+    'InvalidToolApprovalError',
+    () => new InvalidToolApprovalError( { approvalId: 'approval-1' } )
+  ],
+  [
+    'InvalidToolInputError',
+    () => new InvalidToolInputError( {
+      toolName: 'search',
+      toolInput: '{bad json',
+      cause: new Error( 'parse failed' )
+    } )
+  ],
+  [
+    'MessageConversionError',
+    () => new MessageConversionError( {
+      originalMessage: { role: 'critic', content: 'bad role' },
+      message: 'Unsupported role'
+    } )
+  ],
+  [
+    'NoObjectGeneratedError',
+    () => new NoObjectGeneratedError( {
+      text: 'not json',
+      cause: new Error( 'parse failed' )
+    } )
+  ],
+  [
+    'NoOutputGeneratedError',
+    () => new NoOutputGeneratedError()
+  ],
+  [
+    'ToolCallNotFoundForApprovalError',
+    () => new ToolCallNotFoundForApprovalError( {
+      toolCallId: 'tool-call-1',
+      approvalId: 'approval-1'
+    } )
+  ],
+  [
+    'ToolCallRepairError',
+    () => new ToolCallRepairError( {
+      cause: new Error( 'repair failed' ),
+      originalError: new Error( 'invalid tool input' )
+    } )
+  ]
+];
+
+describe( 'findInstanceInCauseChain', () => {
+  class FirstCustomError extends Error {}
+  class SecondCustomError extends Error {}
+
+  it( 'returns the input error when it matches the target constructor', () => {
+    const error = new FirstCustomError( 'first' );
+
+    expect( findInstanceInCauseChain( error, FirstCustomError ) ).toBe( error );
+  } );
+
+  it( 'returns the input error when it matches the target constructor name', () => {
+    const error = new FirstCustomError( 'first' );
+
+    expect( findInstanceInCauseChain( error, 'FirstCustomError' ) ).toBe( error );
+  } );
+
+  it( 'walks the cause chain to find an error by constructor', () => {
+    const target = new SecondCustomError( 'second' );
+    const wrapper = new FirstCustomError( 'first', { cause: target } );
+
+    expect( findInstanceInCauseChain( wrapper, SecondCustomError ) ).toBe( target );
+  } );
+
+  it( 'walks the cause chain to find an error by constructor name', () => {
+    const target = new SecondCustomError( 'second' );
+    const wrapper = new FirstCustomError( 'first', { cause: target } );
+
+    expect( findInstanceInCauseChain( wrapper, 'SecondCustomError' ) ).toBe( target );
+  } );
+
+  it( 'returns null when the target is not found', () => {
+    const error = new FirstCustomError( 'first', { cause: new Error( 'root' ) } );
+
+    expect( findInstanceInCauseChain( error, SecondCustomError ) ).toBeNull();
+  } );
+
+  it( 'returns null for empty or non-object inputs', () => {
+    expect( findInstanceInCauseChain( null, Error ) ).toBeNull();
+    expect( findInstanceInCauseChain( 'not an error', Error ) ).toBeNull();
+  } );
+
+  it( 'stops searching after the depth limit', () => {
+    const makeErrorChain = depth => depth === 0 ?
+      new SecondCustomError( 'target' ) :
+      new FirstCustomError( `level ${depth}`, { cause: makeErrorChain( depth - 1 ) } );
+
+    expect( findInstanceInCauseChain( makeErrorChain( 11 ), SecondCustomError ) ).toBeNull();
+  } );
+} );
+
+describe( 'mapAiError', () => {
+  it( 'preserves existing FatalError instances', () => {
+    const error = new FatalError( 'Already fatal' );
+
+    expect( mapAiError( error ) ).toBe( error );
+  } );
+
+  it( 'adds first schema issue details to NoObjectGeneratedError schema mismatches', () => {
+    class ZodError extends Error {
+      constructor( issues ) {
+        super( 'schema failed' );
+        this.issues = issues;
+      }
+    }
+    const zodError = new ZodError( [
+      {
+        path: [ 'items', 0, 'title' ],
+        message: 'Expected string'
+      }
+    ] );
+    const validationError = new Error( 'validation failed', { cause: zodError } );
+    const error = new NoObjectGeneratedError( {
+      message: 'No object generated: response did not match schema.',
+      text: '{"items":[{}]}',
+      cause: validationError
+    } );
+
+    const result = mapAiError( error );
+
+    expect( result ).not.toBe( error );
+    expect( result.name ).toBe( 'NoObjectGeneratedError' );
+    expect( result.message ).toBe(
+      'No object generated: response did not match schema. First issue is "Expected string" at path [items, 0, title].'
+    );
+    expect( result.cause ).toBe( error );
+  } );
+
+  it( 'preserves NoObjectGeneratedError schema mismatches when no schema issue is available', () => {
+    const error = new NoObjectGeneratedError( {
+      message: 'No object generated: response did not match schema.',
+      text: '{"items":[{}]}',
+      cause: new Error( 'validation failed' )
+    } );
+
+    expect( mapAiError( error ) ).toBe( error );
+  } );
+
+  it( 'maps non-retryable APICallError instances to FatalError', () => {
+    const error = makeApiCallError( {
+      statusCode: 400,
+      isRetryable: false
+    } );
+
+    const result = mapAiError( error );
+
+    expect( result ).toBeInstanceOf( FatalError );
+    expect( result.message ).toBe( 'AI-SDK fatal error (HTTP 400): Provider rejected the request' );
+    expect( result.cause ).toBe( error );
+  } );
+
+  it( 'maps non-retryable APICallError instances without status codes to FatalError', () => {
+    const error = makeApiCallError( {
+      isRetryable: false
+    } );
+
+    const result = mapAiError( error );
+
+    expect( result ).toBeInstanceOf( FatalError );
+    expect( result.message ).toBe( 'AI-SDK fatal error: Provider rejected the request' );
+    expect( result.cause ).toBe( error );
+  } );
+
+  it( 'preserves retryable APICallError instances', () => {
+    const error = makeApiCallError( {
+      statusCode: 429,
+      isRetryable: true
+    } );
+
+    expect( mapAiError( error ) ).toBe( error );
+  } );
+
+  it.each( fatalAiSdkErrors )( 'maps %s to FatalError', ( _name, makeError ) => {
+    const error = makeError();
+
+    const result = mapAiError( error );
+
+    expect( result ).toBeInstanceOf( FatalError );
+    expect( result.message ).toBe( `AI-SDK fatal error: ${error.message}` );
+    expect( result.cause ).toBe( error );
+  } );
+
+  it.each( preservedAiSdkErrors )( 'preserves %s for now', ( _name, makeError ) => {
+    const error = makeError();
+
+    expect( mapAiError( error ) ).toBe( error );
+  } );
+
+  it( 'preserves ordinary errors', () => {
+    const error = new Error( 'Network exploded' );
+
+    expect( mapAiError( error ) ).toBe( error );
+  } );
+} );

package/src/validations.js

@@ -1,3 +1,4 @@
+import { Buffer } from 'node:buffer';
 import { ValidationError, z } from '@outputai/core';
 
 const skillArgSchema = z.object( {
@@ -14,6 +15,44 @@ const generateTextArgsSchema = z.object( {
   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 ) {
   const result = schema.safeParse( input );
   if ( !result.success ) {
@@ -28,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

@@ -1,6 +1,6 @@
 import { describe, it, expect } from 'vitest';
 import { ValidationError } from '@outputai/core';
-import { validateGenerateTextArgs, validateStreamTextArgs } from './validations.js';
+import { validateGenerateTextArgs, validateStreamTextArgs, validateGenerateImageArgs } from './validations.js';
 
 describe( 'validateGenerateTextArgs', () => {
   it( 'accepts a prompt with optional variables', () => {
@@ -88,3 +88,79 @@ describe( 'validateStreamTextArgs', () => {
     } ) ).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/ );
+  } );
+} );