ubiquitous-invention/packages/shared/src/utils/token-bucket.ts

/**
 * Pure, allocation-cheap fixed-window counter used by the in-process rate
 * limiter in `apps/web/server/lib/rate-limit.ts`. Lives in `@tasks/shared`
 * so the algorithm can be unit-tested in isolation — `apps/web` has no
 * vitest harness yet (see Task-bootstrap-vitest-for-apps-web).
 *
 * The shape is a fixed window, not a leaky bucket: a single counter that
 * resets at `resetAt`. This is the simplest correct primitive for "N
 * attempts per window per key" guardrails and matches the limits described
 * in the rate-limit task spec (5/min for sign-in, 10/hr for invites). A
 * sliding-log or leaky-bucket would be more accurate around window
 * boundaries, but the cost is more state per key for no real abuse-control
 * win at these limits.
 *
 * The function is pure: callers pass in the current state (or `undefined`
 * to start fresh) plus the wall clock as `now`. The returned `state` field
 * is the value the caller should store back into its key→state map. This
 * keeps the bucket trivially testable with a fake clock and avoids any
 * hidden global mutable state inside `@tasks/shared`.
 */

export interface TokenBucketState {
  /** Attempts consumed within the current window. */
  count: number;
  /** Epoch milliseconds at which the current window expires. */
  resetAt: number;
}

export interface TokenBucketOptions {
  /** Max allowed attempts per window. Must be >= 1. */
  limit: number;
  /** Window length in milliseconds. Must be >= 1. */
  windowMs: number;
}

export type TokenBucketResult =
  | {
      allowed: true;
      remaining: number;
      resetAt: number;
      state: TokenBucketState;
    }
  | {
      allowed: false;
      remaining: 0;
      resetAt: number;
      retryAfterMs: number;
      state: TokenBucketState;
    };

/**
 * Charge one attempt against `state` for the bucket described by `opts`.
 *
 * - If `state` is undefined or its window has expired (`now >= resetAt`),
 *   a fresh window starts: count = 1, resetAt = now + windowMs.
 * - Otherwise the count is incremented. If the post-increment count
 *   exceeds `limit`, the call is denied; the returned `state` does NOT
 *   include the over-limit increment, so a flood of denied calls cannot
 *   inflate the counter past `limit + 1`.
 *
 * Returning the next-state alongside the decision means the call site
 * doesn't need to special-case "denied" vs "allowed" when writing back —
 * `result.state` is always the value to persist.
 */
export function consumeTokenBucket(
  state: TokenBucketState | undefined,
  now: number,
  opts: TokenBucketOptions,
): TokenBucketResult {
  if (!Number.isFinite(opts.limit) || opts.limit < 1) {
    throw new Error("token-bucket: limit must be >= 1");
  }
  if (!Number.isFinite(opts.windowMs) || opts.windowMs < 1) {
    throw new Error("token-bucket: windowMs must be >= 1");
  }

  const windowExpired = !state || now >= state.resetAt;
  if (windowExpired) {
    const next: TokenBucketState = { count: 1, resetAt: now + opts.windowMs };
    return {
      allowed: true,
      remaining: opts.limit - 1,
      resetAt: next.resetAt,
      state: next,
    };
  }

  const nextCount = state.count + 1;
  if (nextCount > opts.limit) {
    // Keep count pegged at limit so repeated denied calls don't inflate
    // the counter unboundedly. retryAfterMs shrinks as time advances
    // within the same denied window.
    const pegged: TokenBucketState = { count: opts.limit, resetAt: state.resetAt };
    return {
      allowed: false,
      remaining: 0,
      resetAt: state.resetAt,
      retryAfterMs: Math.max(0, state.resetAt - now),
      state: pegged,
    };
  }

  const next: TokenBucketState = { count: nextCount, resetAt: state.resetAt };
  return {
    allowed: true,
    remaining: opts.limit - nextCount,
    resetAt: next.resetAt,
    state: next,
  };
}