← back to Norma

agents/shared/rate-limiter.js

243 lines

/**
 * Norma Rate Limiter
 *
 * Token-bucket rate limiter backed by the sdcc_agent_rate_limits table.
 * Prevents platform bans by enforcing per-agent, per-platform action limits.
 */

const { query } = require('./db');

// Default rate limit windows per platform
// Format: { actionType: { windowSeconds, maxActions } }
const DEFAULT_LIMITS = {
  moveon: {
    post: { windowSeconds: 86400, maxActions: 1 },
    monitor: { windowSeconds: 3600, maxActions: 2 },
  },
  reddit: {
    post: { windowSeconds: 600, maxActions: 1 },
    reply: { windowSeconds: 300, maxActions: 1 },
    search: { windowSeconds: 60, maxActions: 60 },
  },
  discord: {
    post: { windowSeconds: 5, maxActions: 5 },
    monitor: { windowSeconds: 1, maxActions: 999999 }, // effectively always allowed
  },
  twitter: {
    post: { windowSeconds: 86400, maxActions: 8 },
    reply: { windowSeconds: 86400, maxActions: 15 },
    search: { windowSeconds: 900, maxActions: 60 },
  },
  twitch: {
    post: { windowSeconds: 30, maxActions: 20 },
    monitor: { windowSeconds: 1800, maxActions: 30 },
  },
};

class RateLimiter {
  /**
   * @param {string} agent - Agent name (e.g. 'reddit-agent')
   * @param {string} platform - Platform name (e.g. 'reddit')
   */
  constructor(agent, platform) {
    this.agent = agent;
    this.platform = platform.toLowerCase();
    this.limits = DEFAULT_LIMITS[this.platform] || {};
  }

  /**
   * Get or create the rate limit row for an action type.
   * If the window has expired, auto-resets the counter.
   *
   * @param {string} actionType - Action type (e.g. 'post', 'reply', 'search')
   * @returns {Promise<Object>} The rate limit row
   */
  async _getOrCreateLimit(actionType) {
    const defaults = this.limits[actionType] || { windowSeconds: 3600, maxActions: 10 };

    // Try to get existing row
    const { rows } = await query(
      `SELECT * FROM sdcc_agent_rate_limits
       WHERE agent = $1 AND platform = $2 AND action_type = $3`,
      [this.agent, this.platform, actionType]
    );

    if (rows.length === 0) {
      // Create new row with defaults
      const { rows: created } = await query(
        `INSERT INTO sdcc_agent_rate_limits
          (agent, platform, action_type, window_start, window_seconds, max_actions, current_count)
         VALUES ($1, $2, $3, NOW(), $4, $5, 0)
         ON CONFLICT (agent, platform, action_type)
         DO UPDATE SET updated_at = NOW()
         RETURNING *`,
        [this.agent, this.platform, actionType, defaults.windowSeconds, defaults.maxActions]
      );
      return created[0];
    }

    const row = rows[0];

    // Check if window has expired — auto-reset
    const windowEnd = new Date(row.window_start);
    windowEnd.setSeconds(windowEnd.getSeconds() + row.window_seconds);

    if (new Date() > windowEnd) {
      const { rows: reset } = await query(
        `UPDATE sdcc_agent_rate_limits
         SET window_start = NOW(), current_count = 0, is_blocked = false, blocked_until = NULL, updated_at = NOW()
         WHERE id = $1
         RETURNING *`,
        [row.id]
      );
      return reset[0];
    }

    return row;
  }

  /**
   * Check if an action is allowed under current rate limits.
   *
   * @param {string} actionType - Action type to check
   * @returns {Promise<{ allowed: boolean, retryAfterMs: number, current: number, max: number }>}
   */
  async checkLimit(actionType) {
    try {
      const limit = await this._getOrCreateLimit(actionType);

      // Check if explicitly blocked (e.g. platform ban)
      if (limit.is_blocked) {
        const blockedUntil = limit.blocked_until ? new Date(limit.blocked_until) : null;
        const now = new Date();

        if (blockedUntil && now > blockedUntil) {
          // Block expired, clear it
          await query(
            `UPDATE sdcc_agent_rate_limits
             SET is_blocked = false, blocked_until = NULL, updated_at = NOW()
             WHERE id = $1`,
            [limit.id]
          );
        } else {
          const retryAfterMs = blockedUntil
            ? blockedUntil.getTime() - now.getTime()
            : 3600000; // default 1 hour if no end time
          return {
            allowed: false,
            retryAfterMs,
            current: limit.current_count,
            max: limit.max_actions,
            reason: 'blocked',
          };
        }
      }

      // Check count against max
      if (limit.current_count >= limit.max_actions) {
        const windowEnd = new Date(limit.window_start);
        windowEnd.setSeconds(windowEnd.getSeconds() + limit.window_seconds);
        const retryAfterMs = Math.max(0, windowEnd.getTime() - Date.now());

        return {
          allowed: false,
          retryAfterMs,
          current: limit.current_count,
          max: limit.max_actions,
          reason: 'rate_limited',
        };
      }

      return {
        allowed: true,
        retryAfterMs: 0,
        current: limit.current_count,
        max: limit.max_actions,
      };
    } catch (err) {
      console.error(`[${new Date().toISOString()}] [rate-limiter] checkLimit error:`, err.message);
      // Fail open — allow the action if we can't check limits
      return { allowed: true, retryAfterMs: 0, current: 0, max: 0 };
    }
  }

  /**
   * Record that an action was taken. Increments the counter.
   *
   * @param {string} actionType - Action type to record
   * @returns {Promise<Object>} Updated rate limit row
   */
  async recordAction(actionType) {
    try {
      // Ensure the row exists and window is current
      await this._getOrCreateLimit(actionType);

      const { rows } = await query(
        `UPDATE sdcc_agent_rate_limits
         SET current_count = current_count + 1, updated_at = NOW()
         WHERE agent = $1 AND platform = $2 AND action_type = $3
         RETURNING *`,
        [this.agent, this.platform, actionType]
      );
      return rows[0];
    } catch (err) {
      console.error(`[${new Date().toISOString()}] [rate-limiter] recordAction error:`, err.message);
      return null;
    }
  }

  /**
   * Mark this agent/platform/action as blocked (e.g. platform ban detected).
   *
   * @param {string} actionType - Action type to block
   * @param {Date|string} blockedUntil - When the block expires
   * @returns {Promise<Object>} Updated rate limit row
   */
  async markBlocked(actionType, blockedUntil) {
    try {
      await this._getOrCreateLimit(actionType);

      const { rows } = await query(
        `UPDATE sdcc_agent_rate_limits
         SET is_blocked = true, blocked_until = $4, updated_at = NOW()
         WHERE agent = $1 AND platform = $2 AND action_type = $3
         RETURNING *`,
        [this.agent, this.platform, actionType, blockedUntil]
      );

      console.warn(
        `[${new Date().toISOString()}] [rate-limiter] BLOCKED: ${this.agent}/${this.platform}/${actionType} until ${blockedUntil}`
      );

      return rows[0];
    } catch (err) {
      console.error(`[${new Date().toISOString()}] [rate-limiter] markBlocked error:`, err.message);
      return null;
    }
  }

  /**
   * Reset the rate limit window for an action type.
   *
   * @param {string} actionType - Action type to reset
   * @returns {Promise<Object>} Updated rate limit row
   */
  async resetWindow(actionType) {
    try {
      const { rows } = await query(
        `UPDATE sdcc_agent_rate_limits
         SET window_start = NOW(), current_count = 0, is_blocked = false, blocked_until = NULL, updated_at = NOW()
         WHERE agent = $1 AND platform = $2 AND action_type = $3
         RETURNING *`,
        [this.agent, this.platform, actionType]
      );
      return rows[0] || null;
    } catch (err) {
      console.error(`[${new Date().toISOString()}] [rate-limiter] resetWindow error:`, err.message);
      return null;
    }
  }
}

module.exports = { RateLimiter, DEFAULT_LIMITS };