Add BYOK support for user AI provider keys

Introduce Bring-Your-Own-Key (BYOK) support so users can save encrypted AI provider API keys and use them for requests. Adds AES-256-GCM encryption utilities and ENCRYPTION_KEY env var; a new UserProviderKey entity, DTOs, controller, service and ProviderKeysModule; and a migration to create the user_provider_keys table. Integrates BYOK into the stack: Public API will attach decrypted keys for free-tier users, ai-service accepts ephemeral api_key and ProviderRegistry.get_ephemeral creates one-off provider instances, and QuotaInterceptor bypasses platform quota for users with BYOK keys. Billing now records metered overage usage to Stripe (new subscription fields + migration) for paid tiers. Frontend settings UI updated to manage/save/validate/remove provider keys and show hints; minor UI cleanup and exports adjusted.

Author: dewitt4

.env.example

@@ -46,6 +46,10 @@ GEMINI_API_KEY=
 HUGGINGFACE_API_KEY=
 DEFAULT_AI_PROVIDER=openai
 
+# --- BYOK Encryption (AES-256-GCM key for user-provided API key storage) ---
+# Generate with: node -e "console.log(require('crypto').randomBytes(32).toString('hex'))"
+ENCRYPTION_KEY=
+
 # --- AI Service ---
 AI_SERVICE_URL=http://localhost:8000
 AI_SERVICE_PORT=8000

packages/ai-service/app/routers/conversations.py

@@ -1,8 +1,8 @@
 """Conversation management and AI chat endpoints."""
 
-from fastapi import APIRouter, HTTPException
-from pydantic import BaseModel
-from typing import Optional
+from fastapi import APIRouter, HTTPException  # type: ignore
+from pydantic import BaseModel  # type: ignore
+from typing import Any, Dict, Optional
 from datetime import datetime
 
 from app.core.database import get_db
@@ -17,10 +17,10 @@
 
 # Initialize providers on import
 ProviderRegistry.initialize(
-    openai_key=settings.OPENAI_API_KEY,
-    anthropic_key=settings.ANTHROPIC_API_KEY,
-    gemini_key=settings.GEMINI_API_KEY,
-    huggingface_key=settings.HUGGINGFACE_API_KEY,
+    openai_key=settings.OPENAI_API_KEY or "",
+    anthropic_key=settings.ANTHROPIC_API_KEY or "",
+    gemini_key=settings.GEMINI_API_KEY or "",
+    huggingface_key=settings.HUGGINGFACE_API_KEY or "",
 )
 
 
@@ -37,13 +37,16 @@ class SendMessageRequest(BaseModel):
     temperature: Optional[float] = 0.7
     max_tokens: Optional[int] = 2048
     system_prompt: Optional[str] = None
+    # BYOK: caller-supplied key (decrypted by the core service,
+    # transmitted over the internal network — never exposed to browsers).
+    api_key: Optional[str] = None
 
 
 @router.post("/conversations")
 async def create_conversation(req: CreateConversationRequest):
     """Create a new AI conversation."""
     db = get_db()
-    conversation = {
+    conversation: Dict[str, Any] = {
         "applicationId": req.application_id,
         "userId": req.user_id,
         "title": req.title,
@@ -61,10 +64,12 @@ async def create_conversation(req: CreateConversationRequest):
 @router.get("/conversations/{conversation_id}")
 async def get_conversation(conversation_id: str):
     """Get conversation by ID with message history."""
-    from bson import ObjectId
+    from bson import ObjectId  # type: ignore
 
     db = get_db()
-    conv = await db.ai_conversations.find_one({"_id": ObjectId(conversation_id)})
+    conv = await db.ai_conversations.find_one(
+        {"_id": ObjectId(conversation_id)}
+    )
     if not conv:
         raise HTTPException(status_code=404, detail="Conversation not found")
 
@@ -75,21 +80,42 @@ async def get_conversation(conversation_id: str):
 @router.post("/conversations/{conversation_id}/messages")
 async def send_message(conversation_id: str, req: SendMessageRequest):
     """Send a message and get an AI response."""
-    from bson import ObjectId
+    from bson import ObjectId  # type: ignore
 
     db = get_db()
-    conv = await db.ai_conversations.find_one({"_id": ObjectId(conversation_id)})
+    conv = await db.ai_conversations.find_one(
+        {"_id": ObjectId(conversation_id)}
+    )
     if not conv:
         raise HTTPException(status_code=404, detail="Conversation not found")
 
     # Determine provider
     provider_name = req.provider or settings.DEFAULT_AI_PROVIDER
-    provider = ProviderRegistry.get(provider_name)
-    if not provider:
-        raise HTTPException(
-            status_code=400,
-            detail=f"Provider '{provider_name}' not available. Configure API key.",
-        )
+
+    # If a BYOK key was supplied by the core service, create a short-lived
+    # ephemeral provider instance (not cached — avoids key leaks
+    # across requests).
+    if req.api_key:
+        provider = ProviderRegistry.get_ephemeral(provider_name, req.api_key)
+        if not provider:
+            raise HTTPException(
+                status_code=400,
+                detail=(
+                    f"Provider '{provider_name}' is not"
+                    " supported for BYOK."
+                ),
+            )
+    else:
+        provider = ProviderRegistry.get(provider_name)
+        if not provider:
+            raise HTTPException(
+                status_code=400,
+                detail=(
+                    f"Provider '{provider_name}' not available."
+                    " Configure API key in Settings → AI"
+                    " Providers."
+                ),
+            )
 
     # Build message history
     messages = []
@@ -113,7 +139,10 @@ async def send_message(conversation_id: str, req: SendMessageRequest):
     try:
         response = await provider.chat(chat_request)
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"AI provider error: {str(e)}")
+        raise HTTPException(
+            status_code=500,
+            detail=f"AI provider error: {str(e)}",
+        )
 
     # Store messages
     user_msg = {

packages/ai-service/app/services/ai_providers.py

@@ -43,7 +43,9 @@ async def chat(self, request: ChatRequest) -> ChatResponse:
         pass
 
     @abstractmethod
-    async def chat_stream(self, request: ChatRequest) -> AsyncGenerator[str, None]:
+    def chat_stream(
+        self, request: ChatRequest
+    ) -> AsyncGenerator[str, None]:
         pass
 
 
@@ -60,13 +62,19 @@ def name(self) -> str:
 
     @property
     def available_models(self) -> list[str]:
-        return ["gpt-4", "gpt-4-turbo", "gpt-4o", "gpt-4o-mini", "gpt-3.5-turbo"]
+        return [
+            "gpt-4", "gpt-4-turbo", "gpt-4o",
+            "gpt-4o-mini", "gpt-3.5-turbo",
+        ]
 
     async def chat(self, request: ChatRequest) -> ChatResponse:
         model = request.model or "gpt-4"
         response = await self.client.chat.completions.create(
             model=model,
-            messages=[{"role": m.role, "content": m.content} for m in request.messages],
+            messages=[  # type: ignore
+                {"role": m.role, "content": m.content}  # type: ignore
+                for m in request.messages
+            ],
             temperature=request.temperature,
             max_tokens=request.max_tokens,
         )
@@ -75,22 +83,35 @@ async def chat(self, request: ChatRequest) -> ChatResponse:
             model=model,
             provider=self.name,
             usage={
-                "prompt_tokens": response.usage.prompt_tokens if response.usage else 0,
-                "completion_tokens": response.usage.completion_tokens if response.usage else 0,
-                "total_tokens": response.usage.total_tokens if response.usage else 0,
+                "prompt_tokens": (
+                    response.usage.prompt_tokens if response.usage else 0
+                ),
+                "completion_tokens": (
+                    response.usage.completion_tokens
+                    if response.usage
+                    else 0
+                ),
+                "total_tokens": (
+                    response.usage.total_tokens if response.usage else 0
+                ),
             },
         )
 
-    async def chat_stream(self, request: ChatRequest) -> AsyncGenerator[str, None]:
+    async def chat_stream(
+        self, request: ChatRequest
+    ) -> AsyncGenerator[str, None]:
         model = request.model or "gpt-4"
         stream = await self.client.chat.completions.create(
             model=model,
-            messages=[{"role": m.role, "content": m.content} for m in request.messages],
+            messages=[  # type: ignore
+                {"role": m.role, "content": m.content}  # type: ignore
+                for m in request.messages
+            ],
             temperature=request.temperature,
             max_tokens=request.max_tokens,
             stream=True,
         )
-        async for chunk in stream:
+        async for chunk in stream:  # type: ignore[union-attr]
             if chunk.choices[0].delta.content:
                 yield chunk.choices[0].delta.content
 
@@ -126,20 +147,24 @@ async def chat(self, request: ChatRequest) -> ChatResponse:
             model=model,
             max_tokens=request.max_tokens,
             system=system if system else "You are a helpful assistant.",
-            messages=messages,
+            messages=messages,  # type: ignore[arg-type]
         )
         return ChatResponse(
-            content=response.content[0].text,
+            content=response.content[0].text,  # type: ignore[union-attr]
             model=model,
             provider=self.name,
             usage={
                 "prompt_tokens": response.usage.input_tokens,
                 "completion_tokens": response.usage.output_tokens,
-                "total_tokens": response.usage.input_tokens + response.usage.output_tokens,
+                "total_tokens": (
+                    response.usage.input_tokens + response.usage.output_tokens
+                ),
             },
         )
 
-    async def chat_stream(self, request: ChatRequest) -> AsyncGenerator[str, None]:
+    async def chat_stream(
+        self, request: ChatRequest
+    ) -> AsyncGenerator[str, None]:
         model = request.model or "claude-sonnet-4-5-20250929"
         system = ""
         messages = []
@@ -153,7 +178,7 @@ async def chat_stream(self, request: ChatRequest) -> AsyncGenerator[str, None]:
             model=model,
             max_tokens=request.max_tokens,
             system=system if system else "You are a helpful assistant.",
-            messages=messages,
+            messages=messages,  # type: ignore[arg-type]
         ) as stream:
             async for text in stream.text_stream:
                 yield text
@@ -163,7 +188,7 @@ class GeminiProvider(AIProvider):
     """Google Gemini provider."""
 
     def __init__(self, api_key: str):
-        import google.generativeai as genai
+        import google.generativeai as genai  # type: ignore
         self._genai = genai
         genai.configure(api_key=api_key)
 
@@ -175,8 +200,12 @@ def name(self) -> str:
     def available_models(self) -> list[str]:
         return ["gemini-2.0-flash", "gemini-1.5-pro", "gemini-1.5-flash"]
 
-    def _build_contents(self, messages: list[ChatMessage]) -> tuple[list[dict], str]:
-        """Convert ChatMessages to Gemini format, extracting system instruction."""
+    def _build_contents(
+        self, messages: list[ChatMessage]
+    ) -> tuple[list[dict], str]:
+        """Convert ChatMessages to Gemini format, extracting system
+        instruction.
+        """
         system_instruction = ""
         contents = []
         for m in messages:
@@ -226,7 +255,9 @@ async def chat(self, request: ChatRequest) -> ChatResponse:
             usage=usage,
         )
 
-    async def chat_stream(self, request: ChatRequest) -> AsyncGenerator[str, None]:
+    async def chat_stream(
+        self, request: ChatRequest
+    ) -> AsyncGenerator[str, None]:
         model_name = request.model or "gemini-2.0-flash"
         contents, system_instruction = self._build_contents(request.messages)
 
@@ -322,7 +353,9 @@ async def chat(self, request: ChatRequest) -> ChatResponse:
             usage={},
         )
 
-    async def chat_stream(self, request: ChatRequest) -> AsyncGenerator[str, None]:
+    async def chat_stream(
+        self, request: ChatRequest
+    ) -> AsyncGenerator[str, None]:
         model = request.model or "mistralai/Mistral-7B-Instruct-v0.3"
         prompt = self._build_prompt(request.messages)
 
@@ -342,7 +375,6 @@ async def chat_stream(self, request: ChatRequest) -> AsyncGenerator[str, None]:
             json=payload,
         ) as response:
             response.raise_for_status()
-            buffer = ""
             async for line in response.aiter_lines():
                 if line.startswith("data:"):
                     import json
@@ -356,7 +388,9 @@ async def chat_stream(self, request: ChatRequest) -> AsyncGenerator[str, None]:
 
 
 class ProviderRegistry:
-    """Registry of available AI providers."""
+    """Registry of available AI providers (platform keys, initialized at
+    startup).
+    """
 
     _providers: dict[str, AIProvider] = {}
 
@@ -378,10 +412,10 @@ def list_providers(cls) -> list[dict]:
     @classmethod
     def initialize(
         cls,
-        openai_key: str = None,
-        anthropic_key: str = None,
-        gemini_key: str = None,
-        huggingface_key: str = None,
+        openai_key: Optional[str] = None,
+        anthropic_key: Optional[str] = None,
+        gemini_key: Optional[str] = None,
+        huggingface_key: Optional[str] = None,
     ):
         if openai_key:
             cls.register(OpenAIProvider(openai_key))
@@ -391,3 +425,21 @@ def initialize(
             cls.register(GeminiProvider(gemini_key))
         if huggingface_key:
             cls.register(HuggingFaceProvider(huggingface_key))
+
+    @classmethod
+    def get_ephemeral(cls, name: str, api_key: str) -> Optional[AIProvider]:
+        """
+        Create a one-time, non-cached provider instance using a
+        caller-supplied key. Used for BYOK (Bring Your Own Key) requests
+        so user keys are never stored in the registry and do not leak
+        across requests.
+        """
+        if name == "openai":
+            return OpenAIProvider(api_key)
+        if name == "anthropic":
+            return AnthropicProvider(api_key)
+        if name == "gemini":
+            return GeminiProvider(api_key)
+        if name == "huggingface":
+            return HuggingFaceProvider(api_key)
+        return None

packages/core/src/app.module.ts

@@ -33,6 +33,7 @@ import { SsoModule } from "./modules/sso/sso.module";
 import { DataExportModule } from "./modules/data-export/data-export.module";
 import { SystemHealthModule } from "./modules/system-health/system-health.module";
 import { StripeModule } from "./modules/stripe/stripe.module";
+import { ProviderKeysModule } from "./modules/provider-keys/provider-keys.module";
 import { ScheduleModule } from "@nestjs/schedule";
 
 @Module({
@@ -131,6 +132,7 @@ import { ScheduleModule } from "@nestjs/schedule";
     SsoModule,
     DataExportModule,
     SystemHealthModule,
+    ProviderKeysModule,
   ],
 })
 export class AppModule {}