diff --git a/src/index.test.ts b/src/index.test.ts index b4e2c7a..3995d9e 100644 --- a/src/index.test.ts +++ b/src/index.test.ts @@ -1,6 +1,6 @@ import { test, expect } from 'vitest' import { z } from 'zod' -import { createFallback } from './index.js' +import { createFallback, TimeToFirstTokenTimeoutError } from './index.js' import { createOpenAI } from '@ai-sdk/openai' import { createGroq } from '@ai-sdk/groq' import { createAnthropic } from '@ai-sdk/anthropic' @@ -390,6 +390,363 @@ function sleep(ms: number) { return new Promise((resolve) => setTimeout(resolve, ms)) } +test('doStream falls back when first token is too slow', async () => { + const slowModel = new MockLanguageModelV4({ + provider: 'mock-slow', + modelId: 'slow-model', + doStream: async () => ({ + stream: new ReadableStream({ + async start(controller) { + controller.enqueue({ type: 'stream-start', warnings: [] }) + // Delay 500ms before emitting the first real token + await sleep(500) + controller.enqueue({ type: 'text-start', id: 't1' }) + controller.enqueue({ type: 'text-delta', id: 't1', delta: 'slow response' }) + controller.enqueue({ + type: 'finish', + finishReason: { unified: 'stop', raw: 'stop' }, + usage: { + inputTokens: { total: 1, noCache: 1, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: 1, text: 1, reasoning: undefined }, + }, + }) + controller.close() + }, + }), + }), + }) + + const fastModel = new MockLanguageModelV4({ + provider: 'mock-fast', + modelId: 'fast-model', + doStream: async () => ({ + stream: new ReadableStream({ + start(controller) { + controller.enqueue({ type: 'stream-start', warnings: [] }) + controller.enqueue({ type: 'text-start', id: 't1' }) + controller.enqueue({ type: 'text-delta', id: 't1', delta: 'fast response' }) + controller.enqueue({ + type: 'finish', + finishReason: { unified: 'stop', raw: 'stop' }, + usage: { + inputTokens: { total: 1, noCache: 1, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: 1, text: 1, reasoning: undefined }, + }, + }) + controller.close() + }, + }), + }), + }) + + let errorCalled = false + const model = createFallback({ + models: [slowModel, fastModel], + timeToFirstTokenTimeout: 100, + onError(error, modelId) { + errorCalled = true + expect(error).toBeInstanceOf(TimeToFirstTokenTimeoutError) + expect(modelId).toBe('slow-model') + }, + }) + + const result = await model.doStream({ prompt: [] }) + const reader = result.stream.getReader() + const chunks: LanguageModelV4StreamPart[] = [] + while (true) { + const { done, value } = await reader.read() + if (done) break + chunks.push(value) + } + + expect(errorCalled).toBe(true) + expect(model.currentModelIndex).toBe(1) + const textChunks = chunks.filter((c) => c.type === 'text-delta') + expect(textChunks[0]).toMatchObject({ delta: 'fast response' }) +}) + +test('doStream does NOT fall back when first token is fast enough', async () => { + const fastModel = new MockLanguageModelV4({ + provider: 'mock-fast', + modelId: 'primary-fast', + doStream: async () => ({ + stream: new ReadableStream({ + async start(controller) { + controller.enqueue({ type: 'stream-start', warnings: [] }) + // Small delay, well within timeout + await sleep(20) + controller.enqueue({ type: 'text-start', id: 't1' }) + controller.enqueue({ type: 'text-delta', id: 't1', delta: 'primary response' }) + controller.enqueue({ + type: 'finish', + finishReason: { unified: 'stop', raw: 'stop' }, + usage: { + inputTokens: { total: 1, noCache: 1, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: 1, text: 1, reasoning: undefined }, + }, + }) + controller.close() + }, + }), + }), + }) + + const fallbackModel = new MockLanguageModelV4({ + provider: 'mock-fallback', + modelId: 'fallback-model', + doStream: async () => ({ + stream: new ReadableStream({ + start(controller) { + controller.enqueue({ type: 'stream-start', warnings: [] }) + controller.enqueue({ type: 'text-start', id: 't1' }) + controller.enqueue({ type: 'text-delta', id: 't1', delta: 'fallback response' }) + controller.close() + }, + }), + }), + }) + + const model = createFallback({ + models: [fastModel, fallbackModel], + timeToFirstTokenTimeout: 2000, + }) + + const result = await model.doStream({ prompt: [] }) + const reader = result.stream.getReader() + const chunks: LanguageModelV4StreamPart[] = [] + while (true) { + const { done, value } = await reader.read() + if (done) break + chunks.push(value) + } + + expect(model.currentModelIndex).toBe(0) + const textChunks = chunks.filter((c) => c.type === 'text-delta') + expect(textChunks[0]).toMatchObject({ delta: 'primary response' }) +}) + +test('doGenerate falls back when response is too slow', async () => { + const generateResult: LanguageModelV4GenerateResult = { + content: [{ type: 'text', text: 'fast result' }], + finishReason: { unified: 'stop', raw: 'stop' }, + usage: { + inputTokens: { total: 1, noCache: 1, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: 1, text: 1, reasoning: undefined }, + }, + warnings: [], + } + + const slowModel = new MockLanguageModelV4({ + provider: 'mock-slow', + modelId: 'slow-generate', + doGenerate: async () => { + await sleep(500) + return generateResult + }, + }) + + const fastModel = new MockLanguageModelV4({ + provider: 'mock-fast', + modelId: 'fast-generate', + doGenerate: generateResult, + }) + + let errorCalled = false + const model = createFallback({ + models: [slowModel, fastModel], + timeToFirstTokenTimeout: 100, + onError(error, modelId) { + errorCalled = true + expect(error).toBeInstanceOf(TimeToFirstTokenTimeoutError) + expect(modelId).toBe('slow-generate') + }, + }) + + const result = await model.doGenerate({ prompt: [] }) + + expect(errorCalled).toBe(true) + expect(model.currentModelIndex).toBe(1) + expect(result.content).toEqual([{ type: 'text', text: 'fast result' }]) +}) + +test('TTFT timeout retries even with custom shouldRetryThisError', async () => { + const generateResult: LanguageModelV4GenerateResult = { + content: [{ type: 'text', text: 'ok' }], + finishReason: { unified: 'stop', raw: 'stop' }, + usage: { + inputTokens: { total: 1, noCache: 1, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: 1, text: 1, reasoning: undefined }, + }, + warnings: [], + } + + const slowModel = new MockLanguageModelV4({ + provider: 'mock-slow', + modelId: 'slow', + doGenerate: async () => { + await sleep(500) + return generateResult + }, + }) + const fastModel = new MockLanguageModelV4({ + provider: 'mock-fast', + modelId: 'fast', + doGenerate: generateResult, + }) + + const model = createFallback({ + models: [slowModel, fastModel], + timeToFirstTokenTimeout: 100, + // Custom predicate that rejects everything — TTFT should still retry + shouldRetryThisError: () => false, + }) + + const result = await model.doGenerate({ prompt: [] }) + expect(model.currentModelIndex).toBe(1) + expect(result.content).toEqual([{ type: 'text', text: 'ok' }]) +}) + +test('doStream TTFT covers slow doStream() startup', async () => { + // The doStream call itself is slow (simulating slow HTTP handshake) + const slowStartupModel = new MockLanguageModelV4({ + provider: 'mock-slow-startup', + modelId: 'slow-startup', + doStream: async () => { + await sleep(500) + return { + stream: new ReadableStream({ + start(controller) { + controller.enqueue({ type: 'stream-start', warnings: [] }) + controller.enqueue({ type: 'text-start', id: 't1' }) + controller.enqueue({ type: 'text-delta', id: 't1', delta: 'slow' }) + controller.enqueue({ + type: 'finish', + finishReason: { unified: 'stop', raw: 'stop' }, + usage: { + inputTokens: { total: 1, noCache: 1, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: 1, text: 1, reasoning: undefined }, + }, + }) + controller.close() + }, + }), + } + }, + }) + + const fastModel = new MockLanguageModelV4({ + provider: 'mock-fast', + modelId: 'fast-startup', + doStream: async () => ({ + stream: new ReadableStream({ + start(controller) { + controller.enqueue({ type: 'stream-start', warnings: [] }) + controller.enqueue({ type: 'text-start', id: 't1' }) + controller.enqueue({ type: 'text-delta', id: 't1', delta: 'fast' }) + controller.enqueue({ + type: 'finish', + finishReason: { unified: 'stop', raw: 'stop' }, + usage: { + inputTokens: { total: 1, noCache: 1, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: 1, text: 1, reasoning: undefined }, + }, + }) + controller.close() + }, + }), + }), + }) + + const model = createFallback({ + models: [slowStartupModel, fastModel], + timeToFirstTokenTimeout: 100, + }) + + const result = await model.doStream({ prompt: [] }) + const reader = result.stream.getReader() + const chunks: LanguageModelV4StreamPart[] = [] + while (true) { + const { done, value } = await reader.read() + if (done) break + chunks.push(value) + } + + expect(model.currentModelIndex).toBe(1) + const textChunks = chunks.filter((c) => c.type === 'text-delta') + expect(textChunks[0]).toMatchObject({ delta: 'fast' }) +}) + +test('doStream TTFT does not clear on text-start, only on text-delta', async () => { + // Model emits text-start immediately but delays before text-delta + const model_with_slow_content = new MockLanguageModelV4({ + provider: 'mock', + modelId: 'slow-content', + doStream: async () => ({ + stream: new ReadableStream({ + async start(controller) { + controller.enqueue({ type: 'stream-start', warnings: [] }) + // text-start arrives fast + controller.enqueue({ type: 'text-start', id: 't1' }) + // but actual content is slow + await sleep(500) + controller.enqueue({ type: 'text-delta', id: 't1', delta: 'slow content' }) + controller.enqueue({ + type: 'finish', + finishReason: { unified: 'stop', raw: 'stop' }, + usage: { + inputTokens: { total: 1, noCache: 1, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: 1, text: 1, reasoning: undefined }, + }, + }) + controller.close() + }, + }), + }), + }) + + const fastModel = new MockLanguageModelV4({ + provider: 'mock-fast', + modelId: 'fast-content', + doStream: async () => ({ + stream: new ReadableStream({ + start(controller) { + controller.enqueue({ type: 'stream-start', warnings: [] }) + controller.enqueue({ type: 'text-start', id: 't1' }) + controller.enqueue({ type: 'text-delta', id: 't1', delta: 'fast content' }) + controller.enqueue({ + type: 'finish', + finishReason: { unified: 'stop', raw: 'stop' }, + usage: { + inputTokens: { total: 1, noCache: 1, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: 1, text: 1, reasoning: undefined }, + }, + }) + controller.close() + }, + }), + }), + }) + + const model = createFallback({ + models: [model_with_slow_content, fastModel], + timeToFirstTokenTimeout: 100, + }) + + const result = await model.doStream({ prompt: [] }) + const reader = result.stream.getReader() + const chunks: LanguageModelV4StreamPart[] = [] + while (true) { + const { done, value } = await reader.read() + if (done) break + chunks.push(value) + } + + // Should have timed out because text-start doesn't count as content + expect(model.currentModelIndex).toBe(1) + const textChunks = chunks.filter((c) => c.type === 'text-delta') + expect(textChunks[0]).toMatchObject({ delta: 'fast content' }) +}) + test( 'handles overloaded_error from reader.read() and retries with fallback model', async () => { diff --git a/src/index.ts b/src/index.ts index 0718e87..5c82bf7 100644 --- a/src/index.ts +++ b/src/index.ts @@ -10,6 +10,8 @@ interface Settings { models: Array retryAfterOutput?: boolean modelResetInterval?: number + /** If the first token (including thinking/reasoning) takes longer than this (ms), abort and try the next model. */ + timeToFirstTokenTimeout?: number shouldRetryThisError?: (error: Error) => boolean onError?: (error: Error, modelId: string) => void | Promise } @@ -58,6 +60,15 @@ const retryableErrors = [ '504', // Gateway Timeout ] +/** Sentinel error class for TTFT timeouts. Always triggers retry regardless of + * the user-provided shouldRetryThisError predicate. */ +export class TimeToFirstTokenTimeoutError extends Error { + readonly name = 'TimeToFirstTokenTimeoutError' + constructor() { + super('Time to first token timeout') + } +} + export function defaultShouldRetryThisError(error: any): boolean { let statusCode = error?.['statusCode'] @@ -79,6 +90,58 @@ export function defaultShouldRetryThisError(error: any): boolean { return false } +/** Returns true if the stream part carries actual generated content. + * Metadata/lifecycle chunks like stream-start, text-start, text-end, + * response-metadata, raw, finish are excluded so the TTFT timer keeps + * running until real output arrives. */ +function isContentChunk(part: LanguageModelV4StreamPart): boolean { + switch (part.type) { + case 'text-delta': + case 'reasoning-delta': + case 'tool-input-delta': + case 'tool-call': + case 'tool-result': + case 'tool-approval-request': + case 'file': + case 'reasoning-file': + case 'custom': + return true + default: + return false + } +} + +/** Creates a TTFT deadline: a promise that rejects after `ms` with a + * TimeToFirstTokenTimeoutError, plus a clear() to cancel the timer. */ +function createTtftDeadline(ms: number) { + let timer: ReturnType | undefined + const promise = new Promise((_, reject) => { + timer = setTimeout(() => reject(new TimeToFirstTokenTimeoutError()), ms) + }) + return { + promise, + clear() { + if (timer !== undefined) { + clearTimeout(timer) + timer = undefined + } + }, + } +} + +/** Combines two AbortSignals so that if either fires the result fires too. */ +function composeAbortSignals(a: AbortSignal, b: AbortSignal): AbortSignal { + const controller = new AbortController() + const onAbort = () => controller.abort() + if (a.aborted || b.aborted) { + controller.abort() + return controller.signal + } + a.addEventListener('abort', onAbort, { once: true }) + b.addEventListener('abort', onAbort, { once: true }) + return controller.signal +} + function getModel(model: LanguageModelV4 | FallbackModelSettings): LanguageModelV4 { return 'model' in model ? model.model : model } @@ -176,12 +239,15 @@ export class FallbackModel implements LanguageModelV4 { return await fn() } catch (error) { lastError = error as Error - // Only retry if it's a server/capacity error - const shouldRetry = - this.settings.shouldRetryThisError || - defaultShouldRetryThisError - if (!shouldRetry(lastError)) { - throw lastError + // TTFT timeout errors always trigger retry regardless of the + // user-provided shouldRetryThisError predicate. + if (!(lastError instanceof TimeToFirstTokenTimeoutError)) { + const shouldRetry = + this.settings.shouldRetryThisError || + defaultShouldRetryThisError + if (!shouldRetry(lastError)) { + throw lastError + } } if (this.settings.onError) { @@ -217,9 +283,33 @@ export class FallbackModel implements LanguageModelV4 { doGenerate(options: LanguageModelV4CallOptions): PromiseLike { this.checkAndResetModel() - return this.retry(() => - this.currentModel.doGenerate(this.optionsForCurrentModel(options)), - ) + const ttftTimeout = this.settings.timeToFirstTokenTimeout + return this.retry(async () => { + if (ttftTimeout == null || ttftTimeout <= 0) { + return this.currentModel.doGenerate( + this.optionsForCurrentModel(options), + ) + } + const deadline = createTtftDeadline(ttftTimeout) + const abortController = new AbortController() + const mergedOptions = this.optionsForCurrentModel(options) + try { + return await Promise.race([ + this.currentModel.doGenerate({ + ...mergedOptions, + abortSignal: mergedOptions.abortSignal + ? composeAbortSignals(mergedOptions.abortSignal, abortController.signal) + : abortController.signal, + }), + deadline.promise, + ]) + } catch (error) { + abortController.abort() + throw error + } finally { + deadline.clear() + } + }) } doStream(options: LanguageModelV4CallOptions): PromiseLike { @@ -228,12 +318,48 @@ export class FallbackModel implements LanguageModelV4 { const shouldRetry = this.settings.shouldRetryThisError || defaultShouldRetryThisError return this.retry(async () => { - const result = - await self.currentModel.doStream( - self.optionsForCurrentModel(options), - ) + const ttftTimeout = self.settings.timeToFirstTokenTimeout + const hasTtft = ttftTimeout != null && ttftTimeout > 0 + + // Start the TTFT deadline *before* calling doStream so the + // initial HTTP round-trip is also covered by the timeout. + const deadline = hasTtft ? createTtftDeadline(ttftTimeout) : undefined + const abortController = hasTtft ? new AbortController() : undefined + + const mergedOptions = self.optionsForCurrentModel(options) + const optionsWithAbort = + abortController + ? { + ...mergedOptions, + abortSignal: mergedOptions.abortSignal + ? composeAbortSignals(mergedOptions.abortSignal, abortController.signal) + : abortController.signal, + } + : mergedOptions + + let result: LanguageModelV4StreamResult + try { + result = deadline + ? await Promise.race([ + self.currentModel.doStream(optionsWithAbort), + deadline.promise, + ]) + : await self.currentModel.doStream(optionsWithAbort) + } catch (error) { + // If doStream() itself times out (slow HTTP handshake), + // abort the provider request and clean up the timer. + deadline?.clear() + abortController?.abort() + throw error + } - let hasStreamedAny = false + // Two separate booleans: + // - hasEmittedOutput: has any non-stream-start chunk been sent to + // the caller? Used for retry safety (don't mix two models' output). + // - hasSatisfiedTtft: has a content-bearing chunk arrived? Used to + // clear the TTFT deadline. + let hasEmittedOutput = false + let hasSatisfiedTtft = false // Wrap the stream to handle errors and switch providers if needed const wrappedStream = new ReadableStream( { @@ -247,30 +373,50 @@ export class FallbackModel implements LanguageModelV4 { reader = result.stream.getReader() while (true) { - const result = await reader.read() + const readPromise = reader.read() + // Race reads against the TTFT deadline until + // actual content arrives. + const { done, value } = + deadline && !hasSatisfiedTtft + ? await Promise.race([readPromise, deadline.promise]) + : await readPromise - const { done, value } = result if ( - !hasStreamedAny && + !hasEmittedOutput && value && typeof value === 'object' && 'error' in value ) { const error = value.error as any if (shouldRetry(error)) { + deadline?.clear() throw error } } - if (done) break + if (done) { + deadline?.clear() + break + } controller.enqueue(value) if (value?.type !== 'stream-start') { - hasStreamedAny = true + hasEmittedOutput = true + } + + // Only clear TTFT on content-bearing chunks. + // Lifecycle/metadata chunks like stream-start, + // text-start, response-metadata don't count. + if (isContentChunk(value)) { + hasSatisfiedTtft = true + deadline?.clear() } } controller.close() } catch (error) { + deadline?.clear() + abortController?.abort() + if (self.settings.onError) { await self.settings.onError( error as Error, @@ -278,7 +424,7 @@ export class FallbackModel implements LanguageModelV4 { ) } - if (!hasStreamedAny || self.retryAfterOutput) { + if (!hasEmittedOutput || self.retryAfterOutput) { self.switchToNextModel() // TODO should be initialModel instead?