Merge pull request #12 from agencyenterprise/staging

Refactor rate limiting and session management

Author: jpatterson933

src/middleware/rateLimit.ts

@@ -1,8 +1,13 @@
 import { ipKeyGenerator, rateLimit } from "express-rate-limit";
+import {
+  MCP_RATE_LIMIT_MAX,
+  OAUTH_RATE_LIMIT_MAX,
+  RATE_LIMIT_WINDOW_MS,
+} from "../shared/constants";
 
 export const userRateLimiter = rateLimit({
-  windowMs: 60 * 1000,
-  max: 60,
+  windowMs: RATE_LIMIT_WINDOW_MS,
+  max: MCP_RATE_LIMIT_MAX,
   keyGenerator: (req) => {
     if (req.userId) {
       return req.userId;
@@ -11,15 +16,15 @@ export const userRateLimiter = rateLimit({
   },
   message: {
     error: "rate_limit_exceeded",
-    error_description: "Too many requests. Limit: 60 requests per minute.",
+    error_description: `Too many requests. Limit: ${MCP_RATE_LIMIT_MAX} requests per minute.`,
   },
   standardHeaders: true,
   legacyHeaders: false,
 });
 
 export const oauthRateLimiter = rateLimit({
-  windowMs: 60 * 1000,
-  max: 10,
+  windowMs: RATE_LIMIT_WINDOW_MS,
+  max: OAUTH_RATE_LIMIT_MAX,
   keyGenerator: (req) => ipKeyGenerator(req.ip ?? "unknown"),
   message: {
     error: "rate_limit_exceeded",

src/modules/sessions/manager.ts

@@ -1,7 +1,12 @@
 import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
 import { randomUUID } from "crypto";
 import { logger } from "../../middleware/logger";
-import { SESSION_CLEANUP_INTERVAL_MS } from "../../shared/constants";
+import {
+  IDLE_TRANSPORT_REAP_INTERVAL_MS,
+  IDLE_TRANSPORT_TTL_MS,
+  MAX_ACTIVE_TRANSPORTS_WARN,
+  SESSION_CLEANUP_INTERVAL_MS,
+} from "../../shared/constants";
 import { AppError } from "../../shared/errors";
 import { ToolServer } from "../../tools/server";
 import { cleanupExpiredMcpServerOAuthData } from "../oauth/service";
@@ -22,6 +27,7 @@ export class SessionManager {
   private activeTransports: Map<string, ActiveTransport>;
   private toolServer: ToolServer;
   private cleanupIntervalId: NodeJS.Timeout | null = null;
+  private reapIntervalId: NodeJS.Timeout | null = null;
 
   constructor() {
     this.activeTransports = new Map();
@@ -79,8 +85,22 @@ export class SessionManager {
       lastAccessedAt: new Date(),
     });
 
+    const activeCount = this.activeTransports.size;
+
+    if (activeCount > MAX_ACTIVE_TRANSPORTS_WARN) {
+      logger.warn(
+        {
+          sessionId,
+          userId,
+          activeCount,
+          threshold: MAX_ACTIVE_TRANSPORTS_WARN,
+        },
+        "Active transport count exceeds warning threshold",
+      );
+    }
+
     logger.info(
-      { sessionId, userId, activeCount: this.activeTransports.size },
+      { sessionId, userId, activeCount },
       "Transport stored in memory",
     );
   }
@@ -132,6 +152,44 @@ export class SessionManager {
     this.activeTransports.delete(sessionId);
   }
 
+  async reapIdleTransports(): Promise<void> {
+    const now = Date.now();
+    const idleSessions: string[] = [];
+
+    for (const [sessionId, entry] of this.activeTransports) {
+      const idleMs = now - entry.lastAccessedAt.getTime();
+      if (idleMs > IDLE_TRANSPORT_TTL_MS) {
+        idleSessions.push(sessionId);
+      }
+    }
+
+    if (idleSessions.length === 0) return;
+
+    const closePromises = idleSessions.map(async (sessionId) => {
+      const entry = this.activeTransports.get(sessionId);
+      if (entry) {
+        try {
+          await entry.transport.close();
+        } catch (error) {
+          logger.error({ sessionId, error }, "Error closing idle transport");
+        }
+        this.activeTransports.delete(sessionId);
+      }
+      try {
+        await markSessionTerminated(sessionId);
+      } catch {
+        // Best-effort DB update
+      }
+    });
+
+    await Promise.all(closePromises);
+
+    logger.info(
+      { reaped: idleSessions.length, remaining: this.activeTransports.size },
+      "Reaped idle transports",
+    );
+  }
+
   async cleanupExpiredData(): Promise<void> {
     try {
       const expiredSessionIds = await findExpiredSessionIds();
@@ -186,18 +244,31 @@ export class SessionManager {
       SESSION_CLEANUP_INTERVAL_MS,
     );
 
+    this.reapIntervalId = setInterval(
+      () => this.reapIdleTransports(),
+      IDLE_TRANSPORT_REAP_INTERVAL_MS,
+    );
+
     logger.info(
-      { intervalMs: SESSION_CLEANUP_INTERVAL_MS },
-      "Data cleanup scheduler started",
+      {
+        cleanupIntervalMs: SESSION_CLEANUP_INTERVAL_MS,
+        reapIntervalMs: IDLE_TRANSPORT_REAP_INTERVAL_MS,
+        idleTtlMs: IDLE_TRANSPORT_TTL_MS,
+      },
+      "Data cleanup and idle reaper schedulers started",
     );
   }
 
   stopCleanupScheduler(): void {
     if (this.cleanupIntervalId) {
       clearInterval(this.cleanupIntervalId);
       this.cleanupIntervalId = null;
-      logger.info("Session cleanup scheduler stopped");
     }
+    if (this.reapIntervalId) {
+      clearInterval(this.reapIntervalId);
+      this.reapIntervalId = null;
+    }
+    logger.info("Cleanup and reaper schedulers stopped");
   }
 
   async shutdown(): Promise<void> {

src/shared/constants.ts

@@ -8,8 +8,16 @@ export const MCP_SERVER_DEFAULT_SCOPE = "fathom:read";
 export const SESSION_TTL_MS = 24 * 60 * 60 * 1000;
 export const SESSION_CLEANUP_INTERVAL_MS = 60 * 60 * 1000;
 export const STALE_SESSION_CUTOFF_MS = 24 * 60 * 60 * 1000;
+export const IDLE_TRANSPORT_TTL_MS = 5 * 60 * 1000;
+export const IDLE_TRANSPORT_REAP_INTERVAL_MS = 60 * 1000;
+export const MAX_ACTIVE_TRANSPORTS_WARN = 100;
 export const GRACEFUL_SHUTDOWN_TIMEOUT_MS = 10 * 1000;
 
+// Rate Limiting
+export const RATE_LIMIT_WINDOW_MS = 60 * 1000;
+export const MCP_RATE_LIMIT_MAX = 200;
+export const OAUTH_RATE_LIMIT_MAX = 10;
+
 // Fathom API
 export const FATHOM_API_TIMEOUT_MS = 30 * 1000;
 export const FATHOM_API_SCOPE = "public_api";

src/test/modules/sessions/manager.test.ts

@@ -159,23 +159,90 @@ describe("SessionManager", () => {
     });
   });
 
+  describe("reapIdleTransports", () => {
+    it("reaps transports idle beyond IDLE_TRANSPORT_TTL_MS", async () => {
+      vi.mocked(insertSession).mockResolvedValue(undefined);
+      vi.mocked(markSessionTerminated).mockResolvedValue(undefined);
+
+      await sessionManager.createSession("user-123");
+      await vi.advanceTimersByTimeAsync(0);
+
+      expect(
+        sessionManager.getActiveTransport("mock-session-id"),
+      ).toBeDefined();
+
+      // Advance past idle TTL (5 minutes)
+      vi.advanceTimersByTime(6 * 60 * 1000);
+
+      await sessionManager.reapIdleTransports();
+
+      expect(
+        sessionManager.getActiveTransport("mock-session-id"),
+      ).toBeUndefined();
+      expect(markSessionTerminated).toHaveBeenCalledWith("mock-session-id");
+    });
+
+    it("leaves recently accessed transports alone", async () => {
+      vi.mocked(insertSession).mockResolvedValue(undefined);
+
+      await sessionManager.createSession("user-123");
+      await vi.advanceTimersByTimeAsync(0);
+
+      // Advance only 2 minutes (under 5 min TTL)
+      vi.advanceTimersByTime(2 * 60 * 1000);
+
+      await sessionManager.reapIdleTransports();
+
+      expect(
+        sessionManager.getActiveTransport("mock-session-id"),
+      ).toBeDefined();
+    });
+
+    it("does nothing when no transports are idle", async () => {
+      await sessionManager.reapIdleTransports();
+
+      expect(markSessionTerminated).not.toHaveBeenCalled();
+    });
+
+    it("handles transport close errors gracefully", async () => {
+      vi.mocked(insertSession).mockResolvedValue(undefined);
+      vi.mocked(markSessionTerminated).mockResolvedValue(undefined);
+
+      const transport = await sessionManager.createSession("user-123");
+      await vi.advanceTimersByTimeAsync(0);
+
+      vi.mocked(transport.close).mockRejectedValueOnce(
+        new Error("close failed"),
+      );
+
+      vi.advanceTimersByTime(6 * 60 * 1000);
+
+      await expect(sessionManager.reapIdleTransports()).resolves.not.toThrow();
+
+      expect(
+        sessionManager.getActiveTransport("mock-session-id"),
+      ).toBeUndefined();
+    });
+  });
+
   describe("startCleanupScheduler", () => {
-    it("starts cleanup interval", () => {
+    it("starts cleanup and reaper intervals", () => {
       sessionManager.startCleanupScheduler();
 
-      expect(vi.getTimerCount()).toBe(1);
+      // 2 intervals: hourly cleanup + 60s reaper
+      expect(vi.getTimerCount()).toBe(2);
     });
 
     it("does not start duplicate scheduler", () => {
       sessionManager.startCleanupScheduler();
       sessionManager.startCleanupScheduler();
 
-      expect(vi.getTimerCount()).toBe(1);
+      expect(vi.getTimerCount()).toBe(2);
     });
   });
 
   describe("stopCleanupScheduler", () => {
-    it("stops cleanup interval", () => {
+    it("stops both cleanup and reaper intervals", () => {
       sessionManager.startCleanupScheduler();
       sessionManager.stopCleanupScheduler();