sentryagent-idp/src/middleware/tierEnforcement.ts

/**
 * Tier enforcement middleware for SentryAgent.ai AgentIdP.
 *
 * Enforces per-tenant daily API call limits based on the tenant's tier.
 * Uses Redis keys `rate:tier:calls:<org_id>` with TTL aligned to UTC midnight.
 *
 * Behaviour:
 * - Skipped entirely when TIER_ENFORCEMENT env var is 'false'.
 * - Skipped for enterprise tenants (no limits apply).
 * - On Redis unavailability: logs the error and proceeds (fail-open).
 * - Sets X-RateLimit-Limit, X-RateLimit-Remaining, X-RateLimit-Reset on every response.
 * - Returns HTTP 429 with Retry-After header when limit is exceeded.
 */

import { Request, Response, NextFunction, RequestHandler } from 'express';
import type { RedisClientType } from 'redis';
import { Pool } from 'pg';
import { TIER_CONFIG, TierName, isTierName } from '../config/tiers.js';
import { TierLimitError } from '../utils/errors.js';

/** Redis key prefix for daily API call counters. */
const CALLS_KEY_PREFIX = 'rate:tier:calls:';

/**
 * Returns the number of seconds remaining until the next UTC midnight.
 * Used as the Redis key TTL and the Retry-After value on rejection.
 *
 * @returns Seconds until next UTC midnight (minimum 1).
 */
function secondsUntilUtcMidnight(): number {
  const now = Date.now();
  const midnight = new Date();
  midnight.setUTCHours(24, 0, 0, 0);
  return Math.max(1, Math.ceil((midnight.getTime() - now) / 1000));
}

/**
 * Returns the Unix timestamp (seconds) of the next UTC midnight.
 * Used for the X-RateLimit-Reset header.
 *
 * @returns Unix timestamp of next UTC midnight.
 */
function nextUtcMidnightTimestamp(): number {
  const midnight = new Date();
  midnight.setUTCHours(24, 0, 0, 0);
  return Math.ceil(midnight.getTime() / 1000);
}

/**
 * Fetches the tenant's current tier from the organizations table.
 * Falls back to 'free' when the tenant row is not found.
 *
 * @param pool - PostgreSQL connection pool.
 * @param orgId - The organization ID.
 * @returns The tenant's current TierName.
 */
async function getTenantTier(pool: Pool, orgId: string): Promise<TierName> {
  const result = await pool.query<{ tier: string }>(
    `SELECT tier FROM organizations WHERE organization_id = $1 LIMIT 1`,
    [orgId],
  );
  if (result.rows.length === 0) return 'free';
  const tier = result.rows[0].tier;
  return isTierName(tier) ? tier : 'free';
}

/**
 * Creates the tier enforcement middleware.
 *
 * Designed to run after auth middleware (req.user must be populated).
 * Unauthenticated requests pass through unaffected.
 *
 * @param pool - PostgreSQL connection pool (used to look up tenant tier).
 * @param redis - Redis client (used for rate counter storage).
 * @returns Express RequestHandler.
 */
export function createTierEnforcementMiddleware(
  pool: Pool,
  redis: RedisClientType,
): RequestHandler {
  return (req: Request, res: Response, next: NextFunction): void => {
    // Feature flag: bypass all tier enforcement when disabled
    if (process.env['TIER_ENFORCEMENT'] === 'false') {
      // Still set headers reflecting unlimited limits
      res.setHeader('X-RateLimit-Limit', 'unlimited');
      res.setHeader('X-RateLimit-Remaining', 'unlimited');
      res.setHeader('X-RateLimit-Reset', nextUtcMidnightTimestamp());
      next();
      return;
    }

    // Only enforce for authenticated requests
    if (!req.user?.organization_id) {
      next();
      return;
    }

    const orgId = req.user.organization_id;

    void (async (): Promise<void> => {
      try {
        const tier = await getTenantTier(pool, orgId);

        // Enterprise tenants bypass all limits
        if (tier === 'enterprise') {
          res.setHeader('X-RateLimit-Limit', 'unlimited');
          res.setHeader('X-RateLimit-Remaining', 'unlimited');
          res.setHeader('X-RateLimit-Reset', nextUtcMidnightTimestamp());
          next();
          return;
        }

        const limit = TIER_CONFIG[tier].maxCallsPerDay;
        const redisKey = `${CALLS_KEY_PREFIX}${orgId}`;
        const ttl = secondsUntilUtcMidnight();
        const resetAt = nextUtcMidnightTimestamp();

        let currentCount: number;

        try {
          // Atomically increment and set TTL aligned to UTC midnight.
          // INCR returns the new value after increment.
          const newCount = await redis.incr(redisKey);
          currentCount = newCount;

          // Set TTL only on the first increment to avoid resetting the window
          // on every request. If the key was brand new, newCount === 1.
          if (newCount === 1) {
            await redis.expire(redisKey, ttl);
          }
        } catch (redisErr) {
          // Redis unavailable — fail-open: log and proceed without rate limiting
          // eslint-disable-next-line no-console
          console.error('[tierEnforcement] Redis error — proceeding fail-open:', redisErr);
          res.setHeader('X-RateLimit-Limit', limit);
          res.setHeader('X-RateLimit-Remaining', 0);
          res.setHeader('X-RateLimit-Reset', resetAt);
          next();
          return;
        }

        const remaining = Math.max(0, limit - currentCount);

        // Set rate limit headers on all responses
        res.setHeader('X-RateLimit-Limit', limit);
        res.setHeader('X-RateLimit-Remaining', remaining);
        res.setHeader('X-RateLimit-Reset', resetAt);

        // Reject if the new count exceeds the limit
        if (currentCount > limit) {
          res.setHeader('Retry-After', ttl);
          next(new TierLimitError('API call', limit, { orgId, tier }));
          return;
        }

        next();
      } catch (err) {
        next(err);
      }
    })();
  };
}