feat: implement BYOK AI key resolution and rate limits

apps/dashboard-api/src/__tests__/routes.projects.storage.test.js

@@ -92,6 +92,7 @@ jest.mock('../controllers/project.controller', () => {
     updateMemberRole: jest.fn(ok),
     removeMember: jest.fn(ok),
     updateCollection: jest.fn(ok),
+    updateProjectByok: jest.fn(ok),
   };
 });
 

apps/dashboard-api/src/__tests__/routes.user.test.js

@@ -40,6 +40,7 @@ jest.mock('../controllers/auth.controller', () => ({
         },
         message: 'Onboarding updated successfully',
     })),
+    updateByok: jest.fn((_req, res) => res.json({ success: true, data: { hasGroqKey: true }, message: '' })),
 }));
 
 jest.mock('../controllers/pat.controller', () => ({

apps/dashboard-api/src/controllers/ai.controller.js

@@ -1,6 +1,8 @@
-const { Project } = require('@urbackend/common/src/models');
+const { Project, Developer } = require('@urbackend/common/src/models');
 const { forwardToPythonService } = require('../utils/internalPythonClient');
-const { AppError, ApiResponse, getProjectAccessQuery } = require('@urbackend/common');
+const { AppError, ApiResponse, getProjectAccessQuery, resolveEffectivePlan } = require('@urbackend/common');
+const { decrypt } = require('@urbackend/common/src/utils/encryption');
+const { encryptForTransit } = require('../utils/transitEncryption');
 
 /**
  * Controller to handle AI Query Builder requests.
@@ -33,18 +35,26 @@ const queryBuilder = async (req, res, next) => {
             throw new AppError(403, "Cannot query the users collection via AI");
         }
 
-        // 1. Fetch the project and specifically the requested collection schema
-        const project = await Project.findOne(
-            { ...getProjectAccessQuery(req.user._id), _id: projectId, "collections.name": safeCollectionName },
-            { "collections.$": 1 }
+        // 1. Load collection and project details
+        const project = await Project.findOne({
+            _id: projectId,
+            ...getProjectAccessQuery(req.user._id)
+        });
+
+        if (!project) {
+            throw new AppError(404, "Project not found or access denied");
+        }
+
+        const collection = project.collections.find(
+            col => col.name.toLowerCase() === safeCollectionName.toLowerCase()
         );
 
-        if (!project || !project.collections || project.collections.length === 0) {
-            throw new AppError(404, "Collection not found or access denied");
+        if (!collection) {
+            throw new AppError(404, `Collection '${safeCollectionName}' not found in this project`);
         }
 
-        const collection = project.collections[0];
-        const allowedFields = new Set([
+        // Add index checking fields
+        const schemaFieldsCheck = new Set([
             ...collection.model.map(field => field.key),
             '_id',
             'createdAt',
@@ -65,10 +75,38 @@ const queryBuilder = async (req, res, next) => {
             { key: "updatedAt", type: "DATE" }
         );
 
+        // ── Hierarchical BYOK Resolution ──
+        let resolvedKey = null;
+
+        // 1. Check project-level BYOK
+        const projectByok = await Project.findById(projectId)
+            .select('+byok.groqKey.encrypted +byok.groqKey.iv +byok.groqKey.tag')
+            .lean();
+        if (projectByok?.byok?.groqKey?.encrypted) {
+            resolvedKey = decrypt(projectByok.byok.groqKey);
+        }
+
+        // 2. Fallback to developer-level BYOK
+        const dev = await Developer.findById(req.user._id).select('+byok');
+        if (!resolvedKey && dev?.byok?.groqKey?.encrypted) {
+            resolvedKey = decrypt(dev.byok.groqKey);
+        }
+
+        // 3. Encrypt for secure transit to Python
+        const encryptedByok = resolvedKey
+            ? { groqKey: encryptForTransit(resolvedKey) }
+            : null;
+
+        // Resolve effective plan dynamically based on DB (handling subscription expiry safely)
+        const effectivePlan = dev ? resolveEffectivePlan(dev) : 'free';
+
         // 3. Forward request to Python Service
         const aiResponse = await forwardToPythonService('/ai/query-builder', {
             prompt: safePrompt,
-            schema_fields: schemaFields
+            schema_fields: schemaFields,
+            developer_id: req.user._id.toString(),
+            plan: effectivePlan,
+            encrypted_byok: encryptedByok
         });
 
         // 4. Return the structured JSON to the frontend

apps/dashboard-api/src/controllers/auth.controller.js

@@ -672,3 +672,44 @@ module.exports.updateOnboarding = async (req, res, next) => {
         next(err);
     }
 };
+
+const { encrypt, resolveEffectivePlan, getPlanLimits } = require("@urbackend/common");
+
+module.exports.updateByok = async (req, res, next) => {
+    try {
+        const { groqKey } = req.body;
+
+        // Allow null to clear the key
+        if (groqKey === null) {
+            await Developer.findByIdAndUpdate(req.user._id, { $set: { byok: null } });
+            return new ApiResponse({ hasGroqKey: false }, "BYOK key cleared").send(res);
+        }
+
+        const effectivePlan = resolveEffectivePlan(req.user);
+        const limits = getPlanLimits({ plan: effectivePlan });
+        if (!limits.aiByokEnabled) {
+            throw new AppError(403, "Bring Your Own AI Key (Groq) is a Pro feature. Please upgrade to continue.");
+        }
+
+        if (typeof groqKey !== 'string' || !groqKey.startsWith('gsk_') || groqKey.length > 200) {
+            throw new AppError(400, "Invalid Groq API key format. Key must start with 'gsk_'.");
+        }
+
+        const encryptedKey = encrypt(groqKey);
+
+        await Developer.findByIdAndUpdate(req.user._id, {
+            $set: {
+                'byok.groqKey': {
+                    iv: encryptedKey.iv,
+                    encrypted: encryptedKey.encrypted,
+                    tag: encryptedKey.tag
+                }
+            }
+        });
+
+        return new ApiResponse({ hasGroqKey: true }, "BYOK key saved successfully").send(res);
+    } catch (err) {
+        if (err instanceof AppError) return next(err);
+        next(new AppError(500, "Internal server error during BYOK configuration"));
+    }
+};

apps/dashboard-api/src/controllers/project.controller.js

@@ -3427,4 +3427,71 @@ module.exports.removeMember = async (req, res, next) => {
   }
 };
 
+module.exports.updateProjectByok = async (req, res, next) => {
+  try {
+    const { projectId } = req.params;
+    const { groqKey } = req.body;
+
+    const mongoose = require('mongoose');
+    if (!mongoose.Types.ObjectId.isValid(projectId)) {
+      return next(new AppError(400, "Invalid project ID"));
+    }
+
+    const projectCheck = await Project.findOne({
+      _id: projectId,
+      ...getProjectAccessQuery(req.user._id)
+    });
+
+    if (!projectCheck) {
+      return next(new AppError(404, "Project not found or access denied"));
+    }
+
+    // Clearing the key should always be allowed (even if plan downgraded)
+    if (groqKey === null) {
+      await Project.findByIdAndUpdate(projectId, { $set: { 'byok.groqKey': null } });
+      await deleteProjectById(projectId);
+      return res.json({
+        success: true,
+        data: { hasGroqKey: false },
+        message: "Project BYOK key cleared"
+      });
+    }
+
+    // Enforce plan gate for saving a key (supports enterprise customLimits overrides)
+    const { resolveEffectivePlan, getPlanLimits } = require('@urbackend/common');
+    const effectivePlan = resolveEffectivePlan(req.developer || req.user);
+    const limits = getPlanLimits({ plan: effectivePlan, customLimits: projectCheck.customLimits || null });
+    if (!limits.aiByokEnabled) {
+      return next(new AppError(403, 'Bring Your Own AI Key (Groq) is a Pro feature. Please upgrade to continue.'));
+    }
+    if (typeof groqKey !== 'string' || !groqKey.startsWith('gsk_') || groqKey.length > 200) {
+      return next(new AppError(400, "Invalid Groq API key format. Key must start with 'gsk_'."));
+    }
+
+    const encryptedKey = encrypt(groqKey);
+
+    await Project.findByIdAndUpdate(projectId, {
+      $set: {
+        'byok.groqKey': {
+          encrypted: encryptedKey.encrypted,
+          iv: encryptedKey.iv,
+          tag: encryptedKey.tag
+        }
+      }
+    });
+
+    await deleteProjectById(projectId);
+
+    return res.json({
+      success: true,
+      data: { hasGroqKey: true },
+      message: "Project BYOK key saved successfully"
+    });
+  } catch (err) {
+    if (err instanceof AppError) return next(err);
+    next(new AppError(500, "Internal server error during BYOK configuration"));
+  }
+};
+
+
 

apps/dashboard-api/src/routes/projects.js

@@ -54,7 +54,8 @@ const {
     inviteMember,
     updateMemberRole,
     removeMember,
-    updateCollection
+    updateCollection,
+    updateProjectByok
 } = require("../controllers/project.controller");
 
 const { createAdminUser, resetPassword, getUserDetails, updateAdminUser, listAdminUsers, deleteAdminUser, listUserSessions, revokeUserSession } = require('../controllers/userAuth.controller');
@@ -108,6 +109,9 @@ router.delete('/:projectId', authMiddleware, authorizeProject('owner'), verifyEm
 // PATCH REQ FOR UPDATE PROJECT
 router.patch('/:projectId', authMiddleware, authorizeProject('admin'), planEnforcement.attachDeveloper, planEnforcement.checkByokGate, updateProject);
 
+// BYOK Project Save Endpoint
+router.put('/:projectId/byok', authMiddleware, authorizeProject('admin'), planEnforcement.attachDeveloper, updateProjectByok);
+
 // MAIL TEMPLATES (Phase 2)
 router.get('/:projectId/mail/templates', authMiddleware, authorizeProject(), listMailTemplates);
 router.get('/:projectId/mail/templates/global', authMiddleware, authorizeProject(), listGlobalMailTemplates);

apps/dashboard-api/src/routes/user.js

@@ -1,14 +1,15 @@
 const express = require('express');
 const router = express.Router();
 const authorization = require('../middlewares/authMiddleware');
-const { getMe, updateOnboarding } = require('../controllers/auth.controller');
+const { getMe, updateOnboarding, updateByok } = require('../controllers/auth.controller');
 const { createPAT, listPATs, revokePAT } = require('../controllers/pat.controller');
 const authenticateCLI = require("../middlewares/authenticateCLI");
 const {getCLIProfile,} = require("../controllers/cli.controller");
 
 router.get("/cli/me", authenticateCLI, getCLIProfile);
 router.get('/me', authorization, getMe);
 router.patch('/onboarding', authorization, updateOnboarding);
+router.put('/me/byok', authorization, updateByok);
 
 router.post('/pats', authorization, createPAT);
 router.get('/pats', authorization, listPATs);

apps/dashboard-api/src/utils/transitEncryption.js

@@ -0,0 +1,27 @@
+const crypto = require('crypto');
+
+/**
+ * Encrypts a plaintext string for secure transit to the Python service.
+ * Uses AES-256-GCM keyed by INTERNAL_PAYLOAD_KEY.
+ * @param {string} plainText - The secret to encrypt
+ * @returns {{ iv: string, encryptedData: string, authTag: string }}
+ */
+function encryptForTransit(plainText) {
+    const key = process.env.INTERNAL_PAYLOAD_KEY;
+    if (!key) throw new Error("INTERNAL_PAYLOAD_KEY is not defined in environment");
+
+    const iv = crypto.randomBytes(16);
+    const cipher = crypto.createCipheriv('aes-256-gcm', Buffer.from(key, 'hex'), iv);
+
+    let encrypted = cipher.update(plainText, 'utf8', 'hex');
+    encrypted += cipher.final('hex');
+    const tag = cipher.getAuthTag();
+
+    return {
+        iv: iv.toString('hex'),
+        encryptedData: encrypted,
+        authTag: tag.toString('hex')
+    };
+}
+
+module.exports = { encryptForTransit };

apps/python-service/.env.example

@@ -7,6 +7,10 @@
 # Must match the INTERNAL_SECRET in the root Dashboard API .env
 INTERNAL_SECRET=generate_a_random_32_char_secret_here
 
+# Must match INTERNAL_PAYLOAD_KEY in Dashboard API .env
+# Used for decrypting BYOK keys in transit (AES-256-GCM)
+INTERNAL_PAYLOAD_KEY=generate_a_64_char_hex_string_here
+
 # ── AI APIs ──────────────────────────────────────────────────────────────────
 # Required for AI Query Builder
 GROQ_API_KEY=gsk_your_groq_api_key_here

apps/python-service/config.py

@@ -3,6 +3,7 @@
 
 class Settings(BaseSettings):
     INTERNAL_SECRET: str = Field(min_length=32)
+    INTERNAL_PAYLOAD_KEY: str = Field(default="", pattern=r"^$|^[0-9a-fA-F]{64}$")  # 64-char hex for transit decryption
     GROQ_API_KEY: str = "" # Default empty, can be supplied via BYOK later
     REDIS_URL: str = "redis://localhost:6379"
 

apps/python-service/requirements.txt

@@ -6,3 +6,4 @@ pydantic==2.7.4
 pydantic-settings==2.2.1
 python-dotenv==1.0.1
 redis==5.0.3
+pycryptodome==3.22.0

apps/python-service/routers/ai.py

@@ -3,10 +3,10 @@
 from fastapi import APIRouter, HTTPException, Depends
 from pydantic import BaseModel, Field
 from typing import List, Union
-from langchain_groq import ChatGroq
 from langchain_core.prompts import ChatPromptTemplate
 from config import settings
 from dependencies import verify_signature
+from services.byok import resolve_ai_client
 
 router = APIRouter(prefix="/ai", tags=["ai"], dependencies=[Depends(verify_signature)])
 logger = logging.getLogger(__name__)
@@ -20,22 +20,34 @@ class QueryResult(BaseModel):
     filters: List[FilterItem] = Field(default_factory=list, description="List of MongoDB filters to apply to the frontend")
     sort: str = Field(default="-createdAt", description="MongoDB sort string, e.g. '-createdAt' or 'name'. Default to '-createdAt'")
 
+class EncryptedKeyPayload(BaseModel):
+    iv: str
+    encryptedData: str
+    authTag: str
+
+class EncryptedByok(BaseModel):
+    groqKey: EncryptedKeyPayload | None = None
+
 class QueryBuilderRequest(BaseModel):
     prompt: str
     schema_fields: List[dict]
+    developer_id: str
+    plan: str = "free"
+    encrypted_byok: EncryptedByok | None = None
 
 @router.post("/query-builder", response_model=QueryResult)
 async def query_builder(request: QueryBuilderRequest):
-    if not settings.GROQ_API_KEY:
-        raise HTTPException(status_code=500, detail="Groq API key not configured")
-
     try:
-        # Initialize Groq LLM
-        llm = ChatGroq(api_key=settings.GROQ_API_KEY, model_name="llama-3.1-8b-instant", temperature=0)
-
+        # Resolve the AI client (BYOK → Pro → Free with limits)
+        llm = await resolve_ai_client(
+            developer_id=request.developer_id,
+            plan=request.plan,
+            encrypted_byok=request.encrypted_byok.model_dump() if request.encrypted_byok else None,
+        )
+
         # Enforce structured output based on our Pydantic schema
         structured_llm = llm.with_structured_output(QueryResult)
-        
+
         # Build the system prompt
         system_prompt = """You are a highly intelligent database query builder for a MongoDB-based BaaS called urBackend.
 Your job is to take the user's natural language request and convert it into a set of structured filters and a sort string.
@@ -45,19 +57,19 @@ async def query_builder(request: QueryBuilderRequest):
 
 CRITICAL INSTRUCTION:
 Your `filters` output MUST be a list (array) of objects matching the FilterItem schema exactly (each having `field`, `operator`, and `value`).
-DO NOT output a raw MongoDB filter dict like `{{"price": {{"$gt": 1000}} }}`. 
-Correct format: `[{{ "field": "price", "operator": "_gt", "value": 1000 }}]`.
+DO NOT output a raw MongoDB filter dict like {{"price": {{"$gt": 1000}} }}. 
+Correct format: [{{ "field": "price", "operator": "_gt", "value": 1000 }}].
 
 Schema Fields: {schema}"""
 
         prompt = ChatPromptTemplate.from_messages([
             ("system", system_prompt),
             ("human", "{user_prompt}")
         ])
-        
+
         # Create the LangChain chain
         chain = prompt | structured_llm
-        
+
         # Invoke the chain with a timeout to prevent hanging requests
         result = await asyncio.wait_for(
             chain.ainvoke({
@@ -66,9 +78,11 @@ async def query_builder(request: QueryBuilderRequest):
             }),
             timeout=15.0
         )
-        
+
         return result
-
+
+    except HTTPException:
+        raise  # Re-raise BYOK/rate-limit errors as-is
     except asyncio.TimeoutError as e:
         logger.error("AI Query Builder timeout", exc_info=True)
         raise HTTPException(status_code=504, detail="AI request timed out") from e