Merge branch 'fix/quota-on-mongodb'

Author: devthejo

services/api/src/budget.ts

@@ -1,33 +1,5 @@
-import type IORedis from 'ioredis';
 import type { SecurityPolicyDoc } from './db';
-import { createRedisClient } from './redis';
-
-const DEFAULT_REDIS_URL = process.env.REDIS_URL || 'redis://127.0.0.1:6379';
-
-let redisClient: IORedis | null = null;
-let redisHealthy = false;
-
-function getRedis(): IORedis | null {
-  if (!redisClient) {
-    try {
-      redisClient = createRedisClient(DEFAULT_REDIS_URL);
-      redisClient.on('error', (err) => {
-        // eslint-disable-next-line no-console
-        console.error('budget: redis error', err);
-        redisHealthy = false;
-      });
-      redisClient.on('connect', () => {
-        redisHealthy = true;
-      });
-    } catch (err) {
-      // eslint-disable-next-line no-console
-      console.error('budget: failed to create redis client', err);
-      redisHealthy = false;
-      redisClient = null;
-    }
-  }
-  return redisHealthy ? redisClient : null;
-}
+import { getOrgDailyBudgetsCollection } from './db';
 
 function currentDateKeyUtc(): string {
   const now = new Date();
@@ -39,30 +11,57 @@ function currentDateKeyUtc(): string {
   return `${y}${mm}${dd}`;
 }
 
-async function incrementBudgetCounter(
-  key: string,
+async function incrementMongoBudgetCounter(
+  orgId: any,
+  date: string,
+  field: 'scan_count' | 'gpt_count',
   limit: number | null
 ): Promise<boolean> {
   if (limit == null) {
-    return true;
-  }
-
-  const client = getRedis();
-  if (!client) {
-    // Redis unavailable: fail open, do not block scans or GPT
+    // Unlimited budget
     return true;
   }
 
   try {
-    const count = await client.incr(key);
-    if (count === 1) {
-      // Keep key a bit longer than one day to tolerate clock skew
-      await client.expire(key, 60 * 60 * 24 + 60 * 5);
+    const col = await getOrgDailyBudgetsCollection();
+    const now = new Date();
+
+    const result = (await col.findOneAndUpdate(
+      { org_id: orgId, date },
+      {
+        $setOnInsert: {
+          org_id: orgId,
+          date,
+          created_at: now
+        },
+        $inc: { [field]: 1 },
+        $set: { updated_at: now }
+      },
+      {
+        upsert: true,
+        returnDocument: 'after'
+      }
+    )) as any;
+
+    const doc = result && (result.value ?? result);
+    if (!doc) {
+      // Should not happen in normal operation; fail open to preserve
+      // existing semantics where infra issues do not hard-block scans/GPT.
+      // eslint-disable-next-line no-console
+      console.error('budget: Mongo findOneAndUpdate returned no document', {
+        orgId,
+        date,
+        field
+      });
+      return true;
     }
-    return count <= limit;
+
+    const current = typeof doc[field] === 'number' ? doc[field] : 0;
+    return current <= limit;
   } catch (err) {
+    // Preserve fail-open semantics from the Redis implementation.
     // eslint-disable-next-line no-console
-    console.error('budget: increment failed, failing open', err);
+    console.error('budget: Mongo increment failed, failing open', err);
     return true;
   }
 }
@@ -74,8 +73,7 @@ export async function checkAndConsumeScanBudget(
   const limit =
     policy.max_scans_per_day == null ? null : policy.max_scans_per_day;
   const date = currentDateKeyUtc();
-  const key = `budget:scan:${String(orgId)}:${date}`;
-  return incrementBudgetCounter(key, limit);
+  return incrementMongoBudgetCounter(orgId, date, 'scan_count', limit);
 }
 
 export async function checkAndConsumeGptBudget(
@@ -85,6 +83,5 @@ export async function checkAndConsumeGptBudget(
   const limit =
     policy.max_gpt_calls_per_day == null ? null : policy.max_gpt_calls_per_day;
   const date = currentDateKeyUtc();
-  const key = `budget:gpt:${String(orgId)}:${date}`;
-  return incrementBudgetCounter(key, limit);
+  return incrementMongoBudgetCounter(orgId, date, 'gpt_count', limit);
 }

services/api/src/db.ts

@@ -125,6 +125,24 @@ export interface SecurityPolicyDoc {
   created_at: Date;
 }
 
+/**
+ * Daily per-org budget counters for scans and GPT calls. One document per
+ * (org_id, date) pair. Used to enforce SecurityPolicyDoc max_*_per_day
+ * budgets without relying on Redis/Valkey.
+ */
+export interface OrgDailyBudgetDoc {
+  _id?: any;
+  org_id: any;
+  // UTC date key in YYYYMMDD format (e.g. 20251210)
+  date: string;
+  // Number of scans consumed on this day
+  scan_count: number;
+  // Number of GPT calls consumed on this day
+  gpt_count: number;
+  created_at: Date;
+  updated_at: Date;
+}
+
 export interface PolicyOverrideDoc {
   _id?: any;
   org_id: any;
@@ -172,6 +190,27 @@ export async function getSecurityPoliciesCollection(): Promise<
   return col;
 }
 
+export async function getOrgDailyBudgetsCollection(): Promise<
+  Collection<OrgDailyBudgetDoc>
+> {
+  const db = await getDb();
+  const col = db.collection<OrgDailyBudgetDoc>('org_daily_budgets');
+
+  await col.createIndex(
+    { org_id: 1, date: 1 },
+    { unique: true, name: 'uniq_org_date' }
+  );
+
+  // Optional TTL index to avoid unbounded growth. Documents will be removed
+  // approximately 90 days after their last update.
+  await col.createIndex(
+    { updated_at: 1 },
+    { name: 'ttl_org_daily_budgets', expireAfterSeconds: 60 * 60 * 24 * 90 }
+  );
+
+  return col;
+}
+
 export async function getPolicyOverridesCollection(): Promise<
   Collection<PolicyOverrideDoc>
 > {