dflike/server/src/llm/__tests__/generationQueue.test.ts

import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import { createGenerationQueue } from '../generationQueue.js';
import type { OpenRouterClient } from '../openRouterClient.js';

function mockClient(response: string | null = 'generated text'): OpenRouterClient {
  return {
    complete: vi.fn().mockResolvedValue(response),
  };
}

describe('generationQueue', () => {
  beforeEach(() => {
    vi.useFakeTimers();
  });

  afterEach(() => {
    vi.useRealTimers();
  });

  it('processes a queued request', async () => {
    const client = mockClient('result');
    const queue = createGenerationQueue(client, { requestsPerMinute: 60 });

    const promise = queue.enqueue({
      system: 'sys',
      user: 'usr',
    });

    // Advance past the rate limit interval
    await vi.advanceTimersByTimeAsync(100);
    const result = await promise;

    expect(result).toBe('result');
    expect(client.complete).toHaveBeenCalledOnce();
  });

  it('processes requests in order', async () => {
    const results: string[] = [];
    const client: OpenRouterClient = {
      complete: vi.fn()
        .mockResolvedValueOnce('first')
        .mockResolvedValueOnce('second'),
    };
    const queue = createGenerationQueue(client, { requestsPerMinute: 60 });

    const p1 = queue.enqueue({ system: 's', user: 'u1' }).then(r => { if (r) results.push(r); });
    const p2 = queue.enqueue({ system: 's', user: 'u2' }).then(r => { if (r) results.push(r); });

    await vi.advanceTimersByTimeAsync(2000);
    await Promise.all([p1, p2]);

    expect(results).toEqual(['first', 'second']);
  });

  it('respects rate limiting', async () => {
    const client = mockClient('ok');
    // 6 per minute = 1 per 10 seconds
    const queue = createGenerationQueue(client, { requestsPerMinute: 6 });

    queue.enqueue({ system: 's', user: '1' });
    queue.enqueue({ system: 's', user: '2' });
    queue.enqueue({ system: 's', user: '3' });

    // After 1 second, only 1 should have been processed
    await vi.advanceTimersByTimeAsync(1000);
    expect(client.complete).toHaveBeenCalledTimes(1);

    // After 11 seconds, 2 should be processed
    await vi.advanceTimersByTimeAsync(10000);
    expect(client.complete).toHaveBeenCalledTimes(2);
  });

  it('returns null when client fails', async () => {
    const client = mockClient(null);
    const queue = createGenerationQueue(client, { requestsPerMinute: 60 });

    const promise = queue.enqueue({ system: 's', user: 'u' });
    await vi.advanceTimersByTimeAsync(100);
    const result = await promise;

    expect(result).toBeNull();
  });

  it('reports queue depth', async () => {
    const client = mockClient('ok');
    const queue = createGenerationQueue(client, { requestsPerMinute: 6 });

    expect(queue.depth()).toBe(0);

    queue.enqueue({ system: 's', user: '1' });
    queue.enqueue({ system: 's', user: '2' });

    // First processes immediately, second is queued
    await vi.advanceTimersByTimeAsync(100);
    expect(queue.depth()).toBe(1);
  });

  it('exposes isDailyLimitReached as false initially', () => {
    const client = mockClient('ok');
    const queue = createGenerationQueue(client, { requestsPerMinute: 60, requestsPerDay: 5 });
    expect(queue.isDailyLimitReached()).toBe(false);
  });

  it('resolves with null when daily limit is reached', async () => {
    const client = mockClient('ok');
    const queue = createGenerationQueue(client, { requestsPerMinute: 60, requestsPerDay: 2 });

    // First two requests should succeed
    const p1 = queue.enqueue({ system: 's', user: '1' });
    await vi.advanceTimersByTimeAsync(100);
    expect(await p1).toBe('ok');

    const p2 = queue.enqueue({ system: 's', user: '2' });
    await vi.advanceTimersByTimeAsync(100);
    expect(await p2).toBe('ok');

    // Third request should be rejected (daily limit reached)
    expect(queue.isDailyLimitReached()).toBe(true);
    const p3 = queue.enqueue({ system: 's', user: '3' });
    const result = await p3;
    expect(result).toBeNull();
    expect(client.complete).toHaveBeenCalledTimes(2);
  });

  it('resets daily count at midnight', async () => {
    // Set time to 23:59:59
    const nearMidnight = new Date();
    nearMidnight.setHours(23, 59, 59, 0);
    vi.setSystemTime(nearMidnight);

    const client = mockClient('ok');
    const queue = createGenerationQueue(client, { requestsPerMinute: 60, requestsPerDay: 1 });

    // Use up the daily limit
    const p1 = queue.enqueue({ system: 's', user: '1' });
    await vi.advanceTimersByTimeAsync(100);
    expect(await p1).toBe('ok');
    expect(queue.isDailyLimitReached()).toBe(true);

    // Advance past midnight
    await vi.advanceTimersByTimeAsync(2000);

    // Should be able to make requests again
    expect(queue.isDailyLimitReached()).toBe(false);
    const p2 = queue.enqueue({ system: 's', user: '2' });
    await vi.advanceTimersByTimeAsync(100);
    expect(await p2).toBe('ok');
  });

  it('works without requestsPerDay (no daily limit)', async () => {
    const client = mockClient('ok');
    const queue = createGenerationQueue(client, { requestsPerMinute: 60 });

    expect(queue.isDailyLimitReached()).toBe(false);

    // Should process requests without daily limit
    const p1 = queue.enqueue({ system: 's', user: '1' });
    await vi.advanceTimersByTimeAsync(100);
    expect(await p1).toBe('ok');
    expect(queue.isDailyLimitReached()).toBe(false);
  });

  it('can be cleared', async () => {
    const client = mockClient('ok');
    const queue = createGenerationQueue(client, { requestsPerMinute: 6 });

    const p1 = queue.enqueue({ system: 's', user: '1' });
    const p2 = queue.enqueue({ system: 's', user: '2' });
    const p3 = queue.enqueue({ system: 's', user: '3' });

    // Process first
    await vi.advanceTimersByTimeAsync(100);

    // Clear remaining
    queue.clear();
    expect(queue.depth()).toBe(0);

    // Cleared items resolve as null
    await vi.advanceTimersByTimeAsync(30000);
    expect(await p2).toBeNull();
    expect(await p3).toBeNull();
  });
});