EJClaw/src/claude-usage-backoff.test.ts

/**
 * Regression tests for the Claude usage-API 429 `Retry-After` backoff.
 *
 * Root cause this guards against: the usage poller used to retry every
 * MIN_FETCH_INTERVAL_MS (60s), but the endpoint hands back a longer rate-limit
 * window (e.g. `retry-after: 93`). Retrying mid-cooldown re-trips the limit, so
 * the 429s never clear. The fix records a `cooldownUntil` from the header and
 * refuses to call the API until that window closes.
 */
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';

vi.mock('./logger.js', () => ({
  logger: { info: vi.fn(), warn: vi.fn(), error: vi.fn(), debug: vi.fn() },
}));

vi.mock('./config.js', () => ({ DATA_DIR: '/tmp/ejclaw-usage-backoff-data' }));

const { getAllTokensMock } = vi.hoisted(() => ({ getAllTokensMock: vi.fn() }));
vi.mock('./token-rotation.js', () => ({
  getAllTokens: getAllTokensMock,
  getCurrentToken: () => null,
  getConfiguredClaudeTokens: () => [],
}));

const { readJsonFileMock } = vi.hoisted(() => ({ readJsonFileMock: vi.fn() }));
vi.mock('./utils.js', async () => {
  const actual =
    await vi.importActual<typeof import('./utils.js')>('./utils.js');
  return { ...actual, readJsonFile: readJsonFileMock, writeJsonFile: vi.fn() };
});

afterEach(() => {
  vi.resetModules();
  vi.clearAllMocks();
  vi.unstubAllGlobals();
  vi.useRealTimers();
});

function makeResponse(
  status: number,
  headers: Record<string, string> = {},
): Response {
  return {
    status,
    ok: status >= 200 && status < 300,
    headers: { get: (k: string) => headers[k.toLowerCase()] ?? null },
    json: async () => ({}),
  } as unknown as Response;
}

describe('parseRetryAfterMs', () => {
  it('parses delta-seconds', async () => {
    const { parseRetryAfterMs } = await import('./claude-usage.js');
    expect(parseRetryAfterMs('93')).toBe(93_000);
    expect(parseRetryAfterMs('0')).toBe(0);
  });

  it('parses an HTTP-date relative to now', async () => {
    vi.useFakeTimers();
    vi.setSystemTime(new Date('2026-06-20T00:00:00Z'));
    const { parseRetryAfterMs } = await import('./claude-usage.js');
    expect(parseRetryAfterMs('Sat, 20 Jun 2026 00:00:30 GMT')).toBe(30_000);
  });

  it('returns null for missing or junk values', async () => {
    const { parseRetryAfterMs } = await import('./claude-usage.js');
    expect(parseRetryAfterMs(null)).toBeNull();
    expect(parseRetryAfterMs('soon')).toBeNull();
  });
});

describe('Claude usage 429 Retry-After backoff', () => {
  beforeEach(() => {
    readJsonFileMock.mockReturnValue({}); // start from an empty disk cache
    getAllTokensMock.mockReturnValue([
      {
        index: 0,
        token: 'sk-ant-oat01-testTokenAAAAAAAA',
        masked: 'sk-ant-oat01-test...AAAA',
        isActive: true,
        isRateLimited: false,
      },
    ]);
  });

  it('records a cooldown on 429 and skips the next fetch within the window', async () => {
    const fetchMock = vi.fn(async () =>
      makeResponse(429, { 'retry-after': '93' }),
    );
    vi.stubGlobal('fetch', fetchMock);

    const { fetchAllClaudeUsage } = await import('./claude-usage.js');

    const first = await fetchAllClaudeUsage();
    expect(fetchMock).toHaveBeenCalledTimes(1);
    expect(first[0].usageRateLimited).toBe(true);
    expect(first[0].usage).toBeNull();

    // Immediate re-call: cooldown gate must hold, no second network call.
    const second = await fetchAllClaudeUsage();
    expect(fetchMock).toHaveBeenCalledTimes(1);
    expect(second[0].usageRateLimited).toBe(true);
  });

  it('falls back to the default cooldown when Retry-After is 0 (does not disable backoff)', async () => {
    // Production regression: the usage endpoint was returning 429 with
    // `retry-after: 0`, which made `retryAfterMs ?? DEFAULT` evaluate to 0 — a
    // ~5s cooldown. The poller then re-hit the endpoint every 60s and re-tripped
    // the 429 indefinitely. A non-positive Retry-After must use the default.
    vi.useFakeTimers();
    vi.setSystemTime(new Date('2026-06-20T00:00:00Z'));

    const fetchMock = vi.fn(async () =>
      makeResponse(429, { 'retry-after': '0' }),
    );
    vi.stubGlobal('fetch', fetchMock);

    const { fetchAllClaudeUsage } = await import('./claude-usage.js');

    await fetchAllClaudeUsage();
    expect(fetchMock).toHaveBeenCalledTimes(1);

    // +70s: past the 60s throttle. With the bug (0 cooldown) this would refetch;
    // with the default 5-min cooldown it must still be held.
    vi.setSystemTime(Date.now() + 70_000);
    const held = await fetchAllClaudeUsage();
    expect(fetchMock).toHaveBeenCalledTimes(1);
    expect(held[0].usageRateLimited).toBe(true);
  });

  it('keeps backing off after the 60s throttle elapses but before the cooldown closes', async () => {
    // This is the core regression: a 93s Retry-After must outlast the 60s
    // MIN_FETCH_INTERVAL throttle. Without the cooldown the poller would refetch
    // at ~60s — mid rate-limit window — and re-trip the 429 forever.
    vi.useFakeTimers();
    vi.setSystemTime(new Date('2026-06-20T00:00:00Z'));

    const fetchMock = vi
      .fn()
      .mockResolvedValueOnce(makeResponse(429, { 'retry-after': '93' }))
      .mockResolvedValueOnce(makeResponse(200));
    vi.stubGlobal('fetch', fetchMock);

    const { fetchAllClaudeUsage } = await import('./claude-usage.js');

    await fetchAllClaudeUsage();
    expect(fetchMock).toHaveBeenCalledTimes(1);

    // +70s: past the 60s throttle, still inside the ~98s cooldown → no refetch.
    vi.setSystemTime(Date.now() + 70_000);
    const held = await fetchAllClaudeUsage();
    expect(fetchMock).toHaveBeenCalledTimes(1);
    expect(held[0].usageRateLimited).toBe(true);

    // +100s total: cooldown window closed → poller is allowed to fetch again.
    vi.setSystemTime(Date.now() + 30_000);
    const resumed = await fetchAllClaudeUsage();
    expect(fetchMock).toHaveBeenCalledTimes(2);
    expect(resumed[0].usageRateLimited).toBe(false);
  });
});