Offloaded persistence tasks to the background

app/agents/adk_mathminds.py

@@ -1,18 +1,8 @@
-"""
-adk_mathminds.py — MathMinds ADK Agent
-Key changes vs original:
-  1. Semaphore removed. Replaced with Redis-backed daily quota via llm_guard.
-  2. Tenacity retries scoped to 429/rate-limit errors ONLY (not all exceptions),
-     so a quota block is not retried.
-  3. ADK event loop now filters for is_final_response() to avoid
-     collecting tool-call intermediate text.
-  4. Redis client injected via constructor so it can be shared with CacheManager.
-"""
-
 import logging
 import asyncio
 import base64
-from typing import Optional
+import json
+from typing import Optional, AsyncGenerator, Dict, Any
 
 from google.adk.agents import Agent
 from google.adk.runners import Runner
@@ -22,12 +12,11 @@ from google.genai.errors import ClientError
 from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type
 
 from app.core.settings import settings
-from app.core.llm_guard import check_and_increment  # ← new import
+from app.core.llm_guard import check_and_increment
 from app.tools.web_scraper import WebScraper
 from app.tools.symbolic_solver import SymbolicSolver
 from app.tools.similarity_search import SimilarProblemFinder
-from app.core.ocr import OCRProcessor
-from app.tools.vision_analyzer import VisionAnalyzer
+from app.tools.python_executor import PythonInterpreter
 from app.core.math_normalizer import MathQueryNormalizer
 
 logger = logging.getLogger(__name__)
@@ -36,12 +25,12 @@ logger = logging.getLogger(__name__)
 class MathMindsADKAgent:
     """
     Agent-based architecture using Google ADK.
-    Uses a Redis-backed daily quota (not a semaphore) to stay within 20 RPD.
+    Supports real-time streaming of reasoning steps and final answers.
     """
 
-    def __init__(self, model_name: str = "gemini-2.5-flash", redis_client=None):
+    def __init__(self, model_name: str = "gemini-2.0-flash", redis_client=None):
         self.api_key = settings.GOOGLE_API_KEY
-        self.redis_client = redis_client  # injected — shared with CacheManager
+        self.redis_client = redis_client
 
         if not self.api_key:
             logger.warning("No Google API Key found. Agent will fail.")
@@ -51,8 +40,7 @@ class MathMindsADKAgent:
         self.symbolic_solver = SymbolicSolver()
         self.normalizer = MathQueryNormalizer()
         self.similar_finder = SimilarProblemFinder()
-        self.ocr = OCRProcessor()
-        self.vision_analyzer = VisionAnalyzer()
+        self.python_executor = PythonInterpreter()
 
         # ── Tool definitions ──────────────────────────────────────────────────
         async def web_search(query: str) -> str:
@@ -79,6 +67,18 @@ class MathMindsADKAgent:
                 return result.get("content", "No solution found.")
             return f"Error solving math: {result.get('error')}"
 
+        async def execute_python(code: str) -> str:
+            """
+            Execute arbitrary Python code for simulations, complex logic, or data analysis.
+            Use this when SymPy is too restrictive or you need to run a simulation.
+            Args:
+                code: The Python code to execute.
+            """
+            result = await self.python_executor.execute(code)
+            if result.get("status") == "success":
+                return f"Output:\n{result.get('content')}\nResult: {result.get('result')}"
+            return f"Error in Python execution: {result.get('content')}"
+
         def find_similar_problems(query: str) -> str:
             """
             Find similar solved problems from the database for reference.
@@ -93,41 +93,18 @@ class MathMindsADKAgent:
                 formatted += f"Problem: {item.get('problem_text')}\nSolution: {item.get('solution_text')}\n---\n"
             return formatted
 
-        def read_image(image_data: str) -> str:
-            """
-            Extract text/equations from an image using OCR.
-            Args:
-                image_data: Base64 string of the image.
-            """
-            try:
-                text = self.ocr.extract_text(image_data=image_data)
-                return text if text else "No text found in image."
-            except Exception as e:
-                return f"Error reading image: {str(e)}"
-
-        async def analyze_image(image_data: str, focus: str = "") -> str:
-            """
-            Analyze an image mathematically: count objects, describe graphs, extract equations.
-            Args:
-                image_data: Base64 string of the image.
-                focus: Optional focus hint (e.g. 'count red balls').
-            """
-            try:
-                result = self.vision_analyzer.analyze(image_data, focus)
-                return str(result)
-            except Exception as e:
-                return f"Image analysis failed: {str(e)}"
-
         # ── Agent & Runner ─────────────────────────────────────────────────��──
         self.agent = Agent(
             name="math_minds_core",
             model=model_name,
-            tools=[web_search, math_solver, find_similar_problems, read_image, analyze_image],
+            tools=[web_search, math_solver, execute_python, find_similar_problems],
             instruction=(
                 "You are MathMinds AI, a precise mathematical assistant. "
-                "When an image is provided, analyze it first — extract equations, "
-                "count objects, or interpret graphs. Then combine image analysis with "
-                "the text prompt. Use tools only when needed. Show your reasoning clearly."
+                "You can see images natively! When an image is provided, examine it "
+                "carefully to extract equations, count objects, or interpret graphs. "
+                "\n\nCRITICAL: Always start by explaining your step-by-step approach "
+                "before using any tools. Your internal monologue should be clear "
+                "and explain the reasoning behind your tool choices."
             )
         )
 
@@ -138,7 +115,7 @@ class MathMindsADKAgent:
             session_service=self.session_service
         )
 
-        logger.info("MathMindsADKAgent initialized.")
+        logger.info(f"MathMindsADKAgent initialized with model: {model_name}")
 
     async def solve(
         self,
@@ -146,22 +123,17 @@ class MathMindsADKAgent:
         image_data: Optional[str] = None,
         session_id: str = "default_session",
         user_id: str = "default_user"
-    ) -> str:
+    ) -> AsyncGenerator[Dict[str, Any], None]:
         """
-        Main entry point. Enforces daily quota before calling the LLM.
-        Returns the agent's answer string, or an error message.
+        Streaming entry point. Yields events as they occur.
         """
 
         # ── 1. Daily quota check ──────────────────────────────────────────────
-        # This is the ONLY gate. One check per user question = one LLM call.
         if self.redis_client:
             allowed, used, limit = check_and_increment(self.redis_client, user_id)
             if not allowed:
-                logger.warning(f"Quota blocked for user={user_id} ({used}/{limit} today)")
-                return (
-                    f"⚠️ Daily limit reached ({limit} questions per day). "
-                    "Please try again tomorrow."
-                )
+                yield {"type": "error", "content": f"⚠️ Daily limit reached ({limit} today)."}
+                return
         else:
             logger.warning("Redis unavailable — skipping quota check (failing open).")
 
@@ -174,66 +146,56 @@ class MathMindsADKAgent:
                 await self.session_service.create_session(
                     app_name="mathminds", user_id=user_id, session_id=session_id
                 )
-        except Exception:
-            try:
-                await self.session_service.create_session(
-                    app_name="mathminds", user_id=user_id, session_id=session_id
-                )
-            except Exception as e:
-                logger.warning(f"Session create warning: {e}")
+        except Exception as e:
+            logger.warning(f"Session setup warning: {e}")
 
         # ── 3. Build message parts ────────────────────────────────────────────
-        parts = [types.Part.from_text(text=problem)]
+        parts = []
+        if problem:
+            parts.append(types.Part.from_text(text=problem))
+        else:
+            parts.append(types.Part.from_text(text="Analyze this image."))
 
         if image_data:
             try:
-                if image_data.startswith("/9j/"):
-                    mime_type = "image/jpeg"
-                elif image_data.startswith("iVBORw"):
-                    mime_type = "image/png"
-                elif image_data.startswith("R0lGOD"):
-                    mime_type = "image/gif"
-                elif image_data.startswith("UklGR"):
-                    mime_type = "image/webp"
-                else:
-                    mime_type = "image/png"
-
                 img_bytes = base64.b64decode(image_data)
+                mime_type = "image/png"  # Default
+                # Basic sniff
+                if image_data.startswith("/9j/"): mime_type = "image/jpeg"
+                elif image_data.startswith("iVBORw"): mime_type = "image/png"
+                
                 parts.append(types.Part.from_bytes(data=img_bytes, mime_type=mime_type))
-                logger.info("Image attached to agent request.")
             except Exception as e:
-                logger.error(f"Failed to process image: {e}")
-                parts.append(types.Part.from_text(text="[Error: image could not be processed]"))
-
-        # ── 4. Run agent (retry on 429 only, not all exceptions) ─────────────
-        @retry(
-            stop=stop_after_attempt(2),          # max 2 attempts total
-            wait=wait_exponential(multiplier=2, min=5, max=30),
-            retry=retry_if_exception_type(ClientError),  # only retry on API errors
-            reraise=True
-        )
-        async def run_agent_safely() -> str:
-            outcome = ""
+                logger.error(f"Image decode failed: {e}")
+
+        # ── 4. Run agent (Streaming) ──────────────────────────────────────────
+        try:
             async for event in self.runner.run_async(
                 user_id=user_id,
                 session_id=session_id,
                 new_message=types.Content(role="user", parts=parts)
             ):
-                # ✅ Only collect the final response, not tool-call intermediates
-                if hasattr(event, "is_final_response") and event.is_final_response():
-                    if event.content and event.content.parts:
-                        for part in event.content.parts:
-                            if part.text:
-                                outcome += part.text
-            return outcome
-
-        try:
-            response_text = await run_agent_safely()
-            if not response_text:
-                logger.warning("Agent returned empty response.")
-                return "The agent completed but returned no text. Please rephrase your question."
-            return response_text
+                # ── Capture Reasoning / Thoughts ──
+                if hasattr(event, "content") and event.content:
+                    for part in event.content.parts:
+                        if part.text:
+                            # We treat intermittent text as reasoning/logic
+                            yield {"type": "thought", "content": part.text}
+                
+                # ── Capture Tool Usage ──
+                if hasattr(event, "tool_call") and event.tool_call:
+                    yield {
+                        "type": "action", 
+                        "content": f"Using tool: {event.tool_call.function_call.name}"
+                    }
+
+                # ── Capture Tool Response ──
+                if hasattr(event, "tool_response") and event.tool_response:
+                     yield {
+                        "type": "observation", 
+                        "content": f"Obtained result from {event.tool_response.function_response.name}"
+                    }
 
         except Exception as e:
-            logger.error(f"ADK Agent execution failed: {e}")
-            return f"Error processing request: {str(e)}"
+            logger.error(f"Streaming execution failed: {e}")
+            yield {"type": "error", "content": str(e)}

app/api/main.py

@@ -1,4 +1,6 @@
-
+import os
+os.environ["DISABLE_MODEL_SOURCE_CHECK"] = "True"
+from typing import Any, Dict, Optional, List
 import sys
 import asyncio
 
@@ -6,18 +8,19 @@ if sys.platform == 'win32':
     asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy())
 
 import logging
-import datetime
-from datetime import datetime
+from datetime import datetime, timezone
 import uuid
 import sys
+import json
 
 from fastapi import FastAPI, HTTPException, status, Depends, Request
-from fastapi.responses import JSONResponse
+from fastapi.responses import JSONResponse, StreamingResponse
 from slowapi import _rate_limit_exceeded_handler
 from slowapi.errors import RateLimitExceeded
 from app.core.limiter import limiter
 from app.core.orchestrator import Orchestrator
-from app.core.schemas import SolveRequest, SolveResponse, HealthResponse
+from app.core.schemas import SolveRequest, SolveResponse, HealthResponse, Message, ChatSession, SessionRename, UserSignup, UserLogin, TokenResponse
+from app.core.auth_utils import hash_password, verify_password, create_access_token
 from app.core.logging_config import configure_logging
 from app.core.errors import AppError, ErrorCodes, ERROR_MESSAGES
 from app.core.settings import settings  # New settings module
@@ -200,7 +203,7 @@ async def health_check():
     
     return health_status
 
-@app.post("/solve", response_model=SolveResponse)
+@app.post("/solve")
 @limiter.limit("5/minute")
 async def solve_problem(
     request: Request,
@@ -243,55 +246,89 @@ async def solve_problem(
         logger.warning(f"Redis dedup failed (failing open): {e}")
         # If Redis fails, we allow the request to proceed (fail open)
 
-    try:
-        result = await orchestrator.process_problem(
-            text=solve_req.effective_text, 
-            image=solve_req.image, 
-            request_id=final_request_id,
-            model_preference=solve_req.model_preference,
-            session_id=solve_req.session_id,
-            user_id=current_user.get("uid")
-        )
-        
-        # Sanitize metadata for public response
-        public_metadata = result["metadata"].copy()
-        public_metadata.pop("_internal_debug", None)
-        
-        # Map internal result to schema
-        return SolveResponse(
-            request_id=result.get("request_id", final_request_id),
-            status=result["status"],
-            problem_type=result.get("problem_type", "unknown"),
-            source=result.get("source", "unknown"),
-            answer=result.get("answer"),
-            steps=result.get("steps", []),
-            explanation=result.get("explanation"),
-            confidence=result.get("confidence", 0.0),
-            cached=result.get("cached", False),
-            error=result.get("error_msg"), # Keep for backward compat if any
-            error_code=result.get("error_code"),
-            metadata=public_metadata
-        )
+    async def event_generator():
+        try:
+            async for chunk in orchestrator.process_problem(
+                text=solve_req.effective_text, 
+                image=solve_req.image, 
+                request_id=final_request_id,
+                model_preference=solve_req.model_preference,
+                session_id=solve_req.session_id,
+                user_id=current_user.get("uid")
+            ):
+                yield json.dumps(chunk) + "\n"
+        except Exception as e:
+            logger.error(f"Streaming error: {e}")
+            yield json.dumps({"type": "error", "content": "Internal processing error"}) + "\n"
+        finally:
+            if redis_client:
+                try:
+                    redis_client.delete(dedup_key)
+                except Exception:
+                    pass
 
-    except Exception as e:
-        logger.error(f"[{req_id}] Unhandled error in /solve: {e}")
-        # Return generic error
-        return JSONResponse(
-            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
-            content={
-                "status": "error",
-                "error": ERROR_MESSAGES[ErrorCodes.INTERNAL_ERROR],
-                "error_code": ErrorCodes.INTERNAL_ERROR,
-                "metadata": {"request_id": req_id}
-            }
-        )
-    finally:
-        # Cleanup deduplication key
-        if redis_client:
-            try:
-                redis_client.delete(dedup_key)
-            except Exception as e:
-                logger.warning(f"Failed to clear dedup key: {e}")
+    return StreamingResponse(event_generator(), media_type="application/x-ndjson")
+
+# --- Chat History Endpoints ---
+
+@app.get("/chat/sessions", response_model=List[ChatSession])
+async def list_chat_sessions(
+    current_user: dict = Depends(get_current_user),
+    db_manager = Depends(get_db_manager)
+):
+    """List all chat sessions for the current user."""
+    return db_manager.list_sessions(current_user["uid"])
+
+@app.post("/chat/sessions", response_model=ChatSession)
+async def create_chat_session(
+    current_user: dict = Depends(get_current_user),
+    db_manager = Depends(get_db_manager)
+):
+    """Create a new chat session."""
+    session_id = str(uuid.uuid4())
+    title = "New Chat"
+    if db_manager.create_session(current_user["uid"], session_id, title):
+        return {
+            "session_id": session_id,
+            "title": title,
+            "created_at": datetime.utcnow()
+        }
+    raise HTTPException(status_code=500, detail="Failed to create session")
+
+@app.get("/chat/sessions/{session_id}/messages", response_model=List[Message])
+async def get_session_history(
+    session_id: str,
+    current_user: dict = Depends(get_current_user),
+    db_manager = Depends(get_db_manager)
+):
+    """Get message history for a specific session."""
+    history = db_manager.get_chat_history(current_user["uid"], session_id)
+    if not history and history != []:
+        raise HTTPException(status_code=404, detail="Session not found")
+    return history
+
+@app.patch("/chat/sessions/{session_id}")
+async def rename_chat_session(
+    session_id: str,
+    rename_data: SessionRename,
+    current_user: dict = Depends(get_current_user),
+    db_manager = Depends(get_db_manager)
+):
+    """Rename a chat session."""
+    if db_manager.rename_session(current_user["uid"], session_id, rename_data.title):
+        return {"status": "success", "title": rename_data.title}
+    raise HTTPException(status_code=404, detail="Session not found or rename failed")
+
+@app.delete("/chat/sessions/{session_id}")
+async def delete_chat_session(
+    session_id: str,
+    current_user: dict = Depends(get_current_user),
+    db_manager = Depends(get_db_manager)
+):
+    """Delete a chat session."""
+    if db_manager.delete_session(current_user["uid"], session_id):
+        return {"status": "success", "message": "Session deleted"}
+    raise HTTPException(status_code=404, detail="Session not found or delete failed")
 
 # --- User Profile Endpoints ---
 from pydantic import BaseModel
@@ -348,3 +385,64 @@ async def update_profile(
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=8000)
+# ── Auth Endpoints ─────────────────────────────────────────────────────────
+
+@app.post("/auth/signup", response_model=TokenResponse)
+async def signup(
+    user_in: UserSignup,
+    db_manager = Depends(get_db_manager)
+):
+    """Register a new user."""
+    # Check if user already exists
+    existing_user = db_manager.get_user_by_email(user_in.email)
+    if existing_user:
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail="User with this email already exists"
+        )
+    
+    user_id = str(uuid.uuid4())
+    hashed_pw = hash_password(user_in.password)
+    
+    user_dict = {
+        "user_id": user_id,
+        "email": user_in.email,
+        "hashed_password": hashed_pw,
+        "full_name": user_in.full_name,
+        "created_at": datetime.now(timezone.utc)
+    }
+    
+    if db_manager.create_user(user_dict):
+        token = create_access_token(data={"sub": user_id, "email": user_in.email})
+        return {
+            "access_token": token,
+            "token_type": "bearer",
+            "user_id": user_id,
+            "email": user_in.email
+        }
+    
+    raise HTTPException(status_code=500, detail="Failed to create user")
+
+@app.post("/auth/login", response_model=TokenResponse)
+async def login(
+    login_in: UserLogin,
+    db_manager = Depends(get_db_manager)
+):
+    """Login and get access token."""
+    user = db_manager.get_user_by_email(login_in.email)
+    if not user or not verify_password(login_in.password, user["hashed_password"]):
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Incorrect email or password",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+    
+    user_id = user["user_id"]
+    token = create_access_token(data={"sub": user_id, "email": user["email"]})
+    
+    return {
+        "access_token": token,
+        "token_type": "bearer",
+        "user_id": user_id,
+        "email": user["email"]
+    }

app/core/orchestrator.py

@@ -3,7 +3,8 @@ import time
 import hashlib
 import json
 import re
-from typing import Any, Dict, Optional
+import asyncio
+from typing import Any, Dict, Optional, AsyncGenerator
 
 import sympy
 from sympy.parsing.sympy_parser import parse_expr, standard_transformations, implicit_multiplication_application
@@ -17,21 +18,12 @@ from app.core.settings import settings
 
 logger = logging.getLogger(__name__)
 
-# Transformations that allow "2x" → "2*x" in SymPy
 _SYMPY_TRANSFORMATIONS = standard_transformations + (implicit_multiplication_application,)
 
 
 class Orchestrator:
     """
-    Pipeline: Cache → Pre-flight (SymPy) → ADK Agent
-
-    Pre-flight layer solves arithmetic, algebra, derivatives, integrals, and
-    equations locally with SymPy — ZERO LLM calls. Only queries that cannot
-    be resolved symbolically are forwarded to the ADK Agent.
-
-    With a 20 RPD quota this means:
-      - Simple math  → instant, free, no quota used
-      - Complex math → ADK agent, 1-2 LLM calls as needed
+    Evolved Pipeline: Cache → Pre-flight (SymPy) → Agentic Streaming Loop
     """
 
     def __init__(
@@ -51,9 +43,6 @@ class Orchestrator:
             logger.critical(f"Failed to initialize Orchestrator: {e}")
             raise
 
-    # ──────────────────────────────────────────────────────────────────────────
-    # Public entry point
-    # ──────────────────────────────────────────────────────────────────────────
     async def process_problem(
         self,
         text: Optional[str] = None,
@@ -62,222 +51,173 @@ class Orchestrator:
         model_preference: str = "fast",
         session_id: Optional[str] = None,
         user_id: Optional[str] = None,
-    ) -> Dict[str, Any]:
+    ) -> AsyncGenerator[Dict[str, Any], None]:
 
         start_time = time.time()
         request_id = request_id or "unknown"
 
         result_schema: Dict[str, Any] = {
             "request_id": request_id,
-            "status":     "error",
-            "source":     "google_adk_agent",
-            "answer":     None,
-            "steps":      [],
-            "explanation": None,
-            "confidence": 0.0,
-            "cached":     False,
-            "metadata":   {"latency_ms": 0, "model": "sympy_preflight", "tools_used": []},
+            "status":     "success",
+            "source":     "agent",
+            "answer":     "",
+            "metadata":   {"latency_ms": 0, "model": "gemini-2.0-flash", "tools_used": []},
         }
 
         try:
-            # ── Step 1: Input processing ───────────────────────────────────────
+            # ── 1. Input processing ───────────────────────────────────────────
             processed = self.input_processor.process_compound(text_input=text, image_input=image)
             if not processed.is_valid:
-                result_schema["explanation"] = processed.error_message
-                return self._finalize(result_schema, start_time)
+                yield {"type": "error", "content": processed.error_message}
+                return
+
+            query = processed.cleaned_content
+            image_data = processed.metadata.get("image_data")
 
-            query        = processed.cleaned_content
-            image_data   = processed.metadata.get("image_data")
+            # Background: Persist user message
+            if user_id and session_id:
+                asyncio.create_task(self._persist_message(
+                    user_id=user_id, session_id=session_id, role="user", 
+                    content=text or "Uploaded an image", image_data=image_data
+                ))
 
-            # ── Step 2: Cache lookup ───────────────────────────────────────────
+            # ── 2. Cache lookup ───────────────────────────────────────────────
             if settings.ENABLE_CACHE and not image_data:
-                cache_key    = self._make_cache_key(query)
-                cached       = self.cache_manager.get_cached_answer(cache_key)
+                cache_key = self._make_cache_key(query)
+                cached = self.cache_manager.get_cached_answer(cache_key)
                 if cached:
-                    logger.info(f"Cache hit for query: {query[:60]}")
-                    cached["cached"] = True
-                    cached["request_id"] = request_id
-                    return self._finalize(cached, start_time)
+                    yield {"type": "thought", "content": "Retrieving answer from memory..."}
+                    yield {"type": "answer", "content": cached.get("answer")}
+                    # Background: Persist assistant response
+                    if user_id and session_id:
+                        asyncio.create_task(self._persist_message(
+                            user_id=user_id, session_id=session_id, role="assistant",
+                            content=cached.get("answer"), metadata=cached.get("metadata")
+                        ))
+                    return
             else:
                 cache_key = None
 
-            # ── Step 3: Pre-flight — try SymPy before touching the LLM ────────
-            # Only attempted when there is no image (images need vision model).
+            # ── 3. Pre-flight (SymPy) ─────────────────────────────────────────
             if not image_data:
                 preflight_result = self._try_sympy(query)
                 if preflight_result is not None:
+                    yield {"type": "thought", "content": "Calculating result symbolically..."}
+                    yield {"type": "answer", "content": preflight_result}
+                    
                     result_schema.update({
-                        "status":      "success",
-                        "source":      "sympy_preflight",
-                        "answer":      preflight_result,
-                        "explanation": "Solved locally by SymPy — no LLM call needed.",
-                        "confidence":  1.0,
-                        "metadata":    {"latency_ms": 0, "model": "sympy_preflight", "tools_used": ["sympy"]},
+                        "source": "sympy_preflight",
+                        "answer": preflight_result,
+                        "metadata": {"model": "sympy", "tools_used": ["sympy"]}
                     })
-                    logger.info(f"Pre-flight solved: {query[:60]} → {preflight_result[:80]}")
-
-                    # Cache and persist
-                    if settings.ENABLE_CACHE and cache_key:
-                        self.cache_manager.set_cached_answer(cache_key, result_schema)
-                    self.db_manager.save_problem({"content": query}, result_schema)
-                    return self._finalize(result_schema, start_time)
-
-            # ── Step 4: ADK Agent (LLM) ────────────────────────────────────────
-            logger.info("Pre-flight could not solve — routing to ADK Agent")
-            result_schema["metadata"]["model"] = "gemini-flash-adk"
-
-            try:
-                agent_response = await self.adk_agent.solve(
-                    problem=query,
-                    image_data=image_data,
-                    session_id=session_id or "default_session",
-                    user_id=user_id,
-                )
-                result_schema.update({
-                    "status":      "success",
-                    "source":      "google_adk_agent",
-                    "answer":      agent_response,
-                    "explanation": "Processed by MathMinds ADK Agent.",
-                    "confidence":  1.0,
-                })
-            except Exception as e:
-                logger.error(f"ADK Agent execution failed: {e}")
-                result_schema["explanation"] = f"Agent Error: {str(e)}"
-                return self._finalize(result_schema, start_time)
+                    
+                    self._background_log(query, result_schema, user_id, session_id, cache_key)
+                    return
+
+            # ── 4. Agentic Streaming Loop ─────────────────────────────────────
+            full_answer = ""
+            async for event in self.adk_agent.solve(
+                problem=query, image_data=image_data, 
+                session_id=session_id, user_id=user_id
+            ):
+                if event["type"] == "thought":
+                    yield event
+                elif event["type"] in ("action", "observation"):
+                    yield event
+                elif event["type"] == "error":
+                    yield event
+                else:
+                    # Treat everything else as part of the answer
+                    full_answer += event["content"]
+                    yield {"type": "answer", "content": event["content"]}
+
+            # ── 5. Finalize ───────────────────────────────────────────────────
+            result_schema["answer"] = full_answer
+            result_schema["metadata"]["latency_ms"] = int((time.time() - start_time) * 1000)
+            
+            if full_answer:
+                self._background_log(query, result_schema, user_id, session_id, cache_key)
 
-            # ── Step 5: Persist ───────────────────────────────���────────────────
-            if result_schema["status"] == "success":
-                if settings.ENABLE_CACHE and cache_key:
-                    self.cache_manager.set_cached_answer(cache_key, result_schema)
-                self.db_manager.save_problem({"content": query}, result_schema)
-
-            return self._finalize(result_schema, start_time)
+        except Exception as e:
+            logger.error(f"Orchestrator Error: {e}")
+            yield {"type": "error", "content": f"Internal Error: {str(e)}"}
 
+    async def _persist_message(self, user_id, session_id, role, content, **kwargs):
+        try:
+            self.db_manager.create_session(user_id, session_id)
+            self.db_manager.save_chat_message(user_id, session_id, role, content, **kwargs)
         except Exception as e:
-            logger.error(f"Orchestrator Critical Error: {e}")
-            result_schema["explanation"] = f"Internal Error: {str(e)}"
-            return self._finalize(result_schema, start_time)
+            logger.error(f"Failed to persist message: {e}")
+
+    def _background_log(self, query, schema, user_id, session_id, cache_key):
+        """Fire and forget persistence tasks."""
+        asyncio.create_task(self._persist_message(
+            user_id=user_id, session_id=session_id, role="assistant",
+            content=schema["answer"], metadata=schema["metadata"]
+        ))
+        if settings.ENABLE_CACHE and cache_key:
+            self.cache_manager.set_cached_answer(cache_key, schema)
+        self.db_manager.save_problem({"content": query}, schema)
 
-    # ──────────────────────────────────────────────────────────────────────────
-    # Pre-flight: pure SymPy resolution, no LLM
-    # ──────────────────────────────────────────────────────────────────────────
     def _try_sympy(self, query: str) -> Optional[str]:
-        """
-        Attempt to solve the query with the MathQueryNormalizer + SymPy.
-        Returns a human-readable answer string, or None if it can't be solved
-        locally (which means the ADK Agent should handle it).
-        """
         try:
             intent: Optional[MathIntent] = self.normalizer.normalize(query)
-            if intent is None:
-                return None
-
-            expr_str    = self._prep_expr(intent.expression)
-            target_var  = sympy.Symbol(intent.variable or "x")
-
-            if intent.intent == "arithmetic":
-                return self._solve_arithmetic(expr_str)
-
-            if intent.intent == "equation":
-                return self._solve_equation(expr_str, target_var)
-
+            if intent is None: return None
+            expr_str = self._prep_expr(intent.expression)
+            target_var = sympy.Symbol(intent.variable or "x")
+            if intent.intent == "arithmetic": return self._solve_arithmetic(expr_str)
+            if intent.intent == "equation": return self._solve_equation(expr_str, target_var)
             if intent.intent == "derivative":
-                expr   = parse_expr(expr_str, transformations=_SYMPY_TRANSFORMATIONS)
-                result = sympy.diff(expr, target_var)
-                return f"d/d{target_var}({intent.expression}) = {sympy.latex(result)}"
-
+                expr = parse_expr(expr_str, transformations=_SYMPY_TRANSFORMATIONS)
+                return f"d/d{target_var}({intent.expression}) = {sympy.latex(sympy.diff(expr, target_var))}"
             if intent.intent == "integral":
-                expr   = parse_expr(expr_str, transformations=_SYMPY_TRANSFORMATIONS)
-                result = sympy.integrate(expr, target_var)
-                return f"∫({intent.expression}) d{target_var} = {sympy.latex(result)} + C"
-
-            if intent.intent == "limit":
-                return self._solve_limit(intent, query)
-
+                expr = parse_expr(expr_str, transformations=_SYMPY_TRANSFORMATIONS)
+                return f"∫({intent.expression}) d{target_var} = {sympy.latex(sympy.integrate(expr, target_var))} + C"
+            if intent.intent == "limit": return self._solve_limit(intent, query)
             if intent.intent == "simplification":
-                expr   = parse_expr(expr_str, transformations=_SYMPY_TRANSFORMATIONS)
-                result = sympy.simplify(expr)
-                return f"Simplified: {sympy.latex(result)}"
-
-        except Exception as e:
-            # Any parse/evaluation error → fall through to agent
-            logger.debug(f"Pre-flight SymPy failed for '{query}': {e}")
-
+                expr = parse_expr(expr_str, transformations=_SYMPY_TRANSFORMATIONS)
+                return f"Simplified: {sympy.latex(sympy.simplify(expr))}"
+        except Exception: pass
         return None
 
     def _prep_expr(self, expr: str) -> str:
-        """Normalise expression string for SymPy."""
-        expr = expr.replace("^", "**")                       # ^ → **
-        expr = re.sub(r"(\d)([a-zA-Z])", r"\1*\2", expr)    # 2x → 2*x
-        expr = re.sub(r"\)\s*\(", ")*(", expr)               # )( → )*(
+        expr = expr.replace("^", "**")
+        expr = re.sub(r"(\d)([a-zA-Z])", r"\1*\2", expr)
+        expr = re.sub(r"\)\s*\(", ")*(", expr)
         return expr.strip()
 
     def _solve_arithmetic(self, expr_str: str) -> Optional[str]:
-        """Evaluate a pure arithmetic/algebraic expression."""
         try:
-            expr   = parse_expr(expr_str, transformations=_SYMPY_TRANSFORMATIONS)
-            result = sympy.simplify(expr)
-            # If the result is a number, show it plainly; otherwise use LaTeX
+            result = sympy.simplify(parse_expr(expr_str, transformations=_SYMPY_TRANSFORMATIONS))
             if result.is_number:
                 numeric = float(result)
-                # Show as integer if it is one
-                if numeric == int(numeric):
-                    return str(int(numeric))
-                return f"{numeric:.6g}"
+                return str(int(numeric)) if numeric == int(numeric) else f"{numeric:.6g}"
             return sympy.latex(result)
-        except Exception:
-            return None
+        except Exception: return None
 
     def _solve_equation(self, expr_str: str, var: sympy.Symbol) -> Optional[str]:
-        """Solve an equation of the form lhs = rhs, or expr = 0."""
         try:
             parts = expr_str.split("=", 1)
             if len(parts) == 2:
-                lhs      = parse_expr(self._prep_expr(parts[0]), transformations=_SYMPY_TRANSFORMATIONS)
-                rhs      = parse_expr(self._prep_expr(parts[1]), transformations=_SYMPY_TRANSFORMATIONS)
+                lhs = parse_expr(self._prep_expr(parts[0]), transformations=_SYMPY_TRANSFORMATIONS)
+                rhs = parse_expr(self._prep_expr(parts[1]), transformations=_SYMPY_TRANSFORMATIONS)
                 solution = sympy.solve(lhs - rhs, var)
             else:
-                expr     = parse_expr(expr_str, transformations=_SYMPY_TRANSFORMATIONS)
-                solution = sympy.solve(expr, var)
-
-            if not solution:
-                return "No solution found."
-            if len(solution) == 1:
-                return f"{var} = {sympy.latex(solution[0])}"
-            sols = ", ".join(sympy.latex(s) for s in solution)
-            return f"{var} ∈ {{{sols}}}"
-        except Exception:
-            return None
+                solution = sympy.solve(parse_expr(expr_str, transformations=_SYMPY_TRANSFORMATIONS), var)
+            if not solution: return "No solution found."
+            if len(solution) == 1: return f"{var} = {sympy.latex(solution[0])}"
+            return f"{var} ∈ {{{', '.join(sympy.latex(s) for s in solution)}}}"
+        except Exception: return None
 
     def _solve_limit(self, intent: MathIntent, original_query: str) -> Optional[str]:
-        """Parse and evaluate a limit expression."""
         try:
-            # Limit pattern: "limit of X as Y approaches Z"
-            match = re.search(
-                r"limit of\s+(.+?)\s+as\s+(\w+)\s+approaches\s+(.+)",
-                original_query, re.IGNORECASE
-            )
-            if not match:
-                return None
-            expr_raw  = self._prep_expr(match.group(1))
-            var_name  = match.group(2).strip()
-            point_raw = self._prep_expr(match.group(3).strip())
-
-            expr  = parse_expr(expr_raw, transformations=_SYMPY_TRANSFORMATIONS)
-            var   = sympy.Symbol(var_name)
-            point = parse_expr(point_raw, transformations=_SYMPY_TRANSFORMATIONS)
+            match = re.search(r"limit of\s+(.+?)\s+as\s+(\w+)\s+approaches\s+(.+)", original_query, re.IGNORECASE)
+            if not match: return None
+            expr = parse_expr(self._prep_expr(match.group(1)), transformations=_SYMPY_TRANSFORMATIONS)
+            var = sympy.Symbol(match.group(2).strip())
+            point = parse_expr(self._prep_expr(match.group(3).strip()), transformations=_SYMPY_TRANSFORMATIONS)
+            return f"lim({var}→{point}) {sympy.latex(expr)} = {sympy.latex(sympy.limit(expr, var, point))}"
+        except Exception: return None
 
-            result = sympy.limit(expr, var, point)
-            return f"lim({var}→{point}) {sympy.latex(expr)} = {sympy.latex(result)}"
-        except Exception:
-            return None
-
-    # ──────────────────────────────────────────────────────────────────────────
-    # Utilities
-    # ──────────────────────────────────────────────────────────────────────────
     def _make_cache_key(self, query: str) -> str:
         return hashlib.sha256(query.strip().lower().encode()).hexdigest()
-
-    def _finalize(self, schema: Dict, start_time: float) -> Dict:
-        schema["metadata"]["latency_ms"] = int((time.time() - start_time) * 1000)
-        return schema

app/core/schemas.py

@@ -1,3 +1,4 @@
+from datetime import datetime
 from typing import Any, Dict, Optional, List
 from pydantic import BaseModel, Field, model_validator
 
@@ -58,3 +59,39 @@ class HealthResponse(BaseModel):
     """
     status: str
     version: str
+
+# --- Chat History Schemas ---
+
+class Message(BaseModel):
+    role: str
+    content: str
+    timestamp: datetime
+    reasoning: Optional[str] = None
+    metadata: Dict[str, Any] = {}
+    steps: List[str] = []
+
+class ChatSession(BaseModel):
+    session_id: str
+    title: str
+    created_at: datetime
+    # messages: Optional[List[Message]] = None # Optional for listing
+
+class SessionRename(BaseModel):
+    title: str = Field(..., min_length=1, max_length=100)
+
+# --- Auth Schemas ---
+
+class UserSignup(BaseModel):
+    email: str
+    password: str = Field(..., min_length=8, max_length=72)
+    full_name: Optional[str] = None
+
+class UserLogin(BaseModel):
+    email: str
+    password: str = Field(..., max_length=72)
+
+class TokenResponse(BaseModel):
+    access_token: str
+    token_type: str = "bearer"
+    user_id: str
+    email: str

app/core/security.py

@@ -1,31 +1,16 @@
 import logging
-import firebase_admin
-from firebase_admin import credentials, auth
 from fastapi import HTTPException, status, Security
 from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from app.core.settings import settings
+from app.core.auth_utils import decode_access_token
 
 logger = logging.getLogger(__name__)
 
-# Initialize Firebase Admin SDK
-try:
-    if settings.FIREBASE_CREDENTIALS_PATH:
-        if not firebase_admin._apps:
-            cred = credentials.Certificate(settings.FIREBASE_CREDENTIALS_PATH)
-            firebase_admin.initialize_app(cred)
-            logger.info("Firebase Admin SDK initialized successfully.")
-        else:
-            logger.info("Firebase Admin SDK already initialized.")
-    else:
-        logger.warning("FIREBASE_CREDENTIALS_PATH not set. Auth will fail if enabled.")
-except Exception as e:
-    logger.error(f"Failed to initialize Firebase: {e}")
-
 security = HTTPBearer()
 
 def verify_token(credentials: HTTPAuthorizationCredentials = Security(security)):
     """
-    Verifies the Firebase ID token.
+    Verifies the Local JWT access token.
     Returns the decoded token dict if valid.
     """
     token = credentials.credentials
@@ -39,16 +24,21 @@ def verify_token(credentials: HTTPAuthorizationCredentials = Security(security))
         logger.info(f"Using MOCK AUTH for token: {token}")
         return {"uid": "dev_user_123", "email": "dev@mathminds.ai"}
 
-    try:
-        decoded_token = auth.verify_id_token(token)
-        return decoded_token
-    except Exception as e:
-        logger.warning(f"Auth failed: {e}")
-        raise HTTPException(
-            status_code=status.HTTP_401_UNAUTHORIZED,
-            detail="Invalid authentication credentials",
-            headers={"WWW-Authenticate": "Bearer"},
-        )
+    # Use local JWT verification
+    payload = decode_access_token(token)
+    if payload:
+        # Map 'sub' from JWT to 'uid' to maintain compatibility with existing code
+        return {
+            "uid": payload.get("sub"),
+            "email": payload.get("email")
+        }
+    
+    logger.warning(f"Invalid or expired token provided.")
+    raise HTTPException(
+        status_code=status.HTTP_401_UNAUTHORIZED,
+        detail="Invalid or expired authentication credentials",
+        headers={"WWW-Authenticate": "Bearer"},
+    )
 
 def get_current_user(token: dict = Security(verify_token)):
     """

app/core/settings.py

@@ -40,6 +40,11 @@ class Settings(BaseSettings):
     SUPABASE_KEY: Optional[str] = None
     WOLFRAM_APP_ID: Optional[str] = None
 
+    # Security
+    JWT_SECRET_KEY: str = "super_secret_key_change_me"
+    JWT_ALGORITHM: str = "HS256"
+    ACCESS_TOKEN_EXPIRE_MINUTES: int = 60 * 24 * 7 # 1 Week
+
     model_config = {
         "env_file": ".env",
         "case_sensitive": True,

app/memory/database.py

@@ -55,10 +55,15 @@ class DatabaseManager:
 
             self.db = self.client[db_name]
             self.collection = self.db["solved_problems"]
+            self.sessions_collection = self.db["chat_sessions"]
+            self.users_collection = self.db["users"]
             
-            # Ensure index
-            index = IndexModel([("hash", ASCENDING)], name="hash_index")
-            self.collection.create_indexes([index])
+            # Ensure indexes
+            self.collection.create_index([("hash", ASCENDING)], name="hash_index")
+            self.sessions_collection.create_index([("user_id", ASCENDING)], name="user_id_index")
+            self.sessions_collection.create_index([("session_id", ASCENDING)], name="session_id_index", unique=True)
+            self.users_collection.create_index([("email", ASCENDING)], name="email_index", unique=True)
+            self.users_collection.create_index([("user_id", ASCENDING)], name="user_id_index", unique=True)
             
             logger.info(f"Successfully connected to MongoDB at {self.mongo_uri} (DB: {db_name})")
             
@@ -134,18 +139,20 @@ class DatabaseManager:
             logger.error(f"Failed to save problem: {e}")
             return False
 
-    def create_session(self, session_id: str, title: str = "New Chat") -> bool:
+    def create_session(self, user_id: str, session_id: str, title: str = "New Chat") -> bool:
         """
-        Initialize a new chat session.
+        Create a new chat session for a user.
+        Uses upsert with $setOnInsert to be idempotent.
         """
-        if self.db is None:
+        if self.sessions_collection is None:
             return False
         try:
-            self.db["chat_sessions"].update_one(
+            self.sessions_collection.update_one(
                 {"session_id": session_id},
                 {
                     "$setOnInsert": {
                         "session_id": session_id,
+                        "user_id": user_id,
                         "title": title,
                         "created_at": datetime.now(timezone.utc),
                         "messages": []
@@ -154,78 +161,151 @@ class DatabaseManager:
                 upsert=True
             )
             return True
-        except PyMongoError as e:
-            logger.error(f"Failed to create session {session_id}: {e}")
+        except Exception as e:
+            # If it's a duplicate key error, it means another thread just inserted it.
+            # That's fine, we consider the session "created" or at least existing.
+            if "E11000" in str(e) or "duplicate key" in str(e).lower():
+                return True
+            logger.error(f"Failed to create session {session_id} for user {user_id}: {e}")
             return False
 
-    def get_chat_history(self, session_id: str, limit: int = 10) -> List[Dict[str, Any]]:
+    def list_sessions(self, user_id: str) -> List[Dict[str, Any]]:
         """
-        Retrieve recent messages for a session.
+        Retrieve all sessions for a specific user.
         """
-        if self.db is None:
+        if self.sessions_collection is None:
             return []
         try:
-            # Get the session document with sliced messages
-            doc = self.db["chat_sessions"].find_one(
-                {"session_id": session_id},
-                {"messages": {"$slice": -limit}}
-            )
-            if doc and "messages" in doc:
-                return doc["messages"]
-            return []
+            cursor = self.sessions_collection.find(
+                {"user_id": user_id},
+                {"session_id": 1, "title": 1, "created_at": 1, "_id": 0}
+            ).sort("created_at", -1)
+            return list(cursor)
         except PyMongoError as e:
-            logger.error(f"Failed to get history for {session_id}: {e}")
+            logger.error(f"Failed to list sessions for user {user_id}: {e}")
             return []
 
-    def save_chat_message(self, session_id: str, role: str, content: str) -> bool:
+    def get_chat_history(self, user_id: str, session_id: str, limit: int = 50) -> Optional[List[Dict[str, Any]]]:
+        """
+        Retrieve recent messages for a session, ensuring it belongs to the user.
+        Returns None if session not found or not owned by user.
+        """
+        if self.sessions_collection is None:
+            return None
+        try:
+            doc = self.sessions_collection.find_one(
+                {"session_id": session_id, "user_id": user_id},
+                {"messages": {"$slice": -limit}, "_id": 0}
+            )
+            if doc is not None:
+                return doc.get("messages", [])
+            return None
+        except PyMongoError as e:
+            logger.error(f"Failed to get history for {session_id} (user: {user_id}): {e}")
+            return None
+
+    def save_chat_message(self, user_id: str, session_id: str, role: str, content: str, **kwargs) -> bool:
         """
         Append a message to the session history.
-        Also updates the session title if it's the first user message.
+        Only succeeds if the session belongs to the user_id.
         """
-        if self.db is None:
+        if self.sessions_collection is None:
             return False
         try:
             # logic to update title if it's currently "New Chat" and this is a user message
             if role == "user":
-                session = self.db["chat_sessions"].find_one({"session_id": session_id})
-                if session and session.get("title") == "New Chat":
-                    # Generate title from content (truncate)
+                session = self.sessions_collection.find_one({"session_id": session_id, "user_id": user_id})
+                if session and (session.get("title") == "New Chat" or session.get("title") == "New Session" or session.get("title") == "Untitled"):
                     new_title = content[:50] + "..." if len(content) > 50 else content
-                    self.db["chat_sessions"].update_one(
-                        {"session_id": session_id},
+                    self.sessions_collection.update_one(
+                        {"session_id": session_id, "user_id": user_id},
                         {"$set": {"title": new_title}}
                     )
 
             # Push the new message
-            self.db["chat_sessions"].update_one(
-                {"session_id": session_id},
-                {
-                    "$push": {
-                        "messages": {
-                            "role": role, 
-                            "content": content, 
-                            "timestamp": datetime.now(timezone.utc)
-                        }
-                    }
-                },
-                upsert=True
+            msg = {
+                "role": role, 
+                "content": content, 
+                "timestamp": datetime.now(timezone.utc)
+            }
+            msg.update(kwargs)
+            
+            result = self.sessions_collection.update_one(
+                {"session_id": session_id, "user_id": user_id},
+                {"$push": {"messages": msg}}
             )
-            return True
+            # return True if we found the document to update
+            return result.matched_count > 0
+        except PyMongoError as e:
+            logger.error(f"Failed to save message to {session_id} for user {user_id}: {e}")
+            return False
+
+    def delete_session(self, user_id: str, session_id: str) -> bool:
+        """
+        Delete a session belonging to a user.
+        """
+        if self.sessions_collection is None:
+            return False
+        try:
+            result = self.sessions_collection.delete_one({"session_id": session_id, "user_id": user_id})
+            return result.deleted_count > 0
         except PyMongoError as e:
-            logger.error(f"Failed to save message to {session_id}: {e}")
+            logger.error(f"Failed to delete session {session_id} for user {user_id}: {e}")
+            return False
+
+    def rename_session(self, user_id: str, session_id: str, new_title: str) -> bool:
+        """
+        Rename a session belonging to a user.
+        """
+        if self.sessions_collection is None:
+            return False
+        try:
+            result = self.sessions_collection.update_one(
+                {"session_id": session_id, "user_id": user_id},
+                {"$set": {"title": new_title}}
+            )
+            # return True if we found the document to update (even if title was same)
+            return result.matched_count > 0
+        except PyMongoError as e:
+            logger.error(f"Failed to rename session {session_id} for user {user_id}: {e}")
             return False
 
     # -------------------------------------------------------------------------
     # User Profile Management
     # -------------------------------------------------------------------------
+    def create_user(self, user_data: Dict[str, Any]) -> bool:
+        """
+        Create a new user in the database.
+        """
+        if self.users_collection is None:
+            return False
+        try:
+            self.users_collection.insert_one(user_data)
+            return True
+        except PyMongoError as e:
+            logger.error(f"Failed to create user: {e}")
+            return False
+
+    def get_user_by_email(self, email: str) -> Optional[Dict[str, Any]]:
+        """
+        Retrieve a user by email.
+        """
+        if self.users_collection is None:
+            return None
+        try:
+            return self.users_collection.find_one({"email": email})
+        except PyMongoError as e:
+            logger.error(f"Failed to get user by email {email}: {e}")
+            return None
+
     def get_user_profile(self, user_id: str) -> Optional[Dict[str, Any]]:
         """
-        Retrieve user profile by ID (Firebase UID).
+        Retrieve user profile by ID (Firebase UID or local UID).
         """
-        if self.db is None:
+        if self.users_collection is None:
             return None
         try:
-            return self.db["users"].find_one({"user_id": user_id})
+            return self.users_collection.find_one({"user_id": user_id})
         except PyMongoError as e:
             logger.error(f"Failed to get profile for {user_id}: {e}")
             return None
@@ -234,10 +314,10 @@ class DatabaseManager:
         """
         Update or create user profile.
         """
-        if self.db is None:
+        if self.users_collection is None:
             return False
         try:
-            self.db["users"].update_one(
+            self.users_collection.update_one(
                 {"user_id": user_id},
                 {"$set": {**data, "updated_at": datetime.now(timezone.utc)}},
                 upsert=True

frontend/app.py

@@ -1,6 +1,5 @@
 import streamlit as st
 import requests
-import json
 import base64
 from PIL import Image
 import io
@@ -13,21 +12,35 @@ from dotenv import load_dotenv
 load_dotenv()
 
 # ── Session state: ALL keys initialized ONCE at the very top ─────────────────
-# BUG 1 WAS HERE: The original code had TWO `if "user" not in st.session_state`
-# blocks — one here and one again at line ~4797 after the CSS/config blocks.
-# Streamlit re-runs the whole script top-to-bottom on every rerun. On the rerun
-# triggered after login, the second block executed and found "user" ALREADY in
-# session_state (because we just set it during login), so it was a no-op —
-# BUT on a hard browser refresh the two blocks ran in the same execution pass
-# and the second one re-initialized user=None, wiping the login state.
-# FIX: One single initialization block here, never again below.
+# CRITICAL: These must be the very first st.session_state accesses, before any
+# st.* UI calls. Streamlit re-runs the entire script on every interaction.
 if "is_processing" not in st.session_state:
     st.session_state.is_processing = False
 if "user" not in st.session_state:
-    st.session_state.user = None
+    st.session_state.user = None          # None = logged out
 if "current_view" not in st.session_state:
     st.session_state.current_view = "Chat"
 
+# MULTIUSER FIX ─ these three keys must be RESET on logout.
+# They are initialized here so first-run doesn't KeyError.
+if "chat_sessions" not in st.session_state:
+    st.session_state.chat_sessions = []
+if "active_session_id" not in st.session_state:
+    st.session_state.active_session_id = None
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+# MULTIUSER FIX ─ track WHICH user's data is currently loaded.
+# If this doesn't match st.session_state.user["uid"], we know we need to reload.
+if "loaded_for_user" not in st.session_state:
+    st.session_state.loaded_for_user = None
+
+if "renaming_session_id" not in st.session_state:
+    st.session_state.renaming_session_id = None
+
+if "canvas_key" not in st.session_state:
+    st.session_state.canvas_key = "main_canvas"
+
 # ====================================================
 # Page Config — must come before any st.* calls
 # ====================================================
@@ -91,59 +104,177 @@ st.markdown("""
 """, unsafe_allow_html=True)
 
 # ====================================================
-# Config & State
+# Config
 # ====================================================
-API_URL = "http://localhost:8000/solve"
-HISTORY_FILE = "chat_history.json"
+BASE_API_URL = os.getenv("API_BASE_URL", "http://localhost:8000")
+API_URL = f"{BASE_API_URL}/solve"
 
-if "chat_sessions" not in st.session_state:
-    if os.path.exists(HISTORY_FILE):
-        with open(HISTORY_FILE, "r") as f:
-            st.session_state.chat_sessions = json.load(f)
-    else:
-        st.session_state.chat_sessions = {}
 
-if "active_session_id" not in st.session_state:
-    sid = str(uuid.uuid4())
-    st.session_state.chat_sessions[sid] = {"title": "New Session", "messages": [], "created_at": time.time()}
-    st.session_state.active_session_id = sid
+# ====================================================
+# MULTIUSER ISOLATION — Core helper
+# ====================================================
+def _clear_user_state():
+    """
+    Wipe ALL per-user data from Streamlit session state.
+
+    Called on logout and whenever a different user logs in.
+
+    WHY THIS IS THE MOST IMPORTANT FUNCTION FOR MULTIUSER ISOLATION:
+    Streamlit's st.session_state is per browser-tab, not per user. If User A
+    logs in, chats, then User B logs in on the same tab, all of User A's
+    chat_sessions and messages are still sitting in st.session_state. The
+    backend correctly refuses to serve User A's data to User B (every DB query
+    filters by user_id), but the frontend would still DISPLAY User A's messages
+    briefly until the next API call returns. This function prevents that.
+    """
+    st.session_state.chat_sessions       = []
+    st.session_state.active_session_id   = None
+    st.session_state.messages            = []
+    st.session_state.loaded_for_user     = None
+    st.session_state.is_processing       = False
+    st.session_state.current_view        = "Chat"
+    st.session_state.renaming_session_id = None
+    st.session_state.canvas_key          = f"canvas_{uuid.uuid4()}"
+    # Also clear profile cache if it exists
+    if "profile_data" in st.session_state:
+        del st.session_state["profile_data"]
 
-# ── IMPORTANT: No second "user" init block here. See top of file. ─────────────
 
 # ====================================================
 # Helper Functions
 # ====================================================
-def save_history():
-    with open(HISTORY_FILE, "w") as f:
-        json.dump(st.session_state.chat_sessions, f, indent=2)
+def get_auth_headers():
+    if st.session_state.user and "token" in st.session_state.user:
+        return {"Authorization": f"Bearer {st.session_state.user['token']}"}
+    return {}
+
+
+def load_sessions():
+    """Fetch THIS user's chat sessions from the backend and populate state."""
+    try:
+        headers = get_auth_headers()
+        response = requests.get(f"{BASE_API_URL}/chat/sessions", headers=headers, timeout=30)
+        if response.status_code == 200:
+            st.session_state.chat_sessions = response.json()
+            # Mark that we've successfully loaded data for this specific user
+            if st.session_state.user:
+                st.session_state.loaded_for_user = st.session_state.user["uid"]
+            # Auto-select first session if none active
+            if not st.session_state.active_session_id and st.session_state.chat_sessions:
+                st.session_state.active_session_id = st.session_state.chat_sessions[0]["session_id"]
+                load_messages(st.session_state.active_session_id)
+            elif st.session_state.active_session_id and not any(
+                s["session_id"] == st.session_state.active_session_id
+                for s in st.session_state.chat_sessions
+            ):
+                # Active session was deleted — pick first or clear
+                if st.session_state.chat_sessions:
+                    st.session_state.active_session_id = st.session_state.chat_sessions[0]["session_id"]
+                    load_messages(st.session_state.active_session_id)
+                else:
+                    st.session_state.active_session_id = None
+                    st.session_state.messages = []
+        elif response.status_code == 401:
+            # JWT expired — force re-login
+            _clear_user_state()
+            st.session_state.user = None
+            st.error("Session expired. Please log in again.")
+        else:
+            st.error(f"Failed to load sessions: {response.status_code}")
+            st.session_state.chat_sessions = []
+    except Exception as e:
+        st.error(f"Error loading sessions: {e}")
+        st.session_state.chat_sessions = []
+
+
+def load_messages(session_id):
+    """
+    Load messages for a session.
+    The backend enforces user ownership — it will 404 if session_id
+    doesn't belong to the authenticated user, so this is safe.
+    """
+    try:
+        headers = get_auth_headers()
+        response = requests.get(
+            f"{BASE_API_URL}/chat/sessions/{session_id}/messages",
+            headers=headers, timeout=30
+        )
+        if response.status_code == 200:
+            st.session_state.messages = response.json()
+        elif response.status_code == 404:
+            # Session doesn't belong to this user — clear silently
+            st.session_state.messages = []
+            st.session_state.active_session_id = None
+            st.warning("Session not found.")
+        else:
+            st.session_state.messages = []
+            st.error(f"Failed to load messages: {response.status_code}")
+    except Exception as e:
+        st.error(f"Error loading messages: {e}")
+        st.session_state.messages = []
+
 
 def get_active_session():
-    return st.session_state.chat_sessions[st.session_state.active_session_id]
+    for s in st.session_state.chat_sessions:
+        if s["session_id"] == st.session_state.active_session_id:
+            return s
+    return None
+
 
 def add_message(role, content, sent_to_api=False, **kwargs):
-    session = get_active_session()
+    """Optimistic UI update only — persistence happens in the backend via /solve."""
     msg = {"role": role, "content": content, "timestamp": time.time(), "sent_to_api": sent_to_api}
     msg.update(kwargs)
-    session["messages"].append(msg)
-    save_history()
+    st.session_state.messages.append(msg)
+
 
 def new_chat():
-    sid = str(uuid.uuid4())
-    st.session_state.chat_sessions[sid] = {"title": "New Session", "messages": [], "created_at": time.time()}
-    st.session_state.active_session_id = sid
-    save_history()
-    st.rerun()
+    try:
+        headers = get_auth_headers()
+        response = requests.post(f"{BASE_API_URL}/chat/sessions", headers=headers, timeout=30)
+        if response.status_code == 200:
+            new_s = response.json()
+            st.session_state.active_session_id = new_s["session_id"]
+            st.session_state.messages = []
+            load_sessions()
+            st.rerun()
+        else:
+            st.error("Failed to create new chat")
+    except Exception as e:
+        st.error(f"Error: {e}")
+
 
 def delete_chat(sid):
-    if sid in st.session_state.chat_sessions:
-        del st.session_state.chat_sessions[sid]
-        if st.session_state.active_session_id == sid:
-            if st.session_state.chat_sessions:
-                st.session_state.active_session_id = list(st.session_state.chat_sessions.keys())[0]
-            else:
-                new_chat()
-        save_history()
-        st.rerun()
+    try:
+        headers = get_auth_headers()
+        response = requests.delete(f"{BASE_API_URL}/chat/sessions/{sid}", headers=headers, timeout=30)
+        if response.status_code == 200:
+            if st.session_state.active_session_id == sid:
+                st.session_state.active_session_id = None
+                st.session_state.messages = []
+            load_sessions()
+            st.rerun()
+        else:
+            st.error("Failed to delete chat")
+    except Exception as e:
+        st.error(f"Error: {e}")
+
+
+def rename_chat(sid, new_title):
+    try:
+        headers = get_auth_headers()
+        response = requests.patch(
+            f"{BASE_API_URL}/chat/sessions/{sid}",
+            headers=headers, json={"title": new_title}, timeout=30
+        )
+        if response.status_code == 200:
+            load_sessions()
+            st.rerun()
+        else:
+            st.error("Failed to rename chat")
+    except Exception as e:
+        st.error(f"Error: {e}")
+
 
 # ====================================================
 # Login Screen
@@ -168,23 +299,29 @@ def login_screen():
                 password = st.text_input("Password", type="password")
                 if st.form_submit_button("Sign In", use_container_width=True):
                     if email and password:
-                        api_key = os.getenv("FIREBASE_WEB_API_KEY")
-                        if not api_key:
-                            st.error("Missing FIREBASE_WEB_API_KEY in .env")
-                        else:
-                            try:
-                                url = f"https://identitytoolkit.googleapis.com/v1/accounts:signInWithPassword?key={api_key}"
-                                r = requests.post(url, json={"email": email, "password": password, "returnSecureToken": True}, timeout=30)
-                                if r.status_code == 200:
-                                    d = r.json()
-                                    st.session_state.user = {"email": d["email"], "token": d["idToken"], "uid": d["localId"]}
-                                    st.success(f"Welcome back, {d['email']}!")
-                                    time.sleep(0.5)
-                                    st.rerun()
-                                else:
-                                    st.error(f"Login Failed: {r.json().get('error',{}).get('message','Unknown error')}")
-                            except Exception as e:
-                                st.error(f"Connection Error: {e}")
+                        try:
+                            r = requests.post(
+                                f"{BASE_API_URL}/auth/login",
+                                json={"email": email, "password": password},
+                                timeout=30
+                            )
+                            if r.status_code == 200:
+                                d = r.json()
+                                # ✅ MULTIUSER FIX: Clear ALL previous user data
+                                # BEFORE setting the new user identity.
+                                _clear_user_state()
+                                st.session_state.user = {
+                                    "email": d["email"],
+                                    "token": d["access_token"],
+                                    "uid":   d["user_id"]
+                                }
+                                st.success(f"Welcome back, {d['email']}!")
+                                time.sleep(0.5)
+                                st.rerun()
+                            else:
+                                st.error(f"Login Failed: {r.json().get('detail', 'Unknown error')}")
+                        except Exception as e:
+                            st.error(f"Connection Error: {e}")
                     else:
                         st.error("Please enter email and password.")
 
@@ -193,49 +330,78 @@ def login_screen():
                 new_email        = st.text_input("New Email", placeholder="new@student.edu")
                 new_password     = st.text_input("New Password", type="password")
                 confirm_password = st.text_input("Confirm Password", type="password")
+                full_name        = st.text_input("Full Name", placeholder="Optional")
                 if st.form_submit_button("Create Account", use_container_width=True):
                     if new_email and new_password:
                         if new_password != confirm_password:
                             st.error("Passwords do not match!")
                         else:
-                            api_key = os.getenv("FIREBASE_WEB_API_KEY")
-                            if not api_key:
-                                st.error("Missing FIREBASE_WEB_API_KEY in .env")
-                            else:
-                                try:
-                                    url = f"https://identitytoolkit.googleapis.com/v1/accounts:signUp?key={api_key}"
-                                    r = requests.post(url, json={"email": new_email, "password": new_password, "returnSecureToken": True}, timeout=30)
-                                    if r.status_code == 200:
-                                        d = r.json()
-                                        st.session_state.user = {"email": d["email"], "token": d["idToken"], "uid": d["localId"]}
-                                        st.success(f"Account Created! Welcome, {d['email']}!")
-                                        time.sleep(0.5)
-                                        st.rerun()
-                                    else:
-                                        st.error(f"Sign Up Failed: {r.json().get('error',{}).get('message','Unknown error')}")
-                                except Exception as e:
-                                    st.error(f"Connection Error: {e}")
+                            try:
+                                r = requests.post(
+                                    f"{BASE_API_URL}/auth/signup",
+                                    json={
+                                        "email":     new_email,
+                                        "password":  new_password,
+                                        "full_name": full_name
+                                    },
+                                    timeout=30
+                                )
+                                if r.status_code == 200:
+                                    d = r.json()
+                                    # ✅ MULTIUSER FIX: Same as login — clear first
+                                    _clear_user_state()
+                                    st.session_state.user = {
+                                        "email": d["email"],
+                                        "token": d["access_token"],
+                                        "uid":   d["user_id"]
+                                    }
+                                    st.success(f"Account Created! Welcome, {d['email']}!")
+                                    time.sleep(0.5)
+                                    st.rerun()
+                                else:
+                                    st.error(f"Sign Up Failed: {r.json().get('detail', 'Unknown error')}")
+                            except Exception as e:
+                                st.error(f"Connection Error: {e}")
                     else:
                         st.error("Please fill all fields.")
 
-        st.markdown("<p style='text-align:center;font-size:0.8rem;color:#6b7280;'>Powered by Gemini & SymPy</p>", unsafe_allow_html=True)
+        st.markdown(
+            "<p style='text-align:center;font-size:0.8rem;color:#6b7280;'>Powered by Gemini & SymPy</p>",
+            unsafe_allow_html=True
+        )
+
 
 # ── Auth gate ─────────────────────────────────────────────────────────────────
 if not st.session_state.user:
     login_screen()
     st.stop()
 
+# ====================================================
+# ✅ MULTIUSER FIX — Per-rerun data isolation check
+# ====================================================
+# At this point we know a user IS logged in.
+# Check: is the data currently in state actually for THIS user?
+# This handles the scenario where User A's browser tab is reused by User B
+# (e.g. token swap, shared kiosk, etc.) without a full page reload.
+_current_uid = st.session_state.user["uid"]
+if st.session_state.loaded_for_user != _current_uid:
+    # Data in state belongs to a different user (or nobody) — reload for current user
+    _clear_user_state()
+    load_sessions()
+    # loaded_for_user is set inside load_sessions() on success
+
+
 # ====================================================
 # Profile Interface
 # ====================================================
 def profile_interface():
     st.title("👤 User Profile")
     st.markdown("Customize your MathMinds experience.")
-    headers = {"Authorization": f"Bearer {st.session_state.user['token']}"}
+    headers = get_auth_headers()
 
     if "profile_data" not in st.session_state:
         try:
-            r = requests.get(f"{API_URL.replace('/solve','')}/users/profile", headers=headers, timeout=30)
+            r = requests.get(f"{BASE_API_URL}/users/profile", headers=headers, timeout=30)
             st.session_state.profile_data = r.json() if r.status_code == 200 else {}
         except Exception:
             st.session_state.profile_data = {}
@@ -246,43 +412,57 @@ def profile_interface():
 
     with st.form("profile_form"):
         display_name = st.text_input("Display Name", value=data.get("display_name", ""))
-        math_level   = st.selectbox("Math Proficiency Level", levels,
-                                    index=levels.index(data.get("math_level","Undergraduate"))
-                                    if data.get("math_level") in levels else 1)
-        interests    = st.multiselect("Areas of Interest", interests_all,
-                                      default=[i for i in data.get("interests",[]) if i in interests_all])
+        math_level   = st.selectbox(
+            "Math Proficiency Level", levels,
+            index=levels.index(data.get("math_level", "Undergraduate"))
+            if data.get("math_level") in levels else 1
+        )
+        interests = st.multiselect(
+            "Areas of Interest", interests_all,
+            default=[i for i in data.get("interests", []) if i in interests_all]
+        )
         if st.form_submit_button("Save Profile", use_container_width=True, type="primary"):
             payload = {"display_name": display_name, "math_level": math_level, "interests": interests}
             try:
-                r = requests.post(f"{API_URL.replace('/solve','')}/users/profile", json=payload, headers=headers)
+                r = requests.post(f"{BASE_API_URL}/users/profile", json=payload, headers=headers)
                 if r.status_code == 200:
                     st.success("Profile updated!")
                     st.session_state.profile_data = payload
-                    time.sleep(1); st.rerun()
+                    time.sleep(1)
+                    st.rerun()
                 else:
                     st.error(f"Update failed: {r.text}")
             except Exception as e:
                 st.error(f"Error saving: {e}")
 
+
 # ====================================================
 # Chat Interface
 # ====================================================
 def chat_interface():
-    if st.session_state.active_session_id not in st.session_state.chat_sessions:
-        new_chat()
+    if not st.session_state.active_session_id:
+        if st.session_state.chat_sessions:
+            st.session_state.active_session_id = st.session_state.chat_sessions[0]["session_id"]
+            load_messages(st.session_state.active_session_id)
+        else:
+            new_chat()
+            return
 
-    st.title(st.session_state.chat_sessions[st.session_state.active_session_id]["title"])
-    session = get_active_session()
+    active_sess = get_active_session()
+    st.title(active_sess["title"] if active_sess else "Chat")
 
     # ── 1. Render history ─────────────────────────────────────────────────────
-    for msg in session["messages"]:
-        if msg["role"] == "user":
-            with st.chat_message("user", avatar="👤"):
+    for msg in st.session_state.messages:
+        role = msg["role"]
+        with st.chat_message(role, avatar="👤" if role == "user" else "🤖"):
+            if role == "user":
                 if msg.get("image_data"):
-                    st.image(base64.b64decode(msg["image_data"]), width=300)
+                    try:
+                        st.image(base64.b64decode(msg["image_data"]), width=300)
+                    except Exception:
+                        pass
                 st.write(msg["content"])
-        else:
-            with st.chat_message("assistant", avatar="🤖"):
+            else:
                 meta = msg.get("metadata", {})
                 if meta:
                     badges = ""
@@ -293,16 +473,17 @@ def chat_interface():
                         badges += '<span class="badge badge-blue">💾 CACHED</span>'
                     elif src in ("google_adk_agent", "agent"):
                         badges += '<span class="badge badge-purple">🤖 AGENT</span>'
-                    model = meta.get("model_used")
+                    model = meta.get("model_used") or meta.get("model")
                     if model:
                         badges += f'<span class="badge" style="background:rgba(255,255,255,0.1);">{model}</span>'
                     if badges:
                         st.markdown(badges, unsafe_allow_html=True)
 
-                content = msg["content"]
-                if msg.get("reasoning"):
+                if msg.get("reasoning") or msg.get("explanation"):
                     with st.expander("Show Reasoning Steps"):
-                        st.markdown(msg["reasoning"])
+                        st.markdown(msg.get("reasoning") or msg.get("explanation"))
+
+                content = msg["content"]
                 if isinstance(content, dict) and "final_answer" in content:
                     st.markdown(f"**Answer:**\n\n> {content['final_answer']}")
                 else:
@@ -316,24 +497,29 @@ def chat_interface():
     is_processing = st.session_state.get("is_processing", False)
 
     with tab_text:
-        prompt = st.chat_input("Ask a math question...", disabled=is_processing)
+        text_prompt = st.chat_input("Ask a math question...", disabled=is_processing)
+        if text_prompt:
+            prompt = text_prompt
 
     with tab_draw:
         col_canvas, col_controls = st.columns([3, 1])
         with col_canvas:
-            if "canvas_key" not in st.session_state:
-                st.session_state.canvas_key = "main_canvas"
             canvas_result = st_canvas(
                 stroke_width=3, stroke_color="#FFFFFF", background_color="#000000",
-                height=300, width=600, drawing_mode="freedraw", key=st.session_state.canvas_key,
+                height=300, width=600, drawing_mode="freedraw",
+                key=st.session_state.canvas_key,
+            )
+            draw_prompt_input = st.text_input(
+                "Question about drawing (optional)",
+                placeholder="Solve this handwritten problem...",
+                key="draw_prompt_input"
             )
-            draw_prompt = st.text_input("Question about drawing (optional)", placeholder="Solve this handwritten problem...")
         with col_controls:
             st.caption("Controls")
-            if st.button("Clear Canvas"):
+            if st.button("Clear"):
                 st.session_state.canvas_key = f"canvas_{uuid.uuid4()}"
                 st.rerun()
-            if st.button("Solve Drawing", type="primary", disabled=is_processing):
+            if st.button("Solve", type="primary", disabled=is_processing):
                 if canvas_result.image_data is not None:
                     img = Image.fromarray(canvas_result.image_data.astype("uint8"), "RGBA")
                     bg  = Image.new("RGB", img.size, (0, 0, 0))
@@ -341,102 +527,99 @@ def chat_interface():
                     buf = io.BytesIO()
                     bg.save(buf, format="PNG")
                     image_b64 = base64.b64encode(buf.getvalue()).decode()
-                    prompt = draw_prompt or "Solve this handwritten math problem."
+                    prompt = draw_prompt_input or "Solve this handwritten math problem."
 
     with tab_upload:
-        uploaded     = st.file_uploader("Upload Image", type=["png","jpg"], disabled=is_processing)
-        upload_prompt = st.text_input("Question about image (optional)", placeholder="Analyze this image...", disabled=is_processing)
-        if uploaded and st.button("Analyze Image", disabled=is_processing):
-            image_b64 = base64.b64encode(uploaded.getvalue()).decode()
-            prompt    = upload_prompt or "Analyze this image."
+        uploaded_file = st.file_uploader("Upload", type=["png", "jpg"], disabled=is_processing)
+        upload_prompt_input = st.text_input("Question", placeholder="Analyze...", disabled=is_processing, key="upload_prompt_input")
+        if uploaded_file and st.button("Analyze", disabled=is_processing):
+            image_b64 = base64.b64encode(uploaded_file.getvalue()).decode()
+            prompt    = upload_prompt_input or "Analyze this image."
 
-    # ── 3. New user message → optimistic write + rerun ────────────────────────
+    # ── 3. New user message → optimistic UI update + rerun ────────────────────
     if prompt:
         req_id = str(uuid.uuid4())
         add_message("user", prompt, image_data=image_b64, request_id=req_id, sent_to_api=False)
         st.session_state.is_processing = True
         st.rerun()
 
-    # ── 4. Recovery: if we restarted mid-flight, allow retry ──────────────────
-    if session["messages"] and session["messages"][-1]["role"] == "user":
-        last = session["messages"][-1]
-        if last.get("sent_to_api") and not st.session_state.is_processing:
-            last["sent_to_api"] = False
-            save_history()
-
-    # ── 5. Fire API call if last message is unsent user message ───────────────
+    # ── 4. Fire API call if last message is an unsent user message ────────────
     if (
-        session["messages"]
-        and session["messages"][-1]["role"] == "user"
-        and not session["messages"][-1].get("sent_to_api", False)
+        st.session_state.messages
+        and st.session_state.messages[-1]["role"] == "user"
+        and not st.session_state.messages[-1].get("sent_to_api", False)
     ):
-        last = session["messages"][-1]
+        last = st.session_state.messages[-1]
         current_request_id = last.get("request_id") or str(uuid.uuid4())
-        last["request_id"] = current_request_id
+        last["request_id"]  = current_request_id
 
         with st.chat_message("assistant", avatar="🤖"):
-            with st.spinner("Agent is thinking..."):
-                try:
-                    last["sent_to_api"] = True
-                    save_history()
-
-                    payload = {
-                        "text":             last["content"],
-                        "image":            last.get("image_data"),
-                        "model_preference": "agent",
-                        "session_id":       st.session_state.active_session_id,
-                        "request_id":       current_request_id,
-                    }
-                    headers = {}
-                    if st.session_state.user:
-                        headers["Authorization"] = f"Bearer {st.session_state.user['token']}"
-
-                    response = requests.post(API_URL, json=payload, headers=headers, timeout=360)
-
-                    if response.status_code == 200:
-                        data = response.json()
-                        if data.get("status") == "success":
-                            answer_raw = data.get("answer") or data.get("explanation") or "⚠️ No answer returned."
-                            meta       = data.get("metadata", {})
+            status_msg = st.status("Thinking...", expanded=True)
+            logic_placeholder = status_msg.empty()
+            answer_placeholder = st.empty()
+            
+            full_answer = ""
+            logic_trace = []
+            
+            try:
+                last["sent_to_api"] = True
+                payload = {
+                    "text":       last["content"],
+                    "image":      last.get("image_data"),
+                    "session_id": st.session_state.active_session_id,
+                    "request_id": current_request_id,
+                }
+                headers = get_auth_headers()
+                
+                with requests.post(API_URL, json=payload, headers=headers, stream=True, timeout=360) as r:
+                    if r.status_code == 200:
+                        for line in r.iter_lines():
+                            if line:
+                                try:
+                                    data = json.loads(line)
+                                    if data["type"] == "thought":
+                                        logic_trace.append(data["content"])
+                                    elif data["type"] == "action":
+                                        logic_trace.append(f"⚙️ {data['content']}")
+                                    elif data["type"] == "observation":
+                                        logic_trace.append(f"👁️ {data['content']}")
+                                    elif data["type"] == "answer":
+                                        full_answer += data["content"]
+                                        answer_placeholder.markdown(full_answer)
+                                    elif data["type"] == "error":
+                                        st.error(data["content"])
+                                        
+                                    # Update logic trace UI
+                                    logic_placeholder.markdown("\n".join(logic_trace))
+                                except Exception:
+                                    continue
+                        
+                        status_msg.update(label="Solved!", state="complete", expanded=False)
+                        st.session_state.is_processing = False
+                        
+                        if full_answer:
                             add_message(
-                                "assistant",
-                                answer_raw,
-                                reasoning=data.get("explanation"),
-                                metadata={
-                                    "source":     data.get("source", "agent"),
-                                    "model_used": meta.get("model", "gemini-2.5-flash"),
-                                    "latency":    f"{meta.get('latency_ms',0)/1000:.2f}s",
-                                    "tools":      meta.get("tools_used", []),
-                                },
-                                steps=data.get("steps", [])
+                                "assistant", 
+                                full_answer, 
+                                reasoning="\n".join(logic_trace),
+                                metadata={"source": "agent"}
                             )
-                            st.session_state.is_processing = False
-                            st.rerun()   # ← BUG 2 FIX: missing rerun caused blank UI
-
-                        else:
-                            error_msg = data.get("error", "Unknown error")
-                            add_message("assistant", f"⚠️ Error: {error_msg}")
-                            st.session_state.is_processing = False
+                            load_sessions() # Update titles
                             st.rerun()
-
-                    elif response.status_code in [202, 409]:
-                        st.info("ℹ️ Request already processing.")
-                        st.session_state.is_processing = False
+                    elif r.status_code == 401:
+                        _clear_user_state()
+                        st.session_state.user = None
+                        st.error("Session expired. Please log in again.")
                         st.rerun()
-
                     else:
-                        try:
-                            err_msg = response.json().get("error", f"HTTP {response.status_code}")
-                        except Exception:
-                            err_msg = f"HTTP {response.status_code}"
-                        add_message("assistant", f"❌ Server Error: {err_msg}")
+                        st.error(f"Error: {r.status_code}")
                         st.session_state.is_processing = False
-                        st.rerun()
 
-                except Exception as e:
-                    add_message("assistant", f"❌ Connection Failed: {str(e)}")
-                    st.session_state.is_processing = False
-                    st.rerun()
+            except Exception as e:
+                st.error(f"Connection error: {e}")
+                st.session_state.is_processing = False
+                st.rerun()
+
 
 # ====================================================
 # Sidebar
@@ -445,16 +628,31 @@ with st.sidebar:
     st.markdown("### 🧠 MathMinds")
     st.write(f"Logged in as **{st.session_state.user['email']}**")
 
-    view = st.radio("Navigation", ["Chat", "Profile"],
-                    index=0 if st.session_state.current_view == "Chat" else 1)
+    view = st.radio(
+        "Navigation", ["Chat", "Profile"],
+        index=0 if st.session_state.current_view == "Chat" else 1
+    )
     if view != st.session_state.current_view:
         st.session_state.current_view = view
         st.rerun()
 
     if st.button("Sign Out", type="secondary"):
+        # ✅ MULTIUSER FIX: Wipe ALL user-specific state first, THEN clear identity.
+        # Without _clear_user_state() here, the next user to log in on the same
+        # browser tab would see User A's chat history briefly before load_sessions
+        # returns, because st.session_state persists across logins within a tab.
+        _clear_user_state()
         st.session_state.user = None
         st.rerun()
 
+    if st.session_state.is_processing:
+        if st.button(
+            "🔓 Reset Processing Lock", type="primary",
+            help="Use if UI is stuck despite answer finishing."
+        ):
+            st.session_state.is_processing = False
+            st.rerun()
+
     st.divider()
 
     if st.session_state.current_view == "Chat":
@@ -462,26 +660,45 @@ with st.sidebar:
             new_chat()
 
         st.markdown("#### History")
-        sorted_sids = sorted(
-            st.session_state.chat_sessions.keys(),
-            key=lambda k: st.session_state.chat_sessions[k].get("created_at", 0),
-            reverse=True
-        )
-        for sid in sorted_sids:
-            sess    = st.session_state.chat_sessions[sid]
-            title   = sess.get("title", "Untitled")
-            isActive = (sid == st.session_state.active_session_id)
-            col_nav, col_del = st.columns([0.85, 0.15])
-            with col_nav:
-                if st.button(f"{'📍 ' if isActive else ''}{title}", key=sid, use_container_width=True):
+
+        for session in st.session_state.chat_sessions:
+            sid   = session["session_id"]
+            title = session["title"]
+
+            cols = st.columns([0.8, 0.1, 0.1])
+            with cols[0]:
+                is_active = (st.session_state.active_session_id == sid)
+                btn_type  = "primary" if is_active else "secondary"
+                if st.button(title, key=f"sel_{sid}", use_container_width=True, type=btn_type):
                     st.session_state.active_session_id = sid
+                    load_messages(sid)
+                    st.rerun()
+            with cols[1]:
+                if st.button("🖊️", key=f"ren_{sid}", help="Rename"):
+                    st.session_state.renaming_session_id = (
+                        sid if st.session_state.renaming_session_id != sid else None
+                    )
                     st.rerun()
-            with col_del:
-                if isActive and st.button("🗑️", key=f"del_{sid}"):
+            with cols[2]:
+                if st.button("🗑️", key=f"del_{sid}", help="Delete"):
                     delete_chat(sid)
 
-# ── Router ────────────────────────────────────────────────────────────────────
-if st.session_state.current_view == "Profile":
-    profile_interface()
-else:
-    chat_interface()
+            if st.session_state.renaming_session_id == sid:
+                with st.container():
+                    new_title = st.text_input(
+                        "New title", value=title,
+                        key=f"in_{sid}", label_visibility="collapsed"
+                    )
+                    if st.button("Save", key=f"save_{sid}", use_container_width=True):
+                        rename_chat(sid, new_title)
+                        st.session_state.renaming_session_id = None
+                        st.rerun()
+
+
+# ====================================================
+# Main Content Area
+# ====================================================
+if st.session_state.current_view == "Chat":
+    chat_interface()
+elif st.session_state.current_view == "Profile":
+    profile_interface()

requirements.txt

@@ -33,6 +33,9 @@ pandas
 gradio
 transformers
 torch
+passlib[bcrypt]
+bcrypt
+PyJWT
 
 selenium
 webdriver-manager