fix(api): enforce WIKI_AUTH_MODE everywhere and cap task admission

api/routers/auth.py

@@ -1,11 +1,25 @@
-from fastapi import APIRouter
+from fastapi import APIRouter, HTTPException
 
 from api.config import WIKI_AUTH_CODE, WIKI_AUTH_MODE
 from api.schemas import AuthorizationConfig
 
 router = APIRouter(prefix="/auth", tags=["auth"])
 
 
+def check_wiki_auth_code(authorization_code: str | None) -> None:
+    """Enforce WIKI_AUTH_MODE consistently across every protected route.
+
+    Previously only DELETE /api/wiki_cache checked this inline; every other
+    sensitive/costly route (wiki generation, repo indexing, chat, codemap,
+    cache reads) performed no such check at all, so enabling WIKI_AUTH_MODE
+    gave callers no actual protection anywhere but that one endpoint.
+    """
+    if not WIKI_AUTH_MODE:
+        return
+    if not authorization_code or WIKI_AUTH_CODE != authorization_code:
+        raise HTTPException(status_code=401, detail="Authorization code is invalid")
+
+
 @router.get("/status")
 async def get_auth_status():
     """

api/routers/chat.py

@@ -3,6 +3,7 @@
 from fastapi.websockets import WebSocketState
 
 from api.logger import get_logger
+from api.routers.auth import check_wiki_auth_code
 from api.schemas import ChatCompletionRequest
 from api.services.research import RepoNotIndexedError, research_chat
 
@@ -26,6 +27,11 @@ async def handle_websocket_chat(websocket: WebSocket):
     await websocket.accept()
     try:
         request = ChatCompletionRequest(**await websocket.receive_json())
+        try:
+            check_wiki_auth_code(request.authorization_code)
+        except HTTPException as e:
+            await websocket.send_text(f"Error: {e.detail}")
+            return
         if not request.messages or len(request.messages) == 0:
             await websocket.send_text("Error: No messages provided")
             return
@@ -69,6 +75,7 @@ async def handle_websocket_chat(websocket: WebSocket):
 @router.post("/chat/completions/stream")
 async def chat_completions_stream(request: ChatCompletionRequest):
     """Stream a chat completion response directly using Google Generative AI"""  # Validate request
+    check_wiki_auth_code(request.authorization_code)
     if not request.messages or len(request.messages) == 0:
         raise HTTPException(status_code=400, detail="No messages provided")
 

api/routers/codemap.py

@@ -3,6 +3,7 @@
 from fastapi.websockets import WebSocketState
 
 from api.logger import get_logger
+from api.routers.auth import check_wiki_auth_code
 from api.schemas import CodeMapRequest
 from api.services.codemap import generate_codemap, read_repo_file
 
@@ -17,6 +18,11 @@ async def handle_websocket_codemap(websocket: WebSocket):
     await websocket.accept()
     try:
         request = CodeMapRequest(**await websocket.receive_json())
+        try:
+            check_wiki_auth_code(request.authorization_code)
+        except HTTPException as e:
+            await websocket.send_text(str(e.detail))
+            return
         async for event in generate_codemap(request):
             if websocket.application_state != WebSocketState.CONNECTED:
                 break
@@ -39,6 +45,7 @@ async def handle_websocket_codemap(websocket: WebSocket):
 @router.post("/codemap/stream")
 async def codemap_stream(request: CodeMapRequest):
     """HTTP fallback: stream codemap generation events as NDJSON."""
+    check_wiki_auth_code(request.authorization_code)
     try:
         return StreamingResponse(
             generate_codemap(request),
@@ -54,8 +61,10 @@ async def codemap_file(
     repo_url: str = Query(..., description="Repository URL or local path"),
     file_path: str = Query(..., description="Repository-relative file path"),
     type: str = Query("github", description="Repository type"),
+    authorization_code: str | None = Query(None, description="Authorization code"),
 ):
     """Return the full content of a file from the cloned/local repository."""
+    check_wiki_auth_code(authorization_code)
     try:
         content = read_repo_file(repo_url, type, file_path)
         return {"file_path": file_path, "content": content}

api/routers/repo.py

@@ -1,12 +1,14 @@
 import asyncio
 import json
+import os
 
 from fastapi import APIRouter, Query
 from fastapi.responses import StreamingResponse
 
 from api.logger import get_logger
 from api.rag import repo_index_exist
 from api.repository import Repo
+from api.routers.auth import check_wiki_auth_code
 from api.schemas import RepoPrepareRequest
 from api.services.research import prepare_repo_index as prepare_index
 
@@ -19,42 +21,82 @@
 # (undici default = 300s). 10s keeps the connection alive with wide margin.
 _HEARTBEAT_INTERVAL_SEC = 10
 
+# Admission cap: distinct in-flight /repo/prepare indexing jobs allowed at
+# once. This endpoint has no registry of its own (unlike wiki task
+# submission); every request used to get an unconditional asyncio.create_task
+# with no limit at all, so a caller varying repo_url could queue unbounded
+# clone/embedding work.
+try:
+    _MAX_CONCURRENT_PREPARE_TASKS = int(
+        os.environ.get("DEEPWIKI_MAX_CONCURRENT_PREPARE_TASKS", "10")
+    )
+except (TypeError, ValueError):
+    _MAX_CONCURRENT_PREPARE_TASKS = 10
+
+_in_flight_prepare_keys: set[str] = set()
+
 
 @router.post("/prepare")
 async def prepare_repo_index(request: RepoPrepareRequest):
+    check_wiki_auth_code(request.authorization_code)
+
+    key = f"{request.type}_{request.repo_url}"
+
     async def event_stream():
         # Fast path: cache already warm -> return immediately.
         if repo_index_exist(Repo(repo_url=request.repo_url, repo_type=request.type)):
             yield "event: ready\ndata: already indexed\n\n"
             yield "event: done\ndata: ok\n\n"
             return
 
+        if key in _in_flight_prepare_keys:
+            yield "event: ready\ndata: already in progress\n\n"
+            yield "event: done\ndata: ok\n\n"
+            return
+
+        if len(_in_flight_prepare_keys) >= _MAX_CONCURRENT_PREPARE_TASKS:
+            yield (
+                "event: error\ndata: "
+                + json.dumps(
+                    {
+                        "error": f"{_MAX_CONCURRENT_PREPARE_TASKS} indexing jobs are "
+                        "already in flight; try again later"
+                    }
+                )
+                + "\n\n"
+            )
+            return
+
         # First byte -> response headers are flushed now (kills the timeout).
         yield ": indexing-start\n\n"
 
-        task = asyncio.create_task(prepare_index(request))
-        elapsed = 0
-        # Wait on the task, but wake up every interval to emit a heartbeat.
-        while not task.done():
-            try:
-                # shield: a heartbeat timeout must NOT cancel the indexing task.
-                await asyncio.wait_for(
-                    asyncio.shield(task), timeout=_HEARTBEAT_INTERVAL_SEC
-                )
-            except asyncio.TimeoutError:
-                elapsed += _HEARTBEAT_INTERVAL_SEC
-                yield f"event: progress\ndata: {json.dumps({'elapsed_sec': elapsed})}\n\n"
-            except Exception:
-                # Task raised; break out and report via task.exception() below.
-                break
-
-        exc = task.exception()
-        if exc is not None:
-            logger.error("Repo indexing failed for %s: %s", request.repo_url, exc)
-            yield f"event: error\ndata: {json.dumps({'error': str(exc)})}\n\n"
-        else:
-            logger.info("Repo indexing complete for %s", request.repo_url)
-            yield "event: done\ndata: ok\n\n"
+        _in_flight_prepare_keys.add(key)
+        try:
+            task = asyncio.create_task(prepare_index(request))
+            elapsed = 0
+            # Wait on the task, but wake up every interval to emit a heartbeat.
+            while not task.done():
+                try:
+                    # shield: a heartbeat timeout must NOT cancel the indexing task.
+                    await asyncio.wait_for(
+                        asyncio.shield(task), timeout=_HEARTBEAT_INTERVAL_SEC
+                    )
+                except asyncio.TimeoutError:
+                    elapsed += _HEARTBEAT_INTERVAL_SEC
+                    yield f"event: progress\ndata: {json.dumps({'elapsed_sec': elapsed})}\n\n"
+                except Exception:
+                    # Task raised; break out and report via task.exception() below.
+                    break
+
+            exc = task.exception()
+            if exc is not None:
+                logger.error("Repo indexing failed for %s: %s", request.repo_url, exc)
+                yield f"event: error\ndata: {json.dumps({'error': str(exc)})}\n\n"
+            else:
+                logger.info("Repo indexing complete for %s", request.repo_url)
+                yield "event: done\ndata: ok\n\n"
+        finally:
+            _in_flight_prepare_keys.discard(key)
 
     return StreamingResponse(
         event_stream(),

api/routers/wiki.py

@@ -6,8 +6,9 @@
 from fastapi import APIRouter, HTTPException, Query
 from fastapi.responses import JSONResponse, Response, StreamingResponse
 
-from api.config import WIKI_AUTH_CODE, WIKI_AUTH_MODE, configs
+from api.config import configs
 from api.logger import get_logger
+from api.routers.auth import check_wiki_auth_code
 from api.schemas import (
     ProcessedProjectEntry,
     WikiCacheData,
@@ -26,6 +27,7 @@
     list_wiki_cache,
     read_wiki_cache,
     registry,
+    TaskQueueFullError,
     WikiTask,
 )
 
@@ -142,8 +144,10 @@ async def read_wiki(
     repo: str = Query(..., description="Repository name"),
     repo_type: str = Query(..., description="Repository type (e.g., github, gitlab)"),
     language: str = Query(..., description="Language of the wiki content"),
+    authorization_code: Optional[str] = Query(None, description="Authorization code"),
 ):
     """Retrieve cached wiki data (structure and generated pages) for a repository."""
+    check_wiki_auth_code(authorization_code)
     supported_langs = configs["lang_config"]["supported_languages"]
     if language not in supported_langs:
         language = configs["lang_config"]["default"]
@@ -172,16 +176,13 @@ async def delete_wiki(
     """
     Deletes a specific wiki cache from the file system.
     """
+    check_wiki_auth_code(authorization_code)
+
     # Language validation
     supported_langs = configs["lang_config"]["supported_languages"]
     if language not in supported_langs:
         raise HTTPException(status_code=400, detail="Language is not supported")
 
-    if WIKI_AUTH_MODE:
-        logger.info("check the authorization code")
-        if not authorization_code or WIKI_AUTH_CODE != authorization_code:
-            raise HTTPException(status_code=401, detail="Authorization code is invalid")
-
     logger.info(
         f"Attempting to delete wiki cache for {owner}/{repo} ({repo_type}), lang: {language}"
     )
@@ -222,10 +223,14 @@ async def submit_wiki_task(request: WikiTaskRequest):
     Returns one of: created (new task), joined (an active task for the repo
     already exists), or from_cache (this variant is already generated).
     """
+    check_wiki_auth_code(request.authorization_code)
 
-    return await registry.submit(
-        WikiTask.from_wiki_request(request), async_func=generate_repo_wiki
-    )
+    try:
+        return await registry.submit(
+            WikiTask.from_wiki_request(request), async_func=generate_repo_wiki
+        )
+    except TaskQueueFullError as e:
+        raise HTTPException(status_code=429, detail=str(e))
 
 
 @router.get(

api/schemas/base.py

@@ -29,6 +29,10 @@ class RepoRequestBase(BaseModel):
         default_factory=list,
         description="List or newline-separated string of file patterns to include exclusively",
     )
+    authorization_code: str | None = Field(
+        None,
+        description="Authorization code, required on every request when WIKI_AUTH_MODE is enabled",
+    )
 
     @field_validator(
         "excluded_dirs",

api/services/wiki/__init__.py

@@ -13,6 +13,7 @@
     WikiTask,
     registry,
     generate_repo_wiki,
+    TaskQueueFullError,
 )
 
 __all__ = [
@@ -27,4 +28,5 @@
     "WikiTask",
     "registry",
     "generate_repo_wiki",
+    "TaskQueueFullError",
 ]

api/services/wiki/tasks.py

@@ -68,6 +68,19 @@ def _env_int(name, default: int) -> int:
 WIKI_PAGE_RETRIES = _env_int("DEEPWIKI_WIKI_PAGE_RETRIES", 2)
 # How long a terminal (COMPLETED/FAILED) task lingers in the registry.
 WIKI_TASK_TTL_SECONDS = _env_int("DEEPWIKI_WIKI_TASK_TTL_SECONDS", 300)
+# Admission cap: distinct non-terminal tasks allowed in the registry at once.
+# Unlike MAX_CONCURRENT_WIKI_TASKS (the execution semaphore), this bounds how
+# many tasks can be WAITING for that semaphore, since each submitted repo_key
+# is admitted into the registry (and gets its own asyncio.Task) regardless of
+# whether a slot is free. Without this, a caller varying repo_url could queue
+# unlimited pending work.
+WIKI_MAX_PENDING_TASKS = _env_int(
+    "DEEPWIKI_WIKI_MAX_PENDING_TASKS", MAX_CONCURRENT_WIKI_TASKS * 10
+)
+
+
+class TaskQueueFullError(Exception):
+    """Raised when the registry already holds WIKI_MAX_PENDING_TASKS non-terminal tasks."""
 
 
 class WikiTask(BaseModel):
@@ -142,10 +155,15 @@ class TaskRegistry:
     _lock: asyncio.Lock
     _semaphore: asyncio.Semaphore
 
-    def __init__(self, max_concurrent: int = MAX_CONCURRENT_WIKI_TASKS):
+    def __init__(
+        self,
+        max_concurrent: int = MAX_CONCURRENT_WIKI_TASKS,
+        max_pending: int = WIKI_MAX_PENDING_TASKS,
+    ):
         self._tasks = {}
         self._lock = asyncio.Lock()
         self._semaphore = asyncio.Semaphore(max_concurrent)
+        self._max_pending = max_pending
 
     def get(self, id: str) -> WikiTask | None:
         return self._tasks.get(id)
@@ -173,6 +191,11 @@ async def submit(
                     joined=True,
                 )
 
+            if len(self.active()) >= self._max_pending:
+                raise TaskQueueFullError(
+                    f"{self._max_pending} tasks are already pending; try again later"
+                )
+
             if wiki_cache_exists(
                 owner=task.request.owner,
                 repo=task.request.repo,