session-store.ts

import { Redis } from "ioredis";
import os from "node:os";

/* ------------------------------------------------------------------ */
/*  Types                                                              */
/* ------------------------------------------------------------------ */

export interface UserContext {
  userId: number;
  consumerId: string;
  orgId: string;
  email: string;
}

export type ProxyMode = "actions";

/** Serialisable session metadata (no timers, no WebSocket handles). */
export interface SessionRecord {
  id: string;
  owner: UserContext;
  apiKey: string;
  upstreamSessionId: string;
  upstreamCdpUrl: string;
  mode: ProxyMode;
  createdAtMs: number;
  expiresAtMs: number;
  idleTimeoutMs: number;
  maxDurationMs: number;
  connectedClients: number;
  lastActivityMs: number;
  /** Which replica owns the live WebSocket connections for this session. */
  replicaId: string;
}

/* ------------------------------------------------------------------ */
/*  Redis client                                                       */
/* ------------------------------------------------------------------ */

const REDIS_URL = process.env.REDIS_URL?.trim();

function createRedisClient(): Redis {
  if (!REDIS_URL) {
    throw new Error("REDIS_URL is required for multi-replica mode");
  }
  const client = new Redis(REDIS_URL, {
    maxRetriesPerRequest: 3,
    retryStrategy(times) {
      return Math.min(times * 200, 3000);
    },
    lazyConnect: true,
  });
  client.on("error", (err) => {
    console.error(`[session-store] redis error: ${err.message}`);
  });
  return client;
}

let redis: Redis | null = null;

export function getRedis(): Redis {
  if (!redis) {
    redis = createRedisClient();
  }
  return redis;
}

export async function connectRedis(): Promise<void> {
  const r = getRedis();
  await r.connect();
  console.log("[session-store] redis connected");
}

export async function disconnectRedis(): Promise<void> {
  if (redis) {
    await redis.quit();
    redis = null;
  }
}

/* ------------------------------------------------------------------ */
/*  Replica identity                                                   */
/* ------------------------------------------------------------------ */

export const REPLICA_ID =
  process.env.RAILWAY_REPLICA_ID?.trim() ||
  process.env.RAILWAY_DEPLOYMENT_ID?.trim() ||
  os.hostname();

/* ------------------------------------------------------------------ */
/*  Key helpers                                                        */
/* ------------------------------------------------------------------ */

const KEY = {
  session: (id: string) => `proxy:session:${id}`,
  orgSessions: (orgId: string) => `proxy:org:${orgId}:sessions`,
  orgLimits: (orgId: string) => `proxy:org:${orgId}:limits`,
} as const;

/* ------------------------------------------------------------------ */
/*  Serialisation                                                      */
/* ------------------------------------------------------------------ */

function serialise(session: SessionRecord): Record<string, string> {
  return {
    id: session.id,
    owner: JSON.stringify(session.owner),
    apiKey: session.apiKey,
    upstreamSessionId: session.upstreamSessionId,
    upstreamCdpUrl: session.upstreamCdpUrl,
    mode: session.mode,
    createdAtMs: String(session.createdAtMs),
    expiresAtMs: String(session.expiresAtMs),
    idleTimeoutMs: String(session.idleTimeoutMs),
    maxDurationMs: String(session.maxDurationMs),
    connectedClients: String(session.connectedClients),
    lastActivityMs: String(session.lastActivityMs),
    replicaId: session.replicaId,
  };
}

function deserialise(hash: Record<string, string>): SessionRecord | null {
  if (!hash || !hash.id) return null;
  try {
    return {
      id: hash.id,
      owner: JSON.parse(hash.owner) as UserContext,
      apiKey: hash.apiKey,
      upstreamSessionId: hash.upstreamSessionId,
      upstreamCdpUrl: hash.upstreamCdpUrl,
      mode: hash.mode as ProxyMode,
      createdAtMs: Number(hash.createdAtMs),
      expiresAtMs: Number(hash.expiresAtMs),
      idleTimeoutMs: Number(hash.idleTimeoutMs),
      maxDurationMs: Number(hash.maxDurationMs),
      connectedClients: Number(hash.connectedClients),
      lastActivityMs: Number(hash.lastActivityMs),
      replicaId: hash.replicaId,
    };
  } catch {
    return null;
  }
}

/* ------------------------------------------------------------------ */
/*  CRUD operations                                                    */
/* ------------------------------------------------------------------ */

/** Save a new session. Sets Redis TTL to maxDurationMs + 60s buffer. */
export async function saveSession(session: SessionRecord): Promise<void> {
  const r = getRedis();
  const key = KEY.session(session.id);
  const ttlSec = Math.ceil(session.maxDurationMs / 1000) + 60;
  const pipeline = r.pipeline();
  pipeline.hset(key, serialise(session));
  pipeline.expire(key, ttlSec);
  pipeline.sadd(KEY.orgSessions(session.owner.orgId), session.id);
  await pipeline.exec();
}

/** Get session by ID. Returns null if not found or expired in Redis. */
export async function getSession(sessionId: string): Promise<SessionRecord | null> {
  const r = getRedis();
  const hash = await r.hgetall(KEY.session(sessionId));
  if (!hash || Object.keys(hash).length === 0) return null;
  return deserialise(hash as Record<string, string>);
}

/** Delete session from Redis and remove from org set. */
export async function deleteSession(sessionId: string, orgId: string): Promise<void> {
  const r = getRedis();
  const pipeline = r.pipeline();
  pipeline.del(KEY.session(sessionId));
  pipeline.srem(KEY.orgSessions(orgId), sessionId);
  await pipeline.exec();
}

/** Update specific fields on a session. */
export async function updateSession(
  sessionId: string,
  fields: Partial<Pick<SessionRecord, "connectedClients" | "lastActivityMs">>,
): Promise<void> {
  const r = getRedis();
  const updates: Record<string, string> = {};
  if (fields.connectedClients !== undefined) updates.connectedClients = String(fields.connectedClients);
  if (fields.lastActivityMs !== undefined) updates.lastActivityMs = String(fields.lastActivityMs);
  if (Object.keys(updates).length > 0) {
    await r.hset(KEY.session(sessionId), updates);
  }
}

/** Check if this replica owns the session's WebSocket connections. */
export async function isOwnedByThisReplica(sessionId: string): Promise<boolean> {
  const r = getRedis();
  const rid = await r.hget(KEY.session(sessionId), "replicaId");
  return rid === REPLICA_ID;
}

/** Get the replicaId that owns a session. */
export async function getSessionReplicaId(sessionId: string): Promise<string | null> {
  const r = getRedis();
  return r.hget(KEY.session(sessionId), "replicaId");
}

/* ------------------------------------------------------------------ */
/*  Per-org queries                                                    */
/* ------------------------------------------------------------------ */

/** Count active sessions for an org. Cleans up expired entries. */
export async function countOrgSessions(orgId: string): Promise<number> {
  const r = getRedis();
  const sessionIds = await r.smembers(KEY.orgSessions(orgId));
  if (sessionIds.length === 0) return 0;

  // Verify each session still exists (TTL may have expired)
  const pipeline = r.pipeline();
  for (const sid of sessionIds) {
    pipeline.exists(KEY.session(sid));
  }
  const results = await pipeline.exec();
  if (!results) return 0;

  const stale: string[] = [];
  let count = 0;
  for (let i = 0; i < results.length; i++) {
    const [err, exists] = results[i];
    if (!err && exists === 1) {
      count++;
    } else {
      stale.push(sessionIds[i]);
    }
  }

  // Clean up stale entries
  if (stale.length > 0) {
    await r.srem(KEY.orgSessions(orgId), ...stale);
  }

  return count;
}

/** List all session IDs for an org. */
export async function listOrgSessionIds(orgId: string): Promise<string[]> {
  const r = getRedis();
  return r.smembers(KEY.orgSessions(orgId));
}

/* ------------------------------------------------------------------ */
/*  Per-org limits                                                     */
/* ------------------------------------------------------------------ */

export interface OrgLimits {
  maxConcurrentSessions: number;
  maxSessionCreatesPerHour: number;
}

const DEFAULT_ORG_LIMITS: OrgLimits = {
  maxConcurrentSessions: 10,
  maxSessionCreatesPerHour: 100,
};

export async function getOrgLimits(orgId: string): Promise<OrgLimits> {
  const r = getRedis();
  const raw = await r.hgetall(KEY.orgLimits(orgId));
  if (!raw || Object.keys(raw).length === 0) return DEFAULT_ORG_LIMITS;
  return {
    maxConcurrentSessions: Number(raw.maxConcurrentSessions) || DEFAULT_ORG_LIMITS.maxConcurrentSessions,
    maxSessionCreatesPerHour: Number(raw.maxSessionCreatesPerHour) || DEFAULT_ORG_LIMITS.maxSessionCreatesPerHour,
  };
}

export async function setOrgLimits(orgId: string, limits: Partial<OrgLimits>): Promise<void> {
  const r = getRedis();
  const updates: Record<string, string> = {};
  if (limits.maxConcurrentSessions !== undefined) updates.maxConcurrentSessions = String(limits.maxConcurrentSessions);
  if (limits.maxSessionCreatesPerHour !== undefined) updates.maxSessionCreatesPerHour = String(limits.maxSessionCreatesPerHour);
  if (Object.keys(updates).length > 0) {
    await r.hset(KEY.orgLimits(orgId), updates);
  }
}

/** Check concurrency limit. Returns true if allowed. */
export async function checkConcurrencyLimit(orgId: string): Promise<{ allowed: boolean; current: number; limit: number }> {
  const [limits, current] = await Promise.all([
    getOrgLimits(orgId),
    countOrgSessions(orgId),
  ]);
  return {
    allowed: current < limits.maxConcurrentSessions,
    current,
    limit: limits.maxConcurrentSessions,
  };
}

/** Check hourly rate limit. Returns true if allowed. Uses a sliding window counter. */
export async function checkRateLimit(orgId: string): Promise<{ allowed: boolean; current: number; limit: number }> {
  const limits = await getOrgLimits(orgId);
  const r = getRedis();
  const rateKey = `proxy:org:${orgId}:rate:${Math.floor(Date.now() / 3_600_000)}`;
  const current = await r.incr(rateKey);
  if (current === 1) {
    await r.expire(rateKey, 3600);
  }
  if (current > limits.maxSessionCreatesPerHour) {
    return { allowed: false, current, limit: limits.maxSessionCreatesPerHour };
  }
  return { allowed: true, current, limit: limits.maxSessionCreatesPerHour };
}

/* ------------------------------------------------------------------ */
/*  Stale session cleanup                                              */
/* ------------------------------------------------------------------ */

/**
 * On startup, find sessions that claim to be owned by this replica
 * (from a previous instance) and return them for cleanup.
 */
export async function findStaleSessions(): Promise<SessionRecord[]> {
  const r = getRedis();
  const stale: SessionRecord[] = [];

  // SCAN for all session keys
  let cursor = "0";
  do {
    const [nextCursor, keys] = await r.scan(cursor, "MATCH", "proxy:session:*", "COUNT", 100);
    cursor = nextCursor;
    for (const key of keys) {
      const rid = await r.hget(key, "replicaId");
      if (rid === REPLICA_ID) {
        const hash = await r.hgetall(key);
        if (hash && Object.keys(hash).length > 0) {
          const session = deserialise(hash as Record<string, string>);
          if (session) stale.push(session);
        }
      }
    }
  } while (cursor !== "0");

  return stale;
}

/* ------------------------------------------------------------------ */
/*  Global stats                                                       */
/* ------------------------------------------------------------------ */

/** Count total active sessions across all replicas. */
export async function countAllSessions(): Promise<number> {
  const r = getRedis();
  let count = 0;
  let cursor = "0";
  do {
    const [nextCursor, keys] = await r.scan(cursor, "MATCH", "proxy:session:*", "COUNT", 100);
    cursor = nextCursor;
    count += keys.length;
  } while (cursor !== "0");
  return count;
}