`onError` has no knowledge of individual streams

[BrentFarese]

Great library! One issue we are running across at Aline is that onError has no knowledge of individual streams, which makes it somewhat difficult to work with. Take the retryAfterOutput example from the README below:

import { createFallback } from 'ai-fallback'
import { openai } from '@ai-sdk/openai'
import { anthropic } from '@ai-sdk/anthropic'
import { streamText } from 'ai'

let fullText = ''

const model = createFallback({
    models: [anthropic('claude-3-haiku-20240307'), openai('gpt-3.5-turbo')],
    retryAfterOutput: true, // Enable retrying even after partial output
    onError: (err) => {
        console.error('Error:', err)
        // reset the full text because error happened when some tokens were already streamed in
        fullText = ''
    },
})

const stream = await streamText({
    model,
    system: 'You are a helpful assistant.',
    messages: [{ role: 'user', content: 'Write a long story.' }],
})

for await (const chunk of stream.textStream) {
    fullText += chunk
    console.log('Current text:', fullText)
}

In this example, onError is resetting fullText for just one streamText. What if you use this model in multiple streamTexts throughout your codebase (like Aline does). onError has no notion of those executions, but it would be nice if it did.

How we have worked around this at Aline is to use AsyncLocalStorage to inject some context into the onError callback that is scoped to that stream/request.

Here is some example code:

/**
 * Request context that contains chatClient and other request-specific data
 *
 * HOW THIS WORKS:
 *
 * 1. FallbackModel instances are global/persistent (for timer-based recovery logic).
 * 2. AsyncLocalStorage provides request-scoped context isolation.
 * 3. When onError fires, it runs in the same async context as the original request.
 * 4. This allows the global model's error handler to access the request's chatClient.
 *
 * EXAMPLE FLOW:
 * runInRequestContext(context, () => {
 *   streamText({ model: globalFallbackModel }) // Request A context.
 * });
 *
 * runInRequestContext(context, () => {
 *   streamText({ model: globalFallbackModel }) // Request B context.
 * });
 *
 * When either fails, onError gets the correct chatClient for that request!
 */
export interface RequestContext {
  chatClient: ChatClient;
  threadId: string;
}

/**
 * Global AsyncLocalStorage for request context tracking.
 * This ensures fallback notifications are sent to the correct chatClient.
 */
const requestContextStorage = new AsyncLocalStorage<RequestContext>();

/**
 * Get the current request context (if available).
 */
export function getRequestContext(): RequestContext | undefined {
  return requestContextStorage.getStore();
}

/**
 * Run a function within a specific request context.
 * This establishes the context for the entire async chain.
 *
 * @example
 * // In your orchestrator or any function that needs fallback notifications:
 * const requestContext = {
 *   chatClient: createChatClient({...}),
 *   threadId: 'thread-123'
 * };
 *
 * return runInRequestContext(requestContext, async () => {
 *   // Any fallback model usage within this block will have access
 *   // to the chatClient for sending notifications.
 *   const result = await streamText({
 *     model: aiModels.openai41, // Uses global persistent fallback model.
 *     // ... other options
 *   });
 *   return result;
 * });
 */
export function runInRequestContext<T>(
  context: RequestContext,
  fn: () => Promise<T> | T
): Promise<T> | T {
  return requestContextStorage.run(context, fn);
}

/**
 * Creates a fallback-enabled AI model with the given configuration.
 * This wraps the ai-fallback package with Aline-specific functionality.
 */
export function createFallbackModel({ models }: { models: LanguageModelV2Type[] }): FallbackModel {
  const fallbackModel = createFallback({
    // 5 minutes.
    modelResetInterval: 5 * 60 * 1000,
    models,
    onError: (error, modelId) => {
      Sentry.captureException(error, { extra: { modelId } });
      logger.error(
        `🚨 FALLBACK ERROR HANDLER: ${modelId} failed with:`,
        error.message || 'an unknown error occurred.'
      );

      const requestContext = getRequestContext();
      if (!requestContext) {
        logger.debug(
          `🔔 No request context available for fallback notification from ${modelId}. This can validly occur where we don't run inside AsyncLocalStorage (e.g., for non-chat requests).`
        );
        return;
      }

      const currentModelIndex = models.findIndex(model => model.modelId === modelId);
      if (currentModelIndex === -1 || currentModelIndex === models.length - 1) {
        const originalModelDisplayName = getModelNameFromModel(models[0]);
        // All models in chain have failed...
        requestContext.chatClient.onSendError({
          message: `${originalModelDisplayName} and all its fallbacks are unavailable.`,
          stack: error.stack || 'Stack not available.',
          title: 'Critical AI Model Failure',
        });

        return;
      }

      const currentModel = models[currentModelIndex];
      const nextModel = models[currentModelIndex + 1];
      const nextModelName = getModelNameFromModel(nextModel);
      const currentModelName = getModelNameFromModel(currentModel);
      const toolCallId = requestContext.chatClient.onSendToolCall({
        fromModel: currentModelName,
        id: createToolCallId(),
        toModel: nextModelName,
        type: 'modelFallback',
      });
      requestContext.chatClient.onSendToolCallStatus(toolCallId, 'finished');
    },
    /**
     * Flipping this to true would cause the fallback model to retry even mid-stream, which is likely not what we want.
     * See https://github.com/remorses/ai-fallback?tab=readme-ov-file#retry-after-output.
     */
    retryAfterOutput: false,
  });
  const originalDoStream = fallbackModel.doStream.bind(fallbackModel);
  const originalDoGenerate = fallbackModel.doGenerate.bind(fallbackModel);

  fallbackModel.doGenerate = async (...args: Parameters<typeof originalDoGenerate>) => {
    logger.debug(`🎯 Starting doGenerate with fallback model.`);
    const result = await originalDoGenerate(...args);
    logger.debug(`✅ doGenerate completed successfully (fallback may have occurred).`);

    return result;
  };

  fallbackModel.doStream = async (...args: Parameters<typeof originalDoStream>) => {
    logger.debug(`🎯 Starting doStream with fallback model.`);
    const result = await originalDoStream(...args);
    logger.debug(
      `✅ doStream setup completed, returning stream (fallback may occur during streaming).`
    );

    return result;
  };

  return fallbackModel;
}

This becomes pretty useful to do things in a specific stream based on onError occurring. Maybe some pattern like this could be built into the library...

[remorses]

That's actually an old example. The latest version of ai-fallback no longer retries after the first output token is generated, that's because when a provider has downtime it will 90% of the times error before generating the first token

[BrentFarese]

That's actually an old example. The latest version of ai-fallback no longer retries after the first output token is generated, that's because when a provider has downtime it will 90% of the times error before generating the first token

The example was used really to highlight onError limitations. Not necessarily to highlight retryAfterOutput. Thanks!

[remorses]

Ok now I understand. Why do you want to send requestContext.chatClient.onSendToolCall() and requestContext.chatClient.onSendToolCallStatus(toolCallId, 'finished'); inside the onError handler?

If all fallback models were tried ai-fallback will already throw back the error to the usual ai error handler so there should be no need to handle it in ai-fallback

[BrentFarese]

Ok now I understand. Why do you want to send requestContext.chatClient.onSendToolCall() and requestContext.chatClient.onSendToolCallStatus(toolCallId, 'finished'); inside the onError handler?

If all fallback models were tried ai-fallback will already throw back the error to the usual ai error handler so there should be no need to handle it in ai-fallback

When there has been a fallback, we send a notification to the client/show a tool call in the message history so a user knows a fallback has occurred. We do this b/c models have different quality/outputs and we want to be transparent w/ users if they see different output, they know why.

[remorses]

Ok, using context seems the most reasonable solution. ai-fallback needs to be outside of the handler to preserve state.

I don't think you need to override the ai-fallback methods, you can instead wrap your API handler.

You could also use a WeakMap and use the request as the key to store the chatClient values and use them in onError

[BrentFarese]

Ok, using context seems the most reasonable solution. ai-fallback needs to be outside of the handler to preserve state.

I don't think you need to override the ai-fallback methods, you can instead wrap your API handler.

You could also use a WeakMap and use the request as the key to store the chatClient values and use them in onError

Yup, those are viable options. We went with AsyncLocalStorage b/c it seems elegant. You might consider providing a method on FallbackModel like FallbackModel.runInContext(someContext) to allow developers to "inject values" into onError. Something to consider if this ever comes up again...

We are fine b/c we have our own solution but took me ~1-2 days and if it comes up more often, could be cool to support in library.