From 6cfdb2b1c0af822376d57cc49b525d5641dfdbac Mon Sep 17 00:00:00 2001
From: YurenHao0426 <blackhao0426@gmail.com>
Date: Fri, 13 Feb 2026 19:01:53 +0000
Subject: Add username fallback for API key resolution when JWT token expires

When the JWT token is expired or missing, endpoints could not resolve
user API keys and fell back to environment variables (which are unset).
Added resolve_user() helper that falls back to DB lookup by username
query param, and added ?user= to all frontend API calls as a belt-and-
suspenders approach alongside auth tokens.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 backend/app/main.py | 37 ++++++++++++++++++++++++++++++-------
 1 file changed, 30 insertions(+), 7 deletions(-)

(limited to 'backend/app')

diff --git a/backend/app/main.py b/backend/app/main.py
index 2c625a9..d48ec89 100644
--- a/backend/app/main.py
+++ b/backend/app/main.py
@@ -108,6 +108,23 @@ def get_user_api_key(user: User | None, provider: str) -> str | None:
         return os.getenv("OPENROUTER_API_KEY")
     return None
 
+
+def resolve_user(current_user: User | None, username: str | None) -> User | None:
+    """
+    Resolve a User object: prefer authenticated current_user, fall back to DB
+    lookup by username. This handles cases where the JWT token is expired but the
+    username is still known from query params.
+    """
+    if current_user:
+        return current_user
+    if username and username != DEFAULT_USER:
+        try:
+            db = next(get_db())
+            return db.query(User).filter(User.username == username).first()
+        except Exception:
+            return None
+    return None
+
 def ensure_user_root(user: str) -> str:
     """
     Ensures the new data root structure:
@@ -340,20 +357,22 @@ def extract_image_attachments(user: str, attached_ids: List[str]) -> tuple[List[
 @app.post("/api/run_node_stream")
 async def run_node_stream(
     request: NodeRunRequest,
+    user: str = DEFAULT_USER,
     current_user: User | None = Depends(get_current_user_optional)
 ):
     """
     Stream the response from the LLM.
     """
+    resolved = resolve_user(current_user, user)
     # Get API key from user settings if not provided in request
     provider_name = request.config.provider.value if hasattr(request.config.provider, 'value') else str(request.config.provider)
     if not request.config.api_key:
-        user_key = get_user_api_key(current_user, provider_name.lower())
+        user_key = get_user_api_key(resolved, provider_name.lower())
         if user_key:
             request.config.api_key = user_key
-    
+
     # Get username for file operations
-    username = current_user.username if current_user else DEFAULT_USER
+    username = resolved.username if resolved else DEFAULT_USER
 
     # Extract images from attached files (separate from non-image files)
     images, non_image_file_ids = extract_image_attachments(username, request.attached_file_ids)
@@ -416,7 +435,7 @@ async def run_node_stream(
             llm_config=request.config,
         )
 
-    openrouter_key = get_user_api_key(current_user, "openrouter")
+    openrouter_key = get_user_api_key(resolved, "openrouter")
 
     return StreamingResponse(
         llm_streamer(execution_context, request.user_prompt, request.config, attachments, tools, openrouter_api_key=openrouter_key, images=images),
@@ -433,13 +452,15 @@ class TitleResponse(BaseModel):
 @app.post("/api/generate_title", response_model=TitleResponse)
 async def generate_title_endpoint(
     request: TitleRequest,
+    user: str = DEFAULT_USER,
     current_user: User | None = Depends(get_current_user_optional)
 ):
     """
     Generate a short title for a Q-A pair using gpt-5-nano.
     Returns 3-4 short English words summarizing the topic.
     """
-    api_key = get_user_api_key(current_user, "openai")
+    resolved = resolve_user(current_user, user)
+    api_key = get_user_api_key(resolved, "openai")
     title = await generate_title(request.user_prompt, request.response, api_key)
     return TitleResponse(title=title)
 
@@ -454,14 +475,16 @@ class SummarizeResponse(BaseModel):
 @app.post("/api/summarize", response_model=SummarizeResponse)
 async def summarize_endpoint(
     request: SummarizeRequest,
+    user: str = DEFAULT_USER,
     current_user: User | None = Depends(get_current_user_optional)
 ):
     """
     Summarize the given content using the specified model.
     """
+    resolved = resolve_user(current_user, user)
     from app.services.llm import summarize_content
-    openai_key = get_user_api_key(current_user, "openai")
-    gemini_key = get_user_api_key(current_user, "gemini")
+    openai_key = get_user_api_key(resolved, "openai")
+    gemini_key = get_user_api_key(resolved, "gemini")
     summary = await summarize_content(request.content, request.model, openai_key, gemini_key)
     return SummarizeResponse(summary=summary)
 
-- 
cgit v1.2.3