diff --git a/apps/dashboard-api/src/__tests__/routes.projects.storage.test.js b/apps/dashboard-api/src/__tests__/routes.projects.storage.test.js index 3186f8288..cbd9aba31 100644 --- a/apps/dashboard-api/src/__tests__/routes.projects.storage.test.js +++ b/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), }; }); diff --git a/apps/dashboard-api/src/__tests__/routes.user.test.js b/apps/dashboard-api/src/__tests__/routes.user.test.js index f2dc36f4c..5bd3c75c6 100644 --- a/apps/dashboard-api/src/__tests__/routes.user.test.js +++ b/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', () => ({ diff --git a/apps/dashboard-api/src/controllers/ai.controller.js b/apps/dashboard-api/src/controllers/ai.controller.js index 8acc33c60..46a08bd54 100644 --- a/apps/dashboard-api/src/controllers/ai.controller.js +++ b/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,41 @@ 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.findOne({ + _id: projectId, + ...getProjectAccessQuery(req.user._id) + }) + .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 diff --git a/apps/dashboard-api/src/controllers/auth.controller.js b/apps/dashboard-api/src/controllers/auth.controller.js index 623d5b20d..e2ebefaeb 100644 --- a/apps/dashboard-api/src/controllers/auth.controller.js +++ b/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")); + } +}; diff --git a/apps/dashboard-api/src/controllers/project.controller.js b/apps/dashboard-api/src/controllers/project.controller.js index 4991644e8..89b68794e 100644 --- a/apps/dashboard-api/src/controllers/project.controller.js +++ b/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")); + } +}; + + diff --git a/apps/dashboard-api/src/routes/projects.js b/apps/dashboard-api/src/routes/projects.js index db6043538..655096b97 100644 --- a/apps/dashboard-api/src/routes/projects.js +++ b/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); diff --git a/apps/dashboard-api/src/routes/user.js b/apps/dashboard-api/src/routes/user.js index 43ffe037b..8b05f8f05 100644 --- a/apps/dashboard-api/src/routes/user.js +++ b/apps/dashboard-api/src/routes/user.js @@ -1,7 +1,7 @@ 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"); @@ -9,6 +9,7 @@ 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); diff --git a/apps/dashboard-api/src/utils/transitEncryption.js b/apps/dashboard-api/src/utils/transitEncryption.js new file mode 100644 index 000000000..48c0d8d40 --- /dev/null +++ b/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 }; diff --git a/apps/python-service/.env.example b/apps/python-service/.env.example index 6e0eb3dad..e0d3a190d 100644 --- a/apps/python-service/.env.example +++ b/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 diff --git a/apps/python-service/config.py b/apps/python-service/config.py index 387b2f0aa..d94cbb6f1 100644 --- a/apps/python-service/config.py +++ b/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" diff --git a/apps/python-service/requirements.txt b/apps/python-service/requirements.txt index c69a2573f..231fba337 100644 --- a/apps/python-service/requirements.txt +++ b/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 diff --git a/apps/python-service/routers/ai.py b/apps/python-service/routers/ai.py index 7f28e33ae..5919ccab5 100644 --- a/apps/python-service/routers/ai.py +++ b/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,8 +57,8 @@ 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}""" @@ -54,10 +66,10 @@ async def query_builder(request: QueryBuilderRequest): ("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 diff --git a/apps/python-service/services/__init__.py b/apps/python-service/services/__init__.py new file mode 100644 index 000000000..a70b3029a --- /dev/null +++ b/apps/python-service/services/__init__.py @@ -0,0 +1 @@ +# Services package diff --git a/apps/python-service/services/byok.py b/apps/python-service/services/byok.py new file mode 100644 index 000000000..cb4e64cb3 --- /dev/null +++ b/apps/python-service/services/byok.py @@ -0,0 +1,105 @@ +"""BYOK (Bring Your Own Key) resolution and AI generation rate limiting.""" + +import datetime +import logging + +from fastapi import HTTPException +from langchain_groq import ChatGroq + +from config import settings +from dependencies import redis_client +from services.transit_crypto import decrypt_transit + +logger = logging.getLogger(__name__) + +FREE_TIER_LIMIT = 5 +PRO_TIER_LIMIT = 100 + +LUA_RATE_LIMIT = """ +local count = redis.call('INCR', KEYS[1]) +local ttl = redis.call('TTL', KEYS[1]) +if count == 1 or ttl == -1 then + redis.call('EXPIRE', KEYS[1], ARGV[1]) + ttl = tonumber(ARGV[1]) +end +return {count, ttl} +""" + + +async def resolve_ai_client( + developer_id: str | None, + plan: str | None, + encrypted_byok: dict | None, +) -> ChatGroq: + """Resolve the appropriate AI client based on BYOK and plan. + + Resolution order: + 1. BYOK key present → decrypt, use developer's key, no rate limit + 2. Pro plan, no BYOK → platform key, 100 generations/month limit + 3. Free plan, no BYOK → platform key, 5 generations/month limit + + Args: + developer_id: The developer's MongoDB ObjectId string. + plan: The developer's plan ('free' or 'pro'). + encrypted_byok: Transit-encrypted BYOK payload from Node, or None. + + Returns: + An initialized ChatGroq client. + + Raises: + HTTPException: 401 if BYOK key is invalid, 403 if rate limited, 500 if no platform key. + """ + + # ── 1. Try BYOK ── + if encrypted_byok and encrypted_byok.get("groqKey"): + try: + decrypted_key = decrypt_transit(encrypted_byok["groqKey"]) + return ChatGroq( + api_key=decrypted_key, + model_name="llama-3.1-8b-instant", + temperature=0, + ) + except Exception as e: + logger.warning("BYOK decryption/init failed: %s", e) + if isinstance(e, ValueError) and "INTERNAL_PAYLOAD_KEY" in str(e): + raise HTTPException(status_code=500, detail=str(e)) from e + raise HTTPException(status_code=401, detail="Invalid BYOK Groq key provided") from e + + # ── 2. Platform key guard ── + if not settings.GROQ_API_KEY: + raise HTTPException( + status_code=500, detail="Platform AI key not configured" + ) + + # ── 3. Platform key atomic rate limit ── + if not developer_id: + raise HTTPException( + status_code=400, + detail="developer_id is required for platform key rate limiting" + ) + + limit = PRO_TIER_LIMIT if plan == "pro" else FREE_TIER_LIMIT + month = datetime.datetime.now(datetime.UTC).strftime("%Y-%m") + key = f"ai:gen:count:{developer_id}:{month}" + + # Atomic script execution: INCR + EXPIRE (if new key) + count, ttl = await redis_client.eval(LUA_RATE_LIMIT, 1, key, 32 * 86400) + + if count > limit: + tier_name = "Pro" if plan == "pro" else "Free" + upgrade_msg = ( + "Add your own Groq key in Settings to get unlimited usage." + if plan == "pro" + else "Upgrade to Pro or add your own Groq key in Settings." + ) + raise HTTPException( + status_code=403, + detail=f"{tier_name} tier AI limit reached ({limit}/month). {upgrade_msg}", + ) + + # ── 4. Return platform client ── + return ChatGroq( + api_key=settings.GROQ_API_KEY, + model_name="llama-3.1-8b-instant", + temperature=0, + ) diff --git a/apps/python-service/services/transit_crypto.py b/apps/python-service/services/transit_crypto.py new file mode 100644 index 000000000..a08bc4029 --- /dev/null +++ b/apps/python-service/services/transit_crypto.py @@ -0,0 +1,30 @@ +"""Decrypt payloads encrypted by Node's encryptForTransit() using AES-256-GCM.""" + +from Crypto.Cipher import AES +from config import settings + + +def decrypt_transit(encrypted_obj: dict) -> str: + """Decrypt a transit-encrypted value. + + Args: + encrypted_obj: Dict with keys 'iv', 'encryptedData', 'authTag' (all hex strings). + + Returns: + The decrypted plaintext string. + + Raises: + ValueError: If INTERNAL_PAYLOAD_KEY is not configured. + Exception: If decryption or verification fails. + """ + if not settings.INTERNAL_PAYLOAD_KEY: + raise ValueError("INTERNAL_PAYLOAD_KEY is not configured") + + key = bytes.fromhex(settings.INTERNAL_PAYLOAD_KEY) + iv = bytes.fromhex(encrypted_obj["iv"]) + tag = bytes.fromhex(encrypted_obj["authTag"]) + ciphertext = bytes.fromhex(encrypted_obj["encryptedData"]) + + cipher = AES.new(key, AES.MODE_GCM, nonce=iv) + plaintext = cipher.decrypt_and_verify(ciphertext, tag) + return plaintext.decode("utf-8") diff --git a/apps/python-service/tests/test_auth.py b/apps/python-service/tests/test_auth.py index c017b651b..ad571f2c9 100644 --- a/apps/python-service/tests/test_auth.py +++ b/apps/python-service/tests/test_auth.py @@ -12,7 +12,8 @@ SERVICE_ROOT = Path(__file__).resolve().parents[1] if str(SERVICE_ROOT) not in sys.path: sys.path.insert(0, str(SERVICE_ROOT)) -os.environ.setdefault("INTERNAL_SECRET", "test-internal-secret") +os.environ.setdefault("INTERNAL_SECRET", "a" * 32) +os.environ.setdefault("INTERNAL_PAYLOAD_KEY", "a" * 64) from dependencies import verify_signature from config import settings diff --git a/apps/python-service/tests/test_byok.py b/apps/python-service/tests/test_byok.py new file mode 100644 index 000000000..4454b14da --- /dev/null +++ b/apps/python-service/tests/test_byok.py @@ -0,0 +1,133 @@ +import pytest +import os +import sys +import datetime +from pathlib import Path +from unittest.mock import AsyncMock, patch, MagicMock + +SERVICE_ROOT = Path(__file__).resolve().parents[1] +if str(SERVICE_ROOT) not in sys.path: + sys.path.insert(0, str(SERVICE_ROOT)) + +# Set env vars before importing modules +os.environ.setdefault("INTERNAL_SECRET", "a" * 32) +os.environ.setdefault("INTERNAL_PAYLOAD_KEY", "a" * 64) +os.environ.setdefault("GROQ_API_KEY", "gsk_test_platform_key") + +from services.byok import resolve_ai_client, FREE_TIER_LIMIT, LUA_RATE_LIMIT +from config import settings + + +@pytest.fixture(autouse=True) +def mock_redis(): + """Mock Redis client for all tests.""" + with patch('services.byok.redis_client') as mock_rc: + mock_rc.eval = AsyncMock(return_value=[1, -1]) + yield mock_rc + + +@pytest.fixture +def mock_decrypt(): + """Mock transit decryption.""" + with patch('services.byok.decrypt_transit') as mock: + mock.return_value = "gsk_test_byok_key_12345" + yield mock + + +@pytest.mark.asyncio +async def test_byok_key_decrypted_and_used(mock_decrypt): + """When BYOK key is provided, it should be decrypted and used.""" + encrypted_byok = {"groqKey": {"iv": "aa", "encryptedData": "bb", "authTag": "cc"}} + + with patch('services.byok.ChatGroq') as MockGroq: + MockGroq.return_value = MagicMock() + client = await resolve_ai_client("dev123", "free", encrypted_byok) + + mock_decrypt.assert_called_once_with(encrypted_byok["groqKey"]) + MockGroq.assert_called_once_with( + api_key="gsk_test_byok_key_12345", + model_name="llama-3.1-8b-instant", + temperature=0, + ) + + +@pytest.mark.asyncio +async def test_invalid_byok_returns_401(): + """When BYOK decryption fails, should raise 401.""" + encrypted_byok = {"groqKey": {"iv": "bad", "encryptedData": "bad", "authTag": "bad"}} + + with patch('services.byok.decrypt_transit', side_effect=Exception("decrypt failed")): + with pytest.raises(Exception) as exc_info: + await resolve_ai_client("dev123", "free", encrypted_byok) + assert exc_info.value.status_code == 401 + + +@pytest.mark.asyncio +async def test_free_tier_under_limit(mock_redis): + """Free tier under limit should use platform key.""" + mock_redis.eval = AsyncMock(return_value=[3, 100000]) + + with patch('services.byok.ChatGroq') as MockGroq: + MockGroq.return_value = MagicMock() + client = await resolve_ai_client("dev123", "free", None) + + MockGroq.assert_called_once_with( + api_key=settings.GROQ_API_KEY, + model_name="llama-3.1-8b-instant", + temperature=0, + ) + + +@pytest.mark.asyncio +async def test_free_tier_limit_enforced(mock_redis): + """Free tier over limit should raise 403.""" + mock_redis.eval = AsyncMock(return_value=[6, 100000]) + + with pytest.raises(Exception) as exc_info: + await resolve_ai_client("dev123", "free", None) + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_pro_tier_under_limit(mock_redis): + """Pro tier under limit should use platform key.""" + mock_redis.eval = AsyncMock(return_value=[99, 100000]) + + with patch('services.byok.ChatGroq') as MockGroq: + MockGroq.return_value = MagicMock() + client = await resolve_ai_client("dev123", "pro", None) + + MockGroq.assert_called_once_with( + api_key=settings.GROQ_API_KEY, + model_name="llama-3.1-8b-instant", + temperature=0, + ) + + +@pytest.mark.asyncio +async def test_pro_tier_limit_enforced(mock_redis): + """Pro tier over limit should raise 403.""" + mock_redis.eval = AsyncMock(return_value=[101, 100000]) + + with pytest.raises(Exception) as exc_info: + await resolve_ai_client("dev123", "pro", None) + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_first_use_sets_ttl(mock_redis): + """First usage should run the Lua script with 32-day TTL value.""" + mock_redis.eval = AsyncMock(return_value=[1, 2764800]) + + with patch('services.byok.ChatGroq') as MockGroq: + MockGroq.return_value = MagicMock() + await resolve_ai_client("dev123", "free", None) + + expected_key = f"ai:gen:count:dev123:{datetime.datetime.now(datetime.UTC).strftime('%Y-%m')}" + mock_redis.eval.assert_called_once_with( + LUA_RATE_LIMIT, + 1, + expected_key, + 2764800 + ) + diff --git a/packages/common/src/models/Developer.js b/packages/common/src/models/Developer.js index b77f386cc..5855f463d 100644 --- a/packages/common/src/models/Developer.js +++ b/packages/common/src/models/Developer.js @@ -106,6 +106,17 @@ const developerSchema = new mongoose.Schema({ onboarding: { type: onboardingSchema, default: () => ({}) + }, + byok: { + type: new mongoose.Schema({ + groqKey: { + iv: { type: String, default: null }, + encrypted: { type: String, default: null }, + tag: { type: String, default: null }, + } + }, { _id: false }), + select: false, + default: null } }, { timestamps: true }); diff --git a/packages/common/src/models/Project.js b/packages/common/src/models/Project.js index 89bfe8d5b..4e05fc78e 100755 --- a/packages/common/src/models/Project.js +++ b/packages/common/src/models/Project.js @@ -160,6 +160,11 @@ const projectSchema = new mongoose.Schema( maxCollections: { type: Number, default: null } }, + // BYOK — Bring Your Own Key for AI + byok: { + groqKey: { type: resourceConfigSchema, default: null }, + }, + // Team Members members: { type: [