Add public demo rate limiting

src/mastra/index.ts

@@ -9,6 +9,48 @@ import {
 } from '@mastra/observability';
 import { sqlAgent } from './agents/sql-agent';
 
+const RATE_LIMIT_WINDOW_MS = 10 * 60 * 1000;
+const MAX_AGENT_REQUESTS_PER_WINDOW = 10;
+
+type RateLimitEntry = {
+  count: number;
+  resetAt: number;
+};
+
+const agentRequestLimits = new Map<string, RateLimitEntry>();
+
+function getClientIdentifier(c: {
+  req: {
+    header: (name: string) => string | undefined;
+  };
+}): string {
+  const forwardedFor = c.req.header('x-forwarded-for');
+
+  if (forwardedFor) {
+    return forwardedFor.split(',')[0].trim();
+  }
+
+  return c.req.header('x-real-ip') ?? 'unknown-client';
+}
+
+function isPaidAgentRequest(method: string, path: string): boolean {
+  if (method !== 'POST') {
+    return false;
+  }
+
+  return /^\/api\/agents\/[^/]+\/(generate|stream|stream-until-idle|send-message)$/.test(
+    path,
+  );
+}
+
+function cleanExpiredLimits(now: number): void {
+  for (const [key, entry] of agentRequestLimits.entries()) {
+    if (entry.resetAt <= now) {
+      agentRequestLimits.delete(key);
+    }
+  }
+}
+
 export const mastra = new Mastra({
   agents: { sqlAgent },
   storage: new LibSQLStore({
@@ -28,4 +70,57 @@ export const mastra = new Mastra({
       },
     },
   }),
-});
+  server: {
+    bodySizeLimit: 32 * 1024,
+    middleware: [
+      {
+        path: '/api/agents/*',
+        handler: async (c, next) => {
+          const method = c.req.method;
+          const path = new URL(c.req.url).pathname;
+
+          if (!isPaidAgentRequest(method, path)) {
+            await next();
+            return;
+          }
+
+          const now = Date.now();
+          cleanExpiredLimits(now);
+
+          const clientId = getClientIdentifier(c);
+          const current = agentRequestLimits.get(clientId);
+
+          if (!current || current.resetAt <= now) {
+            agentRequestLimits.set(clientId, {
+              count: 1,
+              resetAt: now + RATE_LIMIT_WINDOW_MS,
+            });
+
+            await next();
+            return;
+          }
+
+          if (current.count >= MAX_AGENT_REQUESTS_PER_WINDOW) {
+            const retryAfterSeconds = Math.ceil((current.resetAt - now) / 1000);
+
+            return c.json(
+              {
+                error: 'Public demo rate limit reached. Please try again later.',
+                retryAfterSeconds,
+              },
+              429,
+              {
+                'Retry-After': String(retryAfterSeconds),
+              },
+            );
+          }
+
+          current.count += 1;
+          agentRequestLimits.set(clientId, current);
+
+          await next();
+        },
+      },
+    ],
+  },
+});