Momento/memento-note/tests/unit/fallback.test.ts

import { describe, it, expect, vi, beforeEach } from 'vitest'
import { QuotaExceededError } from '@/lib/entitlements'
import { APICallError } from 'ai'
import type { AIProvider } from '@/lib/ai/types'

const mockState = vi.hoisted(() => {
  const stubProvider = (): AIProvider =>
    ({
      generateTags: async () => [],
      getEmbeddings: async () => [],
      generateTitles: async () => [],
      generateText: async () => '',
      chat: async () => ({ text: '' }),
      getModel: () => ({}),
    }) as AIProvider
  const primaryProvider = stubProvider()
  const secondaryProvider = stubProvider()
  let instanceCall = 0
  return {
    primaryProvider,
    secondaryProvider,
    reset: () => {
      instanceCall = 0
    },
    nextProvider: () => {
      instanceCall++
      return instanceCall === 1 ? primaryProvider : secondaryProvider
    },
    getCallCount: () => instanceCall,
  }
})

vi.mock('@/lib/ai/factory', () => ({
  getProviderInstance: vi.fn(() => mockState.nextProvider()),
}))

import {
  resolveAiFallbackRoute,
  isRetriableProviderError,
  extractProviderErrorStatus,
  withAiProviderFallback,
  FALLBACK_BUDGET_MS,
} from '@/lib/ai/fallback'
import { getProviderInstance } from '@/lib/ai/factory'

describe('resolveAiFallbackRoute', () => {
  it('returns null when fallback provider is unset', () => {
    expect(resolveAiFallbackRoute('chat', { AI_PROVIDER_CHAT: 'openai' })).toBeNull()
  })

  it('resolves chat fallback with slash model ids', () => {
    const route = resolveAiFallbackRoute('chat', {
      AI_PROVIDER_CHAT: 'openai',
      AI_MODEL_CHAT: 'gpt-4o-mini',
      AI_PROVIDER_CHAT_FALLBACK: 'openrouter',
      AI_MODEL_CHAT_FALLBACK: 'deepseek/deepseek-chat',
      AI_MODEL_EMBEDDING: 'text-embedding-3-small',
    })
    expect(route).not.toBeNull()
    expect(route!.providerType).toBe('openrouter')
    expect(route!.modelName).toBe('deepseek/deepseek-chat')
  })

  it('rejects anthropic embedding fallback', () => {
    expect(() =>
      resolveAiFallbackRoute('embedding', {
        AI_PROVIDER_EMBEDDING: 'openai',
        AI_MODEL_EMBEDDING: 'x',
        AI_MODEL_TAGS: 'y',
        AI_PROVIDER_EMBEDDING_FALLBACK: 'anthropic',
      })
    ).toThrow(/AI_PROVIDER_EMBEDDING_FALLBACK cannot use/)
  })

  it('throws on unknown fallback provider', () => {
    expect(() =>
      resolveAiFallbackRoute('chat', {
        AI_PROVIDER_CHAT: 'openai',
        AI_MODEL_CHAT: 'gpt-4o-mini',
        AI_PROVIDER_CHAT_FALLBACK: 'invalid_provider',
        AI_MODEL_EMBEDDING: 'text-embedding-3-small',
      })
    ).toThrow(/Unknown fallback provider/)
  })

  it('returns null when fallback is same provider as primary', () => {
    const route = resolveAiFallbackRoute('chat', {
      AI_PROVIDER_CHAT: 'openai',
      AI_MODEL_CHAT: 'gpt-4o-mini',
      AI_PROVIDER_CHAT_FALLBACK: 'openai',
      AI_MODEL_CHAT_FALLBACK: 'gpt-4.1-mini',
      AI_MODEL_EMBEDDING: 'text-embedding-3-small',
    })
    expect(route).toBeNull()
  })
})

describe('isRetriableProviderError', () => {
  it('treats 429 and 5xx as retriable', () => {
    expect(isRetriableProviderError(new APICallError({ message: 'rate', statusCode: 429 }))).toBe(true)
    expect(isRetriableProviderError(new APICallError({ message: 'srv', statusCode: 500 }))).toBe(true)
    expect(isRetriableProviderError(new APICallError({ message: 'srv', statusCode: 503 }))).toBe(true)
  })

  it('treats 401, 400, 402, 403 as non-retriable', () => {
    expect(isRetriableProviderError(new APICallError({ message: 'auth', statusCode: 401 }))).toBe(false)
    expect(isRetriableProviderError(new APICallError({ message: 'bad', statusCode: 400 }))).toBe(false)
    expect(isRetriableProviderError(new APICallError({ message: 'forbidden', statusCode: 403 }))).toBe(false)
    expect(
      isRetriableProviderError(new QuotaExceededError('PRO', 'chat', 10, 10, false))
    ).toBe(false)
    expect(extractProviderErrorStatus(new QuotaExceededError('PRO', 'chat', 10, 10, false))).toBe(402)
  })

  it('handles circular cause chain without stack overflow', () => {
    const err: { cause: unknown } = { cause: null }
    err.cause = err
    expect(extractProviderErrorStatus(err)).toBeUndefined()
    expect(isRetriableProviderError(err)).toBe(false)
  })

  it('handles nested cause chain', () => {
    const deep = { statusCode: 429 }
    const mid = { cause: deep }
    const top = { cause: mid }
    expect(extractProviderErrorStatus(top)).toBe(429)
  })
})

describe('withAiProviderFallback', () => {
  beforeEach(() => {
    mockState.reset()
    vi.mocked(getProviderInstance).mockImplementation(() => mockState.nextProvider())
  })

  it('falls back on retriable primary failure within NFR-R1 budget', async () => {
    const cfg = {
      AI_PROVIDER_CHAT: 'openai',
      AI_MODEL_CHAT: 'gpt-4o-mini',
      AI_PROVIDER_CHAT_FALLBACK: 'deepseek',
      AI_MODEL_CHAT_FALLBACK: 'deepseek-chat',
      AI_MODEL_EMBEDDING: 'text-embedding-3-small',
    }

    const t0 = performance.now()
    const result = await withAiProviderFallback('chat', cfg, async (provider) => {
      if (provider === mockState.primaryProvider) {
        throw new APICallError({ message: 'rate limited', statusCode: 429 })
      }
      return 'secondary-ok'
    })
    const elapsed = performance.now() - t0

    expect(result).toBe('secondary-ok')
    expect(mockState.getCallCount()).toBe(2)
    expect(elapsed).toBeLessThan(FALLBACK_BUDGET_MS + 250)
  })

  it('rethrows primary error when no fallback configured', async () => {
    const cfg = {
      AI_PROVIDER_CHAT: 'openai',
      AI_MODEL_CHAT: 'gpt-4o-mini',
      AI_MODEL_EMBEDDING: 'text-embedding-3-small',
    }
    await expect(
      withAiProviderFallback('chat', cfg, async () => {
        throw new APICallError({ message: 'down', statusCode: 503 })
      })
    ).rejects.toMatchObject({ statusCode: 503 })
    expect(mockState.getCallCount()).toBe(1)
  })

  it('rethrows primary error when fallback is same provider', async () => {
    const cfg = {
      AI_PROVIDER_CHAT: 'openai',
      AI_MODEL_CHAT: 'gpt-4o-mini',
      AI_PROVIDER_CHAT_FALLBACK: 'openai',
      AI_MODEL_CHAT_FALLBACK: 'gpt-4.1-mini',
      AI_MODEL_EMBEDDING: 'text-embedding-3-small',
    }
    await expect(
      withAiProviderFallback('chat', cfg, async () => {
        throw new APICallError({ message: 'rate', statusCode: 429 })
      })
    ).rejects.toMatchObject({ statusCode: 429 })
    expect(mockState.getCallCount()).toBe(1)
  })

  it('skipSystemFallback runs primary only', async () => {
    const cfg = {
      AI_PROVIDER_CHAT: 'openai',
      AI_MODEL_CHAT: 'gpt-4o-mini',
      AI_MODEL_EMBEDDING: 'text-embedding-3-small',
      AI_PROVIDER_CHAT_FALLBACK: 'deepseek',
    }
    await expect(
      withAiProviderFallback(
        'chat',
        cfg,
        async () => {
          throw new APICallError({ message: 'rate', statusCode: 429 })
        },
        { skipSystemFallback: true }
      )
    ).rejects.toMatchObject({ statusCode: 429 })
    expect(mockState.getCallCount()).toBe(1)
  })

  it('rethrows primary error when resolveAiFallbackRoute throws (config error)', async () => {
    const cfg = {
      AI_PROVIDER_EMBEDDING: 'openai',
      AI_MODEL_EMBEDDING: 'text-embedding-3-small',
      AI_MODEL_TAGS: 'y',
      AI_PROVIDER_EMBEDDING_FALLBACK: 'anthropic',
    }
    await expect(
      withAiProviderFallback('embedding', cfg, async () => {
        throw new APICallError({ message: 'rate', statusCode: 429 })
      })
    ).rejects.toMatchObject({ statusCode: 429 })
    expect(mockState.getCallCount()).toBe(1)
  })
})