deploy on HF

Dockerfile

@@ -0,0 +1,10 @@
+FROM python:3.10
+
+WORKDIR /app
+
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY . .
+
+CMD ["uvicorn", "api.main:app", "--host", "0.0.0.0", "--port", "7860"]

api/main.py

@@ -1,23 +1,11 @@
 # main.py
 import os
-from pathlib import Path
 from contextlib import asynccontextmanager
 import logging
-
-# Load .env from backend root first (so OPENROUTER_API_KEY etc. are set regardless of cwd)
-_backend_root = Path(__file__).resolve().parent.parent
-_env_file = _backend_root / ".env"
-if _env_file.exists():
-    from dotenv import load_dotenv
-    load_dotenv(_env_file)
-    logging.getLogger(__name__).info(f"Loaded env from {_env_file}")
-else:
-    logging.getLogger(__name__).warning(f"No .env file at {_env_file}")
-
-from fastapi import FastAPI
+from fastapi import FastAPI, Request
 from fastapi.middleware.cors import CORSMiddleware
 
-from utils.config import get_settings, get_environment
+from utils.config import get_settings_for_environment, get_environment
 from api.routes import router
 from services.langchain_service import LangChainService
 from services.memory_service import MemoryService
@@ -30,8 +18,8 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
-# Get settings
-settings = get_settings()
+# Get settings based on environment
+settings = get_settings_for_environment()
 
 # Service instances (will be initialized in lifespan)
 langchain_service: LangChainService = None
@@ -82,8 +70,8 @@ async def lifespan(app: FastAPI):
     
     logger.info("=" * 60)
     logger.info("✅ All services initialized successfully")
-    logger.info(f"📡 Server running on http://{settings.host}:{settings.port}")
-    logger.info(f"📚 API docs available at http://{settings.host}:{settings.port}/docs")
+    logger.info("Server running on Hugging Face Space")
+    logger.info("API docs available at /docs")
     logger.info("=" * 60)
     
     yield
@@ -109,15 +97,13 @@ app = FastAPI(
 app.add_middleware(
     CORSMiddleware,
     allow_origins=[
-        "http://localhost:8000",
-        "http://127.0.0.1:8000",
         "vscode-webview://*",
-        "*"  # Allow all origins for development (restrict in production)
+        "https://*.hf.space"
     ],
     allow_credentials=True,
     allow_methods=["*"],  # Allow all HTTP methods (GET, POST, PUT, DELETE, etc.)
     allow_headers=["*"],  # Allow all headers including Authorization
-    expose_headers=["*"],  # Expose all response headers
+    
 )
 
 # Log CORS configuration
@@ -126,6 +112,14 @@ logger.info("✅ CORS middleware configured for VS Code extension compatibility"
 # Include routes with /api prefix
 app.include_router(router, prefix="/api")
 
+@app.middleware("http")
+async def add_security_headers(request: Request, call_next):
+    response = await call_next(request)
+    response.headers["X-Content-Type-Options"] = "nosniff"
+    response.headers["X-Frame-Options"] = "DENY"
+    response.headers["X-XSS-Protection"] = "1; mode=block"
+    return response
+
 
 @app.get("/")
 async def root():
@@ -156,14 +150,14 @@ if __name__ == "__main__":
     import uvicorn
     
     logger.info("Starting Augmas Backend...")
-    logger.info(f"Host: {settings.host}")
-    logger.info(f"Port: {settings.port}")
+    logger.info("Server running on Hugging Face Space")
+    logger.info("API docs available at /docs")
     logger.info(f"Reload: {settings.reload}")
     
     uvicorn.run(
-        "api.main:app", 
-        host=settings.host,
-        port=settings.port,
-        reload=settings.reload,
-        log_level=settings.log_level.lower()
-    )
+    "api.main:app",
+    host="0.0.0.0",
+    port=7860,
+    reload=settings.reload,
+    log_level=settings.log_level.lower()
+)

api/routes.py

@@ -64,16 +64,19 @@ class FileReference(BaseModel):
 class RAGQueryRequest(BaseModel):
     query: str
     workspace_id: str
+    workspace_path: str  # Required for ownership validation
     max_chunks: int = 5
 
 
 class IndexWorkspaceRequest(BaseModel):
     workspace_id: str
+    workspace_path: str  # Required for ownership validation
     files: List[Dict[str, str]]  # List of {path: str, content: str}
 
 
 class IndexFileRequest(BaseModel):
     workspace_id: str
+    workspace_path: str  # Required for ownership validation
     file_path: str
     content: str
 
@@ -257,11 +260,15 @@ async def get_current_user(
 @router.post("/chat", response_model=ChatResponse)
 async def chat(
     request: ChatRequest,
+    user_id: str = Depends(get_current_user_id),
     langchain: LangChainService = Depends(get_langchain_service),
-    rag: RAGServiceSupabase = Depends(get_rag_service)
+    rag: RAGServiceSupabase = Depends(get_rag_service),
+    memory: MemoryService = Depends(get_memory_service)
 ):
     """Process chat message"""
     try:
+        # Set the user in memory service from JWT
+        memory.set_current_user(user_id)
         # Build code context
         context = CodeContext()
         
@@ -329,6 +336,7 @@ async def chat(
 @router.post("/agent/step", response_model=AgentStepResponse)
 async def agent_step(
     request: AgentStepRequest,
+    user_id: str = Depends(get_current_user_id),
     langchain: LangChainService = Depends(get_langchain_service),
 ):
     """Agent mode: get next JSON action from the model (thought, action, input)."""
@@ -350,6 +358,7 @@ async def agent_step(
 @router.post("/rag/query")
 async def rag_query(
     request: RAGQueryRequest,
+    user_id: str = Depends(get_current_user_id),
     rag: RAGServiceSupabase = Depends(get_rag_service)
 ):
     """Query RAG for relevant context (workspace-scoped)"""
@@ -360,6 +369,13 @@ async def rag_query(
                 detail="RAG not ready. Please index workspace first."
             )
         
+        # Validate workspace ownership
+        if not rag._validate_workspace_ownership(request.workspace_id, user_id, request.workspace_path):
+            raise HTTPException(
+                status_code=403,
+                detail="Access denied: Workspace does not belong to user"
+            )
+        
         context = await rag.get_relevant_context(
             request.workspace_id,
             request.query,
@@ -376,12 +392,22 @@ async def rag_query(
 @router.post("/rag/index/workspace")
 async def index_workspace(
     request: IndexWorkspaceRequest,
+    user_id: str = Depends(get_current_user_id),
     rag: RAGServiceSupabase = Depends(get_rag_service)
 ):
     """Index workspace files"""
     try:
+        # Validate workspace ownership
+        if not rag._validate_workspace_ownership(request.workspace_id, user_id, request.workspace_path):
+            raise HTTPException(
+                status_code=403,
+                detail="Access denied: Workspace does not belong to user"
+            )
+        
         result = await rag.index_workspace(request.workspace_id, request.files)
         return result
+    except HTTPException:
+        raise
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
@@ -389,15 +415,25 @@ async def index_workspace(
 @router.post("/rag/index/file")
 async def index_file(
     request: IndexFileRequest,
+    user_id: str = Depends(get_current_user_id),
     rag: RAGServiceSupabase = Depends(get_rag_service)
 ):
     """Index a single file"""
     try:
+        # Validate workspace ownership
+        if not rag._validate_workspace_ownership(request.workspace_id, user_id, request.workspace_path):
+            raise HTTPException(
+                status_code=403,
+                detail="Access denied: Workspace does not belong to user"
+            )
+        
         result = await rag.index_workspace(
             request.workspace_id,
             [{'path': request.file_path, 'content': request.content}]
         )
         return result
+    except HTTPException:
+        raise
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
@@ -405,13 +441,24 @@ async def index_file(
 @router.delete("/rag/index/file")
 async def delete_file(
     workspace_id: str,
+    workspace_path: str,  # Required for ownership validation
     file_path: str,
+    user_id: str = Depends(get_current_user_id),
     rag: RAGServiceSupabase = Depends(get_rag_service)
 ):
     """Delete embeddings for a file"""
     try:
+        # Validate workspace ownership
+        if not rag._validate_workspace_ownership(workspace_id, user_id, workspace_path):
+            raise HTTPException(
+                status_code=403,
+                detail="Access denied: Workspace does not belong to user"
+            )
+        
         result = await rag.delete_file(workspace_id, file_path)
         return result
+    except HTTPException:
+        raise
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
@@ -419,17 +466,30 @@ async def delete_file(
 @router.get("/rag/stats")
 async def get_rag_stats(
     workspace_id: str,
+    workspace_path: str,  # Required for ownership validation
+    user_id: str = Depends(get_current_user_id),
     rag: RAGServiceSupabase = Depends(get_rag_service)
 ):
     """Get RAG indexing statistics for a workspace"""
     try:
+        # Validate workspace ownership
+        if not rag._validate_workspace_ownership(workspace_id, user_id, workspace_path):
+            raise HTTPException(
+                status_code=403,
+                detail="Access denied: Workspace does not belong to user"
+            )
+        
         return await rag.get_index_stats(workspace_id)
+    except HTTPException:
+        raise
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
 
 @router.get("/rag/status")
-async def get_rag_status(rag: RAGServiceSupabase = Depends(get_rag_service)):
+async def get_rag_status(
+    user_id: str = Depends(get_current_user_id),
+    rag: RAGServiceSupabase = Depends(get_rag_service)):
     """Get RAG service status"""
     return {
         "ready": rag.is_ready()
@@ -440,6 +500,7 @@ async def get_rag_status(rag: RAGServiceSupabase = Depends(get_rag_service)):
 
 @router.get("/models")
 async def get_available_models(
+    user_id: str = Depends(get_current_user_id),
     langchain: LangChainService = Depends(get_langchain_service)
 ):
     """Get available models"""
@@ -463,6 +524,7 @@ async def get_available_models(
 @router.post("/models/switch")
 async def switch_model(
     request: ModelSwitchRequest,
+    user_id: str = Depends(get_current_user_id),
     langchain: LangChainService = Depends(get_langchain_service)
 ):
     """Switch to a different model"""
@@ -478,6 +540,7 @@ async def switch_model(
 
 @router.get("/models/ollama/status")
 async def check_ollama_status(
+    user_id: str = Depends(get_current_user_id),
     langchain: LangChainService = Depends(get_langchain_service)
 ):
     """Check Ollama service status"""
@@ -494,6 +557,7 @@ async def check_ollama_status(
 
 @router.get("/models/test")
 async def test_model_connection(
+    user_id: str = Depends(get_current_user_id),
     langchain: LangChainService = Depends(get_langchain_service)
 ):
     """Test LLM connection"""
@@ -507,6 +571,7 @@ async def test_model_connection(
 @router.post("/workspace/set")
 async def set_workspace(
     request: SetWorkspaceRequest,
+    user_id: str = Depends(get_current_user_id),
     langchain: LangChainService = Depends(get_langchain_service)
 ):
     """Set the workspace root path for LangChainService"""
@@ -663,6 +728,14 @@ async def delete_chat_session(
 ):
     """Delete a chat session (requires JWT authentication)"""
     try:
+        memory.set_current_user(user_id)
+
+        # Verify the chat belongs to the user (security check)
+        sessions = await memory.get_chat_sessions()
+        chat_exists = any(s["id"] == chat_id for s in sessions)
+        if not chat_exists:
+            raise HTTPException(status_code=403, detail="Chat not found or access denied")
+        
         await memory.delete_chat_session(chat_id)
         return {"message": "Chat deleted successfully"}
     except ValueError as e:
@@ -678,6 +751,7 @@ async def get_user_stats(
 ):
     """Get user statistics (requires JWT authentication)"""
     try:
+        memory.set_current_user(user_id)
         stats = await memory.get_user_stats()
         return stats
     except HTTPException:
@@ -686,29 +760,4 @@ async def get_user_stats(
         raise HTTPException(status_code=500, detail=str(e))
 
 
-@router.get("/test-openrouter")
-async def test_openrouter():
-    """Test OpenRouter connectivity from backend (diagnostic endpoint)"""
-    import os
-    import httpx
-    
-    try:
-        api_key = os.getenv("OPENROUTER_API_KEY")
-        if not api_key:
-            return {"status": "error", "message": "OPENROUTER_API_KEY not found in environment"}
-        
-        async with httpx.AsyncClient(verify=False) as client:
-            response = await client.get(
-                "https://openrouter.ai/api/v1/models",
-                headers={"Authorization": f"Bearer {api_key}"},
-                timeout=10.0
-            )
-            
-            return {
-                "status": "success" if response.status_code == 200 else "error",
-                "status_code": response.status_code,
-                "message": "OpenRouter is reachable and API key is valid" if response.status_code == 200 else "OpenRouter returned an error"
-            }
-    except Exception as e:
-        logger.error(f"OpenRouter test error: {e}")
-        return {"status": "error", "message": str(e)}
+# Diagnostic endpoint removed for security - use /api/models/test instead

services/langchain_service.py

@@ -377,9 +377,10 @@ class MCPModelServer:
         
         try:
             async with httpx.AsyncClient() as client:
+                app_url = os.getenv("APP_URL", "https://huggingface.co")
                 headers = {
                     "Authorization": f"Bearer {self.openrouter_api_key}",
-                    "HTTP-Referer": "http://localhost:8000",
+                    "HTTP-Referer": app_url,
                     "X-Title": "Code Assistant"
                 }
                 response = await client.get(
@@ -483,7 +484,7 @@ class LangChainService:
                     client=openai_client.chat.completions,
                     async_client=async_openai_client.chat.completions,
                     default_headers={
-                        "HTTP-Referer": "http://localhost:8000",
+                        "HTTP-Referer": os.getenv("APP_URL", "https://huggingface.co"),
                         "X-Title": "Code Assistant"
                     },
                 )
@@ -538,24 +539,14 @@ class LangChainService:
     
     async def check_ollama_status(self) -> bool:
         """Check if Ollama service is running"""
-        try:
-            async with httpx.AsyncClient() as client:
-                response = await client.get('http://localhost:11434/api/tags', timeout=5.0)
-                return response.status_code == 200
-        except Exception as e:
-            logger.warning(f"Ollama status check failed: {e}")
-            return False
-    
+        # Ollama is not available in Hugging Face Spaces container
+        # Return False to indicate Ollama is not available
+        return False
+
     async def get_ollama_models(self) -> List[str]:
         """Get available Ollama models"""
-        try:
-            async with httpx.AsyncClient() as client:
-                response = await client.get('http://localhost:11434/api/tags', timeout=5.0)
-                if response.status_code == 200:
-                    data = response.json()
-                    return [model['name'] for model in data.get('models', [])]
-        except Exception as e:
-            logger.warning(f"Error fetching Ollama models: {e}")
+        # Ollama is not available in Hugging Face Spaces container
+        # Return empty list
         return []
     
     async def test_connection(self) -> bool:

services/rag_service_supabase.py

@@ -116,6 +116,19 @@ class RAGServiceSupabase:
     def _hash_workspace_path(path: str) -> str:
         """Create stable hash of workspace path"""
         return hashlib.md5(path.encode()).hexdigest()
+    
+    @staticmethod
+    def _generate_user_scoped_workspace_id(user_id: str, workspace_path: str) -> str:
+        """Generate user-scoped workspace ID to ensure isolation"""
+        # Combine user_id and workspace_path, then hash
+        combined = f"{user_id}:{workspace_path}"
+        return hashlib.sha256(combined.encode('utf-8')).hexdigest()
+    
+    @staticmethod
+    def _validate_workspace_ownership(workspace_id: str, user_id: str, workspace_path: str) -> bool:
+        """Validate that workspace_id belongs to the user"""
+        expected_id = RAGServiceSupabase._generate_user_scoped_workspace_id(user_id, workspace_path)
+        return workspace_id == expected_id
 
     @staticmethod
     def _content_hash(content: str) -> str:

utils/config.py

@@ -31,7 +31,7 @@ class Settings(BaseSettings):
     
     # Workspace Settings
     workspace_root: str = Field(
-        default_factory=lambda: os.getcwd(),
+        default_factory=lambda: os.getenv("WORKSPACE_ROOT", "/app"),
         description="Path to workspace root directory"
     )
     
@@ -169,6 +169,18 @@ class Settings(BaseSettings):
         description="Secret key for JWT tokens"
     )
     
+    @validator('secret_key')
+    def validate_secret_key(cls, v):
+        """Ensure secret key is changed from default in production/staging"""
+        if v == "your-secret-key-change-in-production":
+            env = os.getenv("ENVIRONMENT", "development").lower()
+            if env in ("production", "staging"):
+                raise ValueError(
+                    "SECRET_KEY environment variable must be set in production/staging! "
+                    "Set SECRET_KEY environment variable to a secure random string."
+                )
+        return v
+    
     access_token_expire_minutes: int = Field(
         default=60 * 24 * 7,  # 7 days
         description="Access token expiration time in minutes",