mcpengine/servers/google-console/src/lib/rate-limit.ts

/**
 * Token bucket rate limiter for Google Search Console API quotas
 * Implements per-API-type tracking with automatic token refill
 */

export type APIType = 'gsc' | 'inspection' | 'indexing';

interface BucketConfig {
  capacity: number;
  refillRate: number; // tokens per millisecond
  refillInterval: number; // interval in ms
}

interface TokenBucket {
  tokens: number;
  lastRefill: number;
  config: BucketConfig;
  queue: Array<{ resolve: () => void; tokens: number }>;
}

export interface RateLimitConfig {
  /** GSC API: 1200 queries/minute (default) */
  gscQueries?: number;
  /** URL Inspection: 600/day per property, 2000/day total */
  inspectionQueries?: number;
  /** Indexing API: 200/day */
  indexingQueries?: number;
}

export class RateLimiter {
  private buckets = new Map<APIType, TokenBucket>();

  constructor(config: RateLimitConfig = {}) {
    // GSC API: 1200 queries/minute
    const gscQueriesPerMin = config.gscQueries ?? 1200;
    this.buckets.set('gsc', {
      tokens: gscQueriesPerMin,
      lastRefill: Date.now(),
      config: {
        capacity: gscQueriesPerMin,
        refillRate: gscQueriesPerMin / (60 * 1000), // per ms
        refillInterval: 60 * 1000 // 1 minute
      },
      queue: []
    });

    // URL Inspection: 600/day (simplified - not tracking per-property)
    const inspectionQueriesPerDay = config.inspectionQueries ?? 600;
    this.buckets.set('inspection', {
      tokens: inspectionQueriesPerDay,
      lastRefill: Date.now(),
      config: {
        capacity: inspectionQueriesPerDay,
        refillRate: inspectionQueriesPerDay / (24 * 60 * 60 * 1000), // per ms
        refillInterval: 24 * 60 * 60 * 1000 // 1 day
      },
      queue: []
    });

    // Indexing API: 200/day
    const indexingQueriesPerDay = config.indexingQueries ?? 200;
    this.buckets.set('indexing', {
      tokens: indexingQueriesPerDay,
      lastRefill: Date.now(),
      config: {
        capacity: indexingQueriesPerDay,
        refillRate: indexingQueriesPerDay / (24 * 60 * 60 * 1000), // per ms
        refillInterval: 24 * 60 * 60 * 1000 // 1 day
      },
      queue: []
    });
  }

  /**
   * Refill tokens in a bucket based on elapsed time
   */
  private refillBucket(bucket: TokenBucket): void {
    const now = Date.now();
    const elapsed = now - bucket.lastRefill;
    const tokensToAdd = elapsed * bucket.config.refillRate;
    
    bucket.tokens = Math.min(
      bucket.config.capacity,
      bucket.tokens + tokensToAdd
    );
    bucket.lastRefill = now;
  }

  /**
   * Process queued requests if tokens available
   */
  private processQueue(apiType: APIType): void {
    const bucket = this.buckets.get(apiType);
    if (!bucket) return;

    this.refillBucket(bucket);

    while (bucket.queue.length > 0 && bucket.tokens >= bucket.queue[0].tokens) {
      const item = bucket.queue.shift()!;
      bucket.tokens -= item.tokens;
      item.resolve();
    }
  }

  /**
   * Acquire tokens for an API call
   * Waits if insufficient tokens are available
   */
  async acquire(apiType: APIType, tokens: number = 1): Promise<void> {
    const bucket = this.buckets.get(apiType);
    if (!bucket) {
      throw new Error(`Unknown API type: ${apiType}`);
    }

    this.refillBucket(bucket);

    // If enough tokens available, consume immediately
    if (bucket.tokens >= tokens) {
      bucket.tokens -= tokens;
      return;
    }

    // Otherwise, queue the request
    return new Promise<void>((resolve) => {
      bucket.queue.push({ resolve, tokens });
      
      // Set up periodic processing
      const interval = setInterval(() => {
        this.processQueue(apiType);
        
        // Stop if this request was processed
        if (!bucket.queue.find(item => item.resolve === resolve)) {
          clearInterval(interval);
        }
      }, 100); // Check every 100ms
    });
  }

  /**
   * Get current token counts for monitoring
   */
  getStatus(): Record<APIType, { tokens: number; capacity: number; queued: number }> {
    const status = {} as Record<APIType, { tokens: number; capacity: number; queued: number }>;
    
    for (const [apiType, bucket] of this.buckets.entries()) {
      this.refillBucket(bucket);
      status[apiType] = {
        tokens: Math.floor(bucket.tokens),
        capacity: bucket.config.capacity,
        queued: bucket.queue.length
      };
    }
    
    return status;
  }

  /**
   * Reset all buckets (useful for testing)
   */
  reset(): void {
    for (const bucket of this.buckets.values()) {
      bucket.tokens = bucket.config.capacity;
      bucket.lastRefill = Date.now();
      bucket.queue = [];
    }
  }
}