Rate Limiting Extension

The @bantai-dev/with-rate-limit package adds rate limiting capabilities to your Bantai contexts with support for multiple rate limiting strategies including fixed-window, sliding-window, and token-bucket algorithms.

Installation

npm install @bantai-dev/with-rate-limit @bantai-dev/core @bantai-dev/with-storage zod

Note: @bantai-dev/core, @bantai-dev/with-storage, and zod are peer dependencies and must be installed separately.

Quick Start

import { z } from 'zod';
import { defineContext, definePolicy, evaluatePolicy, allow } from '@bantai-dev/core';
import {
  withRateLimit,
  defineRateLimitRule,
  createMemoryStorage,
  rateLimitSchema,
} from '@bantai-dev/with-rate-limit';

// 1. Define your base context
const apiContext = defineContext(
  z.object({
    userId: z.string(),
    endpoint: z.string(),
  })
);

// 2. Extend context with rate limiting
// generateKey will automatically create keys from your input
const rateLimitedContext = withRateLimit(apiContext, {
  storage: createMemoryStorage(rateLimitSchema),
  generateKey: (input) => `api:${input.userId}:${input.endpoint}`,
  defaultValues: {
    rateLimit: {
      type: 'fixed-window',
      limit: 100,
      period: '1h',
    },
  },
});

// 3. Define a rate limiting rule using defineRateLimitRule
// This automatically handles rate limit checking and incrementing
const rateLimitRule = defineRateLimitRule(
  rateLimitedContext,
  'check-rate-limit',
  async (input) => {
    // Your business logic here
    // Rate limit is already checked and will be incremented on allow
    return allow({ reason: 'Request allowed' });
  },
  {
    config: {
      limit: 100,
      period: '1h',
      type: 'fixed-window',
    },
  }
);

// 4. Define policy
const policy = definePolicy(
  rateLimitedContext,
  'api-rate-limit-policy',
  [rateLimitRule],
  {
    defaultStrategy: 'preemptive',
  }
);

// 5. Evaluate policy
// The generateKey function will create the key automatically
const result = await evaluatePolicy(policy, {
  userId: 'user123',
  endpoint: '/api/search',
});

Rate Limiting Strategies

Fixed Window

Fixed window rate limiting divides time into discrete windows. All requests within a window count toward the limit, and the counter resets at the start of each new window.

Use cases: Simple rate limiting, API quotas, basic throttling

{
  type: 'fixed-window',
  key: 'user:123',
  limit: 100,
  period: '1h', // Supports ms format: '1h', '30m', '5s', etc.
}

Sliding Window

Sliding window rate limiting tracks individual request timestamps. Only requests within the current window count toward the limit, providing smoother rate limiting.

Use cases: More accurate rate limiting, preventing burst traffic

{
  type: 'sliding-window',
  key: 'user:123',
  limit: 100,
  period: '1h',
}

Token Bucket

Token bucket rate limiting uses a bucket that refills at a constant rate. Requests consume tokens, and requests are allowed when tokens are available.

Use cases: Burst handling, smooth rate limiting with refill

{
  type: 'token-bucket',
  key: 'user:123',
  limit: 10000,
  period: '1d', // Refills to full capacity (10k tokens) over 1 day
  cost: 1, // Optional: tokens consumed per request (default: 1)
}

Note: The period is a time period (e.g. "1d", "1h", "30m") representing the time to refill from empty to full capacity. The refill rate is automatically calculated as limit / period. For example, limit: 10_000 and period: '1d' means the bucket can hold 10,000 tokens and refills at a rate of 10,000 tokens per day.

Cost: The cost parameter (optional, default: 1) specifies how many tokens/requests each operation consumes. This allows you to implement variable-cost rate limiting where different operations consume different amounts. This parameter works for all rate limiting strategies.

API Reference

`withRateLimit`

Extends a Bantai context with rate limiting capabilities. Adds rateLimit schema fields and tools to the context.

function withRateLimit<
  TContext extends ContextDefinition<z.ZodRawShape, Record<string, unknown>>
>(
  context: TContext,
  options?: {
    defaultValues?: Partial<z.infer<typeof rateLimitSchema['shape']['rateLimit']>>;
    storage?: RateLimitStorage;
    generateKey?: (input: ExtractContextShape<TContext>) => string;
  }
): ContextDefinition<...>

Parameters:

context: A Bantai context definition
options?: Optional configuration
- storage?: A storage adapter implementing RateLimitStorage interface
- defaultValues?: Default values for rate limit configuration
- generateKey?: Optional function to generate rate limit keys dynamically from context input. If provided, this function will be used when rateLimit.key is not specified in the input.

Returns: Extended context with rate limiting capabilities

Example with generateKey:

const rateLimitedContext = withRateLimit(apiContext, {
  storage: createMemoryStorage(rateLimitSchema),
  generateKey: (input) => `api:${input.userId}:${input.endpoint}`,
  defaultValues: {
    rateLimit: {
      type: 'fixed-window',
      limit: 100,
      period: '1h',
    },
  },
});

When using defineRateLimitRule, if rateLimit.key is not provided in the input, the generateKey function will be used automatically.

`defineRateLimitRule`

A helper function that automatically handles rate limit checking and incrementing. This simplifies rule creation by handling the rate limit logic for you.

function defineRateLimitRule<
  TContext extends ContextDefinition<z.ZodRawShape, Record<string, unknown>>,
  TName extends string = string
>(
  context: TContext,
  name: TName,
  evaluate: RuleEvaluateFnAsync<ExtractContextShape<TContext>, ExtractContextTools<TContext>>,
  options?: {
    onAllow?: RuleHookFnAsync<ExtractContextShape<TContext>, ExtractContextTools<TContext>>;
    onDeny?: RuleHookFnAsync<ExtractContextShape<TContext>, ExtractContextTools<TContext>>;
  }
): RuleDefinition<TContext, TName>

Parameters:

context: A context extended with rate limiting capabilities via withRateLimit
name: Unique name for the rule
evaluate: Your rule evaluation function (rate limit is checked before this runs)
options?: Optional hooks
- onAllow?: Hook called after rate limit is incremented (if your rule allows)
- onDeny?: Hook called if your rule denies

Returns: RuleDefinition<TContext, TName>

How it works:

Checks the rate limit before evaluating your rule
If rate limit is exceeded, returns deny immediately
If rate limit passes, evaluates your rule
On allow, automatically increments the rate limit counter
Calls your optional hooks

Key resolution order:

If rateLimit.key is provided in the input, use it
Otherwise, if generateKey function is available, use it
Otherwise, fall back to 'unknown-key'

Example:

const rateLimitRule = defineRateLimitRule(
  rateLimitedContext,
  'api-rule',
  async (input) => {
    // Your business logic here
    // Rate limit already checked and will be incremented on allow
    return allow({ reason: 'Request processed' });
  },
  {
    onAllow: async (input) => {
      // Optional: Additional logic after rate limit increment
      console.log(`Request allowed for ${input.userId}`);
    },
  }
);

`checkRateLimit`

Checks if a rate limit would be exceeded without incrementing the counter. Use this in your rule's evaluate function.

async function checkRateLimit<TConfig extends RateLimitConfig>(
  storage: RateLimitStorage,
  config: TConfig,
  clock?: () => number
): Promise<RateLimitCheckResult>

Parameters:

storage: Storage adapter implementing RateLimitStorage
config: Rate limit configuration object
clock?: Optional clock function for testing (defaults to Date.now)

Returns: Promise<RateLimitCheckResult>

RateLimitCheckResult:

{
  allowed: boolean;
  remaining: number;
  resetAt: number; // Unix timestamp in milliseconds
  reason?: string;
}

`incrementRateLimit`

Increments the rate limit counter. Use this in your rule's onAllow hook.

async function incrementRateLimit(
  storage: RateLimitStorage,
  config: RateLimitConfig,
  clock?: () => number
): Promise<void>

Parameters:

storage: Storage adapter implementing RateLimitStorage
config: Rate limit configuration object
clock?: Optional clock function for testing (defaults to Date.now)

Returns: Promise<void>

`createMemoryStorage`

Creates an in-memory storage adapter for development and testing. Not suitable for production use.

function createMemoryStorage<T extends z.ZodType>(
  schema: T
): StorageAdapter<z.infer<T>>

Parameters:

schema: Zod schema for validating storage data

Returns: StorageAdapter<T>

Storage Integration

The rate limiting extension requires a storage adapter. You can use:

Memory storage (development/testing): createMemoryStorage from this package
Redis storage (production): createRedisStorage from @bantai-dev/storage-redis
Custom storage: Implement the StorageAdapter interface from @bantai-dev/with-storage

Using Redis Storage

import { createRedisStorage } from '@bantai-dev/storage-redis';
import { rateLimitSchema } from '@bantai-dev/with-rate-limit';

const redisStorage = createRedisStorage(
  { url: process.env.REDIS_URL },
  rateLimitSchema
);

const rateLimitedContext = withRateLimit(apiContext, {
  storage: redisStorage,
});

Examples

Per-User Rate Limiting

Using defineRateLimitRule with generateKey:

// Extend context with per-user rate limiting
const userRateLimitedContext = withRateLimit(userContext, {
  storage: createMemoryStorage(rateLimitSchema),
  generateKey: (input) => `user:${input.userId}`,
  defaultValues: {
    rateLimit: {
      type: 'fixed-window',
      limit: 1000,
      period: '24h',
    },
  },
});

// Define rule - rate limit checking/incrementing is automatic
const userRateLimitRule = defineRateLimitRule(
  userRateLimitedContext,
  'user-rate-limit',
  async (input) => {
    // Your business logic
    return allow({ reason: 'User request allowed' });
  }
);

Endpoint-Specific Rate Limits

Using generateKey for endpoint-specific rate limiting:

const endpointLimits = {
  '/api/auth/login': { limit: 5, period: '15m' },
  '/api/payment': { limit: 10, period: '1m' },
  '/api/search': { limit: 100, period: '1m' },
};

// Extend context with endpoint-specific key generation
const endpointRateLimitedContext = withRateLimit(apiContext, {
  storage: createMemoryStorage(rateLimitSchema),
  generateKey: (input) => `endpoint:${input.endpoint}:${input.userId}`,
    defaultValues: {
      rateLimit: {
        type: 'sliding-window',
        limit: 50,
        period: '1h',
      },
    },
  });

  // Define rule - override config per request if needed
  const endpointRateLimitRule = defineRateLimitRule(
    endpointRateLimitedContext,
    'endpoint-rate-limit',
    async (input) => {
      // Your business logic
      return allow({ reason: 'Endpoint request allowed' });
    }
  );

  // Evaluate with endpoint-specific limits
  const result = await evaluatePolicy(endpointPolicy, {
    userId: 'user123',
    endpoint: '/api/auth/login',
    rateLimit: {
      // Override default config for this endpoint
      type: 'sliding-window',
      limit: endpointLimits['/api/auth/login'].limit,
      period: endpointLimits['/api/auth/login'].period,
    },
  });

Tier-Based Rate Limiting

Using generateKey for tier-based rate limiting:

const tierLimits = {
  free: { limit: 100, period: '1h' },
  premium: { limit: 1000, period: '1h' },
  enterprise: { limit: 10000, period: '1h' },
};

  // Extend context with tier-based key generation
  const tierRateLimitedContext = withRateLimit(tierContext, {
    storage: createMemoryStorage(rateLimitSchema),
    generateKey: (input) => `tier:${input.userTier}:${input.userId}`,
    defaultValues: {
      rateLimit: {
        type: 'token-bucket',
        limit: 100,
        period: '1h',
      },
    },
  });

  // Define rule
  const tierRateLimitRule = defineRateLimitRule(
    tierRateLimitedContext,
    'tier-rate-limit',
    async (input) => {
      // Your business logic
      return allow({ reason: 'Tier request allowed' });
    }
  );

  // Evaluate with tier-specific limits
  const result = await evaluatePolicy(tierPolicy, {
    userId: 'user123',
    userTier: 'premium',
    rateLimit: {
      // Override with tier-specific config
      type: 'token-bucket',
      limit: tierLimits.premium.limit,
      period: tierLimits.premium.period,
    },
  });

Type Safety

The package provides full TypeScript type safety:

Context extension: Type-safe context merging with rate limit fields
Config validation: Zod schemas validate rate limit configurations
Storage types: Type-safe storage adapter interface
Result types: Typed rate limit check results

Requirements

Node.js >= 18
TypeScript >= 5.0
Zod >= 4.3.5
@bantai-dev/core
@bantai-dev/with-storage

Storage Extension - Learn about storage adapters
Redis Storage - Production-ready Redis storage
Core Concepts - Understand contexts, rules, and policies

Rate Limiting Extension

On this page