deepseek-harness/packages/llm/openrouter-usage/tests/openrouter.spec.ts

import { afterEach, describe, expect, it, vi } from 'vitest'
import {
  DEFAULT_BASE_URL,
  fetchAccountBalance,
  fetchModelPricing,
} from '../src/openrouter.ts'

afterEach(() => {
  vi.restoreAllMocks()
})

function mockFetch(status: number, body: unknown): void {
  vi.stubGlobal('fetch', vi.fn(async () => new Response(JSON.stringify(body), { status })))
}

const KEY = 'sk-test'

describe('fetchModelPricing', () => {
  it('parses per-token USD pricing keyed by model id', async () => {
    mockFetch(200, {
      data: [
        {
          id: 'deepseek/deepseek-chat',
          pricing: { prompt: '0.0000014', completion: '0.0000028', request: '0' },
        },
        // A model with disclosed cache rates and a flat per-request fee.
        {
          id: 'anthropic/claude-3.5-sonnet',
          pricing: {
            prompt: '0.000003', completion: '0.000015', request: '0.0005',
            input_cache_read: '0.0000003', input_cache_write: '0.000003',
          },
        },
      ],
    })
    const pricing = await fetchModelPricing(DEFAULT_BASE_URL, KEY, new AbortController().signal)
    expect(pricing).not.toBeUndefined()
    expect(pricing!.get('deepseek/deepseek-chat')).toEqual({ promptUsd: 1.4e-6, completionUsd: 2.8e-6, requestUsd: 0 })
    expect(pricing!.get('anthropic/claude-3.5-sonnet')).toEqual({
      promptUsd: 3e-6,
      completionUsd: 15e-6,
      requestUsd: 0.0005,
      cacheReadUsd: 3e-7,
      cacheWriteUsd: 3e-6,
    })
  })

  it('skips models with a missing id or no parseable pricing', async () => {
    mockFetch(200, {
      data: [
        { id: '', pricing: { prompt: '0.1', completion: '0.2' } },
        { id: 'no-rates', pricing: {} },
        { id: 'bad-number', pricing: { prompt: 'nope', completion: '0.2' } },
        { id: 'good/model', pricing: { prompt: '0.1', completion: '0.2' } },
      ],
    })
    const pricing = await fetchModelPricing(DEFAULT_BASE_URL, KEY, new AbortController().signal)
    expect([...pricing!.keys()]).toEqual(['good/model'])
  })

  it('returns undefined on a non-OK response', async () => {
    mockFetch(401, {})
    const pricing = await fetchModelPricing(DEFAULT_BASE_URL, KEY, new AbortController().signal)
    expect(pricing).toBeUndefined()
  })

  it('returns undefined on invalid JSON', async () => {
    vi.stubGlobal('fetch', vi.fn(async () => new Response('not json', { status: 200 })))
    const pricing = await fetchModelPricing(DEFAULT_BASE_URL, KEY, new AbortController().signal)
    expect(pricing).toBeUndefined()
  })

  it('sends the bearer token and a user-agent, and rejects redirects', async () => {
    const fetchMock = vi.fn(async () => new Response('{}', { status: 200 }))
    vi.stubGlobal('fetch', fetchMock)
    await fetchModelPricing(DEFAULT_BASE_URL, KEY, new AbortController().signal)
    const call = fetchMock.mock.calls[0] as unknown as [string, RequestInit]
    expect(call[0]).toBe(`${DEFAULT_BASE_URL}/models`)
    expect(call[1].redirect).toBe('error')
    expect((call[1].headers as Record<string, string>).authorization).toBe(`Bearer ${KEY}`)
    expect((call[1].headers as Record<string, string>)['user-agent']).toBeTruthy()
  })
})

describe('fetchAccountBalance', () => {
  function mockAccountFetch(paths: Record<string, unknown>, status = 200): void {
    const calls: string[] = []
    vi.stubGlobal('fetch', vi.fn(async (input: RequestInfo | URL) => {
      const url = String(input)
      calls.push(url)
      if (!(url in paths)) {
        throw new Error(`unexpected fetch: ${url}; seen ${calls.join(', ')}`)
      }
      const body = typeof paths[url] === 'string' ? paths[url] as string : JSON.stringify(paths[url])
      return new Response(body, { status })
    }))
  }

  it('computes the available balance (total minus spent) from /credits', async () => {
    mockAccountFetch({
      [`${DEFAULT_BASE_URL}/credits`]: {
        data: { total_credits: 12.34, total_usage: 2.5, is_free_tier: false },
      },
      [`${DEFAULT_BASE_URL}/auth/key`]: {
        data: { label: 'my key', usage: 5000, limit: 100000 },
      },
    })
    const balance = await fetchAccountBalance(DEFAULT_BASE_URL, KEY, new AbortController().signal)
    expect(balance).toEqual({ balanceUsd: 9.84, label: 'my key', usageTokens: 5000, limitTokens: 100000, isFreeTier: false })
  })

  it('clamps a momentarily under-counted usage to a zero balance', async () => {
    mockAccountFetch({
      [`${DEFAULT_BASE_URL}/credits`]: { data: { total_credits: 5, total_usage: 7 } },
      [`${DEFAULT_BASE_URL}/auth/key`]: { data: {} },
    })
    const balance = await fetchAccountBalance(DEFAULT_BASE_URL, KEY, new AbortController().signal)
    expect(balance).toEqual({ balanceUsd: 0, label: null, usageTokens: null, limitTokens: null, isFreeTier: null })
  })

  it('falls back to the /auth/key free-tier flag when /credits hides it', async () => {
    mockAccountFetch({
      [`${DEFAULT_BASE_URL}/credits`]: { data: { total_credits: 8, total_usage: 5 } },
      [`${DEFAULT_BASE_URL}/auth/key`]: { data: { is_free_tier: true } },
    })
    const balance = await fetchAccountBalance(DEFAULT_BASE_URL, KEY, new AbortController().signal)
    expect(balance).toEqual({ balanceUsd: 3, label: null, usageTokens: null, limitTokens: null, isFreeTier: true })
  })

  it('nulls the available balance when total_usage is absent', async () => {
    mockAccountFetch({
      [`${DEFAULT_BASE_URL}/credits`]: { data: { total_credits: 4 } },
      [`${DEFAULT_BASE_URL}/auth/key`]: { data: {} },
    })
    const balance = await fetchAccountBalance(DEFAULT_BASE_URL, KEY, new AbortController().signal)
    expect(balance).toEqual({ balanceUsd: null, label: null, usageTokens: null, limitTokens: null, isFreeTier: null })
  })

  it('treats a non-JSON response body as an absent envelope', async () => {
    mockAccountFetch({
      [`${DEFAULT_BASE_URL}/credits`]: 'not json {',
      [`${DEFAULT_BASE_URL}/auth/key`]: { data: { label: 'my key' } },
    })
    const balance = await fetchAccountBalance(DEFAULT_BASE_URL, KEY, new AbortController().signal)
    expect(balance).toEqual({ balanceUsd: null, label: 'my key', usageTokens: null, limitTokens: null, isFreeTier: null })
  })

  it('nulls absent or invalid fields', async () => {
    mockAccountFetch({
      [`${DEFAULT_BASE_URL}/credits`]: { data: { total_credits: 'not-a-number', total_usage: 1 } },
      [`${DEFAULT_BASE_URL}/auth/key`]: { data: {} },
    })
    const balance = await fetchAccountBalance(DEFAULT_BASE_URL, KEY, new AbortController().signal)
    expect(balance).toEqual({ balanceUsd: null, label: null, usageTokens: null, limitTokens: null, isFreeTier: null })
  })

  it('returns undefined when both envelopes lack a data object', async () => {
    mockAccountFetch({
      [`${DEFAULT_BASE_URL}/credits`]: {},
      [`${DEFAULT_BASE_URL}/auth/key`]: {},
    })
    const balance = await fetchAccountBalance(DEFAULT_BASE_URL, KEY, new AbortController().signal)
    expect(balance).toBeUndefined()
  })

  it('returns undefined when both responses are non-OK', async () => {
    mockAccountFetch({
      [`${DEFAULT_BASE_URL}/credits`]: {},
      [`${DEFAULT_BASE_URL}/auth/key`]: {},
    }, 500)
    const balance = await fetchAccountBalance(DEFAULT_BASE_URL, KEY, new AbortController().signal)
    expect(balance).toBeUndefined()
  })
})