feat(billing): add generic BillingUsage model for quota tracking (#3273)

* feat(billing): add generic BillingUsage model for quota tracking

* fix(billing): address review feedback on BillingUsage model

- Validate organizationId in increment(), sanitize counter key against injection
- Handle E11000 duplicate key race on concurrent upserts (retry once)
- Add runValidators to reset() for consistency with increment()
- Tighten month regex to reject invalid months (00, 13)
- Use UTC for currentMonth() to avoid timezone-dependent boundaries
- Capture month once in get() to prevent cross-boundary inconsistency
- Remove redundant individual index on organizationId (compound covers it)
- Add missing Jest globals imports in test file
- Add tests for semantically invalid month values

Author: PierreBrisorgueil

modules/billing/models/billing.usage.model.mongoose.js

@@ -0,0 +1,52 @@
+/**
+ * Module dependencies
+ */
+import mongoose from 'mongoose';
+
+const Schema = mongoose.Schema;
+
+/**
+ * Data Model Mongoose
+ */
+const UsageMongoose = new Schema(
+  {
+    organizationId: {
+      type: Schema.ObjectId,
+      ref: 'Organization',
+      required: true,
+    },
+    month: {
+      type: String,
+      required: true,
+      index: true,
+    },
+    counters: {
+      type: Schema.Types.Mixed,
+      default: () => ({}),
+    },
+  },
+  {
+    timestamps: true,
+  },
+);
+
+UsageMongoose.index({ organizationId: 1, month: 1 }, { unique: true });
+
+/**
+ * Returns the hex string representation of the document ObjectId.
+ * @returns {string} Hex string of the ObjectId.
+ */
+function addID() {
+  return this._id.toHexString();
+}
+
+/**
+ * Model configuration
+ */
+UsageMongoose.virtual('id').get(addID);
+// Ensure virtual fields are serialised.
+UsageMongoose.set('toJSON', {
+  virtuals: true,
+});
+
+mongoose.model('BillingUsage', UsageMongoose);

modules/billing/models/billing.usage.schema.js

@@ -0,0 +1,19 @@
+/**
+ * Module dependencies
+ */
+import { z } from 'zod';
+
+/**
+ *  Data Schema
+ */
+const objectIdRegex = /^[a-f\d]{24}$/i;
+
+const BillingUsage = z.object({
+  organizationId: z.string().trim().regex(objectIdRegex, 'organizationId must be a valid ObjectId'),
+  month: z.string().trim().regex(/^\d{4}-(0[1-9]|1[0-2])$/, 'month must be in YYYY-MM format'),
+  counters: z.record(z.string(), z.number()).default(() => ({})),
+});
+
+export default {
+  BillingUsage,
+};

modules/billing/repositories/billing.usage.repository.js

@@ -0,0 +1,72 @@
+/**
+ * Module dependencies
+ */
+import mongoose from 'mongoose';
+
+const BillingUsage = mongoose.model('BillingUsage');
+
+const SAFE_KEY_RE = /^[a-zA-Z0-9_-]+$/;
+
+/**
+ * @function get
+ * @description Fetch a single usage document by organizationId and month.
+ * @param {String} organizationId - The organization ID.
+ * @param {String} month - The month in YYYY-MM format.
+ * @returns {Promise<Object|null>} The usage document or null.
+ */
+const get = (organizationId, month) => {
+  if (!mongoose.Types.ObjectId.isValid(organizationId)) return null;
+  return BillingUsage.findOne({ organizationId, month }).exec();
+};
+
+/**
+ * @function increment
+ * @description Atomically increment a counter key for the given org+month, with upsert.
+ * @param {String} organizationId - The organization ID.
+ * @param {String} month - The month in YYYY-MM format.
+ * @param {String} key - The counter key to increment (e.g. 'executions').
+ * @param {Number} amount - The amount to increment by.
+ * @returns {Promise<Object>} The updated usage document.
+ */
+const increment = async (organizationId, month, key, amount) => {
+  if (!mongoose.Types.ObjectId.isValid(organizationId)) return null;
+  if (!SAFE_KEY_RE.test(key)) throw new Error(`Invalid counter key: ${key}`);
+  try {
+    return await BillingUsage.findOneAndUpdate(
+      { organizationId, month },
+      { $inc: { [`counters.${key}`]: amount } },
+      { upsert: true, returnDocument: 'after', runValidators: true },
+    ).exec();
+  } catch (err) {
+    if (err.code === 11000) {
+      return BillingUsage.findOneAndUpdate(
+        { organizationId, month },
+        { $inc: { [`counters.${key}`]: amount } },
+        { returnDocument: 'after', runValidators: true },
+      ).exec();
+    }
+    throw err;
+  }
+};
+
+/**
+ * @function reset
+ * @description Reset all counters to an empty object for the given org+month.
+ * @param {String} organizationId - The organization ID.
+ * @param {String} month - The month in YYYY-MM format.
+ * @returns {Promise<Object|null>} The updated usage document or null.
+ */
+const reset = (organizationId, month) => {
+  if (!mongoose.Types.ObjectId.isValid(organizationId)) return null;
+  return BillingUsage.findOneAndUpdate(
+    { organizationId, month },
+    { $set: { counters: {} } },
+    { returnDocument: 'after', runValidators: true },
+  ).exec();
+};
+
+export default {
+  get,
+  increment,
+  reset,
+};

modules/billing/services/billing.usage.service.js

@@ -0,0 +1,48 @@
+/**
+ * Module dependencies
+ */
+import UsageRepository from '../repositories/billing.usage.repository.js';
+
+/**
+ * Compute the current month string in YYYY-MM format.
+ * @returns {String} e.g. '2026-03'
+ */
+const currentMonth = () => {
+  const now = new Date();
+  const year = now.getUTCFullYear();
+  const month = String(now.getUTCMonth() + 1).padStart(2, '0');
+  return `${year}-${month}`;
+};
+
+/**
+ * @desc Increment a usage counter for the given organization (current month).
+ * @param {String} organizationId - The organization ID.
+ * @param {String} key - The counter key to increment.
+ * @param {Number} amount - The amount to increment by.
+ * @returns {Promise<Object>} The updated usage document.
+ */
+const increment = (organizationId, key, amount) => UsageRepository.increment(organizationId, currentMonth(), key, amount);
+
+/**
+ * @desc Get usage for the given organization (current month).
+ * @param {String} organizationId - The organization ID.
+ * @returns {Promise<Object>} The usage document or an object with empty counters.
+ */
+const get = async (organizationId) => {
+  const month = currentMonth();
+  const usage = await UsageRepository.get(organizationId, month);
+  return usage || { organizationId, month, counters: {} };
+};
+
+/**
+ * @desc Reset usage counters for the given organization (current month).
+ * @param {String} organizationId - The organization ID.
+ * @returns {Promise<Object|null>} The updated usage document or null.
+ */
+const reset = (organizationId) => UsageRepository.reset(organizationId, currentMonth());
+
+export default {
+  increment,
+  get,
+  reset,
+};

modules/billing/tests/billing.usage.unit.tests.js

@@ -0,0 +1,149 @@
+/**
+ * Module dependencies.
+ */
+import { jest, describe, test, beforeEach, afterEach, expect } from '@jest/globals';
+import schema from '../models/billing.usage.schema.js';
+
+/**
+ * Unit tests
+ */
+describe('BillingUsage unit tests:', () => {
+  describe('Schema validation', () => {
+    let usage;
+
+    beforeEach(() => {
+      usage = {
+        organizationId: '507f1f77bcf86cd799439011',
+        month: '2026-03',
+        counters: { executions: 10 },
+      };
+    });
+
+    test('should be valid with correct data', () => {
+      const result = schema.BillingUsage.safeParse(usage);
+      expect(result.error).toBeFalsy();
+      expect(result.data.organizationId).toBe('507f1f77bcf86cd799439011');
+      expect(result.data.month).toBe('2026-03');
+      expect(result.data.counters.executions).toBe(10);
+    });
+
+    test('should show error when organizationId is missing', () => {
+      usage.organizationId = '';
+      const result = schema.BillingUsage.safeParse(usage);
+      expect(result.error).toBeDefined();
+    });
+
+    test('should show error when organizationId is invalid', () => {
+      usage.organizationId = 'not-valid';
+      const result = schema.BillingUsage.safeParse(usage);
+      expect(result.error).toBeDefined();
+    });
+
+    test('should show error when month format is invalid', () => {
+      usage.month = '2026/03';
+      const result = schema.BillingUsage.safeParse(usage);
+      expect(result.error).toBeDefined();
+    });
+
+    test('should show error when month is missing', () => {
+      delete usage.month;
+      const result = schema.BillingUsage.safeParse(usage);
+      expect(result.error).toBeDefined();
+    });
+
+    test('should reject semantically invalid month 2026-00', () => {
+      usage.month = '2026-00';
+      const result = schema.BillingUsage.safeParse(usage);
+      expect(result.error).toBeDefined();
+    });
+
+    test('should reject semantically invalid month 2026-13', () => {
+      usage.month = '2026-13';
+      const result = schema.BillingUsage.safeParse(usage);
+      expect(result.error).toBeDefined();
+    });
+
+    test('should default counters to empty object', () => {
+      delete usage.counters;
+      const result = schema.BillingUsage.safeParse(usage);
+      expect(result.error).toBeFalsy();
+      expect(result.data.counters).toEqual({});
+    });
+  });
+
+  describe('Service layer', () => {
+    let BillingUsageService;
+    let mockUsageRepository;
+
+    const orgId = '507f1f77bcf86cd799439011';
+
+    beforeEach(async () => {
+      jest.resetModules();
+
+      mockUsageRepository = {
+        get: jest.fn(),
+        increment: jest.fn(),
+        reset: jest.fn(),
+      };
+
+      jest.unstable_mockModule('../repositories/billing.usage.repository.js', () => ({
+        default: mockUsageRepository,
+      }));
+
+      const mod = await import('../services/billing.usage.service.js');
+      BillingUsageService = mod.default;
+    });
+
+    afterEach(() => {
+      jest.restoreAllMocks();
+    });
+
+    test('increment should create document on first call (upsert)', async () => {
+      const created = { organizationId: orgId, month: '2026-03', counters: { executions: 1 } };
+      mockUsageRepository.increment.mockResolvedValue(created);
+
+      const result = await BillingUsageService.increment(orgId, 'executions', 1);
+
+      expect(mockUsageRepository.increment).toHaveBeenCalledWith(orgId, expect.stringMatching(/^\d{4}-\d{2}$/), 'executions', 1);
+      expect(result.counters.executions).toBe(1);
+    });
+
+    test('increment should atomically increase counter', async () => {
+      const updated = { organizationId: orgId, month: '2026-03', counters: { executions: 5 } };
+      mockUsageRepository.increment.mockResolvedValue(updated);
+
+      const result = await BillingUsageService.increment(orgId, 'executions', 3);
+
+      expect(result.counters.executions).toBe(5);
+    });
+
+    test('get should return empty counters for new org/month', async () => {
+      mockUsageRepository.get.mockResolvedValue(null);
+
+      const result = await BillingUsageService.get(orgId);
+
+      expect(result.counters).toEqual({});
+      expect(result.organizationId).toBe(orgId);
+    });
+
+    test('get should return existing usage document', async () => {
+      const existing = { organizationId: orgId, month: '2026-03', counters: { executions: 42, aiCalls: 5 } };
+      mockUsageRepository.get.mockResolvedValue(existing);
+
+      const result = await BillingUsageService.get(orgId);
+
+      expect(result.counters.executions).toBe(42);
+      expect(result.counters.aiCalls).toBe(5);
+    });
+
+    test('reset should clear counters', async () => {
+      const resetDoc = { organizationId: orgId, month: '2026-03', counters: {} };
+      mockUsageRepository.reset.mockResolvedValue(resetDoc);
+
+      const result = await BillingUsageService.reset(orgId);
+
+      expect(mockUsageRepository.reset).toHaveBeenCalledWith(orgId, expect.stringMatching(/^\d{4}-\d{2}$/));
+      expect(result.counters).toEqual({});
+    });
+  });
+});