fix: address secondary PR review feedback

Author: yash-pouranik

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

@@ -10,6 +10,11 @@ const queryBuilder = async (req, res, next) => {
         const { projectId } = req.params;
         const { collectionName, prompt } = req.body;
 
+        const mongoose = require('mongoose');
+        if (!mongoose.Types.ObjectId.isValid(projectId)) {
+            throw new AppError(400, "Invalid project ID");
+        }
+
         if (typeof collectionName !== 'string' || typeof prompt !== 'string') {
             throw new AppError(400, "Collection name and prompt must be strings");
         }

apps/dashboard-api/test_python.js

@@ -4,12 +4,19 @@ const crypto = require('crypto');
 const run = async () => {
     const timestamp = Date.now().toString();
     const payload = JSON.stringify({ prompt: "test", schema_fields: [] });
-    const secret = "test-secret";
+    
+    const secret = process.env.TEST_SECRET || process.env.INTERNAL_SECRET;
+    const apiUrl = process.env.AI_API_URL || 'http://127.0.0.1:8000/ai/query-builder';
+    
+    if (!secret) {
+        console.error("Missing TEST_SECRET or INTERNAL_SECRET in environment variables.");
+        process.exit(1);
+    }
 
     const signature = crypto.createHmac('sha256', secret).update(`${timestamp}.${payload}`).digest('hex');
 
     try {
-        const res = await axios.post('http://127.0.0.1:8000/ai/query-builder', payload, {
+        const res = await axios.post(apiUrl, payload, {
             headers: {
                 'X-Internal-Signature': signature,
                 'X-Timestamp': timestamp,

apps/python-service/config.py

@@ -1,7 +1,8 @@
+from pydantic import Field
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
 class Settings(BaseSettings):
-    INTERNAL_SECRET: str
+    INTERNAL_SECRET: str = Field(min_length=32)
     GROQ_API_KEY: str = "" # Default empty, can be supplied via BYOK later
     REDIS_URL: str = "redis://localhost:6379"
 

apps/python-service/dependencies.py

@@ -17,8 +17,8 @@ async def verify_signature(request: Request):
 
     try:
         timestamp = int(timestamp_str)
-    except ValueError:
-        raise HTTPException(status_code=403, detail="Invalid timestamp format")
+    except ValueError as e:
+        raise HTTPException(status_code=403, detail="Invalid timestamp format") from e
 
     # Replay attack protection — 30s window
     current_time_ms = int(time.time() * 1000)
@@ -38,11 +38,8 @@ async def verify_signature(request: Request):
     if not hmac.compare_digest(expected_mac, signature):
         raise HTTPException(status_code=403, detail="Invalid signature")
 
-    # Redis Nonce check for exact replay attack prevention
+    # Redis Nonce check for exact replay attack prevention (Atomic SET NX)
     nonce_key = f"internal:nonce:{timestamp}:{signature}"
-    exists = await redis_client.get(nonce_key)
-    if exists:
+    success = await redis_client.set(nonce_key, "1", ex=30, nx=True)
+    if not success:
         raise HTTPException(status_code=403, detail="Replay attack detected")
-        
-    # Store the nonce with a 30-second TTL
-    await redis_client.set(nonce_key, "1", ex=30)

apps/python-service/routers/ai.py

@@ -1,3 +1,5 @@
+import asyncio
+import logging
 from fastapi import APIRouter, HTTPException, Depends
 from pydantic import BaseModel, Field
 from typing import List, Union
@@ -7,6 +9,7 @@
 from dependencies import verify_signature
 
 router = APIRouter(prefix="/ai", tags=["ai"], dependencies=[Depends(verify_signature)])
+logger = logging.getLogger(__name__)
 
 class FilterItem(BaseModel):
     field: str = Field(description="The exact field name from the schema")
@@ -19,7 +22,7 @@ class QueryResult(BaseModel):
 
 class QueryBuilderRequest(BaseModel):
     prompt: str
-    schema_fields: list
+    schema_fields: List[dict]
 
 @router.post("/query-builder", response_model=QueryResult)
 async def query_builder(request: QueryBuilderRequest):
@@ -55,13 +58,20 @@ async def query_builder(request: QueryBuilderRequest):
         # Create the LangChain chain
         chain = prompt | structured_llm
 
-        # Invoke the chain
-        result = await chain.ainvoke({
-            "schema": str(request.schema_fields),
-            "user_prompt": request.prompt
-        })
+        # Invoke the chain with a timeout to prevent hanging requests
+        result = await asyncio.wait_for(
+            chain.ainvoke({
+                "schema": str(request.schema_fields),
+                "user_prompt": request.prompt
+            }),
+            timeout=15.0
+        )
 
         return result
 
-    except Exception:
-        raise HTTPException(status_code=500, detail="Failed to generate query")
+    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
+    except Exception as e:
+        logger.error("AI Query Builder failed", exc_info=True)
+        raise HTTPException(status_code=500, detail="Internal server error") from e

apps/python-service/tests/test_auth.py

@@ -38,7 +38,7 @@ def generate_valid_signature(timestamp: int, payload: str) -> str:
 def mock_redis():
     with patch('dependencies.redis_client.get', new_callable=AsyncMock) as mock_get, \
          patch('dependencies.redis_client.set', new_callable=AsyncMock) as mock_set:
-        mock_get.return_value = None  # Simulate nonce does not exist (not a replay attack)
+        mock_set.return_value = True  # Simulate nonce does not exist (SET NX succeeds)
         yield mock_get, mock_set
 
 def test_missing_headers():
@@ -93,9 +93,9 @@ def test_valid_request():
     assert response.json()["message"] == "success"
 
 def test_replay_attack(mock_redis):
-    mock_get, mock_set = mock_redis
-    # Simulate that Redis says this nonce already exists
-    mock_get.return_value = "1" 
+    _, mock_set = mock_redis
+    # Simulate that Redis says this nonce already exists (SET NX fails)
+    mock_set.return_value = None 
 
     timestamp = int(time.time() * 1000)
     payload = '{"foo": "bar"}'

apps/web-dashboard/src/components/Database/AiQueryBar.jsx

@@ -61,6 +61,7 @@ const AiQueryBar = ({ projectId, activeCollection, onFiltersGenerated }) => {
                 value={prompt}
                 onChange={(e) => setPrompt(e.target.value)}
                 placeholder="Ask AI to filter data..."
+                aria-label="Ask AI to filter data"
                 disabled={isLoading || !activeCollection || !projectId}
                 style={{ 
                     background: 'transparent',