From 48e646607ba935cc1bae26b15879de5f13938127 Mon Sep 17 00:00:00 2001
From: Manoj Aggarwal <manoj6543@gmail.com>
Date: Mon, 2 Feb 2026 12:07:53 -0800
Subject: [PATCH 01/40] Fix google calendar and notion erros

---
 .../app/tasks/celery_tasks/connector_tasks.py | 43 +++++++++++++++++++
 .../app/tasks/connector_indexers/base.py      | 14 ++++++
 .../google_calendar_indexer.py                | 14 ++++++
 .../connector_indexers/notion_indexer.py      | 43 ++++++++++++++++---
 .../connector_indexers/webcrawler_indexer.py  | 20 +++++++--
 5 files changed, 124 insertions(+), 10 deletions(-)

diff --git a/surfsense_backend/app/tasks/celery_tasks/connector_tasks.py b/surfsense_backend/app/tasks/celery_tasks/connector_tasks.py
index d0710d246..760651589 100644
--- a/surfsense_backend/app/tasks/celery_tasks/connector_tasks.py
+++ b/surfsense_backend/app/tasks/celery_tasks/connector_tasks.py
@@ -1,6 +1,7 @@
 """Celery tasks for connector indexing."""
 
 import logging
+import traceback
 
 from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
 from sqlalchemy.pool import NullPool
@@ -11,6 +12,36 @@ from app.config import config
 logger = logging.getLogger(__name__)
 
 
+def _handle_greenlet_error(e: Exception, task_name: str, connector_id: int) -> None:
+    """
+    Handle greenlet_spawn errors with detailed logging for debugging.
+    
+    The 'greenlet_spawn has not been called' error occurs when:
+    1. SQLAlchemy lazy-loads a relationship outside of an async context
+    2. A sync operation is called from an async context (or vice versa)
+    3. Session objects are accessed after the session is closed
+    
+    This helper logs detailed context to help identify the root cause.
+    """
+    error_str = str(e)
+    if "greenlet_spawn has not been called" in error_str:
+        logger.error(
+            f"GREENLET ERROR in {task_name} for connector {connector_id}: {error_str}\n"
+            f"This error typically occurs when SQLAlchemy tries to lazy-load a relationship "
+            f"outside of an async context. Check for:\n"
+            f"1. Accessing relationship attributes (e.g., document.chunks, connector.search_space) "
+            f"without using selectinload() or joinedload()\n"
+            f"2. Accessing model attributes after the session is closed\n"
+            f"3. Passing ORM objects between different async contexts\n"
+            f"Stack trace:\n{traceback.format_exc()}"
+        )
+    else:
+        logger.error(
+            f"Error in {task_name} for connector {connector_id}: {error_str}\n"
+            f"Stack trace:\n{traceback.format_exc()}"
+        )
+
+
 def get_celery_session_maker():
     """
     Create a new async session maker for Celery tasks.
@@ -46,6 +77,9 @@ def index_slack_messages_task(
                 connector_id, search_space_id, user_id, start_date, end_date
             )
         )
+    except Exception as e:
+        _handle_greenlet_error(e, "index_slack_messages", connector_id)
+        raise
     finally:
         loop.close()
 
@@ -89,6 +123,9 @@ def index_notion_pages_task(
                 connector_id, search_space_id, user_id, start_date, end_date
             )
         )
+    except Exception as e:
+        _handle_greenlet_error(e, "index_notion_pages", connector_id)
+        raise
     finally:
         loop.close()
 
@@ -347,6 +384,9 @@ def index_google_calendar_events_task(
                 connector_id, search_space_id, user_id, start_date, end_date
             )
         )
+    except Exception as e:
+        _handle_greenlet_error(e, "index_google_calendar_events", connector_id)
+        raise
     finally:
         loop.close()
 
@@ -696,6 +736,9 @@ def index_crawled_urls_task(
                 connector_id, search_space_id, user_id, start_date, end_date
             )
         )
+    except Exception as e:
+        _handle_greenlet_error(e, "index_crawled_urls", connector_id)
+        raise
     finally:
         loop.close()
 
diff --git a/surfsense_backend/app/tasks/connector_indexers/base.py b/surfsense_backend/app/tasks/connector_indexers/base.py
index b390937f0..311fda996 100644
--- a/surfsense_backend/app/tasks/connector_indexers/base.py
+++ b/surfsense_backend/app/tasks/connector_indexers/base.py
@@ -159,6 +159,20 @@ def calculate_date_range(
     )
     end_date_str = end_date if end_date else calculated_end_date.strftime("%Y-%m-%d")
 
+    # FIX: Ensure end_date is at least 1 day after start_date to avoid
+    # "start_date must be strictly before end_date" errors when dates are the same
+    # (e.g., when last_indexed_at is today)
+    if start_date_str == end_date_str:
+        logger.info(
+            f"Start date ({start_date_str}) equals end date ({end_date_str}), "
+            "adjusting end date to next day to ensure valid date range"
+        )
+        # Parse end_date and add 1 day
+        end_dt = datetime.strptime(end_date_str, "%Y-%m-%d")
+        end_dt = end_dt + timedelta(days=1)
+        end_date_str = end_dt.strftime("%Y-%m-%d")
+        logger.info(f"Adjusted end date to {end_date_str}")
+
     return start_date_str, end_date_str
 
 
diff --git a/surfsense_backend/app/tasks/connector_indexers/google_calendar_indexer.py b/surfsense_backend/app/tasks/connector_indexers/google_calendar_indexer.py
index f64a7a5c3..1d8ea32f2 100644
--- a/surfsense_backend/app/tasks/connector_indexers/google_calendar_indexer.py
+++ b/surfsense_backend/app/tasks/connector_indexers/google_calendar_indexer.py
@@ -217,6 +217,20 @@ async def index_google_calendar_events(
             start_date_str = start_date
             end_date_str = end_date
 
+        # FIX: Ensure end_date is at least 1 day after start_date to avoid
+        # "start_date must be strictly before end_date" errors when dates are the same
+        # (e.g., when last_indexed_at is today)
+        if start_date_str == end_date_str:
+            logger.info(
+                f"Start date ({start_date_str}) equals end date ({end_date_str}), "
+                "adjusting end date to next day to ensure valid date range"
+            )
+            # Parse end_date and add 1 day
+            end_dt = datetime.strptime(end_date_str, "%Y-%m-%d")
+            end_dt = end_dt + timedelta(days=1)
+            end_date_str = end_dt.strftime("%Y-%m-%d")
+            logger.info(f"Adjusted end date to {end_date_str}")
+
         await task_logger.log_task_progress(
             log_entry,
             f"Fetching Google Calendar events from {start_date_str} to {end_date_str}",
diff --git a/surfsense_backend/app/tasks/connector_indexers/notion_indexer.py b/surfsense_backend/app/tasks/connector_indexers/notion_indexer.py
index 52622471a..ee5bca5d8 100644
--- a/surfsense_backend/app/tasks/connector_indexers/notion_indexer.py
+++ b/surfsense_backend/app/tasks/connector_indexers/notion_indexer.py
@@ -196,13 +196,44 @@ async def index_notion_pages(
                     "Recommend reconnecting with OAuth."
                 )
         except Exception as e:
-            await task_logger.log_task_failure(
-                log_entry,
-                f"Failed to get Notion pages for connector {connector_id}",
-                str(e),
-                {"error_type": "PageFetchError"},
+            error_str = str(e)
+            # Check if this is an unsupported block type error (transcription, ai_block, etc.)
+            # These are known Notion API limitations and should be logged as warnings, not errors
+            unsupported_block_errors = [
+                "transcription is not supported",
+                "ai_block is not supported",
+                "is not supported via the API",
+            ]
+            is_unsupported_block_error = any(
+                err in error_str.lower() for err in unsupported_block_errors
             )
-            logger.error(f"Error fetching Notion pages: {e!s}", exc_info=True)
+
+            if is_unsupported_block_error:
+                # Log as warning since this is a known Notion API limitation
+                logger.warning(
+                    f"Notion API limitation for connector {connector_id}: {error_str}. "
+                    "This is a known issue with Notion AI blocks (transcription, ai_block) "
+                    "that are not accessible via the Notion API."
+                )
+                await task_logger.log_task_failure(
+                    log_entry,
+                    f"Failed to get Notion pages: Notion API limitation",
+                    f"{error_str} - This page contains Notion AI content (transcription/ai_block) that cannot be accessed via the API.",
+                    {"error_type": "UnsupportedBlockType", "is_known_limitation": True},
+                )
+            else:
+                # Log as error for other failures
+                logger.error(
+                    f"Error fetching Notion pages for connector {connector_id}: {error_str}",
+                    exc_info=True,
+                )
+                await task_logger.log_task_failure(
+                    log_entry,
+                    f"Failed to get Notion pages for connector {connector_id}",
+                    str(e),
+                    {"error_type": "PageFetchError"},
+                )
+
             await notion_client.close()
             return 0, f"Failed to get Notion pages: {e!s}"
 
diff --git a/surfsense_backend/app/tasks/connector_indexers/webcrawler_indexer.py b/surfsense_backend/app/tasks/connector_indexers/webcrawler_indexer.py
index ac16ecde6..b3c24a4e3 100644
--- a/surfsense_backend/app/tasks/connector_indexers/webcrawler_indexer.py
+++ b/surfsense_backend/app/tasks/connector_indexers/webcrawler_indexer.py
@@ -108,10 +108,15 @@ async def index_crawled_urls(
         api_key = connector.config.get("FIRECRAWL_API_KEY")
 
         # Get URLs from connector config
-        urls = parse_webcrawler_urls(connector.config.get("INITIAL_URLS"))
+        raw_initial_urls = connector.config.get("INITIAL_URLS")
+        urls = parse_webcrawler_urls(raw_initial_urls)
 
+        # DEBUG: Log connector config details for troubleshooting empty URL issues
         logger.info(
-            f"Starting crawled web page indexing for connector {connector_id} with {len(urls)} URLs"
+            f"Starting crawled web page indexing for connector {connector_id} with {len(urls)} URLs. "
+            f"Connector name: {connector.name}, "
+            f"INITIAL_URLS type: {type(raw_initial_urls).__name__}, "
+            f"INITIAL_URLS value: {repr(raw_initial_urls)[:200] if raw_initial_urls else 'None'}"
         )
 
         # Initialize webcrawler client
@@ -128,11 +133,18 @@ async def index_crawled_urls(
 
         # Validate URLs
         if not urls:
+            # DEBUG: Log detailed connector config for troubleshooting
+            logger.error(
+                f"No URLs provided for indexing. Connector ID: {connector_id}, "
+                f"Connector name: {connector.name}, "
+                f"Config keys: {list(connector.config.keys()) if connector.config else 'None'}, "
+                f"INITIAL_URLS raw value: {repr(raw_initial_urls)}"
+            )
             await task_logger.log_task_failure(
                 log_entry,
                 "No URLs provided for indexing",
-                "Empty URL list",
-                {"error_type": "ValidationError"},
+                f"Empty URL list. INITIAL_URLS value: {repr(raw_initial_urls)[:100]}",
+                {"error_type": "ValidationError", "connector_name": connector.name},
             )
             return 0, "No URLs provided for indexing"
 

From d761ca199211c1a35c9d8f5d7a9bcf852c1a7a35 Mon Sep 17 00:00:00 2001
From: Eric Lammertsma <e.lammertsma@gmail.com>
Date: Tue, 3 Feb 2026 20:47:18 -0500
Subject: [PATCH 02/40] Added chat renaming and auto naming by the active LLM

---
 surfsense_backend/app/prompts/__init__.py     | 30 +++++++
 .../app/routes/new_chat_routes.py             | 26 +-----
 .../app/services/new_streaming_service.py     | 25 ++++++
 .../app/tasks/chat/stream_new_chat.py         | 75 ++++++++++++++++
 .../new-chat/[[...chat_id]]/page.tsx          | 21 ++++-
 .../layout/providers/LayoutDataProvider.tsx   | 88 ++++++++++++++++++-
 .../layout/ui/shell/LayoutShell.tsx           |  4 +
 .../layout/ui/sidebar/ChatListItem.tsx        | 33 ++++---
 .../layout/ui/sidebar/MobileSidebar.tsx       |  3 +
 .../components/layout/ui/sidebar/Sidebar.tsx  |  4 +
 surfsense_web/messages/en.json                |  7 ++
 surfsense_web/messages/zh.json                |  7 ++
 12 files changed, 287 insertions(+), 36 deletions(-)

diff --git a/surfsense_backend/app/prompts/__init__.py b/surfsense_backend/app/prompts/__init__.py
index 3b21cb9e1..efa31d612 100644
--- a/surfsense_backend/app/prompts/__init__.py
+++ b/surfsense_backend/app/prompts/__init__.py
@@ -104,3 +104,33 @@ SUMMARY_PROMPT = (
 SUMMARY_PROMPT_TEMPLATE = PromptTemplate(
     input_variables=["document"], template=SUMMARY_PROMPT
 )
+
+# =============================================================================
+# Chat Title Generation Prompt
+# =============================================================================
+
+TITLE_GENERATION_PROMPT = """Generate a concise, descriptive title for the following conversation.
+
+<rules>
+- The title MUST be between 1 and 6 words
+- The title MUST be on a single line
+- Capture the main topic or intent of the conversation
+- Do NOT use quotes, punctuation, or formatting
+- Do NOT include words like "Chat about" or "Discussion of"
+- Return ONLY the title, nothing else
+</rules>
+
+<user_query>
+{user_query}
+</user_query>
+
+<assistant_response>
+{assistant_response}
+</assistant_response>
+
+Title:"""
+
+TITLE_GENERATION_PROMPT_TEMPLATE = PromptTemplate(
+    input_variables=["user_query", "assistant_response"],
+    template=TITLE_GENERATION_PROMPT,
+)
diff --git a/surfsense_backend/app/routes/new_chat_routes.py b/surfsense_backend/app/routes/new_chat_routes.py
index 42b8a821b..06e929997 100644
--- a/surfsense_backend/app/routes/new_chat_routes.py
+++ b/surfsense_backend/app/routes/new_chat_routes.py
@@ -886,30 +886,8 @@ async def append_message(
         # Update thread's updated_at timestamp
         thread.updated_at = datetime.now(UTC)
 
-        # Auto-generate title from first user message if title is still default
-        if thread.title == "New Chat" and role_str == "user":
-            # Extract text content for title
-            content = message.content
-            if isinstance(content, str):
-                title_text = content
-            elif isinstance(content, list):
-                # Find first text content
-                title_text = ""
-                for part in content:
-                    if isinstance(part, dict) and part.get("type") == "text":
-                        title_text = part.get("text", "")
-                        break
-                    elif isinstance(part, str):
-                        title_text = part
-                        break
-            else:
-                title_text = str(content)
-
-            # Truncate title
-            if title_text:
-                thread.title = title_text[:100] + (
-                    "..." if len(title_text) > 100 else ""
-                )
+        # Note: Title generation now happens in stream_new_chat.py after the first response
+        # using LLM to generate a descriptive title (with truncation as fallback)
 
         await session.commit()
         await session.refresh(db_message)
diff --git a/surfsense_backend/app/services/new_streaming_service.py b/surfsense_backend/app/services/new_streaming_service.py
index 05dd2d4dd..57fbc9663 100644
--- a/surfsense_backend/app/services/new_streaming_service.py
+++ b/surfsense_backend/app/services/new_streaming_service.py
@@ -479,6 +479,31 @@ class VercelStreamingService:
             },
         )
 
+    def format_thread_title_update(self, thread_id: int, title: str) -> str:
+        """
+        Format a thread title update notification (SurfSense specific).
+
+        This is sent after the first response in a thread to update the
+        auto-generated title based on the conversation content.
+
+        Args:
+            thread_id: The ID of the thread being updated
+            title: The new title for the thread
+
+        Returns:
+            str: SSE formatted thread title update data part
+
+        Example output:
+            data: {"type":"data-thread-title-update","data":{"threadId":123,"title":"New Title"}}
+        """
+        return self.format_data(
+            "thread-title-update",
+            {
+                "threadId": thread_id,
+                "title": title,
+            },
+        )
+
     # =========================================================================
     # Error Part
     # =========================================================================
diff --git a/surfsense_backend/app/tasks/chat/stream_new_chat.py b/surfsense_backend/app/tasks/chat/stream_new_chat.py
index 688777203..47e2a2293 100644
--- a/surfsense_backend/app/tasks/chat/stream_new_chat.py
+++ b/surfsense_backend/app/tasks/chat/stream_new_chat.py
@@ -32,6 +32,7 @@ from app.services.chat_session_state_service import (
     clear_ai_responding,
     set_ai_responding,
 )
+from app.prompts import TITLE_GENERATION_PROMPT_TEMPLATE
 from app.services.connector_service import ConnectorService
 from app.services.new_streaming_service import VercelStreamingService
 from app.utils.content_utils import bootstrap_history_from_db
@@ -1208,6 +1209,80 @@ async def stream_new_chat(
         if completion_event:
             yield completion_event
 
+        # Generate LLM title for new chats after first response
+        # Check if this is the first assistant response by counting existing assistant messages
+        from app.db import NewChatMessage, NewChatThread
+        from sqlalchemy import func
+
+        assistant_count_result = await session.execute(
+            select(func.count(NewChatMessage.id)).filter(
+                NewChatMessage.thread_id == chat_id,
+                NewChatMessage.role == "assistant",
+            )
+        )
+        assistant_message_count = assistant_count_result.scalar() or 0
+
+        # Only generate title on the first response (no prior assistant messages)
+        if assistant_message_count == 0:
+            print(f"[stream_new_chat] First response - generating title for thread {chat_id}")
+            print(f"[stream_new_chat] Query length: {len(user_query)}, Response length: {len(accumulated_text)}")
+
+            generated_title = None
+            try:
+                # Generate title using the same LLM
+                title_chain = TITLE_GENERATION_PROMPT_TEMPLATE | llm
+                # Truncate inputs to avoid context length issues
+                truncated_query = user_query[:500]
+                truncated_response = accumulated_text[:1000]
+                print(f"[stream_new_chat] Calling LLM for title generation...")
+                title_result = await title_chain.ainvoke({
+                    "user_query": truncated_query,
+                    "assistant_response": truncated_response,
+                })
+                print(f"[stream_new_chat] LLM title result type: {type(title_result)}")
+                print(f"[stream_new_chat] LLM title result: {title_result}")
+
+                # Extract and clean the title
+                if title_result and hasattr(title_result, "content"):
+                    raw_title = title_result.content.strip()
+                    print(f"[stream_new_chat] Raw title content: '{raw_title}' (len={len(raw_title)})")
+
+                    # Validate the title (1-6 words, reasonable length)
+                    if raw_title and len(raw_title) <= 100:
+                        # Remove any quotes or extra formatting
+                        generated_title = raw_title.strip('"\'')
+                        print(f"[stream_new_chat] After stripping quotes: '{generated_title}'")
+                    else:
+                        print(f"[stream_new_chat] Title validation failed: empty={not raw_title}, len={len(raw_title)}")
+                        generated_title = None
+                else:
+                    print(f"[stream_new_chat] No content attribute on result")
+            except Exception as title_error:
+                print(f"[stream_new_chat] Title generation failed: {title_error}")
+                import traceback
+                traceback.print_exc()
+                generated_title = None
+
+            # Only update if LLM succeeded (keep truncated prompt title as fallback)
+            if generated_title:
+                print(f"[stream_new_chat] Using LLM-generated title: '{generated_title}'")
+
+                # Fetch thread and update title
+                thread_result = await session.execute(
+                    select(NewChatThread).filter(NewChatThread.id == chat_id)
+                )
+                thread = thread_result.scalars().first()
+                if thread:
+                    thread.title = generated_title
+                    await session.commit()
+
+                    # Notify frontend of the title update
+                    yield streaming_service.format_thread_title_update(
+                        chat_id, generated_title
+                    )
+            else:
+                print(f"[stream_new_chat] LLM title generation failed, keeping truncated prompt title")
+
         # Finish the step and message
         yield streaming_service.format_finish_step()
         yield streaming_service.format_finish()
diff --git a/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx b/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx
index 9b462fcbc..4f99f8021 100644
--- a/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx
+++ b/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx
@@ -437,7 +437,10 @@ export default function NewChatPage() {
 			let isNewThread = false;
 			if (!currentThreadId) {
 				try {
-					const newThread = await createThread(searchSpaceId, "New Chat");
+					// Create thread with truncated prompt as initial title
+					const initialTitle =
+						userQuery.trim().slice(0, 100) + (userQuery.trim().length > 100 ? "..." : "");
+					const newThread = await createThread(searchSpaceId, initialTitle);
 					currentThreadId = newThread.id;
 					setThreadId(currentThreadId);
 					// Set currentThread so ChatHeader can show share button immediately
@@ -827,6 +830,22 @@ export default function NewChatPage() {
 											break;
 										}
 
+										case "data-thread-title-update": {
+											// Handle thread title update from LLM-generated title
+											const titleData = parsed.data as { threadId: number; title: string };
+											if (titleData?.title && titleData?.threadId === currentThreadId) {
+												// Update current thread state with new title
+												setCurrentThread((prev) =>
+													prev ? { ...prev, title: titleData.title } : prev
+												);
+												// Invalidate thread list to refresh sidebar
+												queryClient.invalidateQueries({
+													queryKey: ["threads", String(searchSpaceId)],
+												});
+											}
+											break;
+										}
+
 										case "error":
 											throw new Error(parsed.errorText || "Server error");
 									}
diff --git a/surfsense_web/components/layout/providers/LayoutDataProvider.tsx b/surfsense_web/components/layout/providers/LayoutDataProvider.tsx
index 2f71adad9..9a5c2ed40 100644
--- a/surfsense_web/components/layout/providers/LayoutDataProvider.tsx
+++ b/surfsense_web/components/layout/providers/LayoutDataProvider.tsx
@@ -2,7 +2,7 @@
 
 import { useQuery, useQueryClient } from "@tanstack/react-query";
 import { useAtomValue, useSetAtom } from "jotai";
-import { AlertTriangle, Inbox, LogOut, SquareLibrary, Trash2 } from "lucide-react";
+import { AlertTriangle, Inbox, LogOut, PencilIcon, SquareLibrary, Trash2 } from "lucide-react";
 import { useParams, usePathname, useRouter } from "next/navigation";
 import { useTranslations } from "next-intl";
 import { useTheme } from "next-themes";
@@ -21,6 +21,7 @@ import {
 	DialogHeader,
 	DialogTitle,
 } from "@/components/ui/dialog";
+import { Input } from "@/components/ui/input";
 import { isPageLimitExceededMetadata } from "@/contracts/types/inbox.types";
 import { useInbox } from "@/hooks/use-inbox";
 import { searchSpacesApiService } from "@/lib/apis/search-spaces-api.service";
@@ -207,6 +208,12 @@ export function LayoutDataProvider({
 	const [chatToDelete, setChatToDelete] = useState<{ id: number; name: string } | null>(null);
 	const [isDeletingChat, setIsDeletingChat] = useState(false);
 
+	// Rename dialog state
+	const [showRenameChatDialog, setShowRenameChatDialog] = useState(false);
+	const [chatToRename, setChatToRename] = useState<{ id: number; name: string } | null>(null);
+	const [newChatTitle, setNewChatTitle] = useState("");
+	const [isRenamingChat, setIsRenamingChat] = useState(false);
+
 	// Delete/Leave search space dialog state
 	const [showDeleteSearchSpaceDialog, setShowDeleteSearchSpaceDialog] = useState(false);
 	const [showLeaveSearchSpaceDialog, setShowLeaveSearchSpaceDialog] = useState(false);
@@ -421,6 +428,12 @@ export function LayoutDataProvider({
 		setShowDeleteChatDialog(true);
 	}, []);
 
+	const handleChatRename = useCallback((chat: ChatItem) => {
+		setChatToRename({ id: chat.id, name: chat.name });
+		setNewChatTitle(chat.name);
+		setShowRenameChatDialog(true);
+	}, []);
+
 	const handleChatArchive = useCallback(
 		async (chat: ChatItem) => {
 			const newArchivedState = !chat.archived;
@@ -501,6 +514,27 @@ export function LayoutDataProvider({
 		}
 	}, [chatToDelete, queryClient, searchSpaceId, router, currentChatId]);
 
+	// Rename handler
+	const confirmRenameChat = useCallback(async () => {
+		if (!chatToRename || !newChatTitle.trim()) return;
+		setIsRenamingChat(true);
+		try {
+			await updateThread(chatToRename.id, { title: newChatTitle.trim() });
+			toast.success(tSidebar("chat_renamed") || "Chat renamed");
+			queryClient.invalidateQueries({ queryKey: ["threads", searchSpaceId] });
+			queryClient.invalidateQueries({ queryKey: ["all-threads", searchSpaceId] });
+			queryClient.invalidateQueries({ queryKey: ["search-threads", searchSpaceId] });
+		} catch (error) {
+			console.error("Error renaming thread:", error);
+			toast.error(tSidebar("error_renaming_chat") || "Failed to rename chat");
+		} finally {
+			setIsRenamingChat(false);
+			setShowRenameChatDialog(false);
+			setChatToRename(null);
+			setNewChatTitle("");
+		}
+	}, [chatToRename, newChatTitle, queryClient, searchSpaceId, tSidebar]);
+
 	// Page usage
 	const pageUsage = user
 		? {
@@ -529,6 +563,7 @@ export function LayoutDataProvider({
 				activeChatId={currentChatId}
 				onNewChat={handleNewChat}
 				onChatSelect={handleChatSelect}
+				onChatRename={handleChatRename}
 				onChatDelete={handleChatDelete}
 				onChatArchive={handleChatArchive}
 				onViewAllSharedChats={handleViewAllSharedChats}
@@ -620,6 +655,57 @@ export function LayoutDataProvider({
 				</DialogContent>
 			</Dialog>
 
+			{/* Rename Chat Dialog */}
+			<Dialog open={showRenameChatDialog} onOpenChange={setShowRenameChatDialog}>
+				<DialogContent className="sm:max-w-md">
+					<DialogHeader>
+						<DialogTitle className="flex items-center gap-2">
+							<PencilIcon className="h-5 w-5" />
+							<span>{tSidebar("rename_chat") || "Rename Chat"}</span>
+						</DialogTitle>
+						<DialogDescription>
+							{tSidebar("rename_chat_description") || "Enter a new name for this conversation."}
+						</DialogDescription>
+					</DialogHeader>
+					<Input
+						value={newChatTitle}
+						onChange={(e) => setNewChatTitle(e.target.value)}
+						placeholder={tSidebar("chat_title_placeholder") || "Chat title"}
+						onKeyDown={(e) => {
+							if (e.key === "Enter" && !isRenamingChat && newChatTitle.trim()) {
+								confirmRenameChat();
+							}
+						}}
+					/>
+					<DialogFooter className="flex gap-2 sm:justify-end">
+						<Button
+							variant="outline"
+							onClick={() => setShowRenameChatDialog(false)}
+							disabled={isRenamingChat}
+						>
+							{tCommon("cancel")}
+						</Button>
+						<Button
+							onClick={confirmRenameChat}
+							disabled={isRenamingChat || !newChatTitle.trim()}
+							className="gap-2"
+						>
+							{isRenamingChat ? (
+								<>
+									<span className="h-4 w-4 animate-spin rounded-full border-2 border-current border-t-transparent" />
+									{tSidebar("renaming") || "Renaming..."}
+								</>
+							) : (
+								<>
+									<PencilIcon className="h-4 w-4" />
+									{tSidebar("rename") || "Rename"}
+								</>
+							)}
+						</Button>
+					</DialogFooter>
+				</DialogContent>
+			</Dialog>
+
 			{/* Delete Search Space Dialog */}
 			<Dialog open={showDeleteSearchSpaceDialog} onOpenChange={setShowDeleteSearchSpaceDialog}>
 				<DialogContent className="sm:max-w-md">
diff --git a/surfsense_web/components/layout/ui/shell/LayoutShell.tsx b/surfsense_web/components/layout/ui/shell/LayoutShell.tsx
index 3624c90a3..8eae99b03 100644
--- a/surfsense_web/components/layout/ui/shell/LayoutShell.tsx
+++ b/surfsense_web/components/layout/ui/shell/LayoutShell.tsx
@@ -54,6 +54,7 @@ interface LayoutShellProps {
 	activeChatId?: number | null;
 	onNewChat: () => void;
 	onChatSelect: (chat: ChatItem) => void;
+	onChatRename?: (chat: ChatItem) => void;
 	onChatDelete?: (chat: ChatItem) => void;
 	onChatArchive?: (chat: ChatItem) => void;
 	onViewAllSharedChats?: () => void;
@@ -90,6 +91,7 @@ export function LayoutShell({
 	activeChatId,
 	onNewChat,
 	onChatSelect,
+	onChatRename,
 	onChatDelete,
 	onChatArchive,
 	onViewAllSharedChats,
@@ -147,6 +149,7 @@ export function LayoutShell({
 							activeChatId={activeChatId}
 							onNewChat={onNewChat}
 							onChatSelect={onChatSelect}
+							onChatRename={onChatRename}
 							onChatDelete={onChatDelete}
 							onChatArchive={onChatArchive}
 							onViewAllSharedChats={onViewAllSharedChats}
@@ -215,6 +218,7 @@ export function LayoutShell({
 							activeChatId={activeChatId}
 							onNewChat={onNewChat}
 							onChatSelect={onChatSelect}
+							onChatRename={onChatRename}
 							onChatDelete={onChatDelete}
 							onChatArchive={onChatArchive}
 							onViewAllSharedChats={onViewAllSharedChats}
diff --git a/surfsense_web/components/layout/ui/sidebar/ChatListItem.tsx b/surfsense_web/components/layout/ui/sidebar/ChatListItem.tsx
index 6db6782d0..ba2989145 100644
--- a/surfsense_web/components/layout/ui/sidebar/ChatListItem.tsx
+++ b/surfsense_web/components/layout/ui/sidebar/ChatListItem.tsx
@@ -1,6 +1,6 @@
 "use client";
 
-import { ArchiveIcon, MessageSquare, MoreHorizontal, RotateCcwIcon, Trash2 } from "lucide-react";
+import { ArchiveIcon, MessageSquare, MoreHorizontal, PencilIcon, RotateCcwIcon, Trash2 } from "lucide-react";
 import { useTranslations } from "next-intl";
 import { Button } from "@/components/ui/button";
 import {
@@ -17,6 +17,7 @@ interface ChatListItemProps {
 	isActive?: boolean;
 	archived?: boolean;
 	onClick?: () => void;
+	onRename?: () => void;
 	onArchive?: () => void;
 	onDelete?: () => void;
 }
@@ -26,6 +27,7 @@ export function ChatListItem({
 	isActive,
 	archived,
 	onClick,
+	onRename,
 	onArchive,
 	onDelete,
 }: ChatListItemProps) {
@@ -57,15 +59,26 @@ export function ChatListItem({
 							<span className="sr-only">{t("more_options")}</span>
 						</Button>
 					</DropdownMenuTrigger>
-					<DropdownMenuContent align="end" side="right">
-						{onArchive && (
-							<DropdownMenuItem
-								onClick={(e) => {
-									e.stopPropagation();
-									onArchive();
-								}}
-							>
-								{archived ? (
+				<DropdownMenuContent align="end" side="right">
+					{onRename && (
+						<DropdownMenuItem
+							onClick={(e) => {
+								e.stopPropagation();
+								onRename();
+							}}
+						>
+							<PencilIcon className="mr-2 h-4 w-4" />
+							<span>{t("rename") || "Rename"}</span>
+						</DropdownMenuItem>
+					)}
+					{onArchive && (
+						<DropdownMenuItem
+							onClick={(e) => {
+								e.stopPropagation();
+								onArchive();
+							}}
+						>
+							{archived ? (
 									<>
 										<RotateCcwIcon className="mr-2 h-4 w-4" />
 										<span>{t("unarchive") || "Restore"}</span>
diff --git a/surfsense_web/components/layout/ui/sidebar/MobileSidebar.tsx b/surfsense_web/components/layout/ui/sidebar/MobileSidebar.tsx
index 85f907611..3ed2f9cca 100644
--- a/surfsense_web/components/layout/ui/sidebar/MobileSidebar.tsx
+++ b/surfsense_web/components/layout/ui/sidebar/MobileSidebar.tsx
@@ -24,6 +24,7 @@ interface MobileSidebarProps {
 	activeChatId?: number | null;
 	onNewChat: () => void;
 	onChatSelect: (chat: ChatItem) => void;
+	onChatRename?: (chat: ChatItem) => void;
 	onChatDelete?: (chat: ChatItem) => void;
 	onChatArchive?: (chat: ChatItem) => void;
 	onViewAllSharedChats?: () => void;
@@ -64,6 +65,7 @@ export function MobileSidebar({
 	activeChatId,
 	onNewChat,
 	onChatSelect,
+	onChatRename,
 	onChatDelete,
 	onChatArchive,
 	onViewAllSharedChats,
@@ -142,6 +144,7 @@ export function MobileSidebar({
 							onOpenChange(false);
 						}}
 						onChatSelect={handleChatSelect}
+						onChatRename={onChatRename}
 						onChatDelete={onChatDelete}
 						onChatArchive={onChatArchive}
 						onViewAllSharedChats={onViewAllSharedChats}
diff --git a/surfsense_web/components/layout/ui/sidebar/Sidebar.tsx b/surfsense_web/components/layout/ui/sidebar/Sidebar.tsx
index db04bf6dc..8763056ed 100644
--- a/surfsense_web/components/layout/ui/sidebar/Sidebar.tsx
+++ b/surfsense_web/components/layout/ui/sidebar/Sidebar.tsx
@@ -25,6 +25,7 @@ interface SidebarProps {
 	activeChatId?: number | null;
 	onNewChat: () => void;
 	onChatSelect: (chat: ChatItem) => void;
+	onChatRename?: (chat: ChatItem) => void;
 	onChatDelete?: (chat: ChatItem) => void;
 	onChatArchive?: (chat: ChatItem) => void;
 	onViewAllSharedChats?: () => void;
@@ -51,6 +52,7 @@ export function Sidebar({
 	activeChatId,
 	onNewChat,
 	onChatSelect,
+	onChatRename,
 	onChatDelete,
 	onChatArchive,
 	onViewAllSharedChats,
@@ -163,6 +165,7 @@ export function Sidebar({
 											isActive={chat.id === activeChatId}
 											archived={chat.archived}
 											onClick={() => onChatSelect(chat)}
+											onRename={() => onChatRename?.(chat)}
 											onArchive={() => onChatArchive?.(chat)}
 											onDelete={() => onChatDelete?.(chat)}
 										/>
@@ -215,6 +218,7 @@ export function Sidebar({
 											isActive={chat.id === activeChatId}
 											archived={chat.archived}
 											onClick={() => onChatSelect(chat)}
+											onRename={() => onChatRename?.(chat)}
 											onArchive={() => onChatArchive?.(chat)}
 											onDelete={() => onChatDelete?.(chat)}
 										/>
diff --git a/surfsense_web/messages/en.json b/surfsense_web/messages/en.json
index 75b186420..b44287af5 100644
--- a/surfsense_web/messages/en.json
+++ b/surfsense_web/messages/en.json
@@ -676,6 +676,13 @@
 		"unarchive": "Restore",
 		"chat_archived": "Chat archived",
 		"chat_unarchived": "Chat restored",
+		"chat_renamed": "Chat renamed",
+		"error_renaming_chat": "Failed to rename chat",
+		"rename": "Rename",
+		"rename_chat": "Rename Chat",
+		"rename_chat_description": "Enter a new name for this conversation.",
+		"chat_title_placeholder": "Chat title",
+		"renaming": "Renaming...",
 		"no_archived_chats": "No archived chats",
 		"error_archiving_chat": "Failed to archive chat",
 		"new_chat": "New chat",
diff --git a/surfsense_web/messages/zh.json b/surfsense_web/messages/zh.json
index 81121ef3e..cbcef3691 100644
--- a/surfsense_web/messages/zh.json
+++ b/surfsense_web/messages/zh.json
@@ -661,6 +661,13 @@
 		"unarchive": "恢复",
 		"chat_archived": "对话已归档",
 		"chat_unarchived": "对话已恢复",
+		"chat_renamed": "对话已重命名",
+		"error_renaming_chat": "重命名对话失败",
+		"rename": "重命名",
+		"rename_chat": "重命名对话",
+		"rename_chat_description": "为此对话输入新名称。",
+		"chat_title_placeholder": "对话标题",
+		"renaming": "重命名中...",
 		"no_archived_chats": "暂无已归档对话",
 		"error_archiving_chat": "归档对话失败",
 		"new_chat": "新对话",

From 07f89a426882d663ff9ac2b14a84850a7a7251b0 Mon Sep 17 00:00:00 2001
From: Eric Lammertsma <e.lammertsma@gmail.com>
Date: Tue, 3 Feb 2026 20:54:21 -0500
Subject: [PATCH 03/40] Show the chat title in the breadcrumbs instead of
 "Chat"

---
 .../new-chat/[[...chat_id]]/page.tsx             |  4 ++++
 .../components/dashboard-breadcrumb.tsx          | 16 ++++++++++++++--
 .../layout/providers/LayoutDataProvider.tsx      |  2 ++
 3 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx b/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx
index 4f99f8021..1a535539d 100644
--- a/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx
+++ b/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx
@@ -842,6 +842,10 @@ export default function NewChatPage() {
 												queryClient.invalidateQueries({
 													queryKey: ["threads", String(searchSpaceId)],
 												});
+												// Invalidate thread detail for breadcrumb update
+												queryClient.invalidateQueries({
+													queryKey: ["threads", String(searchSpaceId), "detail", String(titleData.threadId)],
+												});
 											}
 											break;
 										}
diff --git a/surfsense_web/components/dashboard-breadcrumb.tsx b/surfsense_web/components/dashboard-breadcrumb.tsx
index 96bd0ef30..5c6399ce0 100644
--- a/surfsense_web/components/dashboard-breadcrumb.tsx
+++ b/surfsense_web/components/dashboard-breadcrumb.tsx
@@ -14,6 +14,7 @@ import {
 } from "@/components/ui/breadcrumb";
 import { searchSpacesApiService } from "@/lib/apis/search-spaces-api.service";
 import { authenticatedFetch, getBearerToken } from "@/lib/auth-utils";
+import { getThreadFull } from "@/lib/chat/thread-persistence";
 import { cacheKeys } from "@/lib/query-client/cache-keys";
 
 interface BreadcrumbItemInterface {
@@ -34,6 +35,16 @@ export function DashboardBreadcrumb() {
 		enabled: !!searchSpaceId,
 	});
 
+	// Extract chat thread ID from pathname for chat pages
+	const chatThreadId = segments[2] === "new-chat" && segments[3] ? segments[3] : null;
+
+	// Fetch thread details when on a chat page with a thread ID
+	const { data: threadData } = useQuery({
+		queryKey: ["threads", searchSpaceId, "detail", chatThreadId],
+		queryFn: () => getThreadFull(Number(chatThreadId)),
+		enabled: !!chatThreadId && !!searchSpaceId,
+	});
+
 	// State to store document title for editor breadcrumb
 	const [documentTitle, setDocumentTitle] = useState<string | null>(null);
 
@@ -144,10 +155,11 @@ export function DashboardBreadcrumb() {
 					}
 
 					// Handle new-chat sub-sections (thread IDs)
-					// Don't show thread ID in breadcrumb - users identify chats by content, not by ID
+					// Show the chat title if available, otherwise fall back to "Chat"
 					if (section === "new-chat") {
+						const chatLabel = threadData?.title || t("chat") || "Chat";
 						breadcrumbs.push({
-							label: t("chat") || "Chat",
+							label: chatLabel,
 						});
 						return breadcrumbs;
 					}
diff --git a/surfsense_web/components/layout/providers/LayoutDataProvider.tsx b/surfsense_web/components/layout/providers/LayoutDataProvider.tsx
index 9a5c2ed40..4d8e2d23a 100644
--- a/surfsense_web/components/layout/providers/LayoutDataProvider.tsx
+++ b/surfsense_web/components/layout/providers/LayoutDataProvider.tsx
@@ -524,6 +524,8 @@ export function LayoutDataProvider({
 			queryClient.invalidateQueries({ queryKey: ["threads", searchSpaceId] });
 			queryClient.invalidateQueries({ queryKey: ["all-threads", searchSpaceId] });
 			queryClient.invalidateQueries({ queryKey: ["search-threads", searchSpaceId] });
+			// Invalidate thread detail for breadcrumb update
+			queryClient.invalidateQueries({ queryKey: ["threads", searchSpaceId, "detail", String(chatToRename.id)] });
 		} catch (error) {
 			console.error("Error renaming thread:", error);
 			toast.error(tSidebar("error_renaming_chat") || "Failed to rename chat");

From 318ad4a4ba2ff5a6e8edf4837846c649471e9606 Mon Sep 17 00:00:00 2001
From: Eric Lammertsma <e.lammertsma@gmail.com>
Date: Tue, 3 Feb 2026 20:59:50 -0500
Subject: [PATCH 04/40] Removed excessive logging around chat title generation

---
 .../app/tasks/chat/stream_new_chat.py         | 25 ++-----------------
 1 file changed, 2 insertions(+), 23 deletions(-)

diff --git a/surfsense_backend/app/tasks/chat/stream_new_chat.py b/surfsense_backend/app/tasks/chat/stream_new_chat.py
index 47e2a2293..a9751e5d1 100644
--- a/surfsense_backend/app/tasks/chat/stream_new_chat.py
+++ b/surfsense_backend/app/tasks/chat/stream_new_chat.py
@@ -1224,9 +1224,6 @@ async def stream_new_chat(
 
         # Only generate title on the first response (no prior assistant messages)
         if assistant_message_count == 0:
-            print(f"[stream_new_chat] First response - generating title for thread {chat_id}")
-            print(f"[stream_new_chat] Query length: {len(user_query)}, Response length: {len(accumulated_text)}")
-
             generated_title = None
             try:
                 # Generate title using the same LLM
@@ -1234,39 +1231,23 @@ async def stream_new_chat(
                 # Truncate inputs to avoid context length issues
                 truncated_query = user_query[:500]
                 truncated_response = accumulated_text[:1000]
-                print(f"[stream_new_chat] Calling LLM for title generation...")
                 title_result = await title_chain.ainvoke({
                     "user_query": truncated_query,
                     "assistant_response": truncated_response,
                 })
-                print(f"[stream_new_chat] LLM title result type: {type(title_result)}")
-                print(f"[stream_new_chat] LLM title result: {title_result}")
 
                 # Extract and clean the title
                 if title_result and hasattr(title_result, "content"):
                     raw_title = title_result.content.strip()
-                    print(f"[stream_new_chat] Raw title content: '{raw_title}' (len={len(raw_title)})")
-
-                    # Validate the title (1-6 words, reasonable length)
+                    # Validate the title (reasonable length)
                     if raw_title and len(raw_title) <= 100:
                         # Remove any quotes or extra formatting
                         generated_title = raw_title.strip('"\'')
-                        print(f"[stream_new_chat] After stripping quotes: '{generated_title}'")
-                    else:
-                        print(f"[stream_new_chat] Title validation failed: empty={not raw_title}, len={len(raw_title)}")
-                        generated_title = None
-                else:
-                    print(f"[stream_new_chat] No content attribute on result")
-            except Exception as title_error:
-                print(f"[stream_new_chat] Title generation failed: {title_error}")
-                import traceback
-                traceback.print_exc()
+            except Exception:
                 generated_title = None
 
             # Only update if LLM succeeded (keep truncated prompt title as fallback)
             if generated_title:
-                print(f"[stream_new_chat] Using LLM-generated title: '{generated_title}'")
-
                 # Fetch thread and update title
                 thread_result = await session.execute(
                     select(NewChatThread).filter(NewChatThread.id == chat_id)
@@ -1280,8 +1261,6 @@ async def stream_new_chat(
                     yield streaming_service.format_thread_title_update(
                         chat_id, generated_title
                     )
-            else:
-                print(f"[stream_new_chat] LLM title generation failed, keeping truncated prompt title")
 
         # Finish the step and message
         yield streaming_service.format_finish_step()

From 5bee04192c14f516e292f9978f1c85a30ebc9baf Mon Sep 17 00:00:00 2001
From: Eric Lammertsma <e.lammertsma@gmail.com>
Date: Tue, 3 Feb 2026 22:19:10 -0500
Subject: [PATCH 05/40] Overhauled pricing plans.

---
 surfsense_web/components/pricing.tsx          |  5 +-
 .../components/pricing/pricing-section.tsx    | 73 ++++++++++---------
 2 files changed, 42 insertions(+), 36 deletions(-)

diff --git a/surfsense_web/components/pricing.tsx b/surfsense_web/components/pricing.tsx
index 39757e91e..9d05e0262 100644
--- a/surfsense_web/components/pricing.tsx
+++ b/surfsense_web/components/pricing.tsx
@@ -17,6 +17,7 @@ interface PricingPlan {
 	price: string;
 	yearlyPrice: string;
 	period: string;
+	billingText?: string;
 	features: string[];
 	description: string;
 	buttonText: string;
@@ -35,7 +36,7 @@ export function Pricing({
 	title = "Simple, Transparent Pricing",
 	description = "Choose the plan that works for you\nAll plans include access to our SurfSense AI workspace and community support.",
 }: PricingProps) {
-	const [isMonthly, setIsMonthly] = useState(true);
+	const [isMonthly, setIsMonthly] = useState(false);
 	const isDesktop = useMediaQuery("(min-width: 768px)");
 	const switchRef = useRef<HTMLButtonElement>(null);
 
@@ -183,7 +184,7 @@ export function Pricing({
 							</div>
 
 							<p className="text-xs leading-5 text-muted-foreground">
-								{isNaN(Number(plan.price)) ? "" : isMonthly ? "billed monthly" : "billed annually"}
+								{plan.billingText ?? (isNaN(Number(plan.price)) ? "" : isMonthly ? "billed monthly" : "billed annually")}
 							</p>
 
 							<ul className="mt-5 gap-2 flex flex-col">
diff --git a/surfsense_web/components/pricing/pricing-section.tsx b/surfsense_web/components/pricing/pricing-section.tsx
index fdad0796a..5b3513e51 100644
--- a/surfsense_web/components/pricing/pricing-section.tsx
+++ b/surfsense_web/components/pricing/pricing-section.tsx
@@ -4,44 +4,47 @@ import { Pricing } from "@/components/pricing";
 
 const demoPlans = [
 	{
-		name: "COMMUNITY",
+		name: "FREE",
 		price: "0",
 		yearlyPrice: "0",
-		period: "forever",
+		period: "",
+		billingText: "Includes 30 day PRO trial",
 		features: [
-			"Community support",
-			"Supports 100+ LLMs",
-			"Supports OpenAI spec and LiteLLM",
-			"Supports local vLLM or Ollama setups",
-			"6000+ embedding models",
+			"Open source on GitHub",
+			"Upload and chat with up to 1,000 pages of content",
+			"Connects with 8 popular sources, like Drive and Notion.",
+			"Includes limited access to ChatGPT, Claude, and DeepSeek models",
+			"Supports 100+ more LLMs, including Gemini, Llama and many more.",
 			"50+ File extensions supported.",
-			"Podcasts support with local TTS providers.",
-			"Connects with 15+ external sources, like Drive and Notion.",
+			"Generate podcasts in seconds.",
 			"Cross-Browser Extension for dynamic webpages including authenticated content",
-			"Role-based access control (RBAC)",
-			"Collaboration and team features",
+			"Community support on Discord",
 		],
-		description: "Open source version with powerful features",
-		buttonText: "Dive In",
-		href: "/docs",
+		description: "Powerful features with some limitations",
+		buttonText: "Get Started",
+		href: "/",
 		isPopular: false,
 	},
 	{
-		name: "CLOUD",
-		price: "0",
-		yearlyPrice: "0",
-		period: "in beta",
+		name: "PRO",
+		price: "10",
+		yearlyPrice: "10",
+		period: "user / month",
+		billingText: "billed annually",
 		features: [
-			"Everything in Community",
-			"Email support",
-			"Get started in seconds",
-			"Instant access to new features",
-			"Easy access from anywhere",
-			"Remote team management and collaboration",
+			"Everything in Free",
+			"Upload and chat with up to 20,000 pages of content",
+			"Connects with 15+ external sources, like Slack and Airtable.",
+			"Includes extended access to ChatGPT, Claude, and DeepSeek models",
+			"Collaboration and commenting features",
+			"Centralized billing",
+			"Shared BYOK (Bring Your Own Key)",
+			"Team and role management",
+			"Priority support",
 		],
-		description: "Instant access for individuals and teams",
-		buttonText: "Get Started",
-		href: "/",
+		description: "The AIknowledge base for individuals and teams",
+		buttonText: "Upgrade",
+		href: "/contact",
 		isPopular: true,
 	},
 	{
@@ -50,17 +53,19 @@ const demoPlans = [
 		yearlyPrice: "Contact Us",
 		period: "",
 		features: [
-			"Everything in Community",
-			"Priority support",
-			"White-glove setup and deployment",
-			"Monthly managed updates and maintenance",
+			"Everything in Pro",
+			"Connect and chat with virtually unlimited pages of content",
+			"Limit models and/or providers",
 			"On-prem or VPC deployment",
 			"Audit logs and compliance",
 			"SSO, OIDC & SAML",
-			"SLA guarantee",
-			"Uptime guarantee on VPC",
+			"Role-based access control (RBAC)",
+			"White-glove setup and deployment",
+			"Monthly managed updates and maintenance",
+			"SLA commitments",
+			"Dedicated support",
 		],
-		description: "Professional, customized setup for large organizations",
+		description: "Customized setup for large organizations",
 		buttonText: "Contact Sales",
 		href: "/contact",
 		isPopular: false,

From 40a304bd0b15ab834e0ba2523796413727e941eb Mon Sep 17 00:00:00 2001
From: Eric Lammertsma <e.lammertsma@gmail.com>
Date: Tue, 3 Feb 2026 22:49:38 -0500
Subject: [PATCH 06/40] Updated enterprise pricing plan t.

---
 surfsense_web/components/pricing/pricing-section.tsx | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/surfsense_web/components/pricing/pricing-section.tsx b/surfsense_web/components/pricing/pricing-section.tsx
index 5b3513e51..c850852d8 100644
--- a/surfsense_web/components/pricing/pricing-section.tsx
+++ b/surfsense_web/components/pricing/pricing-section.tsx
@@ -49,9 +49,10 @@ const demoPlans = [
 	},
 	{
 		name: "ENTERPRISE",
-		price: "Contact Us",
-		yearlyPrice: "Contact Us",
-		period: "",
+		price: "1000",
+		yearlyPrice: "1000",
+		period: "month",
+		billingText: "billed annually",
 		features: [
 			"Everything in Pro",
 			"Connect and chat with virtually unlimited pages of content",

From fb333fdd541cb37e340759a0630e1e86532dd706 Mon Sep 17 00:00:00 2001
From: CREDO23 <thierrybakera12@gmail.com>
Date: Tue, 3 Feb 2026 21:39:18 +0200
Subject: [PATCH 07/40] Add CommentReplyNotificationHandler for comment thread
 notifications

---
 .../app/services/notification_service.py      | 93 +++++++++++++++++++
 1 file changed, 93 insertions(+)

diff --git a/surfsense_backend/app/services/notification_service.py b/surfsense_backend/app/services/notification_service.py
index 1788d05e1..a759f3536 100644
--- a/surfsense_backend/app/services/notification_service.py
+++ b/surfsense_backend/app/services/notification_service.py
@@ -861,6 +861,98 @@ class MentionNotificationHandler(BaseNotificationHandler):
             raise
 
 
+class CommentReplyNotificationHandler(BaseNotificationHandler):
+    """Handler for comment reply notifications."""
+
+    def __init__(self):
+        super().__init__("comment_reply")
+
+    async def find_notification_by_reply(
+        self,
+        session: AsyncSession,
+        reply_id: int,
+        user_id: UUID,
+    ) -> Notification | None:
+        query = select(Notification).where(
+            Notification.type == self.notification_type,
+            Notification.user_id == user_id,
+            Notification.notification_metadata["reply_id"].astext == str(reply_id),
+        )
+        result = await session.execute(query)
+        return result.scalar_one_or_none()
+
+    async def notify_comment_reply(
+        self,
+        session: AsyncSession,
+        user_id: UUID,
+        reply_id: int,
+        parent_comment_id: int,
+        message_id: int,
+        thread_id: int,
+        thread_title: str,
+        author_id: str,
+        author_name: str,
+        author_avatar_url: str | None,
+        author_email: str,
+        content_preview: str,
+        search_space_id: int,
+    ) -> Notification:
+        existing = await self.find_notification_by_reply(session, reply_id, user_id)
+        if existing:
+            logger.info(
+                f"Notification already exists for reply {reply_id} to user {user_id}"
+            )
+            return existing
+
+        title = f"{author_name} replied in a thread"
+        message = content_preview[:100] + ("..." if len(content_preview) > 100 else "")
+
+        metadata = {
+            "reply_id": reply_id,
+            "parent_comment_id": parent_comment_id,
+            "message_id": message_id,
+            "thread_id": thread_id,
+            "thread_title": thread_title,
+            "author_id": author_id,
+            "author_name": author_name,
+            "author_avatar_url": author_avatar_url,
+            "author_email": author_email,
+            "content_preview": content_preview[:200],
+        }
+
+        try:
+            notification = Notification(
+                user_id=user_id,
+                search_space_id=search_space_id,
+                type=self.notification_type,
+                title=title,
+                message=message,
+                notification_metadata=metadata,
+            )
+            session.add(notification)
+            await session.commit()
+            await session.refresh(notification)
+            logger.info(
+                f"Created comment_reply notification {notification.id} for user {user_id}"
+            )
+            return notification
+        except Exception as e:
+            await session.rollback()
+            if (
+                "duplicate key" in str(e).lower()
+                or "unique constraint" in str(e).lower()
+            ):
+                logger.warning(
+                    f"Duplicate notification for reply {reply_id} to user {user_id}"
+                )
+                existing = await self.find_notification_by_reply(
+                    session, reply_id, user_id
+                )
+                if existing:
+                    return existing
+            raise
+
+
 class PageLimitNotificationHandler(BaseNotificationHandler):
     """Handler for page limit exceeded notifications."""
 
@@ -959,6 +1051,7 @@ class NotificationService:
     connector_indexing = ConnectorIndexingNotificationHandler()
     document_processing = DocumentProcessingNotificationHandler()
     mention = MentionNotificationHandler()
+    comment_reply = CommentReplyNotificationHandler()
     page_limit = PageLimitNotificationHandler()
 
     @staticmethod

From bf91d0c3d21887f909225f50f9b7c55a0523d830 Mon Sep 17 00:00:00 2001
From: CREDO23 <thierrybakera12@gmail.com>
Date: Tue, 3 Feb 2026 21:57:07 +0200
Subject: [PATCH 08/40] Add get_comment_thread_participants helper function

---
 .../app/services/chat_comments_service.py     | 33 ++++++++++++++++++-
 1 file changed, 32 insertions(+), 1 deletion(-)

diff --git a/surfsense_backend/app/services/chat_comments_service.py b/surfsense_backend/app/services/chat_comments_service.py
index dc3b51238..7350dc12f 100644
--- a/surfsense_backend/app/services/chat_comments_service.py
+++ b/surfsense_backend/app/services/chat_comments_service.py
@@ -5,7 +5,7 @@ Service layer for chat comments and mentions.
 from uuid import UUID
 
 from fastapi import HTTPException
-from sqlalchemy import delete, select
+from sqlalchemy import delete, or_, select
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import selectinload
 
@@ -103,6 +103,37 @@ async def process_mentions(
     return mentions_map
 
 
+async def get_comment_thread_participants(
+    session: AsyncSession,
+    parent_comment_id: int,
+    exclude_user_ids: set[UUID],
+) -> list[UUID]:
+    """
+    Get all unique authors in a comment thread (parent + replies), excluding specified users.
+
+    Args:
+        session: Database session
+        parent_comment_id: ID of the parent comment
+        exclude_user_ids: Set of user IDs to exclude (e.g., replier, mentioned users)
+
+    Returns:
+        List of user UUIDs who have participated in the thread
+    """
+    query = select(ChatComment.author_id).where(
+        or_(
+            ChatComment.id == parent_comment_id,
+            ChatComment.parent_id == parent_comment_id,
+        ),
+        ChatComment.author_id.isnot(None),
+    )
+
+    if exclude_user_ids:
+        query = query.where(ChatComment.author_id.notin_(exclude_user_ids))
+
+    result = await session.execute(query.distinct())
+    return [row[0] for row in result.fetchall()]
+
+
 async def get_comments_for_message(
     session: AsyncSession,
     message_id: int,

From cf512153df2fdfccaca577ceeff9ea6cc4e161ed Mon Sep 17 00:00:00 2001
From: CREDO23 <thierrybakera12@gmail.com>
Date: Tue, 3 Feb 2026 22:00:26 +0200
Subject: [PATCH 09/40] Notify thread participants on new reply

---
 .../app/services/chat_comments_service.py     | 22 +++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/surfsense_backend/app/services/chat_comments_service.py b/surfsense_backend/app/services/chat_comments_service.py
index 7350dc12f..209606359 100644
--- a/surfsense_backend/app/services/chat_comments_service.py
+++ b/surfsense_backend/app/services/chat_comments_service.py
@@ -467,6 +467,28 @@ async def create_reply(
             search_space_id=search_space_id,
         )
 
+    # Notify thread participants (excluding replier and mentioned users)
+    exclude_ids = {user.id, *mentions_map.keys()}
+    participants = await get_comment_thread_participants(
+        session, comment_id, exclude_ids
+    )
+    for participant_id in participants:
+        await NotificationService.comment_reply.notify_comment_reply(
+            session=session,
+            user_id=participant_id,
+            reply_id=reply.id,
+            parent_comment_id=comment_id,
+            message_id=parent_comment.message_id,
+            thread_id=thread.id,
+            thread_title=thread.title or "Untitled thread",
+            author_id=str(user.id),
+            author_name=author_name,
+            author_avatar_url=user.avatar_url,
+            author_email=user.email,
+            content_preview=content_preview[:200],
+            search_space_id=search_space_id,
+        )
+
     author = AuthorResponse(
         id=user.id,
         display_name=user.display_name,

From 1c8ec7bbdcd13dc0dfd7158f557a6ee39d33e69c Mon Sep 17 00:00:00 2001
From: CREDO23 <thierrybakera12@gmail.com>
Date: Wed, 4 Feb 2026 14:45:28 +0200
Subject: [PATCH 10/40] Add comment_reply notification type to frontend types

---
 surfsense_web/contracts/types/inbox.types.ts | 31 ++++++++++++++++++++
 1 file changed, 31 insertions(+)

diff --git a/surfsense_web/contracts/types/inbox.types.ts b/surfsense_web/contracts/types/inbox.types.ts
index 8e4b9ae86..ebf1889a1 100644
--- a/surfsense_web/contracts/types/inbox.types.ts
+++ b/surfsense_web/contracts/types/inbox.types.ts
@@ -10,6 +10,7 @@ export const inboxItemTypeEnum = z.enum([
 	"connector_deletion",
 	"document_processing",
 	"new_mention",
+	"comment_reply",
 	"page_limit_exceeded",
 ]);
 
@@ -101,6 +102,19 @@ export const newMentionMetadata = z.object({
 	content_preview: z.string(),
 });
 
+export const commentReplyMetadata = z.object({
+	reply_id: z.number(),
+	parent_comment_id: z.number(),
+	message_id: z.number(),
+	thread_id: z.number(),
+	thread_title: z.string(),
+	author_id: z.string(),
+	author_name: z.string(),
+	author_avatar_url: z.string().nullable().optional(),
+	author_email: z.string().optional(),
+	content_preview: z.string(),
+});
+
 /**
  * Page limit exceeded metadata schema
  */
@@ -125,6 +139,7 @@ export const inboxItemMetadata = z.union([
 	connectorDeletionMetadata,
 	documentProcessingMetadata,
 	newMentionMetadata,
+	commentReplyMetadata,
 	pageLimitExceededMetadata,
 	baseInboxItemMetadata,
 ]);
@@ -168,6 +183,11 @@ export const newMentionInboxItem = inboxItem.extend({
 	metadata: newMentionMetadata,
 });
 
+export const commentReplyInboxItem = inboxItem.extend({
+	type: z.literal("comment_reply"),
+	metadata: commentReplyMetadata,
+});
+
 export const pageLimitExceededInboxItem = inboxItem.extend({
 	type: z.literal("page_limit_exceeded"),
 	metadata: pageLimitExceededMetadata,
@@ -278,6 +298,10 @@ export function isNewMentionMetadata(metadata: unknown): metadata is NewMentionM
 	return newMentionMetadata.safeParse(metadata).success;
 }
 
+export function isCommentReplyMetadata(metadata: unknown): metadata is CommentReplyMetadata {
+	return commentReplyMetadata.safeParse(metadata).success;
+}
+
 /**
  * Type guard for PageLimitExceededMetadata
  */
@@ -298,6 +322,7 @@ export function parseInboxItemMetadata(
 	| ConnectorDeletionMetadata
 	| DocumentProcessingMetadata
 	| NewMentionMetadata
+	| CommentReplyMetadata
 	| PageLimitExceededMetadata
 	| null {
 	switch (type) {
@@ -317,6 +342,10 @@ export function parseInboxItemMetadata(
 			const result = newMentionMetadata.safeParse(metadata);
 			return result.success ? result.data : null;
 		}
+		case "comment_reply": {
+			const result = commentReplyMetadata.safeParse(metadata);
+			return result.success ? result.data : null;
+		}
 		case "page_limit_exceeded": {
 			const result = pageLimitExceededMetadata.safeParse(metadata);
 			return result.success ? result.data : null;
@@ -338,6 +367,7 @@ export type ConnectorIndexingMetadata = z.infer<typeof connectorIndexingMetadata
 export type ConnectorDeletionMetadata = z.infer<typeof connectorDeletionMetadata>;
 export type DocumentProcessingMetadata = z.infer<typeof documentProcessingMetadata>;
 export type NewMentionMetadata = z.infer<typeof newMentionMetadata>;
+export type CommentReplyMetadata = z.infer<typeof commentReplyMetadata>;
 export type PageLimitExceededMetadata = z.infer<typeof pageLimitExceededMetadata>;
 export type InboxItemMetadata = z.infer<typeof inboxItemMetadata>;
 export type InboxItem = z.infer<typeof inboxItem>;
@@ -345,6 +375,7 @@ export type ConnectorIndexingInboxItem = z.infer<typeof connectorIndexingInboxIt
 export type ConnectorDeletionInboxItem = z.infer<typeof connectorDeletionInboxItem>;
 export type DocumentProcessingInboxItem = z.infer<typeof documentProcessingInboxItem>;
 export type NewMentionInboxItem = z.infer<typeof newMentionInboxItem>;
+export type CommentReplyInboxItem = z.infer<typeof commentReplyInboxItem>;
 export type PageLimitExceededInboxItem = z.infer<typeof pageLimitExceededInboxItem>;
 
 // API Request/Response types

From 21a4c254583d2cb604451d226e17c66492b71fe3 Mon Sep 17 00:00:00 2001
From: CREDO23 <thierrybakera12@gmail.com>
Date: Wed, 4 Feb 2026 14:47:03 +0200
Subject: [PATCH 11/40] Handle comment_reply in inbox sidebar (icon + click
 routing)

---
 .../layout/ui/sidebar/InboxSidebar.tsx        | 46 +++++++++++++------
 1 file changed, 33 insertions(+), 13 deletions(-)

diff --git a/surfsense_web/components/layout/ui/sidebar/InboxSidebar.tsx b/surfsense_web/components/layout/ui/sidebar/InboxSidebar.tsx
index 9ef49c0d8..c1f5fcc99 100644
--- a/surfsense_web/components/layout/ui/sidebar/InboxSidebar.tsx
+++ b/surfsense_web/components/layout/ui/sidebar/InboxSidebar.tsx
@@ -46,6 +46,7 @@ import { Tabs, TabsList, TabsTrigger } from "@/components/ui/tabs";
 import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip";
 import { getConnectorIcon } from "@/contracts/enums/connectorIcons";
 import {
+	isCommentReplyMetadata,
 	isConnectorIndexingMetadata,
 	isNewMentionMetadata,
 	isPageLimitExceededMetadata,
@@ -347,19 +348,33 @@ export function InboxSidebar({
 			}
 
 			if (item.type === "new_mention") {
-				// Use type guard for safe metadata access
 				if (isNewMentionMetadata(item.metadata)) {
 					const searchSpaceId = item.search_space_id;
 					const threadId = item.metadata.thread_id;
 					const commentId = item.metadata.comment_id;
 
 					if (searchSpaceId && threadId) {
-						// Pre-set target comment ID before navigation
-						// This also ensures comments panel is not collapsed
 						if (commentId) {
 							setTargetCommentId(commentId);
 						}
+						const url = commentId
+							? `/dashboard/${searchSpaceId}/new-chat/${threadId}?commentId=${commentId}`
+							: `/dashboard/${searchSpaceId}/new-chat/${threadId}`;
+						onOpenChange(false);
+						onCloseMobileSidebar?.();
+						router.push(url);
+					}
+				}
+			} else if (item.type === "comment_reply") {
+				if (isCommentReplyMetadata(item.metadata)) {
+					const searchSpaceId = item.search_space_id;
+					const threadId = item.metadata.thread_id;
+					const commentId = item.metadata.parent_comment_id;
 
+					if (searchSpaceId && threadId) {
+						if (commentId) {
+							setTargetCommentId(commentId);
+						}
 						const url = commentId
 							? `/dashboard/${searchSpaceId}/new-chat/${threadId}?commentId=${commentId}`
 							: `/dashboard/${searchSpaceId}/new-chat/${threadId}`;
@@ -411,24 +426,29 @@ export function InboxSidebar({
 	};
 
 	const getStatusIcon = (item: InboxItem) => {
-		// For mentions, show the author's avatar with initials fallback
-		if (item.type === "new_mention") {
-			// Use type guard for safe metadata access
-			if (isNewMentionMetadata(item.metadata)) {
-				const authorName = item.metadata.author_name;
-				const avatarUrl = item.metadata.author_avatar_url;
-				const authorEmail = item.metadata.author_email;
+		// For mentions and comment replies, show the author's avatar
+		if (item.type === "new_mention" || item.type === "comment_reply") {
+			const metadata =
+				item.type === "new_mention"
+					? isNewMentionMetadata(item.metadata)
+						? item.metadata
+						: null
+					: isCommentReplyMetadata(item.metadata)
+						? item.metadata
+						: null;
 
+			if (metadata) {
 				return (
 					<Avatar className="h-8 w-8">
-						{avatarUrl && <AvatarImage src={avatarUrl} alt={authorName || "User"} />}
+						{metadata.author_avatar_url && (
+							<AvatarImage src={metadata.author_avatar_url} alt={metadata.author_name || "User"} />
+						)}
 						<AvatarFallback className="text-[10px] bg-primary/10 text-primary">
-							{getInitials(authorName, authorEmail)}
+							{getInitials(metadata.author_name, metadata.author_email)}
 						</AvatarFallback>
 					</Avatar>
 				);
 			}
-			// Fallback for invalid metadata
 			return (
 				<Avatar className="h-8 w-8">
 					<AvatarFallback className="text-[10px] bg-primary/10 text-primary">

From d5b75956c36cca3098b1b69aae747a27a9ded01e Mon Sep 17 00:00:00 2001
From: CREDO23 <thierrybakera12@gmail.com>
Date: Wed, 4 Feb 2026 15:13:28 +0200
Subject: [PATCH 12/40] Include comment_reply in status tab filter

---
 surfsense_web/components/layout/ui/sidebar/InboxSidebar.tsx | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/surfsense_web/components/layout/ui/sidebar/InboxSidebar.tsx b/surfsense_web/components/layout/ui/sidebar/InboxSidebar.tsx
index c1f5fcc99..a46ad5478 100644
--- a/surfsense_web/components/layout/ui/sidebar/InboxSidebar.tsx
+++ b/surfsense_web/components/layout/ui/sidebar/InboxSidebar.tsx
@@ -238,7 +238,7 @@ export function InboxSidebar({
 	const currentDataSource = activeTab === "mentions" ? mentions : status;
 	const { loading, loadingMore = false, hasMore = false, loadMore } = currentDataSource;
 
-	// Status tab includes: connector indexing, document processing, page limit exceeded, connector deletion
+	// Status tab includes: connector indexing, document processing, page limit exceeded, connector deletion, comment replies
 	// Filter to only show status notification types
 	const statusItems = useMemo(
 		() =>
@@ -247,7 +247,8 @@ export function InboxSidebar({
 					item.type === "connector_indexing" ||
 					item.type === "document_processing" ||
 					item.type === "page_limit_exceeded" ||
-					item.type === "connector_deletion"
+					item.type === "connector_deletion" ||
+					item.type === "comment_reply"
 			),
 		[status.items]
 	);

From 6d1a3636830d934b76c908e24e755770bece26b6 Mon Sep 17 00:00:00 2001
From: CREDO23 <thierrybakera12@gmail.com>
Date: Wed, 4 Feb 2026 15:18:25 +0200
Subject: [PATCH 13/40] Navigate to reply instead of parent comment

---
 .../components/layout/ui/sidebar/InboxSidebar.tsx      | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/surfsense_web/components/layout/ui/sidebar/InboxSidebar.tsx b/surfsense_web/components/layout/ui/sidebar/InboxSidebar.tsx
index a46ad5478..b63e16317 100644
--- a/surfsense_web/components/layout/ui/sidebar/InboxSidebar.tsx
+++ b/surfsense_web/components/layout/ui/sidebar/InboxSidebar.tsx
@@ -370,14 +370,14 @@ export function InboxSidebar({
 				if (isCommentReplyMetadata(item.metadata)) {
 					const searchSpaceId = item.search_space_id;
 					const threadId = item.metadata.thread_id;
-					const commentId = item.metadata.parent_comment_id;
+					const replyId = item.metadata.reply_id;
 
 					if (searchSpaceId && threadId) {
-						if (commentId) {
-							setTargetCommentId(commentId);
+						if (replyId) {
+							setTargetCommentId(replyId);
 						}
-						const url = commentId
-							? `/dashboard/${searchSpaceId}/new-chat/${threadId}?commentId=${commentId}`
+						const url = replyId
+							? `/dashboard/${searchSpaceId}/new-chat/${threadId}?commentId=${replyId}`
 							: `/dashboard/${searchSpaceId}/new-chat/${threadId}`;
 						onOpenChange(false);
 						onCloseMobileSidebar?.();

From f610e42b9c588e60843e91090540b4abaf6712d3 Mon Sep 17 00:00:00 2001
From: CREDO23 <thierrybakera12@gmail.com>
Date: Wed, 4 Feb 2026 16:37:44 +0200
Subject: [PATCH 14/40] fix: prevent duplicate notifications and fix inbox
 count

---
 .../app/services/chat_comments_service.py     |   7 +-
 .../layout/providers/LayoutDataProvider.tsx   |  11 +-
 .../components/theme/theme-toggle.tsx         | 787 +++++++++---------
 3 files changed, 396 insertions(+), 409 deletions(-)

diff --git a/surfsense_backend/app/services/chat_comments_service.py b/surfsense_backend/app/services/chat_comments_service.py
index 209606359..c9ca920f6 100644
--- a/surfsense_backend/app/services/chat_comments_service.py
+++ b/surfsense_backend/app/services/chat_comments_service.py
@@ -128,7 +128,7 @@ async def get_comment_thread_participants(
     )
 
     if exclude_user_ids:
-        query = query.where(ChatComment.author_id.notin_(exclude_user_ids))
+        query = query.where(ChatComment.author_id.notin_(list(exclude_user_ids)))
 
     result = await session.execute(query.distinct())
     return [row[0] for row in result.fetchall()]
@@ -468,11 +468,14 @@ async def create_reply(
         )
 
     # Notify thread participants (excluding replier and mentioned users)
-    exclude_ids = {user.id, *mentions_map.keys()}
+    mentioned_user_ids = set(mentions_map.keys())
+    exclude_ids = {user.id} | mentioned_user_ids
     participants = await get_comment_thread_participants(
         session, comment_id, exclude_ids
     )
     for participant_id in participants:
+        if participant_id in mentioned_user_ids:
+            continue
         await NotificationService.comment_reply.notify_comment_reply(
             session=session,
             user_id=participant_id,
diff --git a/surfsense_web/components/layout/providers/LayoutDataProvider.tsx b/surfsense_web/components/layout/providers/LayoutDataProvider.tsx
index 2f71adad9..feb34940a 100644
--- a/surfsense_web/components/layout/providers/LayoutDataProvider.tsx
+++ b/surfsense_web/components/layout/providers/LayoutDataProvider.tsx
@@ -109,7 +109,6 @@ export function LayoutDataProvider({
 	// This ensures each tab has independent pagination and data loading
 	const userId = user?.id ? String(user.id) : null;
 
-	// Mentions: Only fetch "new_mention" type notifications
 	const {
 		inboxItems: mentionItems,
 		unreadCount: mentionUnreadCount,
@@ -121,11 +120,9 @@ export function LayoutDataProvider({
 		markAllAsRead: markAllMentionsAsRead,
 	} = useInbox(userId, Number(searchSpaceId) || null, "new_mention");
 
-	// Status: Fetch all types (will be filtered client-side to status types)
-	// We pass null to get all, then InboxSidebar filters to status types
 	const {
 		inboxItems: statusItems,
-		unreadCount: statusUnreadCount,
+		unreadCount: allUnreadCount,
 		loading: statusLoading,
 		loadingMore: statusLoadingMore,
 		hasMore: statusHasMore,
@@ -134,8 +131,8 @@ export function LayoutDataProvider({
 		markAllAsRead: markAllStatusAsRead,
 	} = useInbox(userId, Number(searchSpaceId) || null, null);
 
-	// Combined unread count for nav badge (mentions take priority for visibility)
-	const totalUnreadCount = mentionUnreadCount + statusUnreadCount;
+	const totalUnreadCount = allUnreadCount;
+	const statusOnlyUnreadCount = Math.max(0, allUnreadCount - mentionUnreadCount);
 
 	// Track seen notification IDs to detect new page_limit_exceeded notifications
 	const seenPageLimitNotifications = useRef<Set<number>>(new Set());
@@ -561,7 +558,7 @@ export function LayoutDataProvider({
 					},
 					status: {
 						items: statusItems,
-						unreadCount: statusUnreadCount,
+						unreadCount: statusOnlyUnreadCount,
 						loading: statusLoading,
 						loadingMore: statusLoadingMore,
 						hasMore: statusHasMore,
diff --git a/surfsense_web/components/theme/theme-toggle.tsx b/surfsense_web/components/theme/theme-toggle.tsx
index 382d11087..b9b23656b 100644
--- a/surfsense_web/components/theme/theme-toggle.tsx
+++ b/surfsense_web/components/theme/theme-toggle.tsx
@@ -8,172 +8,167 @@ import { cn } from "@/lib/utils";
 
 // ///////////////////////////////////////////////////////////////////////////
 // Types
-export type AnimationVariant =
-  | "circle"
-  | "rectangle"
-  | "gif"
-  | "polygon"
-  | "circle-blur";
+export type AnimationVariant = "circle" | "rectangle" | "gif" | "polygon" | "circle-blur";
 export type AnimationStart =
-  | "top-left"
-  | "top-right"
-  | "bottom-left"
-  | "bottom-right"
-  | "center"
-  | "top-center"
-  | "bottom-center"
-  | "bottom-up"
-  | "top-down"
-  | "left-right"
-  | "right-left";
+	| "top-left"
+	| "top-right"
+	| "bottom-left"
+	| "bottom-right"
+	| "center"
+	| "top-center"
+	| "bottom-center"
+	| "bottom-up"
+	| "top-down"
+	| "left-right"
+	| "right-left";
 
 interface Animation {
-  name: string;
-  css: string;
+	name: string;
+	css: string;
 }
 
 // ///////////////////////////////////////////////////////////////////////////
 // Helper functions
 
 const getPositionCoords = (position: AnimationStart) => {
-  switch (position) {
-    case "top-left":
-      return { cx: "0", cy: "0" };
-    case "top-right":
-      return { cx: "40", cy: "0" };
-    case "bottom-left":
-      return { cx: "0", cy: "40" };
-    case "bottom-right":
-      return { cx: "40", cy: "40" };
-    case "top-center":
-      return { cx: "20", cy: "0" };
-    case "bottom-center":
-      return { cx: "20", cy: "40" };
-    case "bottom-up":
-    case "top-down":
-    case "left-right":
-    case "right-left":
-      return { cx: "20", cy: "20" };
-  }
+	switch (position) {
+		case "top-left":
+			return { cx: "0", cy: "0" };
+		case "top-right":
+			return { cx: "40", cy: "0" };
+		case "bottom-left":
+			return { cx: "0", cy: "40" };
+		case "bottom-right":
+			return { cx: "40", cy: "40" };
+		case "top-center":
+			return { cx: "20", cy: "0" };
+		case "bottom-center":
+			return { cx: "20", cy: "40" };
+		case "bottom-up":
+		case "top-down":
+		case "left-right":
+		case "right-left":
+			return { cx: "20", cy: "20" };
+	}
 };
 
 const generateSVG = (variant: AnimationVariant, start: AnimationStart) => {
-  if (variant === "circle-blur") {
-    if (start === "center") {
-      return `data:image/svg+xml,<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 40 40"><defs><filter id="blur"><feGaussianBlur stdDeviation="2"/></filter></defs><circle cx="20" cy="20" r="18" fill="white" filter="url(%23blur)"/></svg>`;
-    }
-    const positionCoords = getPositionCoords(start);
-    if (!positionCoords) {
-      throw new Error(`Invalid start position: ${start}`);
-    }
-    const { cx, cy } = positionCoords;
-    return `data:image/svg+xml,<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 40 40"><defs><filter id="blur"><feGaussianBlur stdDeviation="2"/></filter></defs><circle cx="${cx}" cy="${cy}" r="18" fill="white" filter="url(%23blur)"/></svg>`;
-  }
+	if (variant === "circle-blur") {
+		if (start === "center") {
+			return `data:image/svg+xml,<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 40 40"><defs><filter id="blur"><feGaussianBlur stdDeviation="2"/></filter></defs><circle cx="20" cy="20" r="18" fill="white" filter="url(%23blur)"/></svg>`;
+		}
+		const positionCoords = getPositionCoords(start);
+		if (!positionCoords) {
+			throw new Error(`Invalid start position: ${start}`);
+		}
+		const { cx, cy } = positionCoords;
+		return `data:image/svg+xml,<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 40 40"><defs><filter id="blur"><feGaussianBlur stdDeviation="2"/></filter></defs><circle cx="${cx}" cy="${cy}" r="18" fill="white" filter="url(%23blur)"/></svg>`;
+	}
 
-  if (start === "center") return;
+	if (start === "center") return;
 
-  if (variant === "rectangle") return "";
+	if (variant === "rectangle") return "";
 
-  const positionCoords = getPositionCoords(start);
-  if (!positionCoords) {
-    throw new Error(`Invalid start position: ${start}`);
-  }
-  const { cx, cy } = positionCoords;
+	const positionCoords = getPositionCoords(start);
+	if (!positionCoords) {
+		throw new Error(`Invalid start position: ${start}`);
+	}
+	const { cx, cy } = positionCoords;
 
-  if (variant === "circle") {
-    return `data:image/svg+xml,<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 40 40"><circle cx="${cx}" cy="${cy}" r="20" fill="white"/></svg>`;
-  }
+	if (variant === "circle") {
+		return `data:image/svg+xml,<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 40 40"><circle cx="${cx}" cy="${cy}" r="20" fill="white"/></svg>`;
+	}
 
-  return "";
+	return "";
 };
 
 const getTransformOrigin = (start: AnimationStart) => {
-  switch (start) {
-    case "top-left":
-      return "top left";
-    case "top-right":
-      return "top right";
-    case "bottom-left":
-      return "bottom left";
-    case "bottom-right":
-      return "bottom right";
-    case "top-center":
-      return "top center";
-    case "bottom-center":
-      return "bottom center";
-    case "bottom-up":
-    case "top-down":
-    case "left-right":
-    case "right-left":
-      return "center";
-  }
+	switch (start) {
+		case "top-left":
+			return "top left";
+		case "top-right":
+			return "top right";
+		case "bottom-left":
+			return "bottom left";
+		case "bottom-right":
+			return "bottom right";
+		case "top-center":
+			return "top center";
+		case "bottom-center":
+			return "bottom center";
+		case "bottom-up":
+		case "top-down":
+		case "left-right":
+		case "right-left":
+			return "center";
+	}
 };
 
 export const createAnimation = (
-  variant: AnimationVariant,
-  start: AnimationStart = "center",
-  blur = false,
-  url?: string,
+	variant: AnimationVariant,
+	start: AnimationStart = "center",
+	blur = false,
+	url?: string
 ): Animation => {
-  const svg = generateSVG(variant, start);
-  const transformOrigin = getTransformOrigin(start);
+	const svg = generateSVG(variant, start);
+	const transformOrigin = getTransformOrigin(start);
 
-  if (variant === "rectangle") {
-    const getClipPath = (direction: AnimationStart) => {
-      switch (direction) {
-        case "bottom-up":
-          return {
-            from: "polygon(0% 100%, 100% 100%, 100% 100%, 0% 100%)",
-            to: "polygon(0% 0%, 100% 0%, 100% 100%, 0% 100%)",
-          };
-        case "top-down":
-          return {
-            from: "polygon(0% 0%, 100% 0%, 100% 0%, 0% 0%)",
-            to: "polygon(0% 0%, 100% 0%, 100% 100%, 0% 100%)",
-          };
-        case "left-right":
-          return {
-            from: "polygon(0% 0%, 0% 0%, 0% 100%, 0% 100%)",
-            to: "polygon(0% 0%, 100% 0%, 100% 100%, 0% 100%)",
-          };
-        case "right-left":
-          return {
-            from: "polygon(100% 0%, 100% 0%, 100% 100%, 100% 100%)",
-            to: "polygon(0% 0%, 100% 0%, 100% 100%, 0% 100%)",
-          };
-        case "top-left":
-          return {
-            from: "polygon(0% 0%, 0% 0%, 0% 0%, 0% 0%)",
-            to: "polygon(0% 0%, 100% 0%, 100% 100%, 0% 100%)",
-          };
-        case "top-right":
-          return {
-            from: "polygon(100% 0%, 100% 0%, 100% 0%, 100% 0%)",
-            to: "polygon(0% 0%, 100% 0%, 100% 100%, 0% 100%)",
-          };
-        case "bottom-left":
-          return {
-            from: "polygon(0% 100%, 0% 100%, 0% 100%, 0% 100%)",
-            to: "polygon(0% 0%, 100% 0%, 100% 100%, 0% 100%)",
-          };
-        case "bottom-right":
-          return {
-            from: "polygon(100% 100%, 100% 100%, 100% 100%, 100% 100%)",
-            to: "polygon(0% 0%, 100% 0%, 100% 100%, 0% 100%)",
-          };
-        default:
-          return {
-            from: "polygon(0% 100%, 100% 100%, 100% 100%, 0% 100%)",
-            to: "polygon(0% 0%, 100% 0%, 100% 100%, 0% 100%)",
-          };
-      }
-    };
+	if (variant === "rectangle") {
+		const getClipPath = (direction: AnimationStart) => {
+			switch (direction) {
+				case "bottom-up":
+					return {
+						from: "polygon(0% 100%, 100% 100%, 100% 100%, 0% 100%)",
+						to: "polygon(0% 0%, 100% 0%, 100% 100%, 0% 100%)",
+					};
+				case "top-down":
+					return {
+						from: "polygon(0% 0%, 100% 0%, 100% 0%, 0% 0%)",
+						to: "polygon(0% 0%, 100% 0%, 100% 100%, 0% 100%)",
+					};
+				case "left-right":
+					return {
+						from: "polygon(0% 0%, 0% 0%, 0% 100%, 0% 100%)",
+						to: "polygon(0% 0%, 100% 0%, 100% 100%, 0% 100%)",
+					};
+				case "right-left":
+					return {
+						from: "polygon(100% 0%, 100% 0%, 100% 100%, 100% 100%)",
+						to: "polygon(0% 0%, 100% 0%, 100% 100%, 0% 100%)",
+					};
+				case "top-left":
+					return {
+						from: "polygon(0% 0%, 0% 0%, 0% 0%, 0% 0%)",
+						to: "polygon(0% 0%, 100% 0%, 100% 100%, 0% 100%)",
+					};
+				case "top-right":
+					return {
+						from: "polygon(100% 0%, 100% 0%, 100% 0%, 100% 0%)",
+						to: "polygon(0% 0%, 100% 0%, 100% 100%, 0% 100%)",
+					};
+				case "bottom-left":
+					return {
+						from: "polygon(0% 100%, 0% 100%, 0% 100%, 0% 100%)",
+						to: "polygon(0% 0%, 100% 0%, 100% 100%, 0% 100%)",
+					};
+				case "bottom-right":
+					return {
+						from: "polygon(100% 100%, 100% 100%, 100% 100%, 100% 100%)",
+						to: "polygon(0% 0%, 100% 0%, 100% 100%, 0% 100%)",
+					};
+				default:
+					return {
+						from: "polygon(0% 100%, 100% 100%, 100% 100%, 0% 100%)",
+						to: "polygon(0% 0%, 100% 0%, 100% 100%, 0% 100%)",
+					};
+			}
+		};
 
-    const clipPath = getClipPath(start);
+		const clipPath = getClipPath(start);
 
-    return {
-      name: `${variant}-${start}${blur ? "-blur" : ""}`,
-      css: `
+		return {
+			name: `${variant}-${start}${blur ? "-blur" : ""}`,
+			css: `
        ::view-transition-group(root) {
         animation-duration: 0.7s;
         animation-timing-function: var(--expo-out);
@@ -218,12 +213,12 @@ export const createAnimation = (
         }
       }
       `,
-    };
-  }
-  if (variant === "circle" && start == "center") {
-    return {
-      name: `${variant}-${start}${blur ? "-blur" : ""}`,
-      css: `
+		};
+	}
+	if (variant === "circle" && start == "center") {
+		return {
+			name: `${variant}-${start}${blur ? "-blur" : ""}`,
+			css: `
        ::view-transition-group(root) {
         animation-duration: 0.7s;
         animation-timing-function: var(--expo-out);
@@ -268,12 +263,12 @@ export const createAnimation = (
         }
       }
       `,
-    };
-  }
-  if (variant === "gif") {
-    return {
-      name: `${variant}-${start}`,
-      css: `
+		};
+	}
+	if (variant === "gif") {
+		return {
+			name: `${variant}-${start}`,
+			css: `
       ::view-transition-group(root) {
   animation-timing-function: var(--expo-in);
 }
@@ -302,14 +297,14 @@ export const createAnimation = (
     mask-size: 2000vmax;
   }
 }`,
-    };
-  }
+		};
+	}
 
-  if (variant === "circle-blur") {
-    if (start === "center") {
-      return {
-        name: `${variant}-${start}`,
-        css: `
+	if (variant === "circle-blur") {
+		if (start === "center") {
+			return {
+				name: `${variant}-${start}`,
+				css: `
         ::view-transition-group(root) {
           animation-timing-function: var(--expo-out);
         }
@@ -334,12 +329,12 @@ export const createAnimation = (
           }
         }
         `,
-      };
-    }
+			};
+		}
 
-    return {
-      name: `${variant}-${start}`,
-      css: `
+		return {
+			name: `${variant}-${start}`,
+			css: `
       ::view-transition-group(root) {
         animation-timing-function: var(--expo-out);
       }
@@ -364,41 +359,41 @@ export const createAnimation = (
         }
       }
       `,
-    };
-  }
+		};
+	}
 
-  if (variant === "polygon") {
-    const getPolygonClipPaths = (position: AnimationStart) => {
-      switch (position) {
-        case "top-left":
-          return {
-            darkFrom: "polygon(50% -71%, -50% 71%, -50% 71%, 50% -71%)",
-            darkTo: "polygon(50% -71%, -50% 71%, 50% 171%, 171% 50%)",
-            lightFrom: "polygon(171% 50%, 50% 171%, 50% 171%, 171% 50%)",
-            lightTo: "polygon(171% 50%, 50% 171%, -50% 71%, 50% -71%)",
-          };
-        case "top-right":
-          return {
-            darkFrom: "polygon(150% -71%, 250% 71%, 250% 71%, 150% -71%)",
-            darkTo: "polygon(150% -71%, 250% 71%, 50% 171%, -71% 50%)",
-            lightFrom: "polygon(-71% 50%, 50% 171%, 50% 171%, -71% 50%)",
-            lightTo: "polygon(-71% 50%, 50% 171%, 250% 71%, 150% -71%)",
-          };
-        default:
-          return {
-            darkFrom: "polygon(50% -71%, -50% 71%, -50% 71%, 50% -71%)",
-            darkTo: "polygon(50% -71%, -50% 71%, 50% 171%, 171% 50%)",
-            lightFrom: "polygon(171% 50%, 50% 171%, 50% 171%, 171% 50%)",
-            lightTo: "polygon(171% 50%, 50% 171%, -50% 71%, 50% -71%)",
-          };
-      }
-    };
+	if (variant === "polygon") {
+		const getPolygonClipPaths = (position: AnimationStart) => {
+			switch (position) {
+				case "top-left":
+					return {
+						darkFrom: "polygon(50% -71%, -50% 71%, -50% 71%, 50% -71%)",
+						darkTo: "polygon(50% -71%, -50% 71%, 50% 171%, 171% 50%)",
+						lightFrom: "polygon(171% 50%, 50% 171%, 50% 171%, 171% 50%)",
+						lightTo: "polygon(171% 50%, 50% 171%, -50% 71%, 50% -71%)",
+					};
+				case "top-right":
+					return {
+						darkFrom: "polygon(150% -71%, 250% 71%, 250% 71%, 150% -71%)",
+						darkTo: "polygon(150% -71%, 250% 71%, 50% 171%, -71% 50%)",
+						lightFrom: "polygon(-71% 50%, 50% 171%, 50% 171%, -71% 50%)",
+						lightTo: "polygon(-71% 50%, 50% 171%, 250% 71%, 150% -71%)",
+					};
+				default:
+					return {
+						darkFrom: "polygon(50% -71%, -50% 71%, -50% 71%, 50% -71%)",
+						darkTo: "polygon(50% -71%, -50% 71%, 50% 171%, 171% 50%)",
+						lightFrom: "polygon(171% 50%, 50% 171%, 50% 171%, 171% 50%)",
+						lightTo: "polygon(171% 50%, 50% 171%, -50% 71%, 50% -71%)",
+					};
+			}
+		};
 
-    const clipPaths = getPolygonClipPaths(start);
+		const clipPaths = getPolygonClipPaths(start);
 
-    return {
-      name: `${variant}-${start}${blur ? "-blur" : ""}`,
-      css: `
+		return {
+			name: `${variant}-${start}${blur ? "-blur" : ""}`,
+			css: `
       ::view-transition-group(root) {
         animation-duration: 0.7s;
         animation-timing-function: var(--expo-out);
@@ -443,35 +438,35 @@ export const createAnimation = (
         }
       }
       `,
-    };
-  }
+		};
+	}
 
-  // Handle circle variants with start positions using clip-path
-  if (variant === "circle" && start !== "center") {
-    const getClipPathPosition = (position: AnimationStart) => {
-      switch (position) {
-        case "top-left":
-          return "0% 0%";
-        case "top-right":
-          return "100% 0%";
-        case "bottom-left":
-          return "0% 100%";
-        case "bottom-right":
-          return "100% 100%";
-        case "top-center":
-          return "50% 0%";
-        case "bottom-center":
-          return "50% 100%";
-        default:
-          return "50% 50%";
-      }
-    };
+	// Handle circle variants with start positions using clip-path
+	if (variant === "circle" && start !== "center") {
+		const getClipPathPosition = (position: AnimationStart) => {
+			switch (position) {
+				case "top-left":
+					return "0% 0%";
+				case "top-right":
+					return "100% 0%";
+				case "bottom-left":
+					return "0% 100%";
+				case "bottom-right":
+					return "100% 100%";
+				case "top-center":
+					return "50% 0%";
+				case "bottom-center":
+					return "50% 100%";
+				default:
+					return "50% 50%";
+			}
+		};
 
-    const clipPosition = getClipPathPosition(start);
+		const clipPosition = getClipPathPosition(start);
 
-    return {
-      name: `${variant}-${start}${blur ? "-blur" : ""}`,
-      css: `
+		return {
+			name: `${variant}-${start}${blur ? "-blur" : ""}`,
+			css: `
        ::view-transition-group(root) {
         animation-duration: 1s;
         animation-timing-function: var(--expo-out);
@@ -516,12 +511,12 @@ export const createAnimation = (
         }
       }
       `,
-    };
-  }
+		};
+	}
 
-  return {
-    name: `${variant}-${start}${blur ? "-blur" : ""}`,
-    css: `
+	return {
+		name: `${variant}-${start}${blur ? "-blur" : ""}`,
+		css: `
       ::view-transition-group(root) {
         animation-timing-function: var(--expo-in);
       }
@@ -549,237 +544,229 @@ export const createAnimation = (
         }
       }
     `,
-  };
+	};
 };
 
 // ///////////////////////////////////////////////////////////////////////////
 // Custom hook for theme toggle functionality
 export const useThemeToggle = ({
-  variant = "circle",
-  start = "center",
-  blur = false,
-  gifUrl = "",
+	variant = "circle",
+	start = "center",
+	blur = false,
+	gifUrl = "",
 }: {
-  variant?: AnimationVariant;
-  start?: AnimationStart;
-  blur?: boolean;
-  gifUrl?: string;
+	variant?: AnimationVariant;
+	start?: AnimationStart;
+	blur?: boolean;
+	gifUrl?: string;
 } = {}) => {
-  const { theme, setTheme, resolvedTheme } = useTheme();
+	const { theme, setTheme, resolvedTheme } = useTheme();
 
-  const [isDark, setIsDark] = useState(false);
+	const [isDark, setIsDark] = useState(false);
 
-  // Sync isDark state with resolved theme after hydration
-  useEffect(() => {
-    setIsDark(resolvedTheme === "dark");
-  }, [resolvedTheme]);
+	// Sync isDark state with resolved theme after hydration
+	useEffect(() => {
+		setIsDark(resolvedTheme === "dark");
+	}, [resolvedTheme]);
 
-  const styleId = "theme-transition-styles";
+	const styleId = "theme-transition-styles";
 
-  const updateStyles = useCallback((css: string) => {
-    if (typeof window === "undefined") return;
+	const updateStyles = useCallback((css: string) => {
+		if (typeof window === "undefined") return;
 
-    let styleElement = document.getElementById(styleId) as HTMLStyleElement;
+		let styleElement = document.getElementById(styleId) as HTMLStyleElement;
 
-    if (!styleElement) {
-      styleElement = document.createElement("style");
-      styleElement.id = styleId;
-      document.head.appendChild(styleElement);
-    }
+		if (!styleElement) {
+			styleElement = document.createElement("style");
+			styleElement.id = styleId;
+			document.head.appendChild(styleElement);
+		}
 
-    styleElement.textContent = css;
-  }, []);
+		styleElement.textContent = css;
+	}, []);
 
-  const toggleTheme = useCallback(() => {
-    setIsDark(!isDark);
+	const toggleTheme = useCallback(() => {
+		setIsDark(!isDark);
 
-    const animation = createAnimation(variant, start, blur, gifUrl);
+		const animation = createAnimation(variant, start, blur, gifUrl);
 
-    updateStyles(animation.css);
+		updateStyles(animation.css);
 
-    if (typeof window === "undefined") return;
+		if (typeof window === "undefined") return;
 
-    const switchTheme = () => {
-      setTheme(theme === "light" ? "dark" : "light");
-    };
+		const switchTheme = () => {
+			setTheme(theme === "light" ? "dark" : "light");
+		};
 
-    if (!document.startViewTransition) {
-      switchTheme();
-      return;
-    }
+		if (!document.startViewTransition) {
+			switchTheme();
+			return;
+		}
 
-    document.startViewTransition(switchTheme);
-  }, [theme, setTheme, variant, start, blur, gifUrl, updateStyles, isDark]);
+		document.startViewTransition(switchTheme);
+	}, [theme, setTheme, variant, start, blur, gifUrl, updateStyles, isDark]);
 
-  const setCrazyLightTheme = useCallback(() => {
-    setIsDark(false);
+	const setCrazyLightTheme = useCallback(() => {
+		setIsDark(false);
 
-    const animation = createAnimation(variant, start, blur, gifUrl);
+		const animation = createAnimation(variant, start, blur, gifUrl);
 
-    updateStyles(animation.css);
+		updateStyles(animation.css);
 
-    if (typeof window === "undefined") return;
+		if (typeof window === "undefined") return;
 
-    const switchTheme = () => {
-      setTheme("light");
-    };
+		const switchTheme = () => {
+			setTheme("light");
+		};
 
-    if (!document.startViewTransition) {
-      switchTheme();
-      return;
-    }
+		if (!document.startViewTransition) {
+			switchTheme();
+			return;
+		}
 
-    document.startViewTransition(switchTheme);
-  }, [setTheme, variant, start, blur, gifUrl, updateStyles]);
+		document.startViewTransition(switchTheme);
+	}, [setTheme, variant, start, blur, gifUrl, updateStyles]);
 
-  const setCrazyDarkTheme = useCallback(() => {
-    setIsDark(true);
+	const setCrazyDarkTheme = useCallback(() => {
+		setIsDark(true);
 
-    const animation = createAnimation(variant, start, blur, gifUrl);
+		const animation = createAnimation(variant, start, blur, gifUrl);
 
-    updateStyles(animation.css);
+		updateStyles(animation.css);
 
-    if (typeof window === "undefined") return;
+		if (typeof window === "undefined") return;
 
-    const switchTheme = () => {
-      setTheme("dark");
-    };
+		const switchTheme = () => {
+			setTheme("dark");
+		};
 
-    if (!document.startViewTransition) {
-      switchTheme();
-      return;
-    }
+		if (!document.startViewTransition) {
+			switchTheme();
+			return;
+		}
 
-    document.startViewTransition(switchTheme);
-  }, [setTheme, variant, start, blur, gifUrl, updateStyles]);
+		document.startViewTransition(switchTheme);
+	}, [setTheme, variant, start, blur, gifUrl, updateStyles]);
 
-  const setCrazySystemTheme = useCallback(() => {
-    if (typeof window === "undefined") return;
+	const setCrazySystemTheme = useCallback(() => {
+		if (typeof window === "undefined") return;
 
-    const prefersDark = window.matchMedia(
-      "(prefers-color-scheme: dark)",
-    ).matches;
-    setIsDark(prefersDark);
+		const prefersDark = window.matchMedia("(prefers-color-scheme: dark)").matches;
+		setIsDark(prefersDark);
 
-    const animation = createAnimation(variant, start, blur, gifUrl);
+		const animation = createAnimation(variant, start, blur, gifUrl);
 
-    updateStyles(animation.css);
+		updateStyles(animation.css);
 
-    const switchTheme = () => {
-      setTheme("system");
-    };
+		const switchTheme = () => {
+			setTheme("system");
+		};
 
-    if (!document.startViewTransition) {
-      switchTheme();
-      return;
-    }
+		if (!document.startViewTransition) {
+			switchTheme();
+			return;
+		}
 
-    document.startViewTransition(switchTheme);
-  }, [setTheme, variant, start, blur, gifUrl, updateStyles]);
+		document.startViewTransition(switchTheme);
+	}, [setTheme, variant, start, blur, gifUrl, updateStyles]);
 
-  return {
-    isDark,
-    setIsDark,
-    toggleTheme,
-    setCrazyLightTheme,
-    setCrazyDarkTheme,
-    setCrazySystemTheme,
-  };
+	return {
+		isDark,
+		setIsDark,
+		toggleTheme,
+		setCrazyLightTheme,
+		setCrazyDarkTheme,
+		setCrazySystemTheme,
+	};
 };
 
 // ///////////////////////////////////////////////////////////////////////////
 // Theme Toggle Button Component (Sun/Moon Style)
 
 export const ThemeToggleButton = ({
-  className = "",
-  variant = "circle",
-  start = "center",
-  blur = false,
-  gifUrl = "",
+	className = "",
+	variant = "circle",
+	start = "center",
+	blur = false,
+	gifUrl = "",
 }: {
-  className?: string;
-  variant?: AnimationVariant;
-  start?: AnimationStart;
-  blur?: boolean;
-  gifUrl?: string;
+	className?: string;
+	variant?: AnimationVariant;
+	start?: AnimationStart;
+	blur?: boolean;
+	gifUrl?: string;
 }) => {
-  const { isDark, toggleTheme } = useThemeToggle({
-    variant,
-    start,
-    blur,
-    gifUrl,
-  });
-  const clipId = useId();
-  const clipPathId = `theme-toggle-clip-${clipId}`;
+	const { isDark, toggleTheme } = useThemeToggle({
+		variant,
+		start,
+		blur,
+		gifUrl,
+	});
+	const clipId = useId();
+	const clipPathId = `theme-toggle-clip-${clipId}`;
 
-  return (
-    <button
-      type="button"
-      className={cn(
-        "size-10 cursor-pointer rounded-full p-2 transition-all duration-300 active:scale-95 bg-transparent",
-        isDark ? "text-white" : "text-black",
-        className,
-      )}
-      onClick={toggleTheme}
-      aria-label="Toggle theme"
-    >
-      <span className="sr-only">Toggle theme</span>
-      <svg
-        xmlns="http://www.w3.org/2000/svg"
-        aria-hidden="true"
-        fill="currentColor"
-        strokeLinecap="round"
-        viewBox="0 0 32 32"
-      >
-        <clipPath id={clipPathId}>
-          <motion.path
-            animate={{ y: isDark ? 10 : 0, x: isDark ? -12 : 0 }}
-            transition={{ ease: "easeInOut", duration: 0.35 }}
-            d="M0-5h30a1 1 0 0 0 9 13v24H0Z"
-          />
-        </clipPath>
-        <g clipPath={`url(#${clipPathId})`}>
-          <motion.circle
-            animate={{ r: isDark ? 10 : 8 }}
-            transition={{ ease: "easeInOut", duration: 0.35 }}
-            cx="16"
-            cy="16"
-          />
-          <motion.g
-            animate={{
-              rotate: isDark ? -100 : 0,
-              scale: isDark ? 0.5 : 1,
-              opacity: isDark ? 0 : 1,
-            }}
-            transition={{ ease: "easeInOut", duration: 0.35 }}
-            stroke="currentColor"
-            strokeWidth="1.5"
-          >
-            <path d="M16 5.5v-4" />
-            <path d="M16 30.5v-4" />
-            <path d="M1.5 16h4" />
-            <path d="M26.5 16h4" />
-            <path d="m23.4 8.6 2.8-2.8" />
-            <path d="m5.7 26.3 2.9-2.9" />
-            <path d="m5.8 5.8 2.8 2.8" />
-            <path d="m23.4 23.4 2.9 2.9" />
-          </motion.g>
-        </g>
-      </svg>
-    </button>
-  );
+	return (
+		<button
+			type="button"
+			className={cn(
+				"size-10 cursor-pointer rounded-full p-2 transition-all duration-300 active:scale-95 bg-transparent",
+				isDark ? "text-white" : "text-black",
+				className
+			)}
+			onClick={toggleTheme}
+			aria-label="Toggle theme"
+		>
+			<span className="sr-only">Toggle theme</span>
+			<svg
+				xmlns="http://www.w3.org/2000/svg"
+				aria-hidden="true"
+				fill="currentColor"
+				strokeLinecap="round"
+				viewBox="0 0 32 32"
+			>
+				<clipPath id={clipPathId}>
+					<motion.path
+						animate={{ y: isDark ? 10 : 0, x: isDark ? -12 : 0 }}
+						transition={{ ease: "easeInOut", duration: 0.35 }}
+						d="M0-5h30a1 1 0 0 0 9 13v24H0Z"
+					/>
+				</clipPath>
+				<g clipPath={`url(#${clipPathId})`}>
+					<motion.circle
+						animate={{ r: isDark ? 10 : 8 }}
+						transition={{ ease: "easeInOut", duration: 0.35 }}
+						cx="16"
+						cy="16"
+					/>
+					<motion.g
+						animate={{
+							rotate: isDark ? -100 : 0,
+							scale: isDark ? 0.5 : 1,
+							opacity: isDark ? 0 : 1,
+						}}
+						transition={{ ease: "easeInOut", duration: 0.35 }}
+						stroke="currentColor"
+						strokeWidth="1.5"
+					>
+						<path d="M16 5.5v-4" />
+						<path d="M16 30.5v-4" />
+						<path d="M1.5 16h4" />
+						<path d="M26.5 16h4" />
+						<path d="m23.4 8.6 2.8-2.8" />
+						<path d="m5.7 26.3 2.9-2.9" />
+						<path d="m5.8 5.8 2.8 2.8" />
+						<path d="m23.4 23.4 2.9 2.9" />
+					</motion.g>
+				</g>
+			</svg>
+		</button>
+	);
 };
 
 // ///////////////////////////////////////////////////////////////////////////
 // Backwards compatible export (alias for ThemeToggleButton with default settings)
 export function ThemeTogglerComponent() {
-  return (
-    <ThemeToggleButton
-      variant="circle"
-      start="top-right"
-      className="size-8"
-    />
-  );
+	return <ThemeToggleButton variant="circle" start="top-right" className="size-8" />;
 }
 
 /**

From 17eab845d05517d9e5e011f18c5a3d5f7695b08b Mon Sep 17 00:00:00 2001
From: CREDO23 <thierrybakera12@gmail.com>
Date: Wed, 4 Feb 2026 17:17:33 +0200
Subject: [PATCH 15/40] feat: rename Mentions tab to Comments with speech
 bubble icon

---
 .../layout/ui/sidebar/InboxSidebar.tsx        | 58 +++++++++++--------
 surfsense_web/messages/en.json                |  3 +
 surfsense_web/messages/zh.json                |  3 +
 3 files changed, 40 insertions(+), 24 deletions(-)

diff --git a/surfsense_web/components/layout/ui/sidebar/InboxSidebar.tsx b/surfsense_web/components/layout/ui/sidebar/InboxSidebar.tsx
index b63e16317..f313dd6f9 100644
--- a/surfsense_web/components/layout/ui/sidebar/InboxSidebar.tsx
+++ b/surfsense_web/components/layout/ui/sidebar/InboxSidebar.tsx
@@ -4,7 +4,6 @@ import { useAtom } from "jotai";
 import {
 	AlertCircle,
 	AlertTriangle,
-	AtSign,
 	BellDot,
 	Check,
 	CheckCheck,
@@ -15,6 +14,7 @@ import {
 	Inbox,
 	LayoutGrid,
 	ListFilter,
+	MessageSquare,
 	Search,
 	X,
 } from "lucide-react";
@@ -134,7 +134,7 @@ function getConnectorTypeDisplayName(connectorType: string): string {
 	);
 }
 
-type InboxTab = "mentions" | "status";
+type InboxTab = "comments" | "status";
 type InboxFilter = "all" | "unread";
 
 // Tab-specific data source with independent pagination
@@ -187,7 +187,7 @@ export function InboxSidebar({
 	const [, setTargetCommentId] = useAtom(setTargetCommentIdAtom);
 
 	const [searchQuery, setSearchQuery] = useState("");
-	const [activeTab, setActiveTab] = useState<InboxTab>("mentions");
+	const [activeTab, setActiveTab] = useState<InboxTab>("comments");
 	const [activeFilter, setActiveFilter] = useState<InboxFilter>("all");
 	const [selectedConnector, setSelectedConnector] = useState<string | null>(null);
 	const [mounted, setMounted] = useState(false);
@@ -234,12 +234,17 @@ export function InboxSidebar({
 		}
 	}, [activeTab]);
 
-	// Get current tab's data source - each tab has independent data and pagination
-	const currentDataSource = activeTab === "mentions" ? mentions : status;
-	const { loading, loadingMore = false, hasMore = false, loadMore } = currentDataSource;
+	// Both tabs now derive items from status (all types), so use status for pagination
+	const { loading, loadingMore = false, hasMore = false, loadMore } = status;
 
-	// Status tab includes: connector indexing, document processing, page limit exceeded, connector deletion, comment replies
-	// Filter to only show status notification types
+	// Comments tab: mentions and comment replies
+	const commentsItems = useMemo(
+		() =>
+			status.items.filter((item) => item.type === "new_mention" || item.type === "comment_reply"),
+		[status.items]
+	);
+
+	// Status tab: connector indexing, document processing, page limit exceeded, connector deletion
 	const statusItems = useMemo(
 		() =>
 			status.items.filter(
@@ -247,8 +252,7 @@ export function InboxSidebar({
 					item.type === "connector_indexing" ||
 					item.type === "document_processing" ||
 					item.type === "page_limit_exceeded" ||
-					item.type === "connector_deletion" ||
-					item.type === "comment_reply"
+					item.type === "connector_deletion"
 			),
 		[status.items]
 	);
@@ -272,8 +276,8 @@ export function InboxSidebar({
 		}));
 	}, [statusItems]);
 
-	// Get items for current tab - mentions use their source directly, status uses filtered items
-	const displayItems = activeTab === "mentions" ? mentions.items : statusItems;
+	// Get items for current tab
+	const displayItems = activeTab === "comments" ? commentsItems : statusItems;
 
 	// Filter items based on filter type, connector filter, and search query
 	const filteredItems = useMemo(() => {
@@ -336,9 +340,15 @@ export function InboxSidebar({
 		return () => observer.disconnect();
 	}, [loadMore, hasMore, loadingMore, open, searchQuery]);
 
-	// Use unread counts from data sources (more accurate than client-side counting)
-	const unreadMentionsCount = mentions.unreadCount;
-	const unreadStatusCount = status.unreadCount;
+	// Unread counts derived from filtered items
+	const unreadCommentsCount = useMemo(
+		() => commentsItems.filter((item) => !item.read).length,
+		[commentsItems]
+	);
+	const unreadStatusCount = useMemo(
+		() => statusItems.filter((item) => !item.read).length,
+		[statusItems]
+	);
 
 	const handleItemClick = useCallback(
 		async (item: InboxItem) => {
@@ -502,10 +512,10 @@ export function InboxSidebar({
 	};
 
 	const getEmptyStateMessage = () => {
-		if (activeTab === "mentions") {
+		if (activeTab === "comments") {
 			return {
-				title: t("no_mentions") || "No mentions",
-				hint: t("no_mentions_hint") || "You'll see mentions from others here",
+				title: t("no_comments") || "No comments",
+				hint: t("no_comments_hint") || "You'll see mentions and replies here",
 			};
 		}
 		return {
@@ -844,14 +854,14 @@ export function InboxSidebar({
 			>
 				<TabsList className="w-full h-auto p-0 bg-transparent rounded-none border-b">
 					<TabsTrigger
-						value="mentions"
+						value="comments"
 						className="flex-1 rounded-none border-b-2 border-transparent px-1 py-2 text-xs font-medium data-[state=active]:border-primary data-[state=active]:bg-transparent data-[state=active]:shadow-none"
 					>
 						<span className="w-full inline-flex items-center justify-center gap-1.5 px-3 py-1.5 rounded-lg hover:bg-muted transition-colors">
-							<AtSign className="h-4 w-4" />
-							<span>{t("mentions") || "Mentions"}</span>
+							<MessageSquare className="h-4 w-4" />
+							<span>{t("comments") || "Comments"}</span>
 							<span className="inline-flex items-center justify-center min-w-5 h-5 px-1.5 rounded-full bg-primary/20 text-muted-foreground text-xs font-medium">
-								{formatInboxCount(unreadMentionsCount)}
+								{formatInboxCount(unreadCommentsCount)}
 							</span>
 						</span>
 					</TabsTrigger>
@@ -953,8 +963,8 @@ export function InboxSidebar({
 					</div>
 				) : (
 					<div className="text-center py-8">
-						{activeTab === "mentions" ? (
-							<AtSign className="h-12 w-12 mx-auto text-muted-foreground mb-3" />
+						{activeTab === "comments" ? (
+							<MessageSquare className="h-12 w-12 mx-auto text-muted-foreground mb-3" />
 						) : (
 							<History className="h-12 w-12 mx-auto text-muted-foreground mb-3" />
 						)}
diff --git a/surfsense_web/messages/en.json b/surfsense_web/messages/en.json
index 75b186420..0646bdb52 100644
--- a/surfsense_web/messages/en.json
+++ b/surfsense_web/messages/en.json
@@ -698,10 +698,13 @@
 		"mark_all_read": "Mark all as read",
 		"mark_as_read": "Mark as read",
 		"mentions": "Mentions",
+		"comments": "Comments",
 		"status": "Status",
 		"no_results_found": "No results found",
 		"no_mentions": "No mentions",
 		"no_mentions_hint": "You'll see mentions from others here",
+		"no_comments": "No comments",
+		"no_comments_hint": "You'll see mentions and replies here",
 		"no_status_updates": "No status updates",
 		"no_status_updates_hint": "Document and connector updates will appear here",
 		"filter": "Filter",
diff --git a/surfsense_web/messages/zh.json b/surfsense_web/messages/zh.json
index 81121ef3e..67077c348 100644
--- a/surfsense_web/messages/zh.json
+++ b/surfsense_web/messages/zh.json
@@ -683,10 +683,13 @@
 		"mark_all_read": "全部标记为已读",
 		"mark_as_read": "标记为已读",
 		"mentions": "提及",
+		"comments": "评论",
 		"status": "状态",
 		"no_results_found": "未找到结果",
 		"no_mentions": "没有提及",
 		"no_mentions_hint": "您会在这里看到他人的提及",
+		"no_comments": "没有评论",
+		"no_comments_hint": "您会在这里看到提及和回复",
 		"no_status_updates": "没有状态更新",
 		"no_status_updates_hint": "文档和连接器更新将显示在这里",
 		"filter": "筛选",

From 2711563e8be328b39fc58eca3157b5c8048948e6 Mon Sep 17 00:00:00 2001
From: Eric Lammertsma <e.lammertsma@gmail.com>
Date: Wed, 4 Feb 2026 11:03:31 -0500
Subject: [PATCH 16/40] Refined pricing plan descriptions by updating content
 limits and adding "planned" to future features.

---
 surfsense_web/components/pricing/pricing-section.tsx | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/surfsense_web/components/pricing/pricing-section.tsx b/surfsense_web/components/pricing/pricing-section.tsx
index c850852d8..445969ec1 100644
--- a/surfsense_web/components/pricing/pricing-section.tsx
+++ b/surfsense_web/components/pricing/pricing-section.tsx
@@ -11,7 +11,7 @@ const demoPlans = [
 		billingText: "Includes 30 day PRO trial",
 		features: [
 			"Open source on GitHub",
-			"Upload and chat with up to 1,000 pages of content",
+			"Upload and chat with 1,000+ pages of content",
 			"Connects with 8 popular sources, like Drive and Notion.",
 			"Includes limited access to ChatGPT, Claude, and DeepSeek models",
 			"Supports 100+ more LLMs, including Gemini, Llama and many more.",
@@ -33,13 +33,13 @@ const demoPlans = [
 		billingText: "billed annually",
 		features: [
 			"Everything in Free",
-			"Upload and chat with up to 20,000 pages of content",
+			"Upload and chat with up to 20,000+ pages of content",
 			"Connects with 15+ external sources, like Slack and Airtable.",
 			"Includes extended access to ChatGPT, Claude, and DeepSeek models",
 			"Collaboration and commenting features",
-			"Centralized billing",
 			"Shared BYOK (Bring Your Own Key)",
 			"Team and role management",
+			"Planned: Centralized billing",
 			"Priority support",
 		],
 		description: "The AIknowledge base for individuals and teams",
@@ -58,9 +58,9 @@ const demoPlans = [
 			"Connect and chat with virtually unlimited pages of content",
 			"Limit models and/or providers",
 			"On-prem or VPC deployment",
-			"Audit logs and compliance",
-			"SSO, OIDC & SAML",
-			"Role-based access control (RBAC)",
+			"Planned: Audit logs and compliance",
+			"Planned: SSO, OIDC & SAML",
+			"Planned: Role-based access control (RBAC)",
 			"White-glove setup and deployment",
 			"Monthly managed updates and maintenance",
 			"SLA commitments",

From f5aa520743e25d3b600b4af310165eaa09d99696 Mon Sep 17 00:00:00 2001
From: CREDO23 <thierrybakera12@gmail.com>
Date: Wed, 4 Feb 2026 18:04:53 +0200
Subject: [PATCH 17/40] Remove branch picker and restrict edit/reload to last
 message

---
 .../assistant-ui/assistant-message.tsx        | 22 ++++++-----
 .../components/assistant-ui/branch-picker.tsx | 32 ----------------
 .../components/assistant-ui/user-message.tsx  | 38 +++++++++++++++----
 3 files changed, 42 insertions(+), 50 deletions(-)
 delete mode 100644 surfsense_web/components/assistant-ui/branch-picker.tsx

diff --git a/surfsense_web/components/assistant-ui/assistant-message.tsx b/surfsense_web/components/assistant-ui/assistant-message.tsx
index 4fd2446c3..5cdd287de 100644
--- a/surfsense_web/components/assistant-ui/assistant-message.tsx
+++ b/surfsense_web/components/assistant-ui/assistant-message.tsx
@@ -4,20 +4,19 @@ import {
 	ErrorPrimitive,
 	MessagePrimitive,
 	useAssistantState,
+	useMessage,
 } from "@assistant-ui/react";
-import { useAtom, useAtomValue, useSetAtom } from "jotai";
+import { useAtom, useAtomValue } from "jotai";
 import { CheckIcon, CopyIcon, DownloadIcon, MessageSquare, RefreshCwIcon } from "lucide-react";
 import type { FC } from "react";
 import { useContext, useEffect, useMemo, useRef, useState } from "react";
 import {
 	addingCommentToMessageIdAtom,
-	clearTargetCommentIdAtom,
 	commentsCollapsedAtom,
 	commentsEnabledAtom,
 	targetCommentIdAtom,
 } from "@/atoms/chat/current-thread.atom";
 import { activeSearchSpaceIdAtom } from "@/atoms/search-spaces/search-space-query.atoms";
-import { BranchPicker } from "@/components/assistant-ui/branch-picker";
 import { MarkdownText } from "@/components/assistant-ui/markdown-text";
 import {
 	ThinkingStepsContext,
@@ -84,7 +83,6 @@ const AssistantMessageInner: FC = () => {
 			</div>
 
 			<div className="aui-assistant-message-footer mt-1 mb-5 ml-2 flex">
-				<BranchPicker />
 				<AssistantActionBar />
 			</div>
 		</>
@@ -126,7 +124,6 @@ export const AssistantMessage: FC = () => {
 
 	// Target comment navigation - read target from global atom
 	const targetCommentId = useAtomValue(targetCommentIdAtom);
-	const clearTargetCommentId = useSetAtom(clearTargetCommentIdAtom);
 
 	// Check if target comment belongs to this message (including replies)
 	const hasTargetComment = useMemo(() => {
@@ -263,6 +260,8 @@ export const AssistantMessage: FC = () => {
 };
 
 const AssistantActionBar: FC = () => {
+	const { isLast } = useMessage();
+
 	return (
 		<ActionBarPrimitive.Root
 			hideWhenRunning
@@ -285,11 +284,14 @@ const AssistantActionBar: FC = () => {
 					<DownloadIcon />
 				</TooltipIconButton>
 			</ActionBarPrimitive.ExportMarkdown>
-			<ActionBarPrimitive.Reload asChild>
-				<TooltipIconButton tooltip="Refresh">
-					<RefreshCwIcon />
-				</TooltipIconButton>
-			</ActionBarPrimitive.Reload>
+			{/* Only allow regenerating the last assistant message */}
+			{isLast && (
+				<ActionBarPrimitive.Reload asChild>
+					<TooltipIconButton tooltip="Refresh">
+						<RefreshCwIcon />
+					</TooltipIconButton>
+				</ActionBarPrimitive.Reload>
+			)}
 		</ActionBarPrimitive.Root>
 	);
 };
diff --git a/surfsense_web/components/assistant-ui/branch-picker.tsx b/surfsense_web/components/assistant-ui/branch-picker.tsx
deleted file mode 100644
index ee4addd2a..000000000
--- a/surfsense_web/components/assistant-ui/branch-picker.tsx
+++ /dev/null
@@ -1,32 +0,0 @@
-import { BranchPickerPrimitive } from "@assistant-ui/react";
-import { ChevronLeftIcon, ChevronRightIcon } from "lucide-react";
-import type { FC } from "react";
-import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
-import { cn } from "@/lib/utils";
-
-export const BranchPicker: FC<BranchPickerPrimitive.Root.Props> = ({ className, ...rest }) => {
-	return (
-		<BranchPickerPrimitive.Root
-			hideWhenSingleBranch
-			className={cn(
-				"aui-branch-picker-root -ml-2 mr-2 inline-flex items-center text-muted-foreground text-xs",
-				className
-			)}
-			{...rest}
-		>
-			<BranchPickerPrimitive.Previous asChild>
-				<TooltipIconButton tooltip="Previous">
-					<ChevronLeftIcon />
-				</TooltipIconButton>
-			</BranchPickerPrimitive.Previous>
-			<span className="aui-branch-picker-state font-medium">
-				<BranchPickerPrimitive.Number /> / <BranchPickerPrimitive.Count />
-			</span>
-			<BranchPickerPrimitive.Next asChild>
-				<TooltipIconButton tooltip="Next">
-					<ChevronRightIcon />
-				</TooltipIconButton>
-			</BranchPickerPrimitive.Next>
-		</BranchPickerPrimitive.Root>
-	);
-};
diff --git a/surfsense_web/components/assistant-ui/user-message.tsx b/surfsense_web/components/assistant-ui/user-message.tsx
index 896b8c748..1ae8aef3c 100644
--- a/surfsense_web/components/assistant-ui/user-message.tsx
+++ b/surfsense_web/components/assistant-ui/user-message.tsx
@@ -4,7 +4,6 @@ import { FileText, PencilIcon } from "lucide-react";
 import { type FC, useState } from "react";
 import { messageDocumentsMapAtom } from "@/atoms/chat/mentioned-documents.atom";
 import { UserMessageAttachments } from "@/components/assistant-ui/attachment";
-import { BranchPicker } from "@/components/assistant-ui/branch-picker";
 import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
 
 interface AuthorMetadata {
@@ -95,24 +94,47 @@ export const UserMessage: FC = () => {
 					</div>
 				)}
 			</div>
-
-			<BranchPicker className="aui-user-branch-picker -mr-1 col-span-full col-start-1 row-start-3 justify-end" />
 		</MessagePrimitive.Root>
 	);
 };
 
 const UserActionBar: FC = () => {
+	const isThreadRunning = useAssistantState(({ thread }) => thread.isRunning);
+	
+	// Get current message ID
+	const currentMessageId = useAssistantState(({ message }) => message?.id);
+	
+	// Find the last user message ID in the thread (computed once, memoized by selector)
+	const lastUserMessageId = useAssistantState(({ thread }) => {
+		const messages = thread.messages;
+		for (let i = messages.length - 1; i >= 0; i--) {
+			if (messages[i].role === "user") {
+				return messages[i].id;
+			}
+		}
+		return null;
+	});
+
+	// Simple comparison - no iteration needed per message
+	const isLastUserMessage = currentMessageId === lastUserMessageId;
+	
+	// Show edit button only on the last user message and when thread is not running
+	const canEdit = isLastUserMessage && !isThreadRunning;
+
 	return (
 		<ActionBarPrimitive.Root
 			hideWhenRunning
 			autohide="not-last"
 			className="aui-user-action-bar-root flex flex-col items-end"
 		>
-			<ActionBarPrimitive.Edit asChild>
-				<TooltipIconButton tooltip="Edit" className="aui-user-action-edit p-4">
-					<PencilIcon />
-				</TooltipIconButton>
-			</ActionBarPrimitive.Edit>
+			{/* Only allow editing the last user message */}
+			{canEdit && (
+				<ActionBarPrimitive.Edit asChild>
+					<TooltipIconButton tooltip="Edit" className="aui-user-action-edit p-4">
+						<PencilIcon />
+					</TooltipIconButton>
+				</ActionBarPrimitive.Edit>
+			)}
 		</ActionBarPrimitive.Root>
 	);
 };

From fb371d09f51120498de77d7e27940a3b9cf3fa32 Mon Sep 17 00:00:00 2001
From: CREDO23 <thierrybakera12@gmail.com>
Date: Wed, 4 Feb 2026 18:26:38 +0200
Subject: [PATCH 18/40] Add globe indicator for chats with public links

---
 .../components/new-chat/chat-share-button.tsx | 237 ++++++++++--------
 1 file changed, 136 insertions(+), 101 deletions(-)

diff --git a/surfsense_web/components/new-chat/chat-share-button.tsx b/surfsense_web/components/new-chat/chat-share-button.tsx
index fa05f44c1..aae9f18c4 100644
--- a/surfsense_web/components/new-chat/chat-share-button.tsx
+++ b/surfsense_web/components/new-chat/chat-share-button.tsx
@@ -1,8 +1,9 @@
 "use client";
 
-import { useQueryClient } from "@tanstack/react-query";
+import { useQuery, useQueryClient } from "@tanstack/react-query";
 import { useAtomValue, useSetAtom } from "jotai";
 import { Globe, User, Users } from "lucide-react";
+import { useParams, useRouter } from "next/navigation";
 import { useCallback, useMemo, useState } from "react";
 import { toast } from "sonner";
 import { currentThreadAtom, setThreadVisibilityAtom } from "@/atoms/chat/current-thread.atom";
@@ -11,6 +12,7 @@ import { createPublicChatSnapshotMutationAtom } from "@/atoms/public-chat-snapsh
 import { Button } from "@/components/ui/button";
 import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover";
 import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip";
+import { chatThreadsApiService } from "@/lib/apis/chat-threads-api.service";
 import {
 	type ChatVisibility,
 	type ThreadRecord,
@@ -46,6 +48,8 @@ const visibilityOptions: {
 
 export function ChatShareButton({ thread, onVisibilityChange, className }: ChatShareButtonProps) {
 	const queryClient = useQueryClient();
+	const router = useRouter();
+	const params = useParams();
 	const [open, setOpen] = useState(false);
 
 	// Use Jotai atom for visibility (single source of truth)
@@ -65,6 +69,16 @@ export function ChatShareButton({ thread, onVisibilityChange, className }: ChatS
 		return access.permissions?.includes("public_sharing:create") ?? false;
 	}, [access]);
 
+	// Query to check if thread has public snapshots
+	const { data: snapshotsData } = useQuery({
+		queryKey: ["thread-snapshots", thread?.id],
+		queryFn: () => chatThreadsApiService.listPublicChatSnapshots({ thread_id: thread!.id }),
+		enabled: !!thread?.id,
+		staleTime: 30000, // Cache for 30 seconds
+	});
+	const hasPublicSnapshots = (snapshotsData?.snapshots?.length ?? 0) > 0;
+	const snapshotCount = snapshotsData?.snapshots?.length ?? 0;
+
 	// Use Jotai visibility if available (synced from chat page), otherwise fall back to thread prop
 	const currentVisibility = currentThreadState.visibility ?? thread?.visibility ?? "PRIVATE";
 
@@ -106,11 +120,13 @@ export function ChatShareButton({ thread, onVisibilityChange, className }: ChatS
 
 		try {
 			await createSnapshot({ thread_id: thread.id });
+			// Refetch snapshots to show the globe indicator
+			await queryClient.invalidateQueries({ queryKey: ["thread-snapshots", thread.id] });
 			setOpen(false);
 		} catch (error) {
 			console.error("Failed to create public link:", error);
 		}
-	}, [thread, createSnapshot]);
+	}, [thread, createSnapshot, queryClient]);
 
 	// Don't show if no thread (new chat that hasn't been created yet)
 	if (!thread) {
@@ -121,112 +137,131 @@ export function ChatShareButton({ thread, onVisibilityChange, className }: ChatS
 	const buttonLabel = currentVisibility === "PRIVATE" ? "Private" : "Shared";
 
 	return (
-		<Popover open={open} onOpenChange={setOpen}>
-			<Tooltip>
-				<TooltipTrigger asChild>
-					<PopoverTrigger asChild>
-						<Button
-							variant="outline"
-							size="icon"
-							className={cn(
-								"h-8 w-8 md:w-auto md:px-3 md:gap-2 relative bg-muted hover:bg-muted/80 border-0",
-								className
-							)}
-						>
-							<CurrentIcon className="h-4 w-4" />
-							<span className="hidden md:inline text-sm">{buttonLabel}</span>
-						</Button>
-					</PopoverTrigger>
-				</TooltipTrigger>
-				<TooltipContent>Share settings</TooltipContent>
-			</Tooltip>
-
-			<PopoverContent
-				className="w-[280px] md:w-[320px] p-0 rounded-lg shadow-lg border-border/60"
-				align="end"
-				sideOffset={8}
-				onCloseAutoFocus={(e) => e.preventDefault()}
-			>
-				<div className="p-1.5 space-y-1">
-					{/* Visibility Options */}
-					{visibilityOptions.map((option) => {
-						const isSelected = currentVisibility === option.value;
-						const Icon = option.icon;
-
-						return (
-							<button
-								type="button"
-								key={option.value}
-								onClick={() => handleVisibilityChange(option.value)}
-								className={cn(
-									"w-full flex items-center gap-2.5 px-2.5 py-2 rounded-md transition-all",
-									"hover:bg-accent/50 cursor-pointer",
-									"focus:outline-none",
-									isSelected && "bg-accent/80"
-								)}
+		<div className={cn("flex items-center gap-1", className)}>
+			<Popover open={open} onOpenChange={setOpen}>
+				<Tooltip>
+					<TooltipTrigger asChild>
+						<PopoverTrigger asChild>
+							<Button
+								variant="outline"
+								size="icon"
+								className="h-8 w-8 md:w-auto md:px-3 md:gap-2 relative bg-muted hover:bg-muted/80 border-0"
 							>
-								<div
+								<CurrentIcon className="h-4 w-4" />
+								<span className="hidden md:inline text-sm">{buttonLabel}</span>
+							</Button>
+						</PopoverTrigger>
+					</TooltipTrigger>
+					<TooltipContent>Share settings</TooltipContent>
+				</Tooltip>
+
+					<PopoverContent
+					className="w-[280px] md:w-[320px] p-0 rounded-lg shadow-lg border-border/60"
+					align="end"
+					sideOffset={8}
+					onCloseAutoFocus={(e) => e.preventDefault()}
+				>
+					<div className="p-1.5 space-y-1">
+						{/* Visibility Options */}
+						{visibilityOptions.map((option) => {
+							const isSelected = currentVisibility === option.value;
+							const Icon = option.icon;
+
+							return (
+								<button
+									type="button"
+									key={option.value}
+									onClick={() => handleVisibilityChange(option.value)}
 									className={cn(
-										"size-7 rounded-md shrink-0 grid place-items-center",
-										isSelected ? "bg-primary/10" : "bg-muted"
+										"w-full flex items-center gap-2.5 px-2.5 py-2 rounded-md transition-all",
+										"hover:bg-accent/50 cursor-pointer",
+										"focus:outline-none",
+										isSelected && "bg-accent/80"
 									)}
 								>
-									<Icon
+									<div
 										className={cn(
-											"size-4 block",
-											isSelected ? "text-primary" : "text-muted-foreground"
+											"size-7 rounded-md shrink-0 grid place-items-center",
+											isSelected ? "bg-primary/10" : "bg-muted"
 										)}
-									/>
-								</div>
-								<div className="flex-1 text-left min-w-0">
-									<div className="flex items-center gap-1.5">
-										<span className={cn("text-sm font-medium", isSelected && "text-primary")}>
-											{option.label}
-										</span>
+									>
+										<Icon
+											className={cn(
+												"size-4 block",
+												isSelected ? "text-primary" : "text-muted-foreground"
+											)}
+										/>
 									</div>
-									<p className="text-xs text-muted-foreground mt-0.5 leading-snug">
-										{option.description}
-									</p>
-								</div>
-							</button>
-						);
-					})}
-
-					{canCreatePublicLink && (
-						<>
-							{/* Divider */}
-							<div className="border-t border-border my-1" />
-
-							{/* Public Link Option */}
-							<button
-								type="button"
-								onClick={handleCreatePublicLink}
-								disabled={isCreatingSnapshot}
-								className={cn(
-									"w-full flex items-center gap-2.5 px-2.5 py-2 rounded-md transition-all",
-									"hover:bg-accent/50 cursor-pointer",
-									"focus:outline-none",
-									"disabled:opacity-50 disabled:cursor-not-allowed"
-								)}
-							>
-								<div className="size-7 rounded-md shrink-0 grid place-items-center bg-muted">
-									<Globe className="size-4 block text-muted-foreground" />
-								</div>
-								<div className="flex-1 text-left min-w-0">
-									<div className="flex items-center gap-1.5">
-										<span className="text-sm font-medium">
-											{isCreatingSnapshot ? "Creating link..." : "Create public link"}
-										</span>
+									<div className="flex-1 text-left min-w-0">
+										<div className="flex items-center gap-1.5">
+											<span className={cn("text-sm font-medium", isSelected && "text-primary")}>
+												{option.label}
+											</span>
+										</div>
+										<p className="text-xs text-muted-foreground mt-0.5 leading-snug">
+											{option.description}
+										</p>
 									</div>
-									<p className="text-xs text-muted-foreground mt-0.5 leading-snug">
-										Creates a shareable snapshot of this chat
-									</p>
-								</div>
-							</button>
-						</>
-					)}
-				</div>
-			</PopoverContent>
-		</Popover>
+								</button>
+							);
+						})}
+
+						{canCreatePublicLink && (
+							<>
+								{/* Divider */}
+								<div className="border-t border-border my-1" />
+
+								{/* Public Link Option */}
+								<button
+									type="button"
+									onClick={handleCreatePublicLink}
+									disabled={isCreatingSnapshot}
+									className={cn(
+										"w-full flex items-center gap-2.5 px-2.5 py-2 rounded-md transition-all",
+										"hover:bg-accent/50 cursor-pointer",
+										"focus:outline-none",
+										"disabled:opacity-50 disabled:cursor-not-allowed"
+									)}
+								>
+									<div className="size-7 rounded-md shrink-0 grid place-items-center bg-muted">
+										<Globe className="size-4 block text-muted-foreground" />
+									</div>
+									<div className="flex-1 text-left min-w-0">
+										<div className="flex items-center gap-1.5">
+											<span className="text-sm font-medium">
+												{isCreatingSnapshot ? "Creating link..." : "Create public link"}
+											</span>
+										</div>
+										<p className="text-xs text-muted-foreground mt-0.5 leading-snug">
+											Creates a shareable snapshot of this chat
+										</p>
+									</div>
+								</button>
+							</>
+						)}
+					</div>
+				</PopoverContent>
+			</Popover>
+
+				{/* Globe indicator when public snapshots exist - clicks to settings */}
+			{hasPublicSnapshots && (
+				<Tooltip>
+					<TooltipTrigger asChild>
+						<button
+							type="button"
+							onClick={() => router.push(`/dashboard/${params.search_space_id}/settings`)}
+							className="flex items-center justify-center h-8 w-8 rounded-md bg-muted/50 hover:bg-muted transition-colors"
+						>
+							<Globe className="h-4 w-4 text-muted-foreground" />
+						</button>
+					</TooltipTrigger>
+					<TooltipContent>
+						{snapshotCount === 1
+							? "This chat has a public link - Click to manage"
+							: `This chat has ${snapshotCount} public links - Click to manage`}
+					</TooltipContent>
+				</Tooltip>
+			)}
+		</div>
 	);
 }

From 83aa8ef6892507193976c6246dee4f7fbcbff167 Mon Sep 17 00:00:00 2001
From: Eric Lammertsma <e.lammertsma@gmail.com>
Date: Wed, 4 Feb 2026 11:30:30 -0500
Subject: [PATCH 19/40] Added build version for cloud, local and Docker
 installs

---
 .../user/settings/components/UserSettingsSidebar.tsx        | 6 ++++++
 surfsense_web/lib/env-config.ts                             | 6 ++++++
 2 files changed, 12 insertions(+)

diff --git a/surfsense_web/app/dashboard/user/settings/components/UserSettingsSidebar.tsx b/surfsense_web/app/dashboard/user/settings/components/UserSettingsSidebar.tsx
index b7040b4e3..3424113a9 100644
--- a/surfsense_web/app/dashboard/user/settings/components/UserSettingsSidebar.tsx
+++ b/surfsense_web/app/dashboard/user/settings/components/UserSettingsSidebar.tsx
@@ -5,6 +5,7 @@ import { ArrowLeft, ChevronRight, X } from "lucide-react";
 import { AnimatePresence, motion } from "motion/react";
 import { useTranslations } from "next-intl";
 import { Button } from "@/components/ui/button";
+import { APP_VERSION } from "@/lib/env-config";
 import { cn } from "@/lib/utils";
 
 export interface SettingsNavItem {
@@ -148,6 +149,11 @@ export function UserSettingsSidebar({
 						);
 					})}
 				</nav>
+
+				{/* Version display */}
+				<div className="mt-auto border-t px-6 py-3">
+					<p className="text-xs text-muted-foreground/50">v{APP_VERSION}</p>
+				</div>
 			</aside>
 		</>
 	);
diff --git a/surfsense_web/lib/env-config.ts b/surfsense_web/lib/env-config.ts
index 2f9e92357..e36aff10a 100644
--- a/surfsense_web/lib/env-config.ts
+++ b/surfsense_web/lib/env-config.ts
@@ -9,6 +9,8 @@
  * as it may prevent the sed replacement from working correctly.
  */
 
+import packageJson from "../package.json";
+
 // Auth type: "LOCAL" for email/password, "GOOGLE" for OAuth
 // Placeholder: __NEXT_PUBLIC_FASTAPI_BACKEND_AUTH_TYPE__
 export const AUTH_TYPE = process.env.NEXT_PUBLIC_FASTAPI_BACKEND_AUTH_TYPE || "GOOGLE";
@@ -28,6 +30,10 @@ export const ETL_SERVICE = process.env.NEXT_PUBLIC_ETL_SERVICE || "DOCLING";
 // Placeholder: __NEXT_PUBLIC_DEPLOYMENT_MODE__
 export const DEPLOYMENT_MODE = process.env.NEXT_PUBLIC_DEPLOYMENT_MODE || "self-hosted";
 
+// App version - defaults to package.json version
+// Can be overridden at build time with NEXT_PUBLIC_APP_VERSION for full git tag version
+export const APP_VERSION = process.env.NEXT_PUBLIC_APP_VERSION || packageJson.version;
+
 // Helper to check if local auth is enabled
 export const isLocalAuth = () => AUTH_TYPE === "LOCAL";
 

From ab3d99d9e07ebd7aa47cccf929df4d8ee5ca849c Mon Sep 17 00:00:00 2001
From: CREDO23 <thierrybakera12@gmail.com>
Date: Wed, 4 Feb 2026 18:46:12 +0200
Subject: [PATCH 20/40] Use RBAC for listing thread snapshots

---
 .../app/services/public_chat_service.py             | 13 ++++++++-----
 .../components/new-chat/chat-share-button.tsx       |  4 ++--
 2 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/surfsense_backend/app/services/public_chat_service.py b/surfsense_backend/app/services/public_chat_service.py
index 2125dd8ce..4da316240 100644
--- a/surfsense_backend/app/services/public_chat_service.py
+++ b/surfsense_backend/app/services/public_chat_service.py
@@ -366,11 +366,14 @@ async def list_snapshots_for_thread(
     if not thread:
         raise HTTPException(status_code=404, detail="Thread not found")
 
-    if thread.created_by_id != user.id:
-        raise HTTPException(
-            status_code=403,
-            detail="Only the creator can view snapshots",
-        )
+    # Check permission to view public share links
+    await check_permission(
+        session,
+        user,
+        thread.search_space_id,
+        Permission.PUBLIC_SHARING_VIEW.value,
+        "You don't have permission to view public share links",
+    )
 
     result = await session.execute(
         select(PublicChatSnapshot)
diff --git a/surfsense_web/components/new-chat/chat-share-button.tsx b/surfsense_web/components/new-chat/chat-share-button.tsx
index aae9f18c4..2e04fa3ba 100644
--- a/surfsense_web/components/new-chat/chat-share-button.tsx
+++ b/surfsense_web/components/new-chat/chat-share-button.tsx
@@ -257,8 +257,8 @@ export function ChatShareButton({ thread, onVisibilityChange, className }: ChatS
 					</TooltipTrigger>
 					<TooltipContent>
 						{snapshotCount === 1
-							? "This chat has a public link - Click to manage"
-							: `This chat has ${snapshotCount} public links - Click to manage`}
+							? "This chat has a public link"
+							: `This chat has ${snapshotCount} public links`}
 					</TooltipContent>
 				</Tooltip>
 			)}

From 1cf7205a81575f9cdc05c25cc26b653743418b2c Mon Sep 17 00:00:00 2001
From: CREDO23 <thierrybakera12@gmail.com>
Date: Wed, 4 Feb 2026 18:54:59 +0200
Subject: [PATCH 21/40] Add clipboard utility with fallback and show selectable
 URLs

---
 .../public-chat-snapshot-row.tsx              |  7 +++
 surfsense_web/hooks/use-api-key.ts            | 60 +++----------------
 surfsense_web/lib/utils.ts                    | 41 +++++++++++++
 3 files changed, 57 insertions(+), 51 deletions(-)

diff --git a/surfsense_web/components/public-chat-snapshots/public-chat-snapshot-row.tsx b/surfsense_web/components/public-chat-snapshots/public-chat-snapshot-row.tsx
index 696d32466..5f0048100 100644
--- a/surfsense_web/components/public-chat-snapshots/public-chat-snapshot-row.tsx
+++ b/surfsense_web/components/public-chat-snapshots/public-chat-snapshot-row.tsx
@@ -38,6 +38,13 @@ export function PublicChatSnapshotRow({
 						{snapshot.message_count}
 					</span>
 				</div>
+				<input
+					type="text"
+					readOnly
+					value={snapshot.public_url}
+					className="mt-2 w-full text-xs text-muted-foreground bg-muted/50 border rounded px-2 py-1 select-all focus:outline-none focus:ring-1 focus:ring-ring"
+					onClick={(e) => (e.target as HTMLInputElement).select()}
+				/>
 			</div>
 			<div className="flex items-center gap-2">
 				<Button
diff --git a/surfsense_web/hooks/use-api-key.ts b/surfsense_web/hooks/use-api-key.ts
index a5f24d4c6..0c595b420 100644
--- a/surfsense_web/hooks/use-api-key.ts
+++ b/surfsense_web/hooks/use-api-key.ts
@@ -1,6 +1,7 @@
 import { useCallback, useEffect, useState } from "react";
 import { toast } from "sonner";
 import { getBearerToken } from "@/lib/auth-utils";
+import { copyToClipboard as copyToClipboardUtil } from "@/lib/utils";
 
 interface UseApiKeyReturn {
 	apiKey: string | null;
@@ -33,60 +34,17 @@ export function useApiKey(): UseApiKeyReturn {
 		return () => clearTimeout(timer);
 	}, []);
 
-	const fallbackCopyTextToClipboard = (text: string) => {
-		const textArea = document.createElement("textarea");
-		textArea.value = text;
-
-		// Avoid scrolling to bottom
-		textArea.style.top = "0";
-		textArea.style.left = "0";
-		textArea.style.position = "fixed";
-		textArea.style.opacity = "0";
-
-		document.body.appendChild(textArea);
-		textArea.focus();
-		textArea.select();
-
-		try {
-			const successful = document.execCommand("copy");
-			document.body.removeChild(textArea);
-
-			if (successful) {
-				setCopied(true);
-				toast.success("API key copied to clipboard");
-
-				setTimeout(() => {
-					setCopied(false);
-				}, 2000);
-			} else {
-				toast.error("Failed to copy API key");
-			}
-		} catch (err) {
-			console.error("Fallback: Oops, unable to copy", err);
-			document.body.removeChild(textArea);
-			toast.error("Failed to copy API key");
-		}
-	};
-
 	const copyToClipboard = useCallback(async () => {
 		if (!apiKey) return;
 
-		try {
-			if (navigator.clipboard && window.isSecureContext) {
-				// Use Clipboard API if available and in secure context
-				await navigator.clipboard.writeText(apiKey);
-				setCopied(true);
-				toast.success("API key copied to clipboard");
-
-				setTimeout(() => {
-					setCopied(false);
-				}, 2000);
-			} else {
-				// Fallback for non-secure contexts or browsers without clipboard API
-				fallbackCopyTextToClipboard(apiKey);
-			}
-		} catch (err) {
-			console.error("Failed to copy:", err);
+		const success = await copyToClipboardUtil(apiKey);
+		if (success) {
+			setCopied(true);
+			toast.success("API key copied to clipboard");
+			setTimeout(() => {
+				setCopied(false);
+			}, 2000);
+		} else {
 			toast.error("Failed to copy API key");
 		}
 	}, [apiKey]);
diff --git a/surfsense_web/lib/utils.ts b/surfsense_web/lib/utils.ts
index 212ff1259..e7bf8bdbe 100644
--- a/surfsense_web/lib/utils.ts
+++ b/surfsense_web/lib/utils.ts
@@ -12,3 +12,44 @@ export const formatDate = (date: Date): string => {
 		day: "numeric",
 	});
 };
+
+/**
+ * Copy text to clipboard with fallback for older browsers and non-secure contexts.
+ * Returns true if successful, false otherwise.
+ */
+export async function copyToClipboard(text: string): Promise<boolean> {
+	// Use modern Clipboard API if available and in secure context
+	if (navigator.clipboard && window.isSecureContext) {
+		try {
+			await navigator.clipboard.writeText(text);
+			return true;
+		} catch (err) {
+			console.error("Clipboard API failed:", err);
+			return false;
+		}
+	}
+
+	// Fallback for non-secure contexts or browsers without Clipboard API
+	const textArea = document.createElement("textarea");
+	textArea.value = text;
+
+	// Avoid scrolling to bottom
+	textArea.style.top = "0";
+	textArea.style.left = "0";
+	textArea.style.position = "fixed";
+	textArea.style.opacity = "0";
+
+	document.body.appendChild(textArea);
+	textArea.focus();
+	textArea.select();
+
+	try {
+		const successful = document.execCommand("copy");
+		document.body.removeChild(textArea);
+		return successful;
+	} catch (err) {
+		console.error("Fallback copy failed:", err);
+		document.body.removeChild(textArea);
+		return false;
+	}
+}

From a133fb41841f0a694554eb9ffe97258058556c94 Mon Sep 17 00:00:00 2001
From: Eric Lammertsma <e.lammertsma@gmail.com>
Date: Wed, 4 Feb 2026 15:30:15 -0500
Subject: [PATCH 22/40] Updated price tier page limits and removed enterprise
 pricing

---
 surfsense_web/components/pricing/pricing-section.tsx | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/surfsense_web/components/pricing/pricing-section.tsx b/surfsense_web/components/pricing/pricing-section.tsx
index 445969ec1..117be15ec 100644
--- a/surfsense_web/components/pricing/pricing-section.tsx
+++ b/surfsense_web/components/pricing/pricing-section.tsx
@@ -11,7 +11,7 @@ const demoPlans = [
 		billingText: "Includes 30 day PRO trial",
 		features: [
 			"Open source on GitHub",
-			"Upload and chat with 1,000+ pages of content",
+			"Upload and chat with 300+ pages of content",
 			"Connects with 8 popular sources, like Drive and Notion.",
 			"Includes limited access to ChatGPT, Claude, and DeepSeek models",
 			"Supports 100+ more LLMs, including Gemini, Llama and many more.",
@@ -33,7 +33,7 @@ const demoPlans = [
 		billingText: "billed annually",
 		features: [
 			"Everything in Free",
-			"Upload and chat with up to 20,000+ pages of content",
+			"Upload and chat with 5,000+ pages of content",
 			"Connects with 15+ external sources, like Slack and Airtable.",
 			"Includes extended access to ChatGPT, Claude, and DeepSeek models",
 			"Collaboration and commenting features",
@@ -49,10 +49,10 @@ const demoPlans = [
 	},
 	{
 		name: "ENTERPRISE",
-		price: "1000",
-		yearlyPrice: "1000",
-		period: "month",
-		billingText: "billed annually",
+		price: "Contact Us",
+		yearlyPrice: "Contact Us",
+		period: "",
+		billingText: "",
 		features: [
 			"Everything in Pro",
 			"Connect and chat with virtually unlimited pages of content",

From 1c1dcbf47f733f36c31ebad706244f372171b7ce Mon Sep 17 00:00:00 2001
From: Eric Lammertsma <e.lammertsma@gmail.com>
Date: Wed, 4 Feb 2026 15:52:21 -0500
Subject: [PATCH 23/40] Add Gemini 3 Flash and Pro models to LLM_MODELS enum

---
 surfsense_web/contracts/enums/llm-models.ts | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/surfsense_web/contracts/enums/llm-models.ts b/surfsense_web/contracts/enums/llm-models.ts
index c62b2a9d6..5ff15c3df 100644
--- a/surfsense_web/contracts/enums/llm-models.ts
+++ b/surfsense_web/contracts/enums/llm-models.ts
@@ -178,6 +178,18 @@ export const LLM_MODELS: LLMModel[] = [
 	},
 
 	// Google (Gemini)
+	{
+		value: "gemini-3-flash-preview",
+		label: "Gemini 3 Flash",
+		provider: "GOOGLE",
+		contextWindow: "1M",
+	},
+	{
+		value: "gemini-3-pro-preview",
+		label: "Gemini 3 Pro",
+		provider: "GOOGLE",
+		contextWindow: "1M",
+	},
 	{
 		value: "gemini-2.5-flash",
 		label: "Gemini 2.5 Flash",

From 33165830e58820436bb27d915d9e81981ec2f7f0 Mon Sep 17 00:00:00 2001
From: Manoj Aggarwal <manoj6543@gmail.com>
Date: Wed, 4 Feb 2026 13:18:33 -0800
Subject: [PATCH 24/40] add parse date flexible

---
 .../app/tasks/connector_indexers/base.py      | 36 ++++++++++++++++++-
 .../google_calendar_indexer.py                |  9 ++++-
 2 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/surfsense_backend/app/tasks/connector_indexers/base.py b/surfsense_backend/app/tasks/connector_indexers/base.py
index 311fda996..b801b67d6 100644
--- a/surfsense_backend/app/tasks/connector_indexers/base.py
+++ b/surfsense_backend/app/tasks/connector_indexers/base.py
@@ -28,6 +28,34 @@ def get_current_timestamp() -> datetime:
     return datetime.now(UTC)
 
 
+def parse_date_flexible(date_str: str) -> datetime:
+    """
+    Parse date from multiple common formats.
+
+    Args:
+        date_str: Date string to parse
+
+    Returns:
+        Parsed datetime object
+
+    Raises:
+        ValueError: If unable to parse the date string
+    """
+    formats = ["%Y-%m-%d", "%Y-%m-%dT%H:%M:%S"]
+
+    for fmt in formats:
+        try:
+            return datetime.strptime(date_str.rstrip("Z"), fmt)
+        except ValueError:
+            continue
+
+    # Try ISO format as fallback
+    try:
+        return datetime.fromisoformat(date_str.replace("Z", "+00:00"))
+    except ValueError:
+        raise ValueError(f"Unable to parse date: {date_str}")
+
+
 async def check_duplicate_document_by_hash(
     session: AsyncSession, content_hash: str
 ) -> Document | None:
@@ -168,7 +196,13 @@ def calculate_date_range(
             "adjusting end date to next day to ensure valid date range"
         )
         # Parse end_date and add 1 day
-        end_dt = datetime.strptime(end_date_str, "%Y-%m-%d")
+        try:
+            end_dt = parse_date_flexible(end_date_str)
+        except ValueError:
+            logger.warning(
+                f"Could not parse end_date '{end_date_str}', using current date"
+            )
+            end_dt = datetime.now()
         end_dt = end_dt + timedelta(days=1)
         end_date_str = end_dt.strftime("%Y-%m-%d")
         logger.info(f"Adjusted end date to {end_date_str}")
diff --git a/surfsense_backend/app/tasks/connector_indexers/google_calendar_indexer.py b/surfsense_backend/app/tasks/connector_indexers/google_calendar_indexer.py
index 1d8ea32f2..3e1a81356 100644
--- a/surfsense_backend/app/tasks/connector_indexers/google_calendar_indexer.py
+++ b/surfsense_backend/app/tasks/connector_indexers/google_calendar_indexer.py
@@ -27,6 +27,7 @@ from .base import (
     get_connector_by_id,
     get_current_timestamp,
     logger,
+    parse_date_flexible,
     update_connector_last_indexed,
 )
 
@@ -226,7 +227,13 @@ async def index_google_calendar_events(
                 "adjusting end date to next day to ensure valid date range"
             )
             # Parse end_date and add 1 day
-            end_dt = datetime.strptime(end_date_str, "%Y-%m-%d")
+            try:
+                end_dt = parse_date_flexible(end_date_str)
+            except ValueError:
+                logger.warning(
+                    f"Could not parse end_date '{end_date_str}', using current date"
+                )
+                end_dt = datetime.now()
             end_dt = end_dt + timedelta(days=1)
             end_date_str = end_dt.strftime("%Y-%m-%d")
             logger.info(f"Adjusted end date to {end_date_str}")

From ecd0985523b381a37a9cb8222afb08a0aa050810 Mon Sep 17 00:00:00 2001
From: CREDO23 <thierrybakera12@gmail.com>
Date: Thu, 5 Feb 2026 14:39:50 +0200
Subject: [PATCH 25/40] Add access token pre-validation to OAuth connectors

---
 .../app/connectors/airtable_history.py           | 16 ++++++++++++++++
 .../app/connectors/confluence_history.py         | 16 ++++++++++++++++
 surfsense_backend/app/connectors/jira_history.py | 16 ++++++++++++++++
 .../app/connectors/linear_connector.py           | 16 ++++++++++++++++
 4 files changed, 64 insertions(+)

diff --git a/surfsense_backend/app/connectors/airtable_history.py b/surfsense_backend/app/connectors/airtable_history.py
index 64f6465fe..092485f77 100644
--- a/surfsense_backend/app/connectors/airtable_history.py
+++ b/surfsense_backend/app/connectors/airtable_history.py
@@ -71,6 +71,14 @@ class AirtableHistoryConnector:
 
             config_data = connector.config.copy()
 
+            # Check if access_token exists before processing
+            raw_access_token = config_data.get("access_token")
+            if not raw_access_token:
+                raise ValueError(
+                    "Airtable access token not found. "
+                    "Please reconnect your Airtable account."
+                )
+
             # Decrypt credentials if they are encrypted
             token_encrypted = config_data.get("_token_encrypted", False)
             if token_encrypted and config.SECRET_KEY:
@@ -98,6 +106,14 @@ class AirtableHistoryConnector:
                         f"Failed to decrypt Airtable credentials: {e!s}"
                     ) from e
 
+            # Final validation after decryption
+            final_token = config_data.get("access_token")
+            if not final_token or (isinstance(final_token, str) and not final_token.strip()):
+                raise ValueError(
+                    "Airtable access token is invalid or empty. "
+                    "Please reconnect your Airtable account."
+                )
+
             try:
                 self._credentials = AirtableAuthCredentialsBase.from_dict(config_data)
             except Exception as e:
diff --git a/surfsense_backend/app/connectors/confluence_history.py b/surfsense_backend/app/connectors/confluence_history.py
index 9e10ffcf1..908f532db 100644
--- a/surfsense_backend/app/connectors/confluence_history.py
+++ b/surfsense_backend/app/connectors/confluence_history.py
@@ -87,6 +87,14 @@ class ConfluenceHistoryConnector:
 
             if is_oauth:
                 # OAuth 2.0 authentication
+                # Check if access_token exists before processing
+                raw_access_token = config_data.get("access_token")
+                if not raw_access_token:
+                    raise ValueError(
+                        "Confluence access token not found. "
+                        "Please reconnect your Confluence account."
+                    )
+
                 # Decrypt credentials if they are encrypted
                 token_encrypted = config_data.get("_token_encrypted", False)
                 if token_encrypted and config.SECRET_KEY:
@@ -118,6 +126,14 @@ class ConfluenceHistoryConnector:
                             f"Failed to decrypt Confluence credentials: {e!s}"
                         ) from e
 
+                # Final validation after decryption
+                final_token = config_data.get("access_token")
+                if not final_token or (isinstance(final_token, str) and not final_token.strip()):
+                    raise ValueError(
+                        "Confluence access token is invalid or empty. "
+                        "Please reconnect your Confluence account."
+                    )
+
                 try:
                     self._credentials = AtlassianAuthCredentialsBase.from_dict(
                         config_data
diff --git a/surfsense_backend/app/connectors/jira_history.py b/surfsense_backend/app/connectors/jira_history.py
index 6e04ec2a4..46a28324d 100644
--- a/surfsense_backend/app/connectors/jira_history.py
+++ b/surfsense_backend/app/connectors/jira_history.py
@@ -86,6 +86,14 @@ class JiraHistoryConnector:
 
             if is_oauth:
                 # OAuth 2.0 authentication
+                # Check if access_token exists before processing
+                raw_access_token = config_data.get("access_token")
+                if not raw_access_token:
+                    raise ValueError(
+                        "Jira access token not found. "
+                        "Please reconnect your Jira account."
+                    )
+
                 if not config.SECRET_KEY:
                     raise ValueError(
                         "SECRET_KEY not configured but tokens are marked as encrypted"
@@ -119,6 +127,14 @@ class JiraHistoryConnector:
                         f"Failed to decrypt Jira credentials: {e!s}"
                     ) from e
 
+                # Final validation after decryption
+                final_token = config_data.get("access_token")
+                if not final_token or (isinstance(final_token, str) and not final_token.strip()):
+                    raise ValueError(
+                        "Jira access token is invalid or empty. "
+                        "Please reconnect your Jira account."
+                    )
+
                 try:
                     self._credentials = AtlassianAuthCredentialsBase.from_dict(
                         config_data
diff --git a/surfsense_backend/app/connectors/linear_connector.py b/surfsense_backend/app/connectors/linear_connector.py
index b8206a40d..6500b9027 100644
--- a/surfsense_backend/app/connectors/linear_connector.py
+++ b/surfsense_backend/app/connectors/linear_connector.py
@@ -116,6 +116,14 @@ class LinearConnector:
 
             config_data = connector.config.copy()
 
+            # Check if access_token exists before processing
+            raw_access_token = config_data.get("access_token")
+            if not raw_access_token:
+                raise ValueError(
+                    "Linear access token not found. "
+                    "Please reconnect your Linear account."
+                )
+
             # Decrypt credentials if they are encrypted
             token_encrypted = config_data.get("_token_encrypted", False)
             if token_encrypted and config.SECRET_KEY:
@@ -143,6 +151,14 @@ class LinearConnector:
                         f"Failed to decrypt Linear credentials: {e!s}"
                     ) from e
 
+            # Final validation after decryption
+            final_token = config_data.get("access_token")
+            if not final_token or (isinstance(final_token, str) and not final_token.strip()):
+                raise ValueError(
+                    "Linear access token is invalid or empty. "
+                    "Please reconnect your Linear account."
+                )
+
             try:
                 self._credentials = LinearAuthCredentialsBase.from_dict(config_data)
             except Exception as e:

From 04aac379ed3a250b3e5235c6540e633ee5b5d36b Mon Sep 17 00:00:00 2001
From: CREDO23 <thierrybakera12@gmail.com>
Date: Thu, 5 Feb 2026 16:18:45 +0200
Subject: [PATCH 26/40] Add RefreshToken model for user session management

---
 surfsense_backend/app/db.py | 58 +++++++++++++++++++++++++++++++++++++
 1 file changed, 58 insertions(+)

diff --git a/surfsense_backend/app/db.py b/surfsense_backend/app/db.py
index 5cdb712db..a4da1a575 100644
--- a/surfsense_backend/app/db.py
+++ b/surfsense_backend/app/db.py
@@ -1361,6 +1361,13 @@ if config.AUTH_TYPE == "GOOGLE":
         display_name = Column(String, nullable=True)
         avatar_url = Column(String, nullable=True)
 
+        # Refresh tokens for this user
+        refresh_tokens = relationship(
+            "RefreshToken",
+            back_populates="user",
+            cascade="all, delete-orphan",
+        )
+
 else:
 
     class User(SQLAlchemyBaseUserTableUUID, Base):
@@ -1426,6 +1433,57 @@ else:
         display_name = Column(String, nullable=True)
         avatar_url = Column(String, nullable=True)
 
+        # Refresh tokens for this user
+        refresh_tokens = relationship(
+            "RefreshToken",
+            back_populates="user",
+            cascade="all, delete-orphan",
+        )
+
+
+class RefreshToken(Base, TimestampMixin):
+    """
+    Stores refresh tokens for user session management.
+
+    Refresh tokens are long-lived tokens (2 weeks) used to obtain new
+    access tokens without requiring re-authentication.
+    """
+
+    __tablename__ = "refresh_tokens"
+
+    id = Column(Integer, primary_key=True, autoincrement=True)
+
+    # User relationship
+    user_id = Column(
+        UUID(as_uuid=True),
+        ForeignKey("user.id", ondelete="CASCADE"),
+        nullable=False,
+        index=True,
+    )
+    user = relationship("User", back_populates="refresh_tokens")
+
+    # Token hash (stored hashed, not plaintext)
+    token_hash = Column(String(256), unique=True, nullable=False, index=True)
+
+    # Token expiration
+    expires_at = Column(TIMESTAMP(timezone=True), nullable=False, index=True)
+
+    # Revocation flag
+    is_revoked = Column(Boolean, default=False, nullable=False)
+
+    # Token family for rotation tracking (detect reuse attacks)
+    family_id = Column(UUID(as_uuid=True), nullable=False, index=True)
+
+    @property
+    def is_expired(self) -> bool:
+        """Check if the token has expired."""
+        return datetime.now(UTC) >= self.expires_at
+
+    @property
+    def is_valid(self) -> bool:
+        """Check if the token is valid (not expired and not revoked)."""
+        return not self.is_expired and not self.is_revoked
+
 
 engine = create_async_engine(DATABASE_URL)
 async_session_maker = async_sessionmaker(engine, expire_on_commit=False)

From 048ef7024f6ccc08873dd6bbcdcb08f4275eba1f Mon Sep 17 00:00:00 2001
From: CREDO23 <thierrybakera12@gmail.com>
Date: Thu, 5 Feb 2026 16:21:02 +0200
Subject: [PATCH 27/40] Add token lifetime config options

---
 surfsense_backend/app/config/__init__.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/surfsense_backend/app/config/__init__.py b/surfsense_backend/app/config/__init__.py
index 149fedd39..121e5d3b2 100644
--- a/surfsense_backend/app/config/__init__.py
+++ b/surfsense_backend/app/config/__init__.py
@@ -255,6 +255,14 @@ class Config:
     # OAuth JWT
     SECRET_KEY = os.getenv("SECRET_KEY")
 
+    # JWT Token Lifetimes
+    ACCESS_TOKEN_LIFETIME_SECONDS = int(
+        os.getenv("ACCESS_TOKEN_LIFETIME_SECONDS", str(24 * 60 * 60))  # 1 day
+    )
+    REFRESH_TOKEN_LIFETIME_SECONDS = int(
+        os.getenv("REFRESH_TOKEN_LIFETIME_SECONDS", str(14 * 24 * 60 * 60))  # 2 weeks
+    )
+
     # ETL Service
     ETL_SERVICE = os.getenv("ETL_SERVICE")
 

From 9bd7d7475574cb21c37803da9c2dce33167465a1 Mon Sep 17 00:00:00 2001
From: CREDO23 <thierrybakera12@gmail.com>
Date: Thu, 5 Feb 2026 16:49:37 +0200
Subject: [PATCH 28/40] Add RefreshToken model and multi-session refresh token
 logic

---
 surfsense_backend/app/db.py    |  16 +---
 surfsense_backend/app/users.py | 154 ++++++++++++++++++++++++++++++++-
 2 files changed, 152 insertions(+), 18 deletions(-)

diff --git a/surfsense_backend/app/db.py b/surfsense_backend/app/db.py
index a4da1a575..2298e7438 100644
--- a/surfsense_backend/app/db.py
+++ b/surfsense_backend/app/db.py
@@ -1444,16 +1444,12 @@ else:
 class RefreshToken(Base, TimestampMixin):
     """
     Stores refresh tokens for user session management.
-
-    Refresh tokens are long-lived tokens (2 weeks) used to obtain new
-    access tokens without requiring re-authentication.
+    Each row represents one device/session.
     """
 
     __tablename__ = "refresh_tokens"
 
     id = Column(Integer, primary_key=True, autoincrement=True)
-
-    # User relationship
     user_id = Column(
         UUID(as_uuid=True),
         ForeignKey("user.id", ondelete="CASCADE"),
@@ -1461,27 +1457,17 @@ class RefreshToken(Base, TimestampMixin):
         index=True,
     )
     user = relationship("User", back_populates="refresh_tokens")
-
-    # Token hash (stored hashed, not plaintext)
     token_hash = Column(String(256), unique=True, nullable=False, index=True)
-
-    # Token expiration
     expires_at = Column(TIMESTAMP(timezone=True), nullable=False, index=True)
-
-    # Revocation flag
     is_revoked = Column(Boolean, default=False, nullable=False)
-
-    # Token family for rotation tracking (detect reuse attacks)
     family_id = Column(UUID(as_uuid=True), nullable=False, index=True)
 
     @property
     def is_expired(self) -> bool:
-        """Check if the token has expired."""
         return datetime.now(UTC) >= self.expires_at
 
     @property
     def is_valid(self) -> bool:
-        """Check if the token is valid (not expired and not revoked)."""
         return not self.is_expired and not self.is_revoked
 
 
diff --git a/surfsense_backend/app/users.py b/surfsense_backend/app/users.py
index 4be2fe525..cbffd359d 100644
--- a/surfsense_backend/app/users.py
+++ b/surfsense_backend/app/users.py
@@ -1,5 +1,8 @@
+import hashlib
 import logging
+import secrets
 import uuid
+from datetime import UTC, datetime, timedelta
 
 import httpx
 from fastapi import Depends, Request, Response
@@ -12,9 +15,11 @@ from fastapi_users.authentication import (
 )
 from fastapi_users.db import SQLAlchemyUserDatabase
 from pydantic import BaseModel
+from sqlalchemy import select, update
 
 from app.config import config
 from app.db import (
+    RefreshToken,
     SearchSpace,
     SearchSpaceMembership,
     SearchSpaceRole,
@@ -29,11 +34,130 @@ logger = logging.getLogger(__name__)
 
 class BearerResponse(BaseModel):
     access_token: str
+    refresh_token: str
     token_type: str
 
 
 SECRET = config.SECRET_KEY
 
+
+# Refresh token utilities (multi-session)
+def generate_refresh_token() -> str:
+    """Generate a cryptographically secure refresh token."""
+    return secrets.token_urlsafe(32)
+
+
+def hash_token(token: str) -> str:
+    """Hash a token for secure storage."""
+    return hashlib.sha256(token.encode()).hexdigest()
+
+
+async def create_refresh_token(
+    user_id: uuid.UUID,
+    family_id: uuid.UUID | None = None,
+) -> str:
+    """
+    Create and store a new refresh token for a user.
+
+    Args:
+        user_id: The user's ID
+        family_id: Optional family ID for token rotation
+
+    Returns:
+        The plaintext refresh token
+    """
+    token = generate_refresh_token()
+    token_hash = hash_token(token)
+    expires_at = datetime.now(UTC) + timedelta(
+        seconds=config.REFRESH_TOKEN_LIFETIME_SECONDS
+    )
+
+    if family_id is None:
+        family_id = uuid.uuid4()
+
+    async with async_session_maker() as session:
+        refresh_token = RefreshToken(
+            user_id=user_id,
+            token_hash=token_hash,
+            expires_at=expires_at,
+            family_id=family_id,
+        )
+        session.add(refresh_token)
+        await session.commit()
+
+    return token
+
+
+async def validate_refresh_token(token: str) -> RefreshToken | None:
+    """
+    Validate a refresh token. Handles reuse detection.
+
+    Args:
+        token: The plaintext refresh token
+
+    Returns:
+        RefreshToken if valid, None otherwise
+    """
+    token_hash = hash_token(token)
+
+    async with async_session_maker() as session:
+        result = await session.execute(
+            select(RefreshToken).where(RefreshToken.token_hash == token_hash)
+        )
+        refresh_token = result.scalars().first()
+
+        if not refresh_token:
+            return None
+
+        # Reuse detection: revoked token used while family has active tokens
+        if refresh_token.is_revoked:
+            active = await session.execute(
+                select(RefreshToken).where(
+                    RefreshToken.family_id == refresh_token.family_id,
+                    RefreshToken.is_revoked == False,  # noqa: E712
+                    RefreshToken.expires_at > datetime.now(UTC),
+                )
+            )
+            if active.scalars().first():
+                # Revoke entire family
+                await session.execute(
+                    update(RefreshToken)
+                    .where(RefreshToken.family_id == refresh_token.family_id)
+                    .values(is_revoked=True)
+                )
+                await session.commit()
+                logger.warning(f"Token reuse detected for user {refresh_token.user_id}")
+            return None
+
+        if refresh_token.is_expired:
+            return None
+
+        return refresh_token
+
+
+async def rotate_refresh_token(old_token: RefreshToken) -> str:
+    """Revoke old token and create new one in same family."""
+    async with async_session_maker() as session:
+        await session.execute(
+            update(RefreshToken)
+            .where(RefreshToken.id == old_token.id)
+            .values(is_revoked=True)
+        )
+        await session.commit()
+
+    return await create_refresh_token(old_token.user_id, old_token.family_id)
+
+
+async def revoke_all_user_tokens(user_id: uuid.UUID) -> None:
+    """Revoke all refresh tokens for a user (logout all devices)."""
+    async with async_session_maker() as session:
+        await session.execute(
+            update(RefreshToken)
+            .where(RefreshToken.user_id == user_id)
+            .values(is_revoked=True)
+        )
+        await session.commit()
+
 if config.AUTH_TYPE == "GOOGLE":
     from httpx_oauth.clients.google import GoogleOAuth2
 
@@ -183,7 +307,10 @@ async def get_user_manager(user_db: SQLAlchemyUserDatabase = Depends(get_user_db
 
 
 def get_jwt_strategy() -> JWTStrategy[models.UP, models.ID]:
-    return JWTStrategy(secret=SECRET, lifetime_seconds=3600 * 24)
+    return JWTStrategy(
+        secret=SECRET,
+        lifetime_seconds=config.ACCESS_TOKEN_LIFETIME_SECONDS,
+    )
 
 
 # # COOKIE AUTH | Uncomment if you want to use cookie auth.
@@ -209,9 +336,30 @@ def get_jwt_strategy() -> JWTStrategy[models.UP, models.ID]:
 # BEARER AUTH CODE.
 class CustomBearerTransport(BearerTransport):
     async def get_login_response(self, token: str) -> Response:
-        bearer_response = BearerResponse(access_token=token, token_type="bearer")
-        redirect_url = f"{config.NEXT_FRONTEND_URL}/auth/callback?token={bearer_response.access_token}"
+        import jwt
+
+        # Decode JWT to get user_id for refresh token creation
+        try:
+            payload = jwt.decode(token, SECRET, algorithms=["HS256"])
+            user_id = uuid.UUID(payload.get("sub"))
+            refresh_token = await create_refresh_token(user_id)
+        except Exception as e:
+            logger.error(f"Failed to create refresh token: {e}")
+            # Fall back to response without refresh token
+            refresh_token = ""
+
+        bearer_response = BearerResponse(
+            access_token=token,
+            refresh_token=refresh_token,
+            token_type="bearer",
+        )
+
         if config.AUTH_TYPE == "GOOGLE":
+            redirect_url = (
+                f"{config.NEXT_FRONTEND_URL}/auth/callback"
+                f"?token={bearer_response.access_token}"
+                f"&refresh_token={bearer_response.refresh_token}"
+            )
             return RedirectResponse(redirect_url, status_code=302)
         else:
             return JSONResponse(bearer_response.model_dump())

From f3a9922eb96e2651e6d12dcc88063e9e9ae525c3 Mon Sep 17 00:00:00 2001
From: CREDO23 <thierrybakera12@gmail.com>
Date: Thu, 5 Feb 2026 17:29:50 +0200
Subject: [PATCH 29/40] Add refresh token auth routes and utilities

---
 .../versions/92_add_refresh_tokens_table.py   |  92 +++++++++++
 surfsense_backend/app/app.py                  |   4 +
 surfsense_backend/app/routes/auth_routes.py   | 115 +++++++++++++
 surfsense_backend/app/schemas/__init__.py     |   5 +
 surfsense_backend/app/schemas/auth.py         |  23 +++
 surfsense_backend/app/users.py                | 135 ++--------------
 surfsense_backend/app/utils/auth_cookies.py   |  29 ++++
 surfsense_backend/app/utils/refresh_tokens.py | 153 ++++++++++++++++++
 8 files changed, 431 insertions(+), 125 deletions(-)
 create mode 100644 surfsense_backend/alembic/versions/92_add_refresh_tokens_table.py
 create mode 100644 surfsense_backend/app/routes/auth_routes.py
 create mode 100644 surfsense_backend/app/schemas/auth.py
 create mode 100644 surfsense_backend/app/utils/auth_cookies.py
 create mode 100644 surfsense_backend/app/utils/refresh_tokens.py

diff --git a/surfsense_backend/alembic/versions/92_add_refresh_tokens_table.py b/surfsense_backend/alembic/versions/92_add_refresh_tokens_table.py
new file mode 100644
index 000000000..c7e133ae9
--- /dev/null
+++ b/surfsense_backend/alembic/versions/92_add_refresh_tokens_table.py
@@ -0,0 +1,92 @@
+"""Add refresh_tokens table for user session management
+
+Revision ID: 92
+Revises: 91
+
+Changes:
+1. Create refresh_tokens table with columns:
+   - id (primary key)
+   - user_id (foreign key to user)
+   - token_hash (unique, indexed)
+   - expires_at (indexed)
+   - is_revoked
+   - family_id (indexed, for token rotation tracking)
+   - created_at, updated_at (timestamps)
+"""
+
+from collections.abc import Sequence
+
+import sqlalchemy as sa
+from sqlalchemy.dialects.postgresql import UUID
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "92"
+down_revision: str | None = "91"
+branch_labels: str | Sequence[str] | None = None
+depends_on: str | Sequence[str] | None = None
+
+
+def upgrade() -> None:
+    """Create refresh_tokens table (idempotent)."""
+    # Check if table already exists
+    connection = op.get_bind()
+    result = connection.execute(
+        sa.text(
+            "SELECT EXISTS (SELECT FROM information_schema.tables WHERE table_name = 'refresh_tokens')"
+        )
+    )
+    table_exists = result.scalar()
+
+    if not table_exists:
+        op.create_table(
+            "refresh_tokens",
+            sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
+            sa.Column("user_id", UUID(as_uuid=True), nullable=False),
+            sa.Column("token_hash", sa.String(256), nullable=False),
+            sa.Column("expires_at", sa.TIMESTAMP(timezone=True), nullable=False),
+            sa.Column("is_revoked", sa.Boolean(), nullable=False, default=False),
+            sa.Column("family_id", UUID(as_uuid=True), nullable=False),
+            sa.Column(
+                "created_at",
+                sa.TIMESTAMP(timezone=True),
+                server_default=sa.text("now()"),
+                nullable=False,
+            ),
+            sa.Column(
+                "updated_at",
+                sa.TIMESTAMP(timezone=True),
+                server_default=sa.text("now()"),
+                nullable=False,
+            ),
+            sa.PrimaryKeyConstraint("id"),
+            sa.ForeignKeyConstraint(
+                ["user_id"],
+                ["user.id"],
+                ondelete="CASCADE",
+            ),
+        )
+
+    # Create indexes if they don't exist
+    op.execute(
+        "CREATE INDEX IF NOT EXISTS ix_refresh_tokens_user_id ON refresh_tokens (user_id)"
+    )
+    op.execute(
+        "CREATE UNIQUE INDEX IF NOT EXISTS ix_refresh_tokens_token_hash ON refresh_tokens (token_hash)"
+    )
+    op.execute(
+        "CREATE INDEX IF NOT EXISTS ix_refresh_tokens_expires_at ON refresh_tokens (expires_at)"
+    )
+    op.execute(
+        "CREATE INDEX IF NOT EXISTS ix_refresh_tokens_family_id ON refresh_tokens (family_id)"
+    )
+
+
+def downgrade() -> None:
+    """Drop refresh_tokens table (idempotent)."""
+    op.execute("DROP INDEX IF EXISTS ix_refresh_tokens_family_id")
+    op.execute("DROP INDEX IF EXISTS ix_refresh_tokens_expires_at")
+    op.execute("DROP INDEX IF EXISTS ix_refresh_tokens_token_hash")
+    op.execute("DROP INDEX IF EXISTS ix_refresh_tokens_user_id")
+    op.execute("DROP TABLE IF EXISTS refresh_tokens")
diff --git a/surfsense_backend/app/app.py b/surfsense_backend/app/app.py
index 01dd0da3d..63da4e8ad 100644
--- a/surfsense_backend/app/app.py
+++ b/surfsense_backend/app/app.py
@@ -12,6 +12,7 @@ from app.agents.new_chat.checkpointer import (
 from app.config import config, initialize_llm_router
 from app.db import User, create_db_and_tables, get_async_session
 from app.routes import router as crud_router
+from app.routes.auth_routes import router as auth_router
 from app.schemas import UserCreate, UserRead, UserUpdate
 from app.tasks.surfsense_docs_indexer import seed_surfsense_docs
 from app.users import SECRET, auth_backend, current_active_user, fastapi_users
@@ -111,6 +112,9 @@ app.include_router(
     tags=["users"],
 )
 
+# Include custom auth routes (refresh token, logout)
+app.include_router(auth_router)
+
 if config.AUTH_TYPE == "GOOGLE":
     from fastapi.responses import RedirectResponse
 
diff --git a/surfsense_backend/app/routes/auth_routes.py b/surfsense_backend/app/routes/auth_routes.py
new file mode 100644
index 000000000..67abc5482
--- /dev/null
+++ b/surfsense_backend/app/routes/auth_routes.py
@@ -0,0 +1,115 @@
+"""Authentication routes for refresh token management."""
+
+import logging
+
+from fastapi import APIRouter, Cookie, Depends, HTTPException, Response, status
+from sqlalchemy import select
+
+from app.db import User, async_session_maker
+from app.schemas.auth import LogoutAllResponse, LogoutResponse, RefreshTokenResponse
+from app.users import current_active_user, get_jwt_strategy
+from app.utils.auth_cookies import (
+    REFRESH_TOKEN_COOKIE_NAME,
+    delete_refresh_token_cookie,
+    set_refresh_token_cookie,
+)
+from app.utils.refresh_tokens import (
+    revoke_all_user_tokens,
+    revoke_refresh_token,
+    rotate_refresh_token,
+    validate_refresh_token,
+)
+
+logger = logging.getLogger(__name__)
+
+router = APIRouter(prefix="/auth/jwt", tags=["auth"])
+
+
+@router.post("/refresh", response_model=RefreshTokenResponse)
+async def refresh_access_token(
+    response: Response,
+    refresh_token: str | None = Cookie(default=None, alias=REFRESH_TOKEN_COOKIE_NAME),
+):
+    """
+    Exchange a valid refresh token for a new access token and refresh token.
+    Reads refresh token from HTTP-only cookie. Implements token rotation for security.
+    """
+    if not refresh_token:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Refresh token not found",
+        )
+
+    token_record = await validate_refresh_token(refresh_token)
+
+    if not token_record:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid or expired refresh token",
+        )
+
+    # Get user from token record
+    async with async_session_maker() as session:
+        result = await session.execute(
+            select(User).where(User.id == token_record.user_id)
+        )
+        user = result.scalars().first()
+
+    if not user:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="User not found",
+        )
+
+    # Generate new access token
+    strategy = get_jwt_strategy()
+    access_token = await strategy.write_token(user)
+
+    # Rotate refresh token
+    new_refresh_token = await rotate_refresh_token(token_record)
+
+    # Set the new refresh token in cookie
+    set_refresh_token_cookie(response, new_refresh_token)
+
+    logger.info(f"Refreshed token for user {user.id}")
+
+    return RefreshTokenResponse(
+        access_token=access_token,
+        refresh_token=new_refresh_token,
+    )
+
+
+@router.post("/logout", response_model=LogoutResponse)
+async def logout(
+    response: Response,
+    refresh_token: str | None = Cookie(default=None, alias=REFRESH_TOKEN_COOKIE_NAME),
+):
+    """
+    Logout current device by revoking the refresh token from cookie.
+    """
+    if refresh_token:
+        await revoke_refresh_token(refresh_token)
+
+    # Always delete the cookie
+    delete_refresh_token_cookie(response)
+
+    logger.info("User logged out from current device")
+    return LogoutResponse()
+
+
+@router.post("/logout-all", response_model=LogoutAllResponse)
+async def logout_all_devices(
+    response: Response,
+    user: User = Depends(current_active_user),
+):
+    """
+    Logout from all devices by revoking all refresh tokens for the user.
+    Requires valid access token.
+    """
+    await revoke_all_user_tokens(user.id)
+
+    # Delete the cookie on current device
+    delete_refresh_token_cookie(response)
+
+    logger.info(f"User {user.id} logged out from all devices")
+    return LogoutAllResponse()
diff --git a/surfsense_backend/app/schemas/__init__.py b/surfsense_backend/app/schemas/__init__.py
index 6c9577c46..45dba2ba4 100644
--- a/surfsense_backend/app/schemas/__init__.py
+++ b/surfsense_backend/app/schemas/__init__.py
@@ -1,3 +1,4 @@
+from .auth import LogoutAllResponse, LogoutResponse, RefreshTokenResponse
 from .base import IDModel, TimestampModel
 from .chunks import ChunkBase, ChunkCreate, ChunkRead, ChunkUpdate
 from .documents import (
@@ -117,6 +118,9 @@ __all__ = [
     "LogFilter",
     "LogRead",
     "LogUpdate",
+    # Auth schemas
+    "LogoutAllResponse",
+    "LogoutResponse",
     # Search source connector schemas
     "MCPConnectorCreate",
     "MCPConnectorRead",
@@ -146,6 +150,7 @@ __all__ = [
     "PodcastCreate",
     "PodcastRead",
     "PodcastUpdate",
+    "RefreshTokenResponse",
     "RoleCreate",
     "RoleRead",
     "RoleUpdate",
diff --git a/surfsense_backend/app/schemas/auth.py b/surfsense_backend/app/schemas/auth.py
new file mode 100644
index 000000000..77c61de7e
--- /dev/null
+++ b/surfsense_backend/app/schemas/auth.py
@@ -0,0 +1,23 @@
+"""Authentication schemas for refresh token endpoints."""
+
+from pydantic import BaseModel
+
+
+class RefreshTokenResponse(BaseModel):
+    """Response from token refresh endpoint."""
+
+    access_token: str
+    refresh_token: str
+    token_type: str = "bearer"
+
+
+class LogoutResponse(BaseModel):
+    """Response from logout endpoint (current device)."""
+
+    detail: str = "Successfully logged out"
+
+
+class LogoutAllResponse(BaseModel):
+    """Response from logout all devices endpoint."""
+
+    detail: str = "Successfully logged out from all devices"
diff --git a/surfsense_backend/app/users.py b/surfsense_backend/app/users.py
index cbffd359d..ffb6c89e8 100644
--- a/surfsense_backend/app/users.py
+++ b/surfsense_backend/app/users.py
@@ -1,8 +1,5 @@
-import hashlib
 import logging
-import secrets
 import uuid
-from datetime import UTC, datetime, timedelta
 
 import httpx
 from fastapi import Depends, Request, Response
@@ -15,11 +12,9 @@ from fastapi_users.authentication import (
 )
 from fastapi_users.db import SQLAlchemyUserDatabase
 from pydantic import BaseModel
-from sqlalchemy import select, update
 
 from app.config import config
 from app.db import (
-    RefreshToken,
     SearchSpace,
     SearchSpaceMembership,
     SearchSpaceRole,
@@ -28,6 +23,8 @@ from app.db import (
     get_default_roles_config,
     get_user_db,
 )
+from app.utils.auth_cookies import set_refresh_token_cookie
+from app.utils.refresh_tokens import create_refresh_token
 
 logger = logging.getLogger(__name__)
 
@@ -41,123 +38,6 @@ class BearerResponse(BaseModel):
 SECRET = config.SECRET_KEY
 
 
-# Refresh token utilities (multi-session)
-def generate_refresh_token() -> str:
-    """Generate a cryptographically secure refresh token."""
-    return secrets.token_urlsafe(32)
-
-
-def hash_token(token: str) -> str:
-    """Hash a token for secure storage."""
-    return hashlib.sha256(token.encode()).hexdigest()
-
-
-async def create_refresh_token(
-    user_id: uuid.UUID,
-    family_id: uuid.UUID | None = None,
-) -> str:
-    """
-    Create and store a new refresh token for a user.
-
-    Args:
-        user_id: The user's ID
-        family_id: Optional family ID for token rotation
-
-    Returns:
-        The plaintext refresh token
-    """
-    token = generate_refresh_token()
-    token_hash = hash_token(token)
-    expires_at = datetime.now(UTC) + timedelta(
-        seconds=config.REFRESH_TOKEN_LIFETIME_SECONDS
-    )
-
-    if family_id is None:
-        family_id = uuid.uuid4()
-
-    async with async_session_maker() as session:
-        refresh_token = RefreshToken(
-            user_id=user_id,
-            token_hash=token_hash,
-            expires_at=expires_at,
-            family_id=family_id,
-        )
-        session.add(refresh_token)
-        await session.commit()
-
-    return token
-
-
-async def validate_refresh_token(token: str) -> RefreshToken | None:
-    """
-    Validate a refresh token. Handles reuse detection.
-
-    Args:
-        token: The plaintext refresh token
-
-    Returns:
-        RefreshToken if valid, None otherwise
-    """
-    token_hash = hash_token(token)
-
-    async with async_session_maker() as session:
-        result = await session.execute(
-            select(RefreshToken).where(RefreshToken.token_hash == token_hash)
-        )
-        refresh_token = result.scalars().first()
-
-        if not refresh_token:
-            return None
-
-        # Reuse detection: revoked token used while family has active tokens
-        if refresh_token.is_revoked:
-            active = await session.execute(
-                select(RefreshToken).where(
-                    RefreshToken.family_id == refresh_token.family_id,
-                    RefreshToken.is_revoked == False,  # noqa: E712
-                    RefreshToken.expires_at > datetime.now(UTC),
-                )
-            )
-            if active.scalars().first():
-                # Revoke entire family
-                await session.execute(
-                    update(RefreshToken)
-                    .where(RefreshToken.family_id == refresh_token.family_id)
-                    .values(is_revoked=True)
-                )
-                await session.commit()
-                logger.warning(f"Token reuse detected for user {refresh_token.user_id}")
-            return None
-
-        if refresh_token.is_expired:
-            return None
-
-        return refresh_token
-
-
-async def rotate_refresh_token(old_token: RefreshToken) -> str:
-    """Revoke old token and create new one in same family."""
-    async with async_session_maker() as session:
-        await session.execute(
-            update(RefreshToken)
-            .where(RefreshToken.id == old_token.id)
-            .values(is_revoked=True)
-        )
-        await session.commit()
-
-    return await create_refresh_token(old_token.user_id, old_token.family_id)
-
-
-async def revoke_all_user_tokens(user_id: uuid.UUID) -> None:
-    """Revoke all refresh tokens for a user (logout all devices)."""
-    async with async_session_maker() as session:
-        await session.execute(
-            update(RefreshToken)
-            .where(RefreshToken.user_id == user_id)
-            .values(is_revoked=True)
-        )
-        await session.commit()
-
 if config.AUTH_TYPE == "GOOGLE":
     from httpx_oauth.clients.google import GoogleOAuth2
 
@@ -358,11 +238,16 @@ class CustomBearerTransport(BearerTransport):
             redirect_url = (
                 f"{config.NEXT_FRONTEND_URL}/auth/callback"
                 f"?token={bearer_response.access_token}"
-                f"&refresh_token={bearer_response.refresh_token}"
             )
-            return RedirectResponse(redirect_url, status_code=302)
+            response = RedirectResponse(redirect_url, status_code=302)
         else:
-            return JSONResponse(bearer_response.model_dump())
+            response = JSONResponse(bearer_response.model_dump())
+
+        # Set refresh token as HTTP-only cookie
+        if refresh_token:
+            set_refresh_token_cookie(response, refresh_token)
+
+        return response
 
 
 bearer_transport = CustomBearerTransport(tokenUrl="auth/jwt/login")
diff --git a/surfsense_backend/app/utils/auth_cookies.py b/surfsense_backend/app/utils/auth_cookies.py
new file mode 100644
index 000000000..52da80f9d
--- /dev/null
+++ b/surfsense_backend/app/utils/auth_cookies.py
@@ -0,0 +1,29 @@
+"""Utilities for managing authentication cookies."""
+
+from fastapi import Response
+
+from app.config import config
+
+REFRESH_TOKEN_COOKIE_NAME = "refresh_token"
+
+
+def set_refresh_token_cookie(response: Response, token: str) -> None:
+    """Set the refresh token as an HTTP-only cookie."""
+    response.set_cookie(
+        key=REFRESH_TOKEN_COOKIE_NAME,
+        value=token,
+        max_age=config.REFRESH_TOKEN_LIFETIME_SECONDS,
+        httponly=True,
+        secure=True,  # Only send over HTTPS
+        samesite="lax",
+    )
+
+
+def delete_refresh_token_cookie(response: Response) -> None:
+    """Delete the refresh token cookie."""
+    response.delete_cookie(
+        key=REFRESH_TOKEN_COOKIE_NAME,
+        httponly=True,
+        secure=True,
+        samesite="lax",
+    )
diff --git a/surfsense_backend/app/utils/refresh_tokens.py b/surfsense_backend/app/utils/refresh_tokens.py
new file mode 100644
index 000000000..8c0312ba8
--- /dev/null
+++ b/surfsense_backend/app/utils/refresh_tokens.py
@@ -0,0 +1,153 @@
+"""Utilities for managing refresh tokens."""
+
+import hashlib
+import logging
+import secrets
+import uuid
+from datetime import UTC, datetime, timedelta
+
+from sqlalchemy import select, update
+
+from app.config import config
+from app.db import RefreshToken, async_session_maker
+
+logger = logging.getLogger(__name__)
+
+
+def generate_refresh_token() -> str:
+    """Generate a cryptographically secure refresh token."""
+    return secrets.token_urlsafe(32)
+
+
+def hash_token(token: str) -> str:
+    """Hash a token for secure storage."""
+    return hashlib.sha256(token.encode()).hexdigest()
+
+
+async def create_refresh_token(
+    user_id: uuid.UUID,
+    family_id: uuid.UUID | None = None,
+) -> str:
+    """
+    Create and store a new refresh token for a user.
+
+    Args:
+        user_id: The user's ID
+        family_id: Optional family ID for token rotation
+
+    Returns:
+        The plaintext refresh token
+    """
+    token = generate_refresh_token()
+    token_hash = hash_token(token)
+    expires_at = datetime.now(UTC) + timedelta(
+        seconds=config.REFRESH_TOKEN_LIFETIME_SECONDS
+    )
+
+    if family_id is None:
+        family_id = uuid.uuid4()
+
+    async with async_session_maker() as session:
+        refresh_token = RefreshToken(
+            user_id=user_id,
+            token_hash=token_hash,
+            expires_at=expires_at,
+            family_id=family_id,
+        )
+        session.add(refresh_token)
+        await session.commit()
+
+    return token
+
+
+async def validate_refresh_token(token: str) -> RefreshToken | None:
+    """
+    Validate a refresh token. Handles reuse detection.
+
+    Args:
+        token: The plaintext refresh token
+
+    Returns:
+        RefreshToken if valid, None otherwise
+    """
+    token_hash = hash_token(token)
+
+    async with async_session_maker() as session:
+        result = await session.execute(
+            select(RefreshToken).where(RefreshToken.token_hash == token_hash)
+        )
+        refresh_token = result.scalars().first()
+
+        if not refresh_token:
+            return None
+
+        # Reuse detection: revoked token used while family has active tokens
+        if refresh_token.is_revoked:
+            active = await session.execute(
+                select(RefreshToken).where(
+                    RefreshToken.family_id == refresh_token.family_id,
+                    RefreshToken.is_revoked == False,  # noqa: E712
+                    RefreshToken.expires_at > datetime.now(UTC),
+                )
+            )
+            if active.scalars().first():
+                # Revoke entire family
+                await session.execute(
+                    update(RefreshToken)
+                    .where(RefreshToken.family_id == refresh_token.family_id)
+                    .values(is_revoked=True)
+                )
+                await session.commit()
+                logger.warning(f"Token reuse detected for user {refresh_token.user_id}")
+            return None
+
+        if refresh_token.is_expired:
+            return None
+
+        return refresh_token
+
+
+async def rotate_refresh_token(old_token: RefreshToken) -> str:
+    """Revoke old token and create new one in same family."""
+    async with async_session_maker() as session:
+        await session.execute(
+            update(RefreshToken)
+            .where(RefreshToken.id == old_token.id)
+            .values(is_revoked=True)
+        )
+        await session.commit()
+
+    return await create_refresh_token(old_token.user_id, old_token.family_id)
+
+
+async def revoke_refresh_token(token: str) -> bool:
+    """
+    Revoke a single refresh token by its plaintext value.
+
+    Args:
+        token: The plaintext refresh token
+
+    Returns:
+        True if token was found and revoked, False otherwise
+    """
+    token_hash = hash_token(token)
+
+    async with async_session_maker() as session:
+        result = await session.execute(
+            update(RefreshToken)
+            .where(RefreshToken.token_hash == token_hash)
+            .values(is_revoked=True)
+        )
+        await session.commit()
+        return result.rowcount > 0
+
+
+async def revoke_all_user_tokens(user_id: uuid.UUID) -> None:
+    """Revoke all refresh tokens for a user (logout all devices)."""
+    async with async_session_maker() as session:
+        await session.execute(
+            update(RefreshToken)
+            .where(RefreshToken.user_id == user_id)
+            .values(is_revoked=True)
+        )
+        await session.commit()

From 233852b681d2b9f7a76fca6db73b812c0f3264de Mon Sep 17 00:00:00 2001
From: CREDO23 <thierrybakera12@gmail.com>
Date: Thu, 5 Feb 2026 17:55:21 +0200
Subject: [PATCH 30/40] Switch refresh token storage from cookies to
 localStorage

---
 surfsense_backend/app/routes/auth_routes.py |  56 +++-------
 surfsense_backend/app/schemas/__init__.py   |  10 +-
 surfsense_backend/app/schemas/auth.py       |  12 ++
 surfsense_backend/app/users.py              |  12 +-
 surfsense_backend/app/utils/auth_cookies.py |  29 -----
 surfsense_web/components/TokenHandler.tsx   |  12 +-
 surfsense_web/lib/auth-utils.ts             | 117 +++++++++++++++++++-
 7 files changed, 160 insertions(+), 88 deletions(-)
 delete mode 100644 surfsense_backend/app/utils/auth_cookies.py

diff --git a/surfsense_backend/app/routes/auth_routes.py b/surfsense_backend/app/routes/auth_routes.py
index 67abc5482..541dc9d58 100644
--- a/surfsense_backend/app/routes/auth_routes.py
+++ b/surfsense_backend/app/routes/auth_routes.py
@@ -2,17 +2,18 @@
 
 import logging
 
-from fastapi import APIRouter, Cookie, Depends, HTTPException, Response, status
+from fastapi import APIRouter, Depends, HTTPException, status
 from sqlalchemy import select
 
 from app.db import User, async_session_maker
-from app.schemas.auth import LogoutAllResponse, LogoutResponse, RefreshTokenResponse
-from app.users import current_active_user, get_jwt_strategy
-from app.utils.auth_cookies import (
-    REFRESH_TOKEN_COOKIE_NAME,
-    delete_refresh_token_cookie,
-    set_refresh_token_cookie,
+from app.schemas.auth import (
+    LogoutAllResponse,
+    LogoutRequest,
+    LogoutResponse,
+    RefreshTokenRequest,
+    RefreshTokenResponse,
 )
+from app.users import current_active_user, get_jwt_strategy
 from app.utils.refresh_tokens import (
     revoke_all_user_tokens,
     revoke_refresh_token,
@@ -26,21 +27,12 @@ router = APIRouter(prefix="/auth/jwt", tags=["auth"])
 
 
 @router.post("/refresh", response_model=RefreshTokenResponse)
-async def refresh_access_token(
-    response: Response,
-    refresh_token: str | None = Cookie(default=None, alias=REFRESH_TOKEN_COOKIE_NAME),
-):
+async def refresh_access_token(request: RefreshTokenRequest):
     """
     Exchange a valid refresh token for a new access token and refresh token.
-    Reads refresh token from HTTP-only cookie. Implements token rotation for security.
+    Implements token rotation for security.
     """
-    if not refresh_token:
-        raise HTTPException(
-            status_code=status.HTTP_401_UNAUTHORIZED,
-            detail="Refresh token not found",
-        )
-
-    token_record = await validate_refresh_token(refresh_token)
+    token_record = await validate_refresh_token(request.refresh_token)
 
     if not token_record:
         raise HTTPException(
@@ -68,9 +60,6 @@ async def refresh_access_token(
     # Rotate refresh token
     new_refresh_token = await rotate_refresh_token(token_record)
 
-    # Set the new refresh token in cookie
-    set_refresh_token_cookie(response, new_refresh_token)
-
     logger.info(f"Refreshed token for user {user.id}")
 
     return RefreshTokenResponse(
@@ -80,36 +69,21 @@ async def refresh_access_token(
 
 
 @router.post("/logout", response_model=LogoutResponse)
-async def logout(
-    response: Response,
-    refresh_token: str | None = Cookie(default=None, alias=REFRESH_TOKEN_COOKIE_NAME),
-):
+async def logout(request: LogoutRequest):
     """
-    Logout current device by revoking the refresh token from cookie.
+    Logout current device by revoking the provided refresh token.
     """
-    if refresh_token:
-        await revoke_refresh_token(refresh_token)
-
-    # Always delete the cookie
-    delete_refresh_token_cookie(response)
-
+    await revoke_refresh_token(request.refresh_token)
     logger.info("User logged out from current device")
     return LogoutResponse()
 
 
 @router.post("/logout-all", response_model=LogoutAllResponse)
-async def logout_all_devices(
-    response: Response,
-    user: User = Depends(current_active_user),
-):
+async def logout_all_devices(user: User = Depends(current_active_user)):
     """
     Logout from all devices by revoking all refresh tokens for the user.
     Requires valid access token.
     """
     await revoke_all_user_tokens(user.id)
-
-    # Delete the cookie on current device
-    delete_refresh_token_cookie(response)
-
     logger.info(f"User {user.id} logged out from all devices")
     return LogoutAllResponse()
diff --git a/surfsense_backend/app/schemas/__init__.py b/surfsense_backend/app/schemas/__init__.py
index 45dba2ba4..5ff166733 100644
--- a/surfsense_backend/app/schemas/__init__.py
+++ b/surfsense_backend/app/schemas/__init__.py
@@ -1,4 +1,10 @@
-from .auth import LogoutAllResponse, LogoutResponse, RefreshTokenResponse
+from .auth import (
+    LogoutAllResponse,
+    LogoutRequest,
+    LogoutResponse,
+    RefreshTokenRequest,
+    RefreshTokenResponse,
+)
 from .base import IDModel, TimestampModel
 from .chunks import ChunkBase, ChunkCreate, ChunkRead, ChunkUpdate
 from .documents import (
@@ -120,6 +126,7 @@ __all__ = [
     "LogUpdate",
     # Auth schemas
     "LogoutAllResponse",
+    "LogoutRequest",
     "LogoutResponse",
     # Search source connector schemas
     "MCPConnectorCreate",
@@ -150,6 +157,7 @@ __all__ = [
     "PodcastCreate",
     "PodcastRead",
     "PodcastUpdate",
+    "RefreshTokenRequest",
     "RefreshTokenResponse",
     "RoleCreate",
     "RoleRead",
diff --git a/surfsense_backend/app/schemas/auth.py b/surfsense_backend/app/schemas/auth.py
index 77c61de7e..0d958a6d2 100644
--- a/surfsense_backend/app/schemas/auth.py
+++ b/surfsense_backend/app/schemas/auth.py
@@ -3,6 +3,12 @@
 from pydantic import BaseModel
 
 
+class RefreshTokenRequest(BaseModel):
+    """Request body for token refresh endpoint."""
+
+    refresh_token: str
+
+
 class RefreshTokenResponse(BaseModel):
     """Response from token refresh endpoint."""
 
@@ -11,6 +17,12 @@ class RefreshTokenResponse(BaseModel):
     token_type: str = "bearer"
 
 
+class LogoutRequest(BaseModel):
+    """Request body for logout endpoint (current device)."""
+
+    refresh_token: str
+
+
 class LogoutResponse(BaseModel):
     """Response from logout endpoint (current device)."""
 
diff --git a/surfsense_backend/app/users.py b/surfsense_backend/app/users.py
index ffb6c89e8..aef94d558 100644
--- a/surfsense_backend/app/users.py
+++ b/surfsense_backend/app/users.py
@@ -23,7 +23,6 @@ from app.db import (
     get_default_roles_config,
     get_user_db,
 )
-from app.utils.auth_cookies import set_refresh_token_cookie
 from app.utils.refresh_tokens import create_refresh_token
 
 logger = logging.getLogger(__name__)
@@ -238,16 +237,11 @@ class CustomBearerTransport(BearerTransport):
             redirect_url = (
                 f"{config.NEXT_FRONTEND_URL}/auth/callback"
                 f"?token={bearer_response.access_token}"
+                f"&refresh_token={bearer_response.refresh_token}"
             )
-            response = RedirectResponse(redirect_url, status_code=302)
+            return RedirectResponse(redirect_url, status_code=302)
         else:
-            response = JSONResponse(bearer_response.model_dump())
-
-        # Set refresh token as HTTP-only cookie
-        if refresh_token:
-            set_refresh_token_cookie(response, refresh_token)
-
-        return response
+            return JSONResponse(bearer_response.model_dump())
 
 
 bearer_transport = CustomBearerTransport(tokenUrl="auth/jwt/login")
diff --git a/surfsense_backend/app/utils/auth_cookies.py b/surfsense_backend/app/utils/auth_cookies.py
deleted file mode 100644
index 52da80f9d..000000000
--- a/surfsense_backend/app/utils/auth_cookies.py
+++ /dev/null
@@ -1,29 +0,0 @@
-"""Utilities for managing authentication cookies."""
-
-from fastapi import Response
-
-from app.config import config
-
-REFRESH_TOKEN_COOKIE_NAME = "refresh_token"
-
-
-def set_refresh_token_cookie(response: Response, token: str) -> None:
-    """Set the refresh token as an HTTP-only cookie."""
-    response.set_cookie(
-        key=REFRESH_TOKEN_COOKIE_NAME,
-        value=token,
-        max_age=config.REFRESH_TOKEN_LIFETIME_SECONDS,
-        httponly=True,
-        secure=True,  # Only send over HTTPS
-        samesite="lax",
-    )
-
-
-def delete_refresh_token_cookie(response: Response) -> None:
-    """Delete the refresh token cookie."""
-    response.delete_cookie(
-        key=REFRESH_TOKEN_COOKIE_NAME,
-        httponly=True,
-        secure=True,
-        samesite="lax",
-    )
diff --git a/surfsense_web/components/TokenHandler.tsx b/surfsense_web/components/TokenHandler.tsx
index e3295df7c..230cda81a 100644
--- a/surfsense_web/components/TokenHandler.tsx
+++ b/surfsense_web/components/TokenHandler.tsx
@@ -3,7 +3,7 @@
 import { useSearchParams } from "next/navigation";
 import { useEffect } from "react";
 import { useGlobalLoadingEffect } from "@/hooks/use-global-loading";
-import { getAndClearRedirectPath, setBearerToken } from "@/lib/auth-utils";
+import { getAndClearRedirectPath, setBearerToken, setRefreshToken } from "@/lib/auth-utils";
 import { trackLoginSuccess } from "@/lib/posthog/events";
 
 interface TokenHandlerProps {
@@ -35,8 +35,9 @@ const TokenHandler = ({
 		// Only run on client-side
 		if (typeof window === "undefined") return;
 
-		// Get token from URL parameters
+		// Get tokens from URL parameters
 		const token = searchParams.get(tokenParamName);
+		const refreshToken = searchParams.get("refresh_token");
 
 		if (token) {
 			try {
@@ -50,10 +51,15 @@ const TokenHandler = ({
 				// Clear the flag for future logins
 				sessionStorage.removeItem("login_success_tracked");
 
-				// Store token in localStorage using both methods for compatibility
+				// Store access token in localStorage using both methods for compatibility
 				localStorage.setItem(storageKey, token);
 				setBearerToken(token);
 
+				// Store refresh token if provided
+				if (refreshToken) {
+					setRefreshToken(refreshToken);
+				}
+
 				// Check if there's a saved redirect path from before the auth flow
 				const savedRedirectPath = getAndClearRedirectPath();
 
diff --git a/surfsense_web/lib/auth-utils.ts b/surfsense_web/lib/auth-utils.ts
index 604843292..24db377a8 100644
--- a/surfsense_web/lib/auth-utils.ts
+++ b/surfsense_web/lib/auth-utils.ts
@@ -4,6 +4,11 @@
 
 const REDIRECT_PATH_KEY = "surfsense_redirect_path";
 const BEARER_TOKEN_KEY = "surfsense_bearer_token";
+const REFRESH_TOKEN_KEY = "surfsense_refresh_token";
+
+// Flag to prevent multiple simultaneous refresh attempts
+let isRefreshing = false;
+let refreshPromise: Promise<string | null> | null = null;
 
 /**
  * Saves the current path and redirects to login page
@@ -21,8 +26,9 @@ export function handleUnauthorized(): void {
 		localStorage.setItem(REDIRECT_PATH_KEY, currentPath);
 	}
 
-	// Clear the token
+	// Clear both tokens
 	localStorage.removeItem(BEARER_TOKEN_KEY);
+	localStorage.removeItem(REFRESH_TOKEN_KEY);
 
 	// Redirect to home page (which has login options)
 	window.location.href = "/login";
@@ -66,6 +72,38 @@ export function clearBearerToken(): void {
 	localStorage.removeItem(BEARER_TOKEN_KEY);
 }
 
+/**
+ * Gets the refresh token from localStorage
+ */
+export function getRefreshToken(): string | null {
+	if (typeof window === "undefined") return null;
+	return localStorage.getItem(REFRESH_TOKEN_KEY);
+}
+
+/**
+ * Sets the refresh token in localStorage
+ */
+export function setRefreshToken(token: string): void {
+	if (typeof window === "undefined") return;
+	localStorage.setItem(REFRESH_TOKEN_KEY, token);
+}
+
+/**
+ * Clears the refresh token from localStorage
+ */
+export function clearRefreshToken(): void {
+	if (typeof window === "undefined") return;
+	localStorage.removeItem(REFRESH_TOKEN_KEY);
+}
+
+/**
+ * Clears all auth tokens from localStorage
+ */
+export function clearAllTokens(): void {
+	clearBearerToken();
+	clearRefreshToken();
+}
+
 /**
  * Checks if the user is authenticated (has a token)
  */
@@ -106,14 +144,66 @@ export function getAuthHeaders(additionalHeaders?: Record<string, string>): Reco
 }
 
 /**
- * Authenticated fetch wrapper that handles 401 responses uniformly
- * Automatically redirects to login on 401 and saves the current path
+ * Attempts to refresh the access token using the stored refresh token.
+ * Returns the new access token if successful, null otherwise.
+ */
+async function refreshAccessToken(): Promise<string | null> {
+	// If already refreshing, wait for that request to complete
+	if (isRefreshing && refreshPromise) {
+		return refreshPromise;
+	}
+
+	const currentRefreshToken = getRefreshToken();
+	if (!currentRefreshToken) {
+		return null;
+	}
+
+	isRefreshing = true;
+	refreshPromise = (async () => {
+		try {
+			const backendUrl = process.env.NEXT_PUBLIC_BACKEND_URL || "";
+			const response = await fetch(`${backendUrl}/auth/jwt/refresh`, {
+				method: "POST",
+				headers: {
+					"Content-Type": "application/json",
+				},
+				body: JSON.stringify({ refresh_token: currentRefreshToken }),
+			});
+
+			if (!response.ok) {
+				// Refresh failed, clear tokens
+				clearAllTokens();
+				return null;
+			}
+
+			const data = await response.json();
+			if (data.access_token && data.refresh_token) {
+				setBearerToken(data.access_token);
+				setRefreshToken(data.refresh_token);
+				return data.access_token;
+			}
+			return null;
+		} catch {
+			return null;
+		} finally {
+			isRefreshing = false;
+			refreshPromise = null;
+		}
+	})();
+
+	return refreshPromise;
+}
+
+/**
+ * Authenticated fetch wrapper that handles 401 responses uniformly.
+ * On 401, attempts to refresh the token and retry the request.
+ * If refresh fails, redirects to login and saves the current path.
  */
 export async function authenticatedFetch(
 	url: string,
-	options?: RequestInit & { skipAuthRedirect?: boolean }
+	options?: RequestInit & { skipAuthRedirect?: boolean; skipRefresh?: boolean }
 ): Promise<Response> {
-	const { skipAuthRedirect = false, ...fetchOptions } = options || {};
+	const { skipAuthRedirect = false, skipRefresh = false, ...fetchOptions } = options || {};
 
 	const headers = getAuthHeaders(fetchOptions.headers as Record<string, string>);
 
@@ -124,6 +214,23 @@ export async function authenticatedFetch(
 
 	// Handle 401 Unauthorized
 	if (response.status === 401 && !skipAuthRedirect) {
+		// Try to refresh the token (unless skipRefresh is set to prevent infinite loops)
+		if (!skipRefresh) {
+			const newToken = await refreshAccessToken();
+			if (newToken) {
+				// Retry the original request with the new token
+				const retryHeaders = {
+					...(fetchOptions.headers as Record<string, string>),
+					Authorization: `Bearer ${newToken}`,
+				};
+				return fetch(url, {
+					...fetchOptions,
+					headers: retryHeaders,
+				});
+			}
+		}
+
+		// Refresh failed or was skipped, redirect to login
 		handleUnauthorized();
 		throw new Error("Unauthorized: Redirecting to login page");
 	}

From 287e5afbac0c52a309dd70ec49fed539500efa80 Mon Sep 17 00:00:00 2001
From: CREDO23 <thierrybakera12@gmail.com>
Date: Thu, 5 Feb 2026 18:11:33 +0200
Subject: [PATCH 31/40] Fix JWT audience validation when creating refresh token

---
 surfsense_backend/app/users.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/surfsense_backend/app/users.py b/surfsense_backend/app/users.py
index aef94d558..696cdf25e 100644
--- a/surfsense_backend/app/users.py
+++ b/surfsense_backend/app/users.py
@@ -219,7 +219,7 @@ class CustomBearerTransport(BearerTransport):
 
         # Decode JWT to get user_id for refresh token creation
         try:
-            payload = jwt.decode(token, SECRET, algorithms=["HS256"])
+            payload = jwt.decode(token, SECRET, algorithms=["HS256"], options={"verify_aud": False})
             user_id = uuid.UUID(payload.get("sub"))
             refresh_token = await create_refresh_token(user_id)
         except Exception as e:

From f13345b226be4f280c7d0c22185c67f3219837f0 Mon Sep 17 00:00:00 2001
From: CREDO23 <thierrybakera12@gmail.com>
Date: Thu, 5 Feb 2026 18:56:38 +0200
Subject: [PATCH 32/40] Add logout with token revocation and loading states

---
 surfsense_backend/app/routes/auth_routes.py   | 12 ++++---
 surfsense_web/components/UserDropdown.tsx     | 29 +++++++++++----
 .../layout/providers/LayoutDataProvider.tsx   |  6 +++-
 .../layout/ui/sidebar/SidebarUserProfile.tsx  | 34 ++++++++++++++----
 surfsense_web/lib/auth-utils.ts               | 35 ++++++++++++++++++-
 surfsense_web/messages/en.json                |  1 +
 surfsense_web/messages/zh.json                |  1 +
 7 files changed, 98 insertions(+), 20 deletions(-)

diff --git a/surfsense_backend/app/routes/auth_routes.py b/surfsense_backend/app/routes/auth_routes.py
index 541dc9d58..b1cbaf2a5 100644
--- a/surfsense_backend/app/routes/auth_routes.py
+++ b/surfsense_backend/app/routes/auth_routes.py
@@ -68,13 +68,17 @@ async def refresh_access_token(request: RefreshTokenRequest):
     )
 
 
-@router.post("/logout", response_model=LogoutResponse)
-async def logout(request: LogoutRequest):
+@router.post("/revoke", response_model=LogoutResponse)
+async def revoke_token(request: LogoutRequest):
     """
     Logout current device by revoking the provided refresh token.
+    Does not require authentication - just the refresh token.
     """
-    await revoke_refresh_token(request.refresh_token)
-    logger.info("User logged out from current device")
+    revoked = await revoke_refresh_token(request.refresh_token)
+    if revoked:
+        logger.info("User logged out from current device - token revoked")
+    else:
+        logger.warning("Logout called but no matching token found to revoke")
     return LogoutResponse()
 
 
diff --git a/surfsense_web/components/UserDropdown.tsx b/surfsense_web/components/UserDropdown.tsx
index 3dac745cf..233a41a1f 100644
--- a/surfsense_web/components/UserDropdown.tsx
+++ b/surfsense_web/components/UserDropdown.tsx
@@ -1,7 +1,8 @@
 "use client";
 
-import { BadgeCheck, LogOut } from "lucide-react";
+import { BadgeCheck, Loader2, LogOut } from "lucide-react";
 import { useRouter } from "next/navigation";
+import { useState } from "react";
 import { Avatar, AvatarFallback, AvatarImage } from "@/components/ui/avatar";
 import { Button } from "@/components/ui/button";
 import {
@@ -13,6 +14,7 @@ import {
 	DropdownMenuSeparator,
 	DropdownMenuTrigger,
 } from "@/components/ui/dropdown-menu";
+import { logout } from "@/lib/auth-utils";
 import { cleanupElectric } from "@/lib/electric/client";
 import { resetUser, trackLogout } from "@/lib/posthog/events";
 
@@ -26,8 +28,11 @@ export function UserDropdown({
 	};
 }) {
 	const router = useRouter();
+	const [isLoggingOut, setIsLoggingOut] = useState(false);
 
 	const handleLogout = async () => {
+		if (isLoggingOut) return;
+		setIsLoggingOut(true);
 		try {
 			// Track logout event and reset PostHog identity
 			trackLogout();
@@ -41,15 +46,17 @@ export function UserDropdown({
 				console.warn("[Logout] Electric cleanup failed (will be handled on next login):", err);
 			}
 
+			// Revoke refresh token on server and clear all tokens from localStorage
+			await logout();
+
 			if (typeof window !== "undefined") {
-				localStorage.removeItem("surfsense_bearer_token");
 				window.location.href = "/";
 			}
 		} catch (error) {
 			console.error("Error during logout:", error);
-			// Optionally, provide user feedback
+			// Even if there's an error, try to clear tokens and redirect
+			await logout();
 			if (typeof window !== "undefined") {
-				localStorage.removeItem("surfsense_bearer_token");
 				window.location.href = "/";
 			}
 		}
@@ -85,9 +92,17 @@ export function UserDropdown({
 					</DropdownMenuItem>
 				</DropdownMenuGroup>
 				<DropdownMenuSeparator />
-				<DropdownMenuItem onClick={handleLogout} className="text-xs md:text-sm">
-					<LogOut className="mr-2 h-3.5 w-3.5 md:h-4 md:w-4" />
-					Log out
+				<DropdownMenuItem
+					onClick={handleLogout}
+					className="text-xs md:text-sm"
+					disabled={isLoggingOut}
+				>
+					{isLoggingOut ? (
+						<Loader2 className="mr-2 h-3.5 w-3.5 md:h-4 md:w-4 animate-spin" />
+					) : (
+						<LogOut className="mr-2 h-3.5 w-3.5 md:h-4 md:w-4" />
+					)}
+					{isLoggingOut ? "Logging out..." : "Log out"}
 				</DropdownMenuItem>
 			</DropdownMenuContent>
 		</DropdownMenu>
diff --git a/surfsense_web/components/layout/providers/LayoutDataProvider.tsx b/surfsense_web/components/layout/providers/LayoutDataProvider.tsx
index 2f71adad9..68350bce1 100644
--- a/surfsense_web/components/layout/providers/LayoutDataProvider.tsx
+++ b/surfsense_web/components/layout/providers/LayoutDataProvider.tsx
@@ -25,6 +25,7 @@ import { isPageLimitExceededMetadata } from "@/contracts/types/inbox.types";
 import { useInbox } from "@/hooks/use-inbox";
 import { searchSpacesApiService } from "@/lib/apis/search-spaces-api.service";
 import { deleteThread, fetchThreads, updateThread } from "@/lib/chat/thread-persistence";
+import { logout } from "@/lib/auth-utils";
 import { cleanupElectric } from "@/lib/electric/client";
 import { resetUser, trackLogout } from "@/lib/posthog/events";
 import { cacheKeys } from "@/lib/query-client/cache-keys";
@@ -464,12 +465,15 @@ export function LayoutDataProvider({
 				console.warn("[Logout] Electric cleanup failed (will be handled on next login):", err);
 			}
 
+			// Revoke refresh token on server and clear all tokens from localStorage
+			await logout();
+
 			if (typeof window !== "undefined") {
-				localStorage.removeItem("surfsense_bearer_token");
 				router.push("/");
 			}
 		} catch (error) {
 			console.error("Error during logout:", error);
+			await logout();
 			router.push("/");
 		}
 	}, [router]);
diff --git a/surfsense_web/components/layout/ui/sidebar/SidebarUserProfile.tsx b/surfsense_web/components/layout/ui/sidebar/SidebarUserProfile.tsx
index 7c96b1dcb..38b3028d2 100644
--- a/surfsense_web/components/layout/ui/sidebar/SidebarUserProfile.tsx
+++ b/surfsense_web/components/layout/ui/sidebar/SidebarUserProfile.tsx
@@ -1,7 +1,8 @@
 "use client";
 
-import { Check, ChevronUp, Languages, Laptop, LogOut, Moon, Settings, Sun } from "lucide-react";
+import { Check, ChevronUp, Languages, Laptop, Loader2, LogOut, Moon, Settings, Sun } from "lucide-react";
 import { useTranslations } from "next-intl";
+import { useState } from "react";
 import {
 	DropdownMenu,
 	DropdownMenuContent,
@@ -124,6 +125,7 @@ export function SidebarUserProfile({
 }: SidebarUserProfileProps) {
 	const t = useTranslations("sidebar");
 	const { locale, setLocale } = useLocaleContext();
+	const [isLoggingOut, setIsLoggingOut] = useState(false);
 	const bgColor = stringToColor(user.email);
 	const initials = getInitials(user.email);
 	const displayName = user.name || user.email.split("@")[0];
@@ -136,6 +138,16 @@ export function SidebarUserProfile({
 		setTheme?.(newTheme);
 	};
 
+	const handleLogout = async () => {
+		if (isLoggingOut || !onLogout) return;
+		setIsLoggingOut(true);
+		try {
+			await onLogout();
+		} finally {
+			setIsLoggingOut(false);
+		}
+	};
+
 	// Collapsed view - just show avatar with dropdown
 	if (isCollapsed) {
 		return (
@@ -242,9 +254,13 @@ export function SidebarUserProfile({
 
 						<DropdownMenuSeparator />
 
-						<DropdownMenuItem onClick={onLogout}>
-							<LogOut className="mr-2 h-4 w-4" />
-							{t("logout")}
+						<DropdownMenuItem onClick={handleLogout} disabled={isLoggingOut}>
+							{isLoggingOut ? (
+								<Loader2 className="mr-2 h-4 w-4 animate-spin" />
+							) : (
+								<LogOut className="mr-2 h-4 w-4" />
+							)}
+							{isLoggingOut ? t("loggingOut") : t("logout")}
 						</DropdownMenuItem>
 					</DropdownMenuContent>
 				</DropdownMenu>
@@ -360,9 +376,13 @@ export function SidebarUserProfile({
 
 					<DropdownMenuSeparator />
 
-					<DropdownMenuItem onClick={onLogout}>
-						<LogOut className="mr-2 h-4 w-4" />
-						{t("logout")}
+					<DropdownMenuItem onClick={handleLogout} disabled={isLoggingOut}>
+						{isLoggingOut ? (
+							<Loader2 className="mr-2 h-4 w-4 animate-spin" />
+						) : (
+							<LogOut className="mr-2 h-4 w-4" />
+						)}
+						{isLoggingOut ? t("loggingOut") : t("logout")}
 					</DropdownMenuItem>
 				</DropdownMenuContent>
 			</DropdownMenu>
diff --git a/surfsense_web/lib/auth-utils.ts b/surfsense_web/lib/auth-utils.ts
index 24db377a8..409f29d0c 100644
--- a/surfsense_web/lib/auth-utils.ts
+++ b/surfsense_web/lib/auth-utils.ts
@@ -104,6 +104,39 @@ export function clearAllTokens(): void {
 	clearRefreshToken();
 }
 
+/**
+ * Logout the current user by revoking the refresh token and clearing localStorage.
+ * Returns true if logout was successful (or tokens were cleared), false otherwise.
+ */
+export async function logout(): Promise<boolean> {
+	const refreshToken = getRefreshToken();
+
+	// Call backend to revoke the refresh token
+	if (refreshToken) {
+		try {
+			const backendUrl = process.env.NEXT_PUBLIC_FASTAPI_BACKEND_URL || "http://localhost:8000";
+			const response = await fetch(`${backendUrl}/auth/jwt/revoke`, {
+				method: "POST",
+				headers: {
+					"Content-Type": "application/json",
+				},
+				body: JSON.stringify({ refresh_token: refreshToken }),
+			});
+
+			if (!response.ok) {
+				console.warn("Failed to revoke refresh token:", response.status, await response.text());
+			}
+		} catch (error) {
+			console.warn("Failed to revoke refresh token on server:", error);
+			// Continue to clear local tokens even if server call fails
+		}
+	}
+
+	// Clear all tokens from localStorage
+	clearAllTokens();
+	return true;
+}
+
 /**
  * Checks if the user is authenticated (has a token)
  */
@@ -161,7 +194,7 @@ async function refreshAccessToken(): Promise<string | null> {
 	isRefreshing = true;
 	refreshPromise = (async () => {
 		try {
-			const backendUrl = process.env.NEXT_PUBLIC_BACKEND_URL || "";
+			const backendUrl = process.env.NEXT_PUBLIC_FASTAPI_BACKEND_URL || "http://localhost:8000";
 			const response = await fetch(`${backendUrl}/auth/jwt/refresh`, {
 				method: "POST",
 				headers: {
diff --git a/surfsense_web/messages/en.json b/surfsense_web/messages/en.json
index 75b186420..93877096b 100644
--- a/surfsense_web/messages/en.json
+++ b/surfsense_web/messages/en.json
@@ -693,6 +693,7 @@
 		"dark": "Dark",
 		"system": "System",
 		"logout": "Logout",
+		"loggingOut": "Logging out...",
 		"inbox": "Inbox",
 		"search_inbox": "Search inbox",
 		"mark_all_read": "Mark all as read",
diff --git a/surfsense_web/messages/zh.json b/surfsense_web/messages/zh.json
index 81121ef3e..639bf31f4 100644
--- a/surfsense_web/messages/zh.json
+++ b/surfsense_web/messages/zh.json
@@ -678,6 +678,7 @@
 		"dark": "深色",
 		"system": "系统",
 		"logout": "退出登录",
+		"loggingOut": "正在退出...",
 		"inbox": "收件箱",
 		"search_inbox": "搜索收件箱",
 		"mark_all_read": "全部标记为已读",

From 1690ce1891580b80bc816514110b0c3a4036b8bb Mon Sep 17 00:00:00 2001
From: CREDO23 <thierrybakera12@gmail.com>
Date: Thu, 5 Feb 2026 18:59:15 +0200
Subject: [PATCH 33/40] Document token lifetime env vars in .env.example

---
 surfsense_backend/.env.example | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/surfsense_backend/.env.example b/surfsense_backend/.env.example
index 2e10f4e36..628329917 100644
--- a/surfsense_backend/.env.example
+++ b/surfsense_backend/.env.example
@@ -32,6 +32,11 @@ ELECTRIC_DB_PASSWORD=electric_password
 SCHEDULE_CHECKER_INTERVAL=5m
 
 SECRET_KEY=SECRET
+
+# JWT Token Lifetimes (optional, defaults shown)
+# ACCESS_TOKEN_LIFETIME_SECONDS=86400      # 1 day
+# REFRESH_TOKEN_LIFETIME_SECONDS=1209600   # 2 weeks
+
 NEXT_FRONTEND_URL=http://localhost:3000
 
 # Backend URL for OAuth callbacks (optional, set when behind reverse proxy with HTTPS)

From a8e8b0e2fd81720ea5f6245682262d7eaa4c6c95 Mon Sep 17 00:00:00 2001
From: CREDO23 <thierrybakera12@gmail.com>
Date: Thu, 5 Feb 2026 19:29:34 +0200
Subject: [PATCH 34/40] Fix auto-refresh on 401 in base API service

---
 surfsense_web/lib/apis/base-api.service.ts | 20 ++++++++++++++++++--
 surfsense_web/lib/auth-utils.ts            |  3 ++-
 2 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/surfsense_web/lib/apis/base-api.service.ts b/surfsense_web/lib/apis/base-api.service.ts
index a87d4deaf..933e54656 100644
--- a/surfsense_web/lib/apis/base-api.service.ts
+++ b/surfsense_web/lib/apis/base-api.service.ts
@@ -1,5 +1,5 @@
 import type { ZodType } from "zod";
-import { getBearerToken, handleUnauthorized } from "../auth-utils";
+import { getBearerToken, handleUnauthorized, refreshAccessToken } from "../auth-utils";
 import { AppError, AuthenticationError, AuthorizationError, NotFoundError } from "../error";
 
 enum ResponseType {
@@ -17,6 +17,7 @@ export type RequestOptions = {
 	signal?: AbortSignal;
 	body?: any;
 	responseType?: ResponseType;
+	_isRetry?: boolean; // Internal flag to prevent infinite retry loops
 	// Add more options as needed
 };
 
@@ -135,8 +136,23 @@ class BaseApiService {
 					throw new AppError("Failed to parse response", response.status, response.statusText);
 				}
 
-				// Handle 401 first before other error handling - ensures token is cleared and user redirected
+				// Handle 401 - try to refresh token first (only once)
 				if (response.status === 401) {
+					if (!options?._isRetry) {
+						const newToken = await refreshAccessToken();
+						if (newToken) {
+							// Retry the request with the new token
+							return this.request(url, responseSchema, {
+								...mergedOptions,
+								headers: {
+									...mergedOptions.headers,
+									Authorization: `Bearer ${newToken}`,
+								},
+								_isRetry: true,
+							} as RequestOptions & { responseType?: R });
+						}
+					}
+					// Refresh failed or retry failed, redirect to login
 					handleUnauthorized();
 					throw new AuthenticationError(
 						typeof data === "object" && "detail" in data
diff --git a/surfsense_web/lib/auth-utils.ts b/surfsense_web/lib/auth-utils.ts
index 409f29d0c..8c067a4b7 100644
--- a/surfsense_web/lib/auth-utils.ts
+++ b/surfsense_web/lib/auth-utils.ts
@@ -179,8 +179,9 @@ export function getAuthHeaders(additionalHeaders?: Record<string, string>): Reco
 /**
  * Attempts to refresh the access token using the stored refresh token.
  * Returns the new access token if successful, null otherwise.
+ * Exported for use by API services.
  */
-async function refreshAccessToken(): Promise<string | null> {
+export async function refreshAccessToken(): Promise<string | null> {
 	// If already refreshing, wait for that request to complete
 	if (isRefreshing && refreshPromise) {
 		return refreshPromise;

From 19e28573434593dc5c513ee2afd086f01ffa6688 Mon Sep 17 00:00:00 2001
From: "DESKTOP-RTLN3BA\\$punk" <vermarohanfinal@gmail.com>
Date: Thu, 5 Feb 2026 16:43:48 -0800
Subject: [PATCH 35/40] feat: added image gen support

---
 .../93_add_image_generations_table.py         | 292 ++++++++
 ...4_add_access_token_to_image_generations.py |  39 +
 .../app/agents/new_chat/chat_deepagent.py     |   1 +
 .../app/agents/new_chat/system_prompt.py      |  39 +-
 .../app/agents/new_chat/tools/__init__.py     |   3 +
 .../agents/new_chat/tools/display_image.py    |  12 +-
 .../agents/new_chat/tools/generate_image.py   | 254 +++++++
 .../app/agents/new_chat/tools/registry.py     |  11 +
 surfsense_backend/app/app.py                  |   4 +-
 surfsense_backend/app/celery_app.py           |   7 +-
 surfsense_backend/app/config/__init__.py      |  83 +++
 .../app/config/global_llm_config.example.yaml |  74 ++
 surfsense_backend/app/db.py                   | 155 ++++
 surfsense_backend/app/routes/__init__.py      |   2 +
 .../app/routes/image_generation_routes.py     | 646 ++++++++++++++++
 .../app/routes/search_spaces_routes.py        |  74 ++
 surfsense_backend/app/schemas/__init__.py     |  20 +
 .../app/schemas/image_generation.py           | 231 ++++++
 .../app/schemas/new_llm_config.py             |   9 +
 .../app/services/image_gen_router_service.py  | 294 ++++++++
 .../app/utils/signed_image_urls.py            |  44 ++
 .../[search_space_id]/settings/page.tsx       |  15 +-
 .../image-gen-config-mutation.atoms.ts        |  91 +++
 .../image-gen-config-query.atoms.ts           |  33 +
 surfsense_web/components/Logo.tsx             |  26 +-
 surfsense_web/components/homepage/navbar.tsx  |   4 +-
 .../components/new-chat/chat-header.tsx       |  10 +-
 .../new-chat/image-model-selector.tsx         | 335 +++++++++
 .../components/new-chat/model-selector.tsx    |   4 +-
 .../settings/image-model-manager.tsx          | 692 ++++++++++++++++++
 .../components/settings/llm-role-manager.tsx  | 202 ++---
 .../components/tool-ui/display-image.tsx      |   2 +-
 .../components/tool-ui/image/index.tsx        | 128 ++--
 .../contracts/enums/image-gen-providers.ts    | 105 +++
 .../contracts/types/new-llm-config.types.ts   |  98 ++-
 .../lib/apis/image-gen-config-api.service.ts  |  83 +++
 surfsense_web/lib/query-client/cache-keys.ts  |   5 +
 surfsense_web/messages/en.json                |   2 +
 surfsense_web/messages/zh.json                |   2 +
 39 files changed, 3950 insertions(+), 181 deletions(-)
 create mode 100644 surfsense_backend/alembic/versions/93_add_image_generations_table.py
 create mode 100644 surfsense_backend/alembic/versions/94_add_access_token_to_image_generations.py
 create mode 100644 surfsense_backend/app/agents/new_chat/tools/generate_image.py
 create mode 100644 surfsense_backend/app/routes/image_generation_routes.py
 create mode 100644 surfsense_backend/app/schemas/image_generation.py
 create mode 100644 surfsense_backend/app/services/image_gen_router_service.py
 create mode 100644 surfsense_backend/app/utils/signed_image_urls.py
 create mode 100644 surfsense_web/atoms/image-gen-config/image-gen-config-mutation.atoms.ts
 create mode 100644 surfsense_web/atoms/image-gen-config/image-gen-config-query.atoms.ts
 create mode 100644 surfsense_web/components/new-chat/image-model-selector.tsx
 create mode 100644 surfsense_web/components/settings/image-model-manager.tsx
 create mode 100644 surfsense_web/contracts/enums/image-gen-providers.ts
 create mode 100644 surfsense_web/lib/apis/image-gen-config-api.service.ts

diff --git a/surfsense_backend/alembic/versions/93_add_image_generations_table.py b/surfsense_backend/alembic/versions/93_add_image_generations_table.py
new file mode 100644
index 000000000..151208229
--- /dev/null
+++ b/surfsense_backend/alembic/versions/93_add_image_generations_table.py
@@ -0,0 +1,292 @@
+"""Add image generation tables and search space preference
+
+Revision ID: 93
+Revises: 92
+
+Changes:
+1. Create image_generation_configs table (user-created image model configs)
+2. Create image_generations table (stores generation requests/results)
+3. Add image_generation_config_id column to searchspaces table
+4. Add image generation permissions to existing system roles
+"""
+
+from collections.abc import Sequence
+
+import sqlalchemy as sa
+from sqlalchemy.dialects.postgresql import JSONB, UUID
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "93"
+down_revision: str | None = "92"
+branch_labels: str | Sequence[str] | None = None
+depends_on: str | Sequence[str] | None = None
+
+
+def upgrade() -> None:
+    connection = op.get_bind()
+
+    # 1. Create imagegenprovider enum type if it doesn't exist
+    connection.execute(
+        sa.text(
+            """
+            DO $$
+            BEGIN
+                IF NOT EXISTS (SELECT 1 FROM pg_type WHERE typname = 'imagegenprovider') THEN
+                    CREATE TYPE imagegenprovider AS ENUM (
+                        'OPENAI', 'AZURE_OPENAI', 'GOOGLE', 'VERTEX_AI', 'BEDROCK',
+                        'RECRAFT', 'OPENROUTER', 'XINFERENCE', 'NSCALE'
+                    );
+                END IF;
+            END
+            $$;
+            """
+        )
+    )
+
+    # 2. Create image_generation_configs table (uses imagegenprovider enum)
+    result = connection.execute(
+        sa.text(
+            "SELECT EXISTS (SELECT FROM information_schema.tables WHERE table_name = 'image_generation_configs')"
+        )
+    )
+    if not result.scalar():
+        op.create_table(
+            "image_generation_configs",
+            sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
+            sa.Column("name", sa.String(100), nullable=False),
+            sa.Column("description", sa.String(500), nullable=True),
+            sa.Column(
+                "provider",
+                sa.Enum(
+                    "OPENAI", "AZURE_OPENAI", "GOOGLE", "VERTEX_AI", "BEDROCK",
+                    "RECRAFT", "OPENROUTER", "XINFERENCE", "NSCALE",
+                    name="imagegenprovider",
+                    create_type=False,
+                ),
+                nullable=False,
+            ),
+            sa.Column("custom_provider", sa.String(100), nullable=True),
+            sa.Column("model_name", sa.String(100), nullable=False),
+            sa.Column("api_key", sa.String(), nullable=False),
+            sa.Column("api_base", sa.String(500), nullable=True),
+            sa.Column("api_version", sa.String(50), nullable=True),
+            sa.Column("litellm_params", sa.JSON(), nullable=True),
+            sa.Column("search_space_id", sa.Integer(), nullable=False),
+            sa.Column(
+                "created_at",
+                sa.TIMESTAMP(timezone=True),
+                server_default=sa.text("now()"),
+                nullable=False,
+            ),
+            sa.PrimaryKeyConstraint("id"),
+            sa.ForeignKeyConstraint(
+                ["search_space_id"], ["searchspaces.id"], ondelete="CASCADE"
+            ),
+        )
+        op.execute(
+            "CREATE INDEX IF NOT EXISTS ix_image_generation_configs_name "
+            "ON image_generation_configs (name)"
+        )
+        op.execute(
+            "CREATE INDEX IF NOT EXISTS ix_image_generation_configs_search_space_id "
+            "ON image_generation_configs (search_space_id)"
+        )
+
+    # 3. Create image_generations table
+    result = connection.execute(
+        sa.text(
+            "SELECT EXISTS (SELECT FROM information_schema.tables WHERE table_name = 'image_generations')"
+        )
+    )
+    if not result.scalar():
+        op.create_table(
+            "image_generations",
+            sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
+            sa.Column("prompt", sa.Text(), nullable=False),
+            sa.Column("model", sa.String(200), nullable=True),
+            sa.Column("n", sa.Integer(), nullable=True),
+            sa.Column("quality", sa.String(50), nullable=True),
+            sa.Column("size", sa.String(50), nullable=True),
+            sa.Column("style", sa.String(50), nullable=True),
+            sa.Column("response_format", sa.String(50), nullable=True),
+            sa.Column("image_generation_config_id", sa.Integer(), nullable=True),
+            sa.Column("response_data", JSONB(), nullable=True),
+            sa.Column("error_message", sa.Text(), nullable=True),
+            sa.Column("search_space_id", sa.Integer(), nullable=False),
+            sa.Column("created_by_id", UUID(as_uuid=True), nullable=True),
+            sa.Column(
+                "created_at",
+                sa.TIMESTAMP(timezone=True),
+                server_default=sa.text("now()"),
+                nullable=False,
+            ),
+            sa.PrimaryKeyConstraint("id"),
+            sa.ForeignKeyConstraint(
+                ["search_space_id"], ["searchspaces.id"], ondelete="CASCADE"
+            ),
+            sa.ForeignKeyConstraint(
+                ["created_by_id"], ["user.id"], ondelete="SET NULL"
+            ),
+        )
+
+    op.execute(
+        "CREATE INDEX IF NOT EXISTS ix_image_generations_search_space_id "
+        "ON image_generations (search_space_id)"
+    )
+    op.execute(
+        "CREATE INDEX IF NOT EXISTS ix_image_generations_created_by_id "
+        "ON image_generations (created_by_id)"
+    )
+    op.execute(
+        "CREATE INDEX IF NOT EXISTS ix_image_generations_created_at "
+        "ON image_generations (created_at)"
+    )
+
+    # 4. Add image_generation_config_id column to searchspaces
+    result = connection.execute(
+        sa.text(
+            """
+            SELECT EXISTS (
+                SELECT 1 FROM information_schema.columns
+                WHERE table_name = 'searchspaces'
+                  AND column_name = 'image_generation_config_id'
+            )
+            """
+        )
+    )
+    if not result.scalar():
+        op.add_column(
+            "searchspaces",
+            sa.Column(
+                "image_generation_config_id",
+                sa.Integer(),
+                nullable=True,
+                server_default="0",
+            ),
+        )
+
+    # Drop old column name if it exists (from earlier version of this migration)
+    result = connection.execute(
+        sa.text(
+            """
+            SELECT EXISTS (
+                SELECT 1 FROM information_schema.columns
+                WHERE table_name = 'searchspaces'
+                  AND column_name = 'image_generation_llm_id'
+            )
+            """
+        )
+    )
+    if result.scalar():
+        op.drop_column("searchspaces", "image_generation_llm_id")
+
+    # Drop old column name on image_generations if it exists
+    result = connection.execute(
+        sa.text(
+            """
+            SELECT EXISTS (
+                SELECT 1 FROM information_schema.columns
+                WHERE table_name = 'image_generations'
+                  AND column_name = 'llm_config_id'
+            )
+            """
+        )
+    )
+    if result.scalar():
+        op.drop_column("image_generations", "llm_config_id")
+
+    # Drop old api_version column on image_generations if it exists
+    result = connection.execute(
+        sa.text(
+            """
+            SELECT EXISTS (
+                SELECT 1 FROM information_schema.columns
+                WHERE table_name = 'image_generations'
+                  AND column_name = 'api_version'
+            )
+            """
+        )
+    )
+    if result.scalar():
+        op.drop_column("image_generations", "api_version")
+
+    # 5. Add image generation permissions to existing system roles
+    connection.execute(
+        sa.text(
+            """
+            UPDATE search_space_roles
+            SET permissions = array_cat(
+                permissions,
+                ARRAY['image_generations:create', 'image_generations:read']
+            )
+            WHERE is_system_role = true
+              AND name = 'Editor'
+              AND NOT ('image_generations:create' = ANY(permissions))
+            """
+        )
+    )
+    connection.execute(
+        sa.text(
+            """
+            UPDATE search_space_roles
+            SET permissions = array_cat(
+                permissions,
+                ARRAY['image_generations:read']
+            )
+            WHERE is_system_role = true
+              AND name = 'Viewer'
+              AND NOT ('image_generations:read' = ANY(permissions))
+            """
+        )
+    )
+
+
+def downgrade() -> None:
+    connection = op.get_bind()
+
+    # Remove permissions
+    connection.execute(
+        sa.text(
+            """
+            UPDATE search_space_roles
+            SET permissions = array_remove(
+                array_remove(
+                    array_remove(permissions, 'image_generations:create'),
+                    'image_generations:read'
+                ),
+                'image_generations:delete'
+            )
+            WHERE is_system_role = true
+            """
+        )
+    )
+
+    # Remove image_generation_config_id from searchspaces
+    result = connection.execute(
+        sa.text(
+            """
+            SELECT EXISTS (
+                SELECT 1 FROM information_schema.columns
+                WHERE table_name = 'searchspaces'
+                  AND column_name = 'image_generation_config_id'
+            )
+            """
+        )
+    )
+    if result.scalar():
+        op.drop_column("searchspaces", "image_generation_config_id")
+
+    # Drop indexes and tables
+    op.execute("DROP INDEX IF EXISTS ix_image_generations_created_at")
+    op.execute("DROP INDEX IF EXISTS ix_image_generations_created_by_id")
+    op.execute("DROP INDEX IF EXISTS ix_image_generations_search_space_id")
+    op.execute("DROP TABLE IF EXISTS image_generations")
+
+    op.execute("DROP INDEX IF EXISTS ix_image_generation_configs_search_space_id")
+    op.execute("DROP INDEX IF EXISTS ix_image_generation_configs_name")
+    op.execute("DROP TABLE IF EXISTS image_generation_configs")
+
+    # Drop the imagegenprovider enum type
+    op.execute("DROP TYPE IF EXISTS imagegenprovider")
diff --git a/surfsense_backend/alembic/versions/94_add_access_token_to_image_generations.py b/surfsense_backend/alembic/versions/94_add_access_token_to_image_generations.py
new file mode 100644
index 000000000..09bea2c19
--- /dev/null
+++ b/surfsense_backend/alembic/versions/94_add_access_token_to_image_generations.py
@@ -0,0 +1,39 @@
+"""Add access_token column to image_generations
+
+Revision ID: 94
+Revises: 93
+
+Adds an indexed access_token column to the image_generations table.
+This token is stored per-record so that image serving URLs survive
+SECRET_KEY rotation.
+"""
+
+from collections.abc import Sequence
+
+import sqlalchemy as sa
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "94"
+down_revision: str | None = "93"
+branch_labels: str | Sequence[str] | None = None
+depends_on: str | Sequence[str] | None = None
+
+
+def upgrade() -> None:
+    # Add access_token column (nullable so existing rows are unaffected)
+    op.add_column(
+        "image_generations",
+        sa.Column("access_token", sa.String(64), nullable=True),
+    )
+    op.create_index(
+        "ix_image_generations_access_token",
+        "image_generations",
+        ["access_token"],
+    )
+
+
+def downgrade() -> None:
+    op.drop_index("ix_image_generations_access_token", table_name="image_generations")
+    op.drop_column("image_generations", "access_token")
diff --git a/surfsense_backend/app/agents/new_chat/chat_deepagent.py b/surfsense_backend/app/agents/new_chat/chat_deepagent.py
index 9c383c308..9da6ea3c2 100644
--- a/surfsense_backend/app/agents/new_chat/chat_deepagent.py
+++ b/surfsense_backend/app/agents/new_chat/chat_deepagent.py
@@ -133,6 +133,7 @@ async def create_surfsense_deep_agent(
     The agent comes with built-in tools that can be configured:
     - search_knowledge_base: Search the user's personal knowledge base
     - generate_podcast: Generate audio podcasts from content
+    - generate_image: Generate images from text descriptions using AI models
     - link_preview: Fetch rich previews for URLs
     - display_image: Display images in chat
     - scrape_webpage: Extract content from webpages
diff --git a/surfsense_backend/app/agents/new_chat/system_prompt.py b/surfsense_backend/app/agents/new_chat/system_prompt.py
index fdb80acb9..38bae230d 100644
--- a/surfsense_backend/app/agents/new_chat/system_prompt.py
+++ b/surfsense_backend/app/agents/new_chat/system_prompt.py
@@ -83,6 +83,7 @@ You have access to the following tools:
     * Showing an image from a URL the user explicitly mentioned in their message
     * Displaying images found in scraped webpage content (from scrape_webpage tool)
     * Showing a publicly accessible diagram or chart from a known URL
+    * Displaying an AI-generated image after calling the generate_image tool (ALWAYS required)
   
   CRITICAL - NEVER USE THIS TOOL FOR USER-UPLOADED ATTACHMENTS:
   When a user uploads/attaches an image file to their message:
@@ -100,7 +101,23 @@ You have access to the following tools:
   - Returns: An image card with the image, title, and description
   - The image will automatically be displayed in the chat.
 
-5. scrape_webpage: Scrape and extract the main content from a webpage.
+5. generate_image: Generate images from text descriptions using AI image models.
+  - Use this when the user asks you to create, generate, draw, design, or make an image.
+  - Trigger phrases: "generate an image of", "create a picture of", "draw me", "make an image", "design a logo", "create artwork"
+  - Args:
+    - prompt: A detailed text description of the image to generate. Be specific about subject, style, colors, composition, and mood.
+    - size: Image size. Options: "1024x1024" (square, default), "1536x1024" (landscape), "1024x1536" (portrait), "1792x1024" (wide)
+    - quality: Image quality. Options: "auto" (default), "high", "medium", "low"
+    - n: Number of images to generate (1-4, default: 1)
+  - Returns: A dictionary with the generated image URL in the "src" field, along with metadata.
+  - CRITICAL: After calling generate_image, you MUST call `display_image` with the returned "src" URL
+    to actually show the image in the chat. The generate_image tool only generates the image and returns
+    the URL — it does NOT display anything. You must always follow up with display_image.
+  - IMPORTANT: Write a detailed, descriptive prompt for best results. Don't just pass the user's words verbatim -
+    expand and improve the prompt with specific details about style, lighting, composition, and mood.
+  - If the user's request is vague (e.g., "make me an image of a cat"), enhance the prompt with artistic details.
+
+6. scrape_webpage: Scrape and extract the main content from a webpage.
   - Use this when the user wants you to READ and UNDERSTAND the actual content of a webpage.
   - IMPORTANT: This is different from link_preview:
     * link_preview: Only fetches metadata (title, description, thumbnail) for display
@@ -123,7 +140,7 @@ You have access to the following tools:
     * Prioritize showing: diagrams, charts, infographics, key illustrations, or images that help explain the content.
     * Don't show every image - just the most relevant 1-3 images that enhance understanding.
 
-6. save_memory: Save facts, preferences, or context about the user for personalized responses.
+7. save_memory: Save facts, preferences, or context about the user for personalized responses.
   - Use this when the user explicitly or implicitly shares information worth remembering.
   - Trigger scenarios:
     * User says "remember this", "keep this in mind", "note that", or similar
@@ -146,7 +163,7 @@ You have access to the following tools:
   - IMPORTANT: Only save information that would be genuinely useful for future conversations.
     Don't save trivial or temporary information.
 
-7. recall_memory: Retrieve relevant memories about the user for personalized responses.
+8. recall_memory: Retrieve relevant memories about the user for personalized responses.
   - Use this to access stored information about the user.
   - Trigger scenarios:
     * You need user context to give a better, more personalized answer
@@ -281,6 +298,22 @@ You have access to the following tools:
   - Then, if the content contains useful diagrams/images like `![Neural Network Diagram](https://example.com/nn-diagram.png)`:
     - Call: `display_image(src="https://example.com/nn-diagram.png", alt="Neural Network Diagram", title="Neural Network Architecture")`
   - Then provide your explanation, referencing the displayed image
+
+- User: "Generate an image of a cat"
+  - Step 1: `generate_image(prompt="A fluffy orange tabby cat sitting on a windowsill, bathed in warm golden sunlight, soft bokeh background with green houseplants, photorealistic style, cozy atmosphere", size="1024x1024", quality="auto")`
+  - Step 2: Use the returned "src" URL to display it: `display_image(src="<returned_url>", alt="A fluffy orange tabby cat on a windowsill", title="Generated Image")`
+
+- User: "Create a landscape painting of mountains"
+  - Step 1: `generate_image(prompt="Majestic snow-capped mountain range at sunset, dramatic orange and purple sky, alpine meadow with wildflowers in the foreground, oil painting style with visible brushstrokes, inspired by the Hudson River School art movement", size="1536x1024", quality="high")`
+  - Step 2: `display_image(src="<returned_url>", alt="Mountain landscape painting", title="Generated Image")`
+
+- User: "Draw me a logo for a coffee shop called Bean Dream"
+  - Step 1: `generate_image(prompt="Minimalist modern logo design for a coffee shop called 'Bean Dream', featuring a stylized coffee bean with dream-like swirls of steam, clean vector style, warm brown and cream color palette, white background, professional branding", size="1024x1024", quality="high")`
+  - Step 2: `display_image(src="<returned_url>", alt="Bean Dream coffee shop logo", title="Generated Image")`
+
+- User: "Make a wide banner image for my blog about AI"
+  - Step 1: `generate_image(prompt="Wide banner illustration for an AI technology blog, featuring abstract neural network patterns, glowing blue and purple connections, modern futuristic aesthetic, digital art style, clean and professional", size="1792x1024", quality="high")`
+  - Step 2: `display_image(src="<returned_url>", alt="AI blog banner", title="Generated Image")`
 </tool_call_examples>
 """
 
diff --git a/surfsense_backend/app/agents/new_chat/tools/__init__.py b/surfsense_backend/app/agents/new_chat/tools/__init__.py
index 9e1a4f19c..0a11951f0 100644
--- a/surfsense_backend/app/agents/new_chat/tools/__init__.py
+++ b/surfsense_backend/app/agents/new_chat/tools/__init__.py
@@ -8,6 +8,7 @@ Available tools:
 - search_knowledge_base: Search the user's personal knowledge base
 - search_surfsense_docs: Search Surfsense documentation for usage help
 - generate_podcast: Generate audio podcasts from content
+- generate_image: Generate images from text descriptions using AI models
 - link_preview: Fetch rich previews for URLs
 - display_image: Display images in chat
 - scrape_webpage: Extract content from webpages
@@ -18,6 +19,7 @@ Available tools:
 # Registry exports
 # Tool factory exports (for direct use)
 from .display_image import create_display_image_tool
+from .generate_image import create_generate_image_tool
 from .knowledge_base import (
     CONNECTOR_DESCRIPTIONS,
     create_search_knowledge_base_tool,
@@ -47,6 +49,7 @@ __all__ = [
     "build_tools",
     # Tool factories
     "create_display_image_tool",
+    "create_generate_image_tool",
     "create_generate_podcast_tool",
     "create_link_preview_tool",
     "create_recall_memory_tool",
diff --git a/surfsense_backend/app/agents/new_chat/tools/display_image.py b/surfsense_backend/app/agents/new_chat/tools/display_image.py
index 5eb846063..4424cc0d3 100644
--- a/surfsense_backend/app/agents/new_chat/tools/display_image.py
+++ b/surfsense_backend/app/agents/new_chat/tools/display_image.py
@@ -82,14 +82,20 @@ def create_display_image_tool():
 
         domain = extract_domain(src)
 
-        # Determine aspect ratio based on common image sources
-        ratio = "16:9"  # Default
-        if "unsplash.com" in src or "pexels.com" in src:
+        # Determine aspect ratio based on image source
+        # AI-generated images should use "auto" to preserve their native ratio
+        is_generated = "/image-generations/" in src
+        if is_generated:
+            ratio = "auto"
+            domain = "ai-generated"
+        elif "unsplash.com" in src or "pexels.com" in src:
             ratio = "16:9"
         elif (
             "imgur.com" in src or "github.com" in src or "githubusercontent.com" in src
         ):
             ratio = "auto"
+        else:
+            ratio = "auto"
 
         return {
             "id": image_id,
diff --git a/surfsense_backend/app/agents/new_chat/tools/generate_image.py b/surfsense_backend/app/agents/new_chat/tools/generate_image.py
new file mode 100644
index 000000000..091fb122f
--- /dev/null
+++ b/surfsense_backend/app/agents/new_chat/tools/generate_image.py
@@ -0,0 +1,254 @@
+"""
+Image generation tool for the SurfSense agent.
+
+This module provides a tool that generates images using litellm.aimage_generation()
+and returns the result via the existing display_image tool format so the frontend
+renders the generated image inline in the chat.
+
+Config resolution:
+1. Uses the search space's image_generation_config_id preference
+2. Falls back to Auto mode (router load balancing) if available
+3. Supports global YAML configs (negative IDs) and user DB configs (positive IDs)
+"""
+
+import logging
+from typing import Any
+
+from langchain_core.tools import tool
+from litellm import aimage_generation
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from app.config import config
+from app.db import ImageGeneration, ImageGenerationConfig, SearchSpace
+from app.utils.signed_image_urls import generate_image_token
+from app.services.image_gen_router_service import (
+    IMAGE_GEN_AUTO_MODE_ID,
+    ImageGenRouterService,
+    is_image_gen_auto_mode,
+)
+
+logger = logging.getLogger(__name__)
+
+# Provider mapping (same as routes)
+_PROVIDER_MAP = {
+    "OPENAI": "openai",
+    "AZURE_OPENAI": "azure",
+    "GOOGLE": "gemini",
+    "VERTEX_AI": "vertex_ai",
+    "BEDROCK": "bedrock",
+    "RECRAFT": "recraft",
+    "OPENROUTER": "openrouter",
+    "XINFERENCE": "xinference",
+    "NSCALE": "nscale",
+}
+
+
+def _build_model_string(
+    provider: str, model_name: str, custom_provider: str | None
+) -> str:
+    if custom_provider:
+        return f"{custom_provider}/{model_name}"
+    prefix = _PROVIDER_MAP.get(provider.upper(), provider.lower())
+    return f"{prefix}/{model_name}"
+
+
+def _get_global_image_gen_config(config_id: int) -> dict | None:
+    """Get a global image gen config by negative ID."""
+    for cfg in config.GLOBAL_IMAGE_GEN_CONFIGS:
+        if cfg.get("id") == config_id:
+            return cfg
+    return None
+
+
+def create_generate_image_tool(
+    search_space_id: int,
+    db_session: AsyncSession,
+):
+    """
+    Factory function to create the generate_image tool.
+
+    Args:
+        search_space_id: The search space ID (for config resolution)
+        db_session: Async database session
+    """
+
+    @tool
+    async def generate_image(
+        prompt: str,
+        size: str = "1024x1024",
+        quality: str = "auto",
+        n: int = 1,
+    ) -> dict[str, Any]:
+        """
+        Generate an image from a text description using AI image models.
+
+        Use this tool when the user asks you to create, generate, draw, or make an image.
+        The generated image will be displayed directly in the chat.
+
+        Args:
+            prompt: A detailed text description of the image to generate.
+                    Be specific about subject, style, colors, composition, and mood.
+            size: Image size. Options: "1024x1024" (square), "1536x1024" (landscape),
+                  "1024x1536" (portrait), "1792x1024" (wide). Default: "1024x1024"
+            quality: Image quality. Options: "auto" (default), "high", "medium", "low".
+                     Default: "auto"
+            n: Number of images to generate (1-4). Default: 1
+
+        Returns:
+            A dictionary containing the generated image(s) for display in the chat.
+        """
+        try:
+            # Resolve the image generation config from the search space preference
+            result = await db_session.execute(
+                select(SearchSpace).filter(SearchSpace.id == search_space_id)
+            )
+            search_space = result.scalars().first()
+            if not search_space:
+                return {"error": "Search space not found"}
+
+            config_id = (
+                search_space.image_generation_config_id or IMAGE_GEN_AUTO_MODE_ID
+            )
+
+            # Build generation kwargs
+            # NOTE: 'style' is intentionally excluded from gen_kwargs because
+            # it is only supported by DALL-E 3 and causes errors with other
+            # models (e.g. gpt-image-1 rejects it as an unknown parameter).
+            # Since we can't predict which model auto-mode will route to,
+            # it's safest to omit it.
+            gen_kwargs: dict[str, Any] = {}
+            if n is not None and n > 1:
+                gen_kwargs["n"] = n
+            if quality:
+                gen_kwargs["quality"] = quality
+            if size:
+                gen_kwargs["size"] = size
+
+            # Call litellm based on config type
+            if is_image_gen_auto_mode(config_id):
+                if not ImageGenRouterService.is_initialized():
+                    return {
+                        "error": "No image generation models configured. "
+                        "Please add an image model in Settings > Image Models."
+                    }
+                response = await ImageGenRouterService.aimage_generation(
+                    prompt=prompt, model="auto", **gen_kwargs
+                )
+            elif config_id < 0:
+                cfg = _get_global_image_gen_config(config_id)
+                if not cfg:
+                    return {"error": f"Image generation config {config_id} not found"}
+
+                model_string = _build_model_string(
+                    cfg.get("provider", ""),
+                    cfg["model_name"],
+                    cfg.get("custom_provider"),
+                )
+                gen_kwargs["api_key"] = cfg.get("api_key")
+                if cfg.get("api_base"):
+                    gen_kwargs["api_base"] = cfg["api_base"]
+                if cfg.get("api_version"):
+                    gen_kwargs["api_version"] = cfg["api_version"]
+                if cfg.get("litellm_params"):
+                    gen_kwargs.update(cfg["litellm_params"])
+
+                response = await aimage_generation(
+                    prompt=prompt, model=model_string, **gen_kwargs
+                )
+            else:
+                # Positive ID = user-created ImageGenerationConfig
+                cfg_result = await db_session.execute(
+                    select(ImageGenerationConfig).filter(
+                        ImageGenerationConfig.id == config_id
+                    )
+                )
+                db_cfg = cfg_result.scalars().first()
+                if not db_cfg:
+                    return {"error": f"Image generation config {config_id} not found"}
+
+                model_string = _build_model_string(
+                    db_cfg.provider.value,
+                    db_cfg.model_name,
+                    db_cfg.custom_provider,
+                )
+                gen_kwargs["api_key"] = db_cfg.api_key
+                if db_cfg.api_base:
+                    gen_kwargs["api_base"] = db_cfg.api_base
+                if db_cfg.api_version:
+                    gen_kwargs["api_version"] = db_cfg.api_version
+                if db_cfg.litellm_params:
+                    gen_kwargs.update(db_cfg.litellm_params)
+
+                response = await aimage_generation(
+                    prompt=prompt, model=model_string, **gen_kwargs
+                )
+
+            # Parse the response and store in DB
+            response_dict = (
+                response.model_dump()
+                if hasattr(response, "model_dump")
+                else dict(response)
+            )
+
+            # Generate a random access token for this image
+            access_token = generate_image_token()
+
+            # Save to image_generations table for history
+            db_image_gen = ImageGeneration(
+                prompt=prompt,
+                model=getattr(response, "_hidden_params", {}).get("model"),
+                n=n,
+                quality=quality,
+                size=size,
+                image_generation_config_id=config_id,
+                response_data=response_dict,
+                search_space_id=search_space_id,
+                access_token=access_token,
+            )
+            db_session.add(db_image_gen)
+            await db_session.commit()
+            await db_session.refresh(db_image_gen)
+
+            # Extract image URLs from response
+            images = response_dict.get("data", [])
+            if not images:
+                return {"error": "No images were generated"}
+
+            first_image = images[0]
+            revised_prompt = first_image.get("revised_prompt", prompt)
+
+            # Resolve image URL:
+            # - If the API returned a URL, use it directly.
+            # - If the API returned b64_json (e.g. gpt-image-1), serve the
+            #   image through our backend endpoint to avoid bloating the
+            #   LLM context with megabytes of base64 data.
+            if first_image.get("url"):
+                image_url = first_image["url"]
+            elif first_image.get("b64_json"):
+                backend_url = config.BACKEND_URL or "http://localhost:8000"
+                image_url = (
+                    f"{backend_url}/api/v1/image-generations/"
+                    f"{db_image_gen.id}/image?token={access_token}"
+                )
+            else:
+                return {"error": "No displayable image data in the response"}
+
+            return {
+                "src": image_url,
+                "alt": revised_prompt or prompt,
+                "title": "Generated Image",
+                "description": revised_prompt if revised_prompt != prompt else None,
+                "generated": True,
+                "prompt": prompt,
+                "image_count": len(images),
+            }
+
+        except Exception as e:
+            logger.exception("Image generation failed in tool")
+            return {
+                "error": f"Image generation failed: {e!s}",
+                "prompt": prompt,
+            }
+
+    return generate_image
diff --git a/surfsense_backend/app/agents/new_chat/tools/registry.py b/surfsense_backend/app/agents/new_chat/tools/registry.py
index c65445419..2cf43c973 100644
--- a/surfsense_backend/app/agents/new_chat/tools/registry.py
+++ b/surfsense_backend/app/agents/new_chat/tools/registry.py
@@ -44,6 +44,7 @@ from typing import Any
 from langchain_core.tools import BaseTool
 
 from .display_image import create_display_image_tool
+from .generate_image import create_generate_image_tool
 from .knowledge_base import create_search_knowledge_base_tool
 from .link_preview import create_link_preview_tool
 from .mcp_tool import load_mcp_tools
@@ -125,6 +126,16 @@ BUILTIN_TOOLS: list[ToolDefinition] = [
         factory=lambda deps: create_display_image_tool(),
         requires=[],
     ),
+    # Generate image tool - creates images using AI models (DALL-E, GPT Image, etc.)
+    ToolDefinition(
+        name="generate_image",
+        description="Generate images from text descriptions using AI image models",
+        factory=lambda deps: create_generate_image_tool(
+            search_space_id=deps["search_space_id"],
+            db_session=deps["db_session"],
+        ),
+        requires=["search_space_id", "db_session"],
+    ),
     # Web scraping tool - extracts content from webpages
     ToolDefinition(
         name="scrape_webpage",
diff --git a/surfsense_backend/app/app.py b/surfsense_backend/app/app.py
index 63da4e8ad..0f619097e 100644
--- a/surfsense_backend/app/app.py
+++ b/surfsense_backend/app/app.py
@@ -9,7 +9,7 @@ from app.agents.new_chat.checkpointer import (
     close_checkpointer,
     setup_checkpointer_tables,
 )
-from app.config import config, initialize_llm_router
+from app.config import config, initialize_image_gen_router, initialize_llm_router
 from app.db import User, create_db_and_tables, get_async_session
 from app.routes import router as crud_router
 from app.routes.auth_routes import router as auth_router
@@ -26,6 +26,8 @@ async def lifespan(app: FastAPI):
     await setup_checkpointer_tables()
     # Initialize LLM Router for Auto mode load balancing
     initialize_llm_router()
+    # Initialize Image Generation Router for Auto mode load balancing
+    initialize_image_gen_router()
     # Seed Surfsense documentation
     await seed_surfsense_docs()
     yield
diff --git a/surfsense_backend/app/celery_app.py b/surfsense_backend/app/celery_app.py
index 74b21fbf0..af406eab7 100644
--- a/surfsense_backend/app/celery_app.py
+++ b/surfsense_backend/app/celery_app.py
@@ -13,14 +13,15 @@ load_dotenv()
 
 @worker_process_init.connect
 def init_worker(**kwargs):
-    """Initialize the LLM Router when a Celery worker process starts.
+    """Initialize the LLM Router and Image Gen Router when a Celery worker process starts.
 
     This ensures the Auto mode (LiteLLM Router) is available for background tasks
-    like document summarization.
+    like document summarization and image generation.
     """
-    from app.config import initialize_llm_router
+    from app.config import initialize_image_gen_router, initialize_llm_router
 
     initialize_llm_router()
+    initialize_image_gen_router()
 
 
 # Get Celery configuration from environment
diff --git a/surfsense_backend/app/config/__init__.py b/surfsense_backend/app/config/__init__.py
index 121e5d3b2..bb299e583 100644
--- a/surfsense_backend/app/config/__init__.py
+++ b/surfsense_backend/app/config/__init__.py
@@ -81,6 +81,56 @@ def load_router_settings():
         return default_settings
 
 
+def load_global_image_gen_configs():
+    """
+    Load global image generation configurations from YAML file.
+
+    Returns:
+        list: List of global image generation config dictionaries, or empty list
+    """
+    global_config_file = BASE_DIR / "app" / "config" / "global_llm_config.yaml"
+
+    if not global_config_file.exists():
+        return []
+
+    try:
+        with open(global_config_file, encoding="utf-8") as f:
+            data = yaml.safe_load(f)
+            return data.get("global_image_generation_configs", [])
+    except Exception as e:
+        print(f"Warning: Failed to load global image generation configs: {e}")
+        return []
+
+
+def load_image_gen_router_settings():
+    """
+    Load router settings for image generation Auto mode from YAML file.
+
+    Returns:
+        dict: Router settings dictionary
+    """
+    default_settings = {
+        "routing_strategy": "usage-based-routing",
+        "num_retries": 3,
+        "allowed_fails": 3,
+        "cooldown_time": 60,
+    }
+
+    global_config_file = BASE_DIR / "app" / "config" / "global_llm_config.yaml"
+
+    if not global_config_file.exists():
+        return default_settings
+
+    try:
+        with open(global_config_file, encoding="utf-8") as f:
+            data = yaml.safe_load(f)
+            settings = data.get("image_generation_router_settings", {})
+            return {**default_settings, **settings}
+    except Exception as e:
+        print(f"Warning: Failed to load image generation router settings: {e}")
+        return default_settings
+
+
 def initialize_llm_router():
     """
     Initialize the LLM Router service for Auto mode.
@@ -105,6 +155,33 @@ def initialize_llm_router():
         print(f"Warning: Failed to initialize LLM Router: {e}")
 
 
+def initialize_image_gen_router():
+    """
+    Initialize the Image Generation Router service for Auto mode.
+    This should be called during application startup.
+    """
+    image_gen_configs = load_global_image_gen_configs()
+    router_settings = load_image_gen_router_settings()
+
+    if not image_gen_configs:
+        print(
+            "Info: No global image generation configs found, "
+            "Image Generation Auto mode will not be available"
+        )
+        return
+
+    try:
+        from app.services.image_gen_router_service import ImageGenRouterService
+
+        ImageGenRouterService.initialize(image_gen_configs, router_settings)
+        print(
+            f"Info: Image Generation Router initialized with {len(image_gen_configs)} models "
+            f"(strategy: {router_settings.get('routing_strategy', 'usage-based-routing')})"
+        )
+    except Exception as e:
+        print(f"Warning: Failed to initialize Image Generation Router: {e}")
+
+
 class Config:
     # Check if ffmpeg is installed
     if not is_ffmpeg_installed():
@@ -216,6 +293,12 @@ class Config:
     # Router settings for Auto mode (LiteLLM Router load balancing)
     ROUTER_SETTINGS = load_router_settings()
 
+    # Global Image Generation Configurations (optional)
+    GLOBAL_IMAGE_GEN_CONFIGS = load_global_image_gen_configs()
+
+    # Router settings for Image Generation Auto mode
+    IMAGE_GEN_ROUTER_SETTINGS = load_image_gen_router_settings()
+
     # Chonkie Configuration | Edit this to your needs
     EMBEDDING_MODEL = os.getenv("EMBEDDING_MODEL")
     # Azure OpenAI credentials from environment variables
diff --git a/surfsense_backend/app/config/global_llm_config.example.yaml b/surfsense_backend/app/config/global_llm_config.example.yaml
index 9b213aafe..e727b8d56 100644
--- a/surfsense_backend/app/config/global_llm_config.example.yaml
+++ b/surfsense_backend/app/config/global_llm_config.example.yaml
@@ -183,6 +183,73 @@ global_llm_configs:
     use_default_system_instructions: true
     citations_enabled: true
 
+# =============================================================================
+# Image Generation Configuration
+# =============================================================================
+# These configurations power the image generation feature using litellm.aimage_generation().
+# Supported providers: OpenAI, Azure, Google AI Studio, Vertex AI, AWS Bedrock,
+# Recraft, OpenRouter, Xinference, Nscale
+#
+# Auto mode (ID 0) uses LiteLLM Router for load balancing across all image gen configs.
+
+# Router Settings for Image Generation Auto Mode
+image_generation_router_settings:
+  routing_strategy: "usage-based-routing"
+  num_retries: 3
+  allowed_fails: 3
+  cooldown_time: 60
+
+global_image_generation_configs:
+  # Example: OpenAI DALL-E 3
+  - id: -1
+    name: "Global DALL-E 3"
+    description: "OpenAI's DALL-E 3 for high-quality image generation"
+    provider: "OPENAI"
+    model_name: "dall-e-3"
+    api_key: "sk-your-openai-api-key-here"
+    api_base: ""
+    rpm: 50
+    tpm: 100000
+    litellm_params: {}
+
+  # Example: OpenAI GPT Image 1
+  - id: -2
+    name: "Global GPT Image 1"
+    description: "OpenAI's GPT Image 1 model"
+    provider: "OPENAI"
+    model_name: "gpt-image-1"
+    api_key: "sk-your-openai-api-key-here"
+    api_base: ""
+    rpm: 50
+    tpm: 100000
+    litellm_params: {}
+
+  # Example: Azure OpenAI DALL-E 3
+  - id: -3
+    name: "Global Azure DALL-E 3"
+    description: "Azure-hosted DALL-E 3 deployment"
+    provider: "AZURE_OPENAI"
+    model_name: "azure/dall-e-3-deployment"
+    api_key: "your-azure-api-key-here"
+    api_base: "https://your-resource.openai.azure.com"
+    api_version: "2024-02-15-preview"
+    rpm: 50
+    tpm: 100000
+    litellm_params:
+      base_model: "dall-e-3"
+
+  # Example: OpenRouter Gemini Image Generation
+  # - id: -4
+  #   name: "Global Gemini Image Gen"
+  #   description: "Google Gemini image generation via OpenRouter"
+  #   provider: "OPENROUTER"
+  #   model_name: "google/gemini-2.5-flash-image"
+  #   api_key: "your-openrouter-api-key-here"
+  #   api_base: ""
+  #   rpm: 30
+  #   tpm: 50000
+  #   litellm_params: {}
+
 # Notes:
 # - ID 0 is reserved for "Auto" mode - uses LiteLLM Router for load balancing
 # - Use negative IDs to distinguish global configs from user configs (NewLLMConfig in DB)
@@ -202,3 +269,10 @@ global_llm_configs:
 # - model_name format: "azure/<your-deployment-name>"
 # - api_version: Use a recent Azure API version (e.g., "2024-02-15-preview")
 # - See: https://docs.litellm.ai/docs/proxy/cost_tracking#spend-tracking-for-azure-openai-models
+#
+# IMAGE GENERATION NOTES:
+# - Image generation configs use the same ID scheme as LLM configs (negative for global)
+# - Supported models: dall-e-2, dall-e-3, gpt-image-1 (OpenAI), azure/* (Azure), 
+#   bedrock/* (AWS), vertex_ai/* (Google), recraft/* (Recraft), openrouter/* (OpenRouter)
+# - The router uses litellm.aimage_generation() for async image generation
+# - api_version is required for Azure image generation deployments
diff --git a/surfsense_backend/app/db.py b/surfsense_backend/app/db.py
index 2298e7438..320ff6d8d 100644
--- a/surfsense_backend/app/db.py
+++ b/surfsense_backend/app/db.py
@@ -100,6 +100,7 @@ class PodcastStatus(str, Enum):
     FAILED = "failed"
 
 
+
 class LiteLLMProvider(str, Enum):
     """
     Enum for LLM providers supported by LiteLLM.
@@ -137,6 +138,24 @@ class LiteLLMProvider(str, Enum):
     CUSTOM = "CUSTOM"
 
 
+class ImageGenProvider(str, Enum):
+    """
+    Enum for image generation providers supported by LiteLLM.
+    This is a subset of LLM providers — only those that support image generation.
+    See: https://docs.litellm.ai/docs/image_generation#supported-providers
+    """
+
+    OPENAI = "OPENAI"
+    AZURE_OPENAI = "AZURE_OPENAI"
+    GOOGLE = "GOOGLE"  # Google AI Studio
+    VERTEX_AI = "VERTEX_AI"
+    BEDROCK = "BEDROCK"  # AWS Bedrock
+    RECRAFT = "RECRAFT"
+    OPENROUTER = "OPENROUTER"
+    XINFERENCE = "XINFERENCE"
+    NSCALE = "NSCALE"
+
+
 class LogLevel(str, Enum):
     DEBUG = "DEBUG"
     INFO = "INFO"
@@ -237,6 +256,11 @@ class Permission(str, Enum):
     PODCASTS_UPDATE = "podcasts:update"
     PODCASTS_DELETE = "podcasts:delete"
 
+    # Image Generations
+    IMAGE_GENERATIONS_CREATE = "image_generations:create"
+    IMAGE_GENERATIONS_READ = "image_generations:read"
+    IMAGE_GENERATIONS_DELETE = "image_generations:delete"
+
     # Connectors
     CONNECTORS_CREATE = "connectors:create"
     CONNECTORS_READ = "connectors:read"
@@ -298,6 +322,9 @@ DEFAULT_ROLE_PERMISSIONS = {
         Permission.PODCASTS_CREATE.value,
         Permission.PODCASTS_READ.value,
         Permission.PODCASTS_UPDATE.value,
+        # Image Generations (create and read, no delete)
+        Permission.IMAGE_GENERATIONS_CREATE.value,
+        Permission.IMAGE_GENERATIONS_READ.value,
         # Connectors (no delete)
         Permission.CONNECTORS_CREATE.value,
         Permission.CONNECTORS_READ.value,
@@ -327,6 +354,8 @@ DEFAULT_ROLE_PERMISSIONS = {
         Permission.LLM_CONFIGS_READ.value,
         # Podcasts (read only)
         Permission.PODCASTS_READ.value,
+        # Image Generations (read only)
+        Permission.IMAGE_GENERATIONS_READ.value,
         # Connectors (read only)
         Permission.CONNECTORS_READ.value,
         # Logs (read only)
@@ -881,6 +910,103 @@ class Podcast(BaseModel, TimestampMixin):
     thread = relationship("NewChatThread")
 
 
+class ImageGenerationConfig(BaseModel, TimestampMixin):
+    """
+    Dedicated configuration table for image generation models.
+
+    Separate from NewLLMConfig because image generation models don't need
+    system_instructions, citations_enabled, or use_default_system_instructions.
+    They only need provider credentials and model parameters.
+    """
+
+    __tablename__ = "image_generation_configs"
+
+    name = Column(String(100), nullable=False, index=True)
+    description = Column(String(500), nullable=True)
+
+    # Provider & model (uses ImageGenProvider, NOT LiteLLMProvider)
+    provider = Column(SQLAlchemyEnum(ImageGenProvider), nullable=False)
+    custom_provider = Column(String(100), nullable=True)
+    model_name = Column(String(100), nullable=False)
+
+    # Credentials
+    api_key = Column(String, nullable=False)
+    api_base = Column(String(500), nullable=True)
+    api_version = Column(String(50), nullable=True)  # Azure-specific
+
+    # Additional litellm parameters
+    litellm_params = Column(JSON, nullable=True, default={})
+
+    # Relationships
+    search_space_id = Column(
+        Integer, ForeignKey("searchspaces.id", ondelete="CASCADE"), nullable=False
+    )
+    search_space = relationship("SearchSpace", back_populates="image_generation_configs")
+
+
+class ImageGeneration(BaseModel, TimestampMixin):
+    """
+    Stores image generation requests and results using litellm.aimage_generation().
+
+    Since aimage_generation is a single async call (not a background job),
+    there is no status enum. A row with response_data means success;
+    a row with error_message means failure.
+
+    Response data is stored as JSONB matching the litellm output format:
+    {
+        "created": int,
+        "data": [{"b64_json": str|None, "revised_prompt": str|None, "url": str|None}],
+        "usage": {"prompt_tokens": int, "completion_tokens": int, "total_tokens": int}
+    }
+    """
+
+    __tablename__ = "image_generations"
+
+    # Request parameters (matching litellm.aimage_generation() params)
+    prompt = Column(Text, nullable=False)
+    model = Column(String(200), nullable=True)  # e.g., "dall-e-3", "gpt-image-1"
+    n = Column(Integer, nullable=True, default=1)
+    quality = Column(
+        String(50), nullable=True
+    )  # "auto", "high", "medium", "low", "hd", "standard"
+    size = Column(
+        String(50), nullable=True
+    )  # "1024x1024", "1536x1024", "1024x1536", etc.
+    style = Column(String(50), nullable=True)  # Model-specific style parameter
+    response_format = Column(
+        String(50), nullable=True
+    )  # "url" or "b64_json"
+
+    # Image generation config reference
+    # 0 = Auto mode (router), negative IDs = global configs from YAML,
+    # positive IDs = ImageGenerationConfig records in DB
+    image_generation_config_id = Column(Integer, nullable=True)
+
+    # Response data (full litellm response as JSONB) — present on success
+    response_data = Column(JSONB, nullable=True)
+    # Error message — present on failure
+    error_message = Column(Text, nullable=True)
+
+    # Signed access token for serving images via <img> tags.
+    # Stored in DB so it survives SECRET_KEY rotation.
+    access_token = Column(String(64), nullable=True, index=True)
+
+    # Foreign keys
+    search_space_id = Column(
+        Integer, ForeignKey("searchspaces.id", ondelete="CASCADE"), nullable=False
+    )
+    created_by_id = Column(
+        UUID(as_uuid=True),
+        ForeignKey("user.id", ondelete="SET NULL"),
+        nullable=True,
+        index=True,
+    )
+
+    # Relationships
+    search_space = relationship("SearchSpace", back_populates="image_generations")
+    created_by = relationship("User", back_populates="image_generations")
+
+
 class SearchSpace(BaseModel, TimestampMixin):
     __tablename__ = "searchspaces"
 
@@ -905,6 +1031,9 @@ class SearchSpace(BaseModel, TimestampMixin):
     document_summary_llm_id = Column(
         Integer, nullable=True, default=0
     )  # For document summarization, defaults to Auto mode
+    image_generation_config_id = Column(
+        Integer, nullable=True, default=0
+    )  # For image generation, defaults to Auto mode
 
     user_id = Column(
         UUID(as_uuid=True), ForeignKey("user.id", ondelete="CASCADE"), nullable=False
@@ -929,6 +1058,12 @@ class SearchSpace(BaseModel, TimestampMixin):
         order_by="Podcast.id.desc()",
         cascade="all, delete-orphan",
     )
+    image_generations = relationship(
+        "ImageGeneration",
+        back_populates="search_space",
+        order_by="ImageGeneration.id.desc()",
+        cascade="all, delete-orphan",
+    )
     logs = relationship(
         "Log",
         back_populates="search_space",
@@ -953,6 +1088,12 @@ class SearchSpace(BaseModel, TimestampMixin):
         order_by="NewLLMConfig.id",
         cascade="all, delete-orphan",
     )
+    image_generation_configs = relationship(
+        "ImageGenerationConfig",
+        back_populates="search_space",
+        order_by="ImageGenerationConfig.id",
+        cascade="all, delete-orphan",
+    )
 
     # RBAC relationships
     roles = relationship(
@@ -1333,6 +1474,13 @@ if config.AUTH_TYPE == "GOOGLE":
             passive_deletes=True,
         )
 
+        # Image generations created by this user
+        image_generations = relationship(
+            "ImageGeneration",
+            back_populates="created_by",
+            passive_deletes=True,
+        )
+
         # User memories for personalized AI responses
         memories = relationship(
             "UserMemory",
@@ -1405,6 +1553,13 @@ else:
             passive_deletes=True,
         )
 
+        # Image generations created by this user
+        image_generations = relationship(
+            "ImageGeneration",
+            back_populates="created_by",
+            passive_deletes=True,
+        )
+
         # User memories for personalized AI responses
         memories = relationship(
             "UserMemory",
diff --git a/surfsense_backend/app/routes/__init__.py b/surfsense_backend/app/routes/__init__.py
index 746c18c6d..683f3548b 100644
--- a/surfsense_backend/app/routes/__init__.py
+++ b/surfsense_backend/app/routes/__init__.py
@@ -11,6 +11,7 @@ from .confluence_add_connector_route import router as confluence_add_connector_r
 from .discord_add_connector_route import router as discord_add_connector_router
 from .documents_routes import router as documents_router
 from .editor_routes import router as editor_router
+from .image_generation_routes import router as image_generation_router
 from .google_calendar_add_connector_route import (
     router as google_calendar_add_connector_router,
 )
@@ -49,6 +50,7 @@ router.include_router(notes_router)
 router.include_router(new_chat_router)  # Chat with assistant-ui persistence
 router.include_router(chat_comments_router)
 router.include_router(podcasts_router)  # Podcast task status and audio
+router.include_router(image_generation_router)  # Image generation via litellm
 router.include_router(search_source_connectors_router)
 router.include_router(google_calendar_add_connector_router)
 router.include_router(google_gmail_add_connector_router)
diff --git a/surfsense_backend/app/routes/image_generation_routes.py b/surfsense_backend/app/routes/image_generation_routes.py
new file mode 100644
index 000000000..9b79771eb
--- /dev/null
+++ b/surfsense_backend/app/routes/image_generation_routes.py
@@ -0,0 +1,646 @@
+"""
+Image Generation routes:
+- CRUD for ImageGenerationConfig (user-created image model configs)
+- Global image gen configs endpoint (from YAML)
+- Image generation execution (calls litellm.aimage_generation())
+- CRUD for ImageGeneration records (results)
+- Image serving endpoint (serves b64_json images from DB, protected by signed tokens)
+"""
+
+import base64
+import logging
+
+from fastapi import APIRouter, Depends, HTTPException, Query
+from fastapi.responses import Response
+from litellm import aimage_generation
+from sqlalchemy import select
+from sqlalchemy.exc import SQLAlchemyError
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from app.config import config
+from app.db import (
+    ImageGeneration,
+    ImageGenerationConfig,
+    Permission,
+    SearchSpace,
+    SearchSpaceMembership,
+    User,
+    get_async_session,
+)
+from app.schemas import (
+    GlobalImageGenConfigRead,
+    ImageGenerationConfigCreate,
+    ImageGenerationConfigRead,
+    ImageGenerationConfigUpdate,
+    ImageGenerationCreate,
+    ImageGenerationListRead,
+    ImageGenerationRead,
+)
+from app.services.image_gen_router_service import (
+    IMAGE_GEN_AUTO_MODE_ID,
+    ImageGenRouterService,
+    is_image_gen_auto_mode,
+)
+from app.users import current_active_user
+from app.utils.rbac import check_permission
+from app.utils.signed_image_urls import verify_image_token
+
+router = APIRouter()
+logger = logging.getLogger(__name__)
+
+# Provider mapping for building litellm model strings.
+# Only includes providers that support image generation.
+# See: https://docs.litellm.ai/docs/image_generation#supported-providers
+_PROVIDER_MAP = {
+    "OPENAI": "openai",
+    "AZURE_OPENAI": "azure",
+    "GOOGLE": "gemini",       # Google AI Studio
+    "VERTEX_AI": "vertex_ai",
+    "BEDROCK": "bedrock",     # AWS Bedrock
+    "RECRAFT": "recraft",
+    "OPENROUTER": "openrouter",
+    "XINFERENCE": "xinference",
+    "NSCALE": "nscale",
+}
+
+
+def _get_global_image_gen_config(config_id: int) -> dict | None:
+    """Get a global image generation configuration by ID (negative IDs)."""
+    if config_id == IMAGE_GEN_AUTO_MODE_ID:
+        return {
+            "id": IMAGE_GEN_AUTO_MODE_ID,
+            "name": "Auto (Load Balanced)",
+            "provider": "AUTO",
+            "model_name": "auto",
+            "is_auto_mode": True,
+        }
+    if config_id > 0:
+        return None
+    for cfg in config.GLOBAL_IMAGE_GEN_CONFIGS:
+        if cfg.get("id") == config_id:
+            return cfg
+    return None
+
+
+def _build_model_string(provider: str, model_name: str, custom_provider: str | None) -> str:
+    """Build a litellm model string from provider + model_name."""
+    if custom_provider:
+        return f"{custom_provider}/{model_name}"
+    prefix = _PROVIDER_MAP.get(provider.upper(), provider.lower())
+    return f"{prefix}/{model_name}"
+
+
+async def _execute_image_generation(
+    session: AsyncSession,
+    image_gen: ImageGeneration,
+    search_space: SearchSpace,
+) -> None:
+    """
+    Call litellm.aimage_generation() with the appropriate config.
+
+    Resolution order:
+    1. Explicit image_generation_config_id on the request
+    2. Search space's image_generation_config_id preference
+    3. Falls back to Auto mode if available
+    """
+    config_id = image_gen.image_generation_config_id
+    if config_id is None:
+        config_id = search_space.image_generation_config_id or IMAGE_GEN_AUTO_MODE_ID
+        image_gen.image_generation_config_id = config_id
+
+    # Build kwargs
+    gen_kwargs = {}
+    if image_gen.n is not None:
+        gen_kwargs["n"] = image_gen.n
+    if image_gen.quality is not None:
+        gen_kwargs["quality"] = image_gen.quality
+    if image_gen.size is not None:
+        gen_kwargs["size"] = image_gen.size
+    if image_gen.style is not None:
+        gen_kwargs["style"] = image_gen.style
+    if image_gen.response_format is not None:
+        gen_kwargs["response_format"] = image_gen.response_format
+
+    if is_image_gen_auto_mode(config_id):
+        if not ImageGenRouterService.is_initialized():
+            raise ValueError(
+                "Auto mode requested but Image Generation Router not initialized. "
+                "Ensure global_llm_config.yaml has global_image_generation_configs."
+            )
+        response = await ImageGenRouterService.aimage_generation(
+            prompt=image_gen.prompt, model="auto", **gen_kwargs
+        )
+    elif config_id < 0:
+        # Global config from YAML
+        cfg = _get_global_image_gen_config(config_id)
+        if not cfg:
+            raise ValueError(f"Global image generation config {config_id} not found")
+
+        model_string = _build_model_string(
+            cfg.get("provider", ""), cfg["model_name"], cfg.get("custom_provider")
+        )
+        gen_kwargs["api_key"] = cfg.get("api_key")
+        if cfg.get("api_base"):
+            gen_kwargs["api_base"] = cfg["api_base"]
+        if cfg.get("api_version"):
+            gen_kwargs["api_version"] = cfg["api_version"]
+        if cfg.get("litellm_params"):
+            gen_kwargs.update(cfg["litellm_params"])
+
+        # User model override
+        if image_gen.model:
+            model_string = image_gen.model
+
+        response = await aimage_generation(
+            prompt=image_gen.prompt, model=model_string, **gen_kwargs
+        )
+    else:
+        # Positive ID = DB ImageGenerationConfig
+        result = await session.execute(
+            select(ImageGenerationConfig).filter(ImageGenerationConfig.id == config_id)
+        )
+        db_cfg = result.scalars().first()
+        if not db_cfg:
+            raise ValueError(f"Image generation config {config_id} not found")
+
+        model_string = _build_model_string(
+            db_cfg.provider.value, db_cfg.model_name, db_cfg.custom_provider
+        )
+        gen_kwargs["api_key"] = db_cfg.api_key
+        if db_cfg.api_base:
+            gen_kwargs["api_base"] = db_cfg.api_base
+        if db_cfg.api_version:
+            gen_kwargs["api_version"] = db_cfg.api_version
+        if db_cfg.litellm_params:
+            gen_kwargs.update(db_cfg.litellm_params)
+
+        # User model override
+        if image_gen.model:
+            model_string = image_gen.model
+
+        response = await aimage_generation(
+            prompt=image_gen.prompt, model=model_string, **gen_kwargs
+        )
+
+    # Store response
+    image_gen.response_data = (
+        response.model_dump() if hasattr(response, "model_dump") else dict(response)
+    )
+    if not image_gen.model and hasattr(response, "_hidden_params"):
+        hidden = response._hidden_params
+        if isinstance(hidden, dict) and hidden.get("model"):
+            image_gen.model = hidden["model"]
+
+
+# =============================================================================
+# Global Image Generation Configs (from YAML)
+# =============================================================================
+
+
+@router.get(
+    "/global-image-generation-configs",
+    response_model=list[GlobalImageGenConfigRead],
+)
+async def get_global_image_gen_configs(
+    user: User = Depends(current_active_user),
+):
+    """Get all global image generation configs. API keys are hidden."""
+    try:
+        global_configs = config.GLOBAL_IMAGE_GEN_CONFIGS
+        safe_configs = []
+
+        if global_configs and len(global_configs) > 0:
+            safe_configs.append({
+                "id": 0,
+                "name": "Auto (Load Balanced)",
+                "description": "Automatically routes across available image generation providers.",
+                "provider": "AUTO",
+                "custom_provider": None,
+                "model_name": "auto",
+                "api_base": None,
+                "api_version": None,
+                "litellm_params": {},
+                "is_global": True,
+                "is_auto_mode": True,
+            })
+
+        for cfg in global_configs:
+            safe_configs.append({
+                "id": cfg.get("id"),
+                "name": cfg.get("name"),
+                "description": cfg.get("description"),
+                "provider": cfg.get("provider"),
+                "custom_provider": cfg.get("custom_provider"),
+                "model_name": cfg.get("model_name"),
+                "api_base": cfg.get("api_base") or None,
+                "api_version": cfg.get("api_version") or None,
+                "litellm_params": cfg.get("litellm_params", {}),
+                "is_global": True,
+            })
+
+        return safe_configs
+    except Exception as e:
+        logger.exception("Failed to fetch global image generation configs")
+        raise HTTPException(status_code=500, detail=f"Failed to fetch configs: {e!s}") from e
+
+
+# =============================================================================
+# ImageGenerationConfig CRUD
+# =============================================================================
+
+
+@router.post("/image-generation-configs", response_model=ImageGenerationConfigRead)
+async def create_image_gen_config(
+    config_data: ImageGenerationConfigCreate,
+    session: AsyncSession = Depends(get_async_session),
+    user: User = Depends(current_active_user),
+):
+    """Create a new image generation config for a search space."""
+    try:
+        await check_permission(
+            session, user, config_data.search_space_id,
+            Permission.IMAGE_GENERATIONS_CREATE.value,
+            "You don't have permission to create image generation configs in this search space",
+        )
+
+        db_config = ImageGenerationConfig(**config_data.model_dump())
+        session.add(db_config)
+        await session.commit()
+        await session.refresh(db_config)
+        return db_config
+
+    except HTTPException:
+        raise
+    except Exception as e:
+        await session.rollback()
+        logger.exception("Failed to create ImageGenerationConfig")
+        raise HTTPException(status_code=500, detail=f"Failed to create config: {e!s}") from e
+
+
+@router.get("/image-generation-configs", response_model=list[ImageGenerationConfigRead])
+async def list_image_gen_configs(
+    search_space_id: int,
+    skip: int = 0,
+    limit: int = 100,
+    session: AsyncSession = Depends(get_async_session),
+    user: User = Depends(current_active_user),
+):
+    """List image generation configs for a search space."""
+    try:
+        await check_permission(
+            session, user, search_space_id,
+            Permission.IMAGE_GENERATIONS_READ.value,
+            "You don't have permission to view image generation configs in this search space",
+        )
+
+        result = await session.execute(
+            select(ImageGenerationConfig)
+            .filter(ImageGenerationConfig.search_space_id == search_space_id)
+            .order_by(ImageGenerationConfig.created_at.desc())
+            .offset(skip)
+            .limit(limit)
+        )
+        return result.scalars().all()
+
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.exception("Failed to list ImageGenerationConfigs")
+        raise HTTPException(status_code=500, detail=f"Failed to fetch configs: {e!s}") from e
+
+
+@router.get("/image-generation-configs/{config_id}", response_model=ImageGenerationConfigRead)
+async def get_image_gen_config(
+    config_id: int,
+    session: AsyncSession = Depends(get_async_session),
+    user: User = Depends(current_active_user),
+):
+    """Get a specific image generation config by ID."""
+    try:
+        result = await session.execute(
+            select(ImageGenerationConfig).filter(ImageGenerationConfig.id == config_id)
+        )
+        db_config = result.scalars().first()
+        if not db_config:
+            raise HTTPException(status_code=404, detail="Config not found")
+
+        await check_permission(
+            session, user, db_config.search_space_id,
+            Permission.IMAGE_GENERATIONS_READ.value,
+            "You don't have permission to view image generation configs in this search space",
+        )
+        return db_config
+
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.exception("Failed to get ImageGenerationConfig")
+        raise HTTPException(status_code=500, detail=f"Failed to fetch config: {e!s}") from e
+
+
+@router.put("/image-generation-configs/{config_id}", response_model=ImageGenerationConfigRead)
+async def update_image_gen_config(
+    config_id: int,
+    update_data: ImageGenerationConfigUpdate,
+    session: AsyncSession = Depends(get_async_session),
+    user: User = Depends(current_active_user),
+):
+    """Update an existing image generation config."""
+    try:
+        result = await session.execute(
+            select(ImageGenerationConfig).filter(ImageGenerationConfig.id == config_id)
+        )
+        db_config = result.scalars().first()
+        if not db_config:
+            raise HTTPException(status_code=404, detail="Config not found")
+
+        await check_permission(
+            session, user, db_config.search_space_id,
+            Permission.IMAGE_GENERATIONS_CREATE.value,
+            "You don't have permission to update image generation configs in this search space",
+        )
+
+        for key, value in update_data.model_dump(exclude_unset=True).items():
+            setattr(db_config, key, value)
+
+        await session.commit()
+        await session.refresh(db_config)
+        return db_config
+
+    except HTTPException:
+        raise
+    except Exception as e:
+        await session.rollback()
+        logger.exception("Failed to update ImageGenerationConfig")
+        raise HTTPException(status_code=500, detail=f"Failed to update config: {e!s}") from e
+
+
+@router.delete("/image-generation-configs/{config_id}", response_model=dict)
+async def delete_image_gen_config(
+    config_id: int,
+    session: AsyncSession = Depends(get_async_session),
+    user: User = Depends(current_active_user),
+):
+    """Delete an image generation config."""
+    try:
+        result = await session.execute(
+            select(ImageGenerationConfig).filter(ImageGenerationConfig.id == config_id)
+        )
+        db_config = result.scalars().first()
+        if not db_config:
+            raise HTTPException(status_code=404, detail="Config not found")
+
+        await check_permission(
+            session, user, db_config.search_space_id,
+            Permission.IMAGE_GENERATIONS_DELETE.value,
+            "You don't have permission to delete image generation configs in this search space",
+        )
+
+        await session.delete(db_config)
+        await session.commit()
+        return {"message": "Image generation config deleted successfully", "id": config_id}
+
+    except HTTPException:
+        raise
+    except Exception as e:
+        await session.rollback()
+        logger.exception("Failed to delete ImageGenerationConfig")
+        raise HTTPException(status_code=500, detail=f"Failed to delete config: {e!s}") from e
+
+
+# =============================================================================
+# Image Generation Execution + Results CRUD
+# =============================================================================
+
+
+@router.post("/image-generations", response_model=ImageGenerationRead)
+async def create_image_generation(
+    data: ImageGenerationCreate,
+    session: AsyncSession = Depends(get_async_session),
+    user: User = Depends(current_active_user),
+):
+    """Create and execute an image generation request."""
+    try:
+        await check_permission(
+            session, user, data.search_space_id,
+            Permission.IMAGE_GENERATIONS_CREATE.value,
+            "You don't have permission to create image generations in this search space",
+        )
+
+        result = await session.execute(
+            select(SearchSpace).filter(SearchSpace.id == data.search_space_id)
+        )
+        search_space = result.scalars().first()
+        if not search_space:
+            raise HTTPException(status_code=404, detail="Search space not found")
+
+        db_image_gen = ImageGeneration(
+            prompt=data.prompt,
+            model=data.model,
+            n=data.n,
+            quality=data.quality,
+            size=data.size,
+            style=data.style,
+            response_format=data.response_format,
+            image_generation_config_id=data.image_generation_config_id,
+            search_space_id=data.search_space_id,
+            created_by_id=user.id,
+        )
+        session.add(db_image_gen)
+        await session.flush()
+
+        try:
+            await _execute_image_generation(session, db_image_gen, search_space)
+        except Exception as e:
+            logger.exception("Image generation call failed")
+            db_image_gen.error_message = str(e)
+
+        await session.commit()
+        await session.refresh(db_image_gen)
+        return db_image_gen
+
+    except HTTPException:
+        raise
+    except SQLAlchemyError:
+        await session.rollback()
+        raise HTTPException(status_code=500, detail="Database error during image generation") from None
+    except Exception as e:
+        await session.rollback()
+        logger.exception("Failed to create image generation")
+        raise HTTPException(status_code=500, detail=f"Image generation failed: {e!s}") from e
+
+
+@router.get("/image-generations", response_model=list[ImageGenerationListRead])
+async def list_image_generations(
+    search_space_id: int | None = None,
+    skip: int = 0,
+    limit: int = 50,
+    session: AsyncSession = Depends(get_async_session),
+    user: User = Depends(current_active_user),
+):
+    """List image generations."""
+    if skip < 0 or limit < 1:
+        raise HTTPException(status_code=400, detail="Invalid pagination parameters")
+    if limit > 100:
+        limit = 100
+
+    try:
+        if search_space_id is not None:
+            await check_permission(
+                session, user, search_space_id,
+                Permission.IMAGE_GENERATIONS_READ.value,
+                "You don't have permission to read image generations in this search space",
+            )
+            result = await session.execute(
+                select(ImageGeneration)
+                .filter(ImageGeneration.search_space_id == search_space_id)
+                .order_by(ImageGeneration.created_at.desc())
+                .offset(skip).limit(limit)
+            )
+        else:
+            result = await session.execute(
+                select(ImageGeneration)
+                .join(SearchSpace)
+                .join(SearchSpaceMembership)
+                .filter(SearchSpaceMembership.user_id == user.id)
+                .order_by(ImageGeneration.created_at.desc())
+                .offset(skip).limit(limit)
+            )
+
+        return [ImageGenerationListRead.from_orm_with_count(img) for img in result.scalars().all()]
+
+    except HTTPException:
+        raise
+    except SQLAlchemyError:
+        raise HTTPException(status_code=500, detail="Database error fetching image generations") from None
+
+
+@router.get("/image-generations/{image_gen_id}", response_model=ImageGenerationRead)
+async def get_image_generation(
+    image_gen_id: int,
+    session: AsyncSession = Depends(get_async_session),
+    user: User = Depends(current_active_user),
+):
+    """Get a specific image generation by ID."""
+    try:
+        result = await session.execute(
+            select(ImageGeneration).filter(ImageGeneration.id == image_gen_id)
+        )
+        image_gen = result.scalars().first()
+        if not image_gen:
+            raise HTTPException(status_code=404, detail="Image generation not found")
+
+        await check_permission(
+            session, user, image_gen.search_space_id,
+            Permission.IMAGE_GENERATIONS_READ.value,
+            "You don't have permission to read image generations in this search space",
+        )
+        return image_gen
+
+    except HTTPException:
+        raise
+    except SQLAlchemyError:
+        raise HTTPException(status_code=500, detail="Database error fetching image generation") from None
+
+
+@router.delete("/image-generations/{image_gen_id}", response_model=dict)
+async def delete_image_generation(
+    image_gen_id: int,
+    session: AsyncSession = Depends(get_async_session),
+    user: User = Depends(current_active_user),
+):
+    """Delete an image generation record."""
+    try:
+        result = await session.execute(
+            select(ImageGeneration).filter(ImageGeneration.id == image_gen_id)
+        )
+        db_image_gen = result.scalars().first()
+        if not db_image_gen:
+            raise HTTPException(status_code=404, detail="Image generation not found")
+
+        await check_permission(
+            session, user, db_image_gen.search_space_id,
+            Permission.IMAGE_GENERATIONS_DELETE.value,
+            "You don't have permission to delete image generations in this search space",
+        )
+
+        await session.delete(db_image_gen)
+        await session.commit()
+        return {"message": "Image generation deleted successfully"}
+
+    except HTTPException:
+        raise
+    except SQLAlchemyError:
+        await session.rollback()
+        raise HTTPException(status_code=500, detail="Database error deleting image generation") from None
+
+
+# =============================================================================
+# Image Serving (serves generated images from DB, protected by signed tokens)
+# =============================================================================
+
+@router.get("/image-generations/{image_gen_id}/image")
+async def serve_generated_image(
+    image_gen_id: int,
+    token: str = Query(..., description="Signed access token"),
+    index: int = 0,
+    session: AsyncSession = Depends(get_async_session),
+):
+    """
+    Serve a generated image by ID, protected by a signed token.
+
+    The token is generated when the image URL is created by the generate_image
+    tool and encodes the image_gen_id, search_space_id, and an expiry timestamp.
+    This ensures only users with access to the search space can view images,
+    without requiring auth headers (which <img> tags cannot pass).
+
+    Args:
+        image_gen_id: The image generation record ID
+        token: HMAC-signed access token (included as query parameter)
+        index: Which image to serve if multiple were generated (default: 0)
+    """
+    try:
+        result = await session.execute(
+            select(ImageGeneration).filter(ImageGeneration.id == image_gen_id)
+        )
+        image_gen = result.scalars().first()
+        if not image_gen:
+            raise HTTPException(status_code=404, detail="Image generation not found")
+
+        # Verify the access token against the one stored on the record
+        if not verify_image_token(image_gen.access_token, token):
+            raise HTTPException(status_code=403, detail="Invalid image access token")
+
+        if not image_gen.response_data:
+            raise HTTPException(status_code=404, detail="No image data available")
+
+        images = image_gen.response_data.get("data", [])
+        if not images or index >= len(images):
+            raise HTTPException(status_code=404, detail="Image not found at the specified index")
+
+        image_entry = images[index]
+
+        # If there's a URL, redirect to it
+        if image_entry.get("url"):
+            from fastapi.responses import RedirectResponse
+            return RedirectResponse(url=image_entry["url"])
+
+        # If there's b64_json data, decode and serve it
+        if image_entry.get("b64_json"):
+            image_bytes = base64.b64decode(image_entry["b64_json"])
+            return Response(
+                content=image_bytes,
+                media_type="image/png",
+                headers={
+                    "Cache-Control": "public, max-age=86400",
+                    "Content-Disposition": f'inline; filename="generated-{image_gen_id}-{index}.png"',
+                },
+            )
+
+        raise HTTPException(status_code=404, detail="No displayable image data")
+
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.exception("Failed to serve generated image")
+        raise HTTPException(status_code=500, detail=f"Failed to serve image: {e!s}") from e
diff --git a/surfsense_backend/app/routes/search_spaces_routes.py b/surfsense_backend/app/routes/search_spaces_routes.py
index a8916f2ea..fd84c0f45 100644
--- a/surfsense_backend/app/routes/search_spaces_routes.py
+++ b/surfsense_backend/app/routes/search_spaces_routes.py
@@ -7,6 +7,7 @@ from sqlalchemy.future import select
 
 from app.config import config
 from app.db import (
+    ImageGenerationConfig,
     NewLLMConfig,
     Permission,
     SearchSpace,
@@ -387,6 +388,69 @@ async def _get_llm_config_by_id(
         return None
 
 
+async def _get_image_gen_config_by_id(
+    session: AsyncSession, config_id: int | None
+) -> dict | None:
+    """
+    Get an image generation config by ID as a dictionary.
+    Returns Auto mode for ID 0, global config for negative IDs,
+    DB ImageGenerationConfig for positive IDs, or None.
+    """
+    if config_id is None:
+        return None
+
+    if config_id == 0:
+        return {
+            "id": 0,
+            "name": "Auto (Load Balanced)",
+            "description": "Automatically routes requests across available image generation providers",
+            "provider": "AUTO",
+            "model_name": "auto",
+            "is_global": True,
+            "is_auto_mode": True,
+        }
+
+    if config_id < 0:
+        for cfg in config.GLOBAL_IMAGE_GEN_CONFIGS:
+            if cfg.get("id") == config_id:
+                return {
+                    "id": cfg.get("id"),
+                    "name": cfg.get("name"),
+                    "description": cfg.get("description"),
+                    "provider": cfg.get("provider"),
+                    "custom_provider": cfg.get("custom_provider"),
+                    "model_name": cfg.get("model_name"),
+                    "api_base": cfg.get("api_base") or None,
+                    "api_version": cfg.get("api_version") or None,
+                    "litellm_params": cfg.get("litellm_params", {}),
+                    "is_global": True,
+                }
+        return None
+
+    # Positive ID: query ImageGenerationConfig table
+    result = await session.execute(
+        select(ImageGenerationConfig).filter(ImageGenerationConfig.id == config_id)
+    )
+    db_config = result.scalars().first()
+    if db_config:
+        return {
+            "id": db_config.id,
+            "name": db_config.name,
+            "description": db_config.description,
+            "provider": db_config.provider.value if db_config.provider else None,
+            "custom_provider": db_config.custom_provider,
+            "model_name": db_config.model_name,
+            "api_base": db_config.api_base,
+            "api_version": db_config.api_version,
+            "litellm_params": db_config.litellm_params or {},
+            "created_at": db_config.created_at.isoformat()
+            if db_config.created_at
+            else None,
+            "search_space_id": db_config.search_space_id,
+        }
+    return None
+
+
 @router.get(
     "/search-spaces/{search_space_id}/llm-preferences",
     response_model=LLMPreferencesRead,
@@ -423,12 +487,17 @@ async def get_llm_preferences(
         document_summary_llm = await _get_llm_config_by_id(
             session, search_space.document_summary_llm_id
         )
+        image_generation_config = await _get_image_gen_config_by_id(
+            session, search_space.image_generation_config_id
+        )
 
         return LLMPreferencesRead(
             agent_llm_id=search_space.agent_llm_id,
             document_summary_llm_id=search_space.document_summary_llm_id,
+            image_generation_config_id=search_space.image_generation_config_id,
             agent_llm=agent_llm,
             document_summary_llm=document_summary_llm,
+            image_generation_config=image_generation_config,
         )
 
     except HTTPException:
@@ -485,12 +554,17 @@ async def update_llm_preferences(
         document_summary_llm = await _get_llm_config_by_id(
             session, search_space.document_summary_llm_id
         )
+        image_generation_config = await _get_image_gen_config_by_id(
+            session, search_space.image_generation_config_id
+        )
 
         return LLMPreferencesRead(
             agent_llm_id=search_space.agent_llm_id,
             document_summary_llm_id=search_space.document_summary_llm_id,
+            image_generation_config_id=search_space.image_generation_config_id,
             agent_llm=agent_llm,
             document_summary_llm=document_summary_llm,
+            image_generation_config=image_generation_config,
         )
 
     except HTTPException:
diff --git a/surfsense_backend/app/schemas/__init__.py b/surfsense_backend/app/schemas/__init__.py
index 5ff166733..332af55fd 100644
--- a/surfsense_backend/app/schemas/__init__.py
+++ b/surfsense_backend/app/schemas/__init__.py
@@ -20,6 +20,16 @@ from .documents import (
     PaginatedResponse,
 )
 from .google_drive import DriveItem, GoogleDriveIndexingOptions, GoogleDriveIndexRequest
+from .image_generation import (
+    GlobalImageGenConfigRead,
+    ImageGenerationConfigCreate,
+    ImageGenerationConfigPublic,
+    ImageGenerationConfigRead,
+    ImageGenerationConfigUpdate,
+    ImageGenerationCreate,
+    ImageGenerationListRead,
+    ImageGenerationRead,
+)
 from .logs import LogBase, LogCreate, LogFilter, LogRead, LogUpdate
 from .new_chat import (
     ChatMessage,
@@ -106,6 +116,16 @@ __all__ = [
     "GlobalNewLLMConfigRead",
     "GoogleDriveIndexRequest",
     "GoogleDriveIndexingOptions",
+    "GlobalImageGenConfigRead",
+    # Image Generation Config schemas
+    "ImageGenerationConfigCreate",
+    "ImageGenerationConfigPublic",
+    "ImageGenerationConfigRead",
+    "ImageGenerationConfigUpdate",
+    # Image Generation schemas
+    "ImageGenerationCreate",
+    "ImageGenerationListRead",
+    "ImageGenerationRead",
     # Base schemas
     "IDModel",
     # RBAC schemas
diff --git a/surfsense_backend/app/schemas/image_generation.py b/surfsense_backend/app/schemas/image_generation.py
new file mode 100644
index 000000000..367a35a77
--- /dev/null
+++ b/surfsense_backend/app/schemas/image_generation.py
@@ -0,0 +1,231 @@
+"""
+Pydantic schemas for Image Generation configs and generation requests.
+
+ImageGenerationConfig: CRUD schemas for user-created image gen model configs.
+ImageGeneration: Schemas for the actual image generation requests/results.
+GlobalImageGenConfigRead: Schema for admin-configured YAML configs.
+"""
+
+from datetime import datetime
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from app.db import ImageGenProvider
+
+
+# =============================================================================
+# ImageGenerationConfig CRUD Schemas
+# =============================================================================
+
+
+class ImageGenerationConfigBase(BaseModel):
+    """Base schema with fields for ImageGenerationConfig."""
+
+    name: str = Field(
+        ..., max_length=100, description="User-friendly name for the config"
+    )
+    description: str | None = Field(
+        None, max_length=500, description="Optional description"
+    )
+    provider: ImageGenProvider = Field(
+        ...,
+        description="Image generation provider (OpenAI, Azure, Google AI Studio, Vertex AI, Bedrock, Recraft, OpenRouter, Xinference, Nscale)",
+    )
+    custom_provider: str | None = Field(
+        None, max_length=100, description="Custom provider name"
+    )
+    model_name: str = Field(
+        ..., max_length=100, description="Model name (e.g., dall-e-3, gpt-image-1)"
+    )
+    api_key: str = Field(..., description="API key for the provider")
+    api_base: str | None = Field(
+        None, max_length=500, description="Optional API base URL"
+    )
+    api_version: str | None = Field(
+        None,
+        max_length=50,
+        description="Azure-specific API version (e.g., '2024-02-15-preview')",
+    )
+    litellm_params: dict[str, Any] | None = Field(
+        default=None, description="Additional LiteLLM parameters"
+    )
+
+
+class ImageGenerationConfigCreate(ImageGenerationConfigBase):
+    """Schema for creating a new ImageGenerationConfig."""
+
+    search_space_id: int = Field(
+        ..., description="Search space ID to associate the config with"
+    )
+
+
+class ImageGenerationConfigUpdate(BaseModel):
+    """Schema for updating an existing ImageGenerationConfig. All fields optional."""
+
+    name: str | None = Field(None, max_length=100)
+    description: str | None = Field(None, max_length=500)
+    provider: ImageGenProvider | None = None
+    custom_provider: str | None = Field(None, max_length=100)
+    model_name: str | None = Field(None, max_length=100)
+    api_key: str | None = None
+    api_base: str | None = Field(None, max_length=500)
+    api_version: str | None = Field(None, max_length=50)
+    litellm_params: dict[str, Any] | None = None
+
+
+class ImageGenerationConfigRead(ImageGenerationConfigBase):
+    """Schema for reading an ImageGenerationConfig (includes id and timestamps)."""
+
+    id: int
+    created_at: datetime
+    search_space_id: int
+
+    model_config = ConfigDict(from_attributes=True)
+
+
+class ImageGenerationConfigPublic(BaseModel):
+    """Public schema that hides the API key (for list views)."""
+
+    id: int
+    name: str
+    description: str | None = None
+    provider: ImageGenProvider
+    custom_provider: str | None = None
+    model_name: str
+    api_base: str | None = None
+    api_version: str | None = None
+    litellm_params: dict[str, Any] | None = None
+    created_at: datetime
+    search_space_id: int
+
+    model_config = ConfigDict(from_attributes=True)
+
+
+# =============================================================================
+# ImageGeneration (request/result) Schemas
+# =============================================================================
+
+
+class ImageGenerationCreate(BaseModel):
+    """Schema for creating an image generation request."""
+
+    prompt: str = Field(
+        ...,
+        min_length=1,
+        max_length=4000,
+        description="A text description of the desired image(s)",
+    )
+    model: str | None = Field(
+        None,
+        max_length=200,
+        description="The model to use (e.g., 'dall-e-3', 'gpt-image-1'). Overrides the config model.",
+    )
+    n: int | None = Field(
+        None,
+        ge=1,
+        le=10,
+        description="Number of images to generate (1-10).",
+    )
+    quality: str | None = Field(None, max_length=50)
+    size: str | None = Field(None, max_length=50)
+    style: str | None = Field(None, max_length=50)
+    response_format: str | None = Field(None, max_length=50)
+    search_space_id: int = Field(
+        ..., description="Search space ID to associate the generation with"
+    )
+    image_generation_config_id: int | None = Field(
+        None,
+        description=(
+            "Image generation config ID. "
+            "0 = Auto mode (router), negative = global YAML config, positive = DB config. "
+            "If not provided, uses the search space's image_generation_config_id preference."
+        ),
+    )
+
+
+class ImageGenerationRead(BaseModel):
+    """Schema for reading an image generation record."""
+
+    id: int
+    prompt: str
+    model: str | None = None
+    n: int | None = None
+    quality: str | None = None
+    size: str | None = None
+    style: str | None = None
+    response_format: str | None = None
+    image_generation_config_id: int | None = None
+    response_data: dict[str, Any] | None = None
+    error_message: str | None = None
+    search_space_id: int
+    created_at: datetime
+
+    model_config = ConfigDict(from_attributes=True)
+
+
+class ImageGenerationListRead(BaseModel):
+    """Lightweight schema for listing image generations (without full response_data)."""
+
+    id: int
+    prompt: str
+    model: str | None = None
+    n: int | None = None
+    quality: str | None = None
+    size: str | None = None
+    search_space_id: int
+    created_at: datetime
+    is_success: bool
+    image_count: int | None = None
+
+    model_config = ConfigDict(from_attributes=True)
+
+    @classmethod
+    def from_orm_with_count(cls, obj: Any) -> "ImageGenerationListRead":
+        """Create ImageGenerationListRead with computed fields."""
+        image_count = None
+        if obj.response_data and isinstance(obj.response_data, dict):
+            data = obj.response_data.get("data")
+            if isinstance(data, list):
+                image_count = len(data)
+
+        return cls(
+            id=obj.id,
+            prompt=obj.prompt,
+            model=obj.model,
+            n=obj.n,
+            quality=obj.quality,
+            size=obj.size,
+            search_space_id=obj.search_space_id,
+            created_at=obj.created_at,
+            is_success=obj.response_data is not None,
+            image_count=image_count,
+        )
+
+
+# =============================================================================
+# Global Image Gen Config (from YAML)
+# =============================================================================
+
+
+class GlobalImageGenConfigRead(BaseModel):
+    """
+    Schema for reading global image generation configs from YAML.
+    Global configs have negative IDs. API key is hidden.
+    ID 0 is reserved for Auto mode (LiteLLM Router load balancing).
+    """
+
+    id: int = Field(
+        ...,
+        description="Config ID: 0 for Auto mode, negative for global configs",
+    )
+    name: str
+    description: str | None = None
+    provider: str
+    custom_provider: str | None = None
+    model_name: str
+    api_base: str | None = None
+    api_version: str | None = None
+    litellm_params: dict[str, Any] | None = None
+    is_global: bool = True
+    is_auto_mode: bool = False
diff --git a/surfsense_backend/app/schemas/new_llm_config.py b/surfsense_backend/app/schemas/new_llm_config.py
index 286c07843..a6294fba2 100644
--- a/surfsense_backend/app/schemas/new_llm_config.py
+++ b/surfsense_backend/app/schemas/new_llm_config.py
@@ -176,12 +176,18 @@ class LLMPreferencesRead(BaseModel):
     document_summary_llm_id: int | None = Field(
         None, description="ID of the LLM config to use for document summarization"
     )
+    image_generation_config_id: int | None = Field(
+        None, description="ID of the image generation config to use"
+    )
     agent_llm: dict[str, Any] | None = Field(
         None, description="Full config for agent LLM"
     )
     document_summary_llm: dict[str, Any] | None = Field(
         None, description="Full config for document summary LLM"
     )
+    image_generation_config: dict[str, Any] | None = Field(
+        None, description="Full config for image generation"
+    )
 
     model_config = ConfigDict(from_attributes=True)
 
@@ -195,3 +201,6 @@ class LLMPreferencesUpdate(BaseModel):
     document_summary_llm_id: int | None = Field(
         None, description="ID of the LLM config to use for document summarization"
     )
+    image_generation_config_id: int | None = Field(
+        None, description="ID of the image generation config to use"
+    )
diff --git a/surfsense_backend/app/services/image_gen_router_service.py b/surfsense_backend/app/services/image_gen_router_service.py
new file mode 100644
index 000000000..3b8a15d2a
--- /dev/null
+++ b/surfsense_backend/app/services/image_gen_router_service.py
@@ -0,0 +1,294 @@
+"""
+Image Generation Router Service for Load Balancing
+
+This module provides a singleton LiteLLM Router for automatic load balancing
+across multiple image generation deployments. It uses litellm.Router which
+natively supports aimage_generation() for async image generation.
+
+The router handles:
+- Rate limit management with automatic cooldowns
+- Automatic failover and retries
+- Usage-based routing to distribute load evenly
+
+Supported providers: OpenAI, Azure, Google AI Studio, Vertex AI,
+AWS Bedrock, Recraft, OpenRouter, Xinference, Nscale.
+"""
+
+import logging
+from typing import Any
+
+from litellm import Router
+from litellm.utils import ImageResponse
+
+logger = logging.getLogger(__name__)
+
+# Special ID for Auto mode - uses router for load balancing
+IMAGE_GEN_AUTO_MODE_ID = 0
+
+# Provider mapping for LiteLLM model string construction.
+# Only includes providers that support image generation.
+# See: https://docs.litellm.ai/docs/image_generation#supported-providers
+IMAGE_GEN_PROVIDER_MAP = {
+    "OPENAI": "openai",
+    "AZURE_OPENAI": "azure",
+    "GOOGLE": "gemini",       # Google AI Studio
+    "VERTEX_AI": "vertex_ai",
+    "BEDROCK": "bedrock",     # AWS Bedrock
+    "RECRAFT": "recraft",
+    "OPENROUTER": "openrouter",
+    "XINFERENCE": "xinference",
+    "NSCALE": "nscale",
+}
+
+
+class ImageGenRouterService:
+    """
+    Singleton service for managing LiteLLM Router for image generation.
+
+    The router provides automatic load balancing, failover, and rate limit
+    handling across multiple image generation deployments.
+    Uses Router.aimage_generation() for async image generation calls.
+    """
+
+    _instance = None
+    _router: Router | None = None
+    _model_list: list[dict] = []
+    _router_settings: dict = {}
+    _initialized: bool = False
+
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    @classmethod
+    def get_instance(cls) -> "ImageGenRouterService":
+        """Get the singleton instance of the router service."""
+        if cls._instance is None:
+            cls._instance = cls()
+        return cls._instance
+
+    @classmethod
+    def initialize(
+        cls,
+        global_configs: list[dict],
+        router_settings: dict | None = None,
+    ) -> None:
+        """
+        Initialize the router with global image generation configurations.
+
+        Args:
+            global_configs: List of global image gen config dictionaries from YAML
+            router_settings: Optional router settings (routing_strategy, num_retries, etc.)
+        """
+        instance = cls.get_instance()
+
+        if instance._initialized:
+            logger.debug("Image Generation Router already initialized, skipping")
+            return
+
+        # Build model list from global configs
+        model_list = []
+        for config in global_configs:
+            deployment = cls._config_to_deployment(config)
+            if deployment:
+                model_list.append(deployment)
+
+        if not model_list:
+            logger.warning(
+                "No valid image generation configs found for router initialization"
+            )
+            return
+
+        instance._model_list = model_list
+        instance._router_settings = router_settings or {}
+
+        # Default router settings optimized for rate limit handling
+        default_settings = {
+            "routing_strategy": "usage-based-routing",
+            "num_retries": 3,
+            "allowed_fails": 3,
+            "cooldown_time": 60,
+            "retry_after": 5,
+        }
+
+        # Merge with provided settings
+        final_settings = {**default_settings, **instance._router_settings}
+
+        try:
+            instance._router = Router(
+                model_list=model_list,
+                routing_strategy=final_settings.get(
+                    "routing_strategy", "usage-based-routing"
+                ),
+                num_retries=final_settings.get("num_retries", 3),
+                allowed_fails=final_settings.get("allowed_fails", 3),
+                cooldown_time=final_settings.get("cooldown_time", 60),
+                set_verbose=False,
+            )
+            instance._initialized = True
+            logger.info(
+                f"Image Generation Router initialized with {len(model_list)} deployments, "
+                f"strategy: {final_settings.get('routing_strategy')}"
+            )
+        except Exception as e:
+            logger.error(f"Failed to initialize Image Generation Router: {e}")
+            instance._router = None
+
+    @classmethod
+    def _config_to_deployment(cls, config: dict) -> dict | None:
+        """
+        Convert a global image gen config to a router deployment entry.
+
+        Args:
+            config: Global image gen config dictionary
+
+        Returns:
+            Router deployment dictionary or None if invalid
+        """
+        try:
+            # Skip if essential fields are missing
+            if not config.get("model_name") or not config.get("api_key"):
+                return None
+
+            # Build model string
+            if config.get("custom_provider"):
+                model_string = f"{config['custom_provider']}/{config['model_name']}"
+            else:
+                provider = config.get("provider", "").upper()
+                provider_prefix = IMAGE_GEN_PROVIDER_MAP.get(
+                    provider, provider.lower()
+                )
+                model_string = f"{provider_prefix}/{config['model_name']}"
+
+            # Build litellm params
+            litellm_params: dict[str, Any] = {
+                "model": model_string,
+                "api_key": config.get("api_key"),
+            }
+
+            # Add optional api_base
+            if config.get("api_base"):
+                litellm_params["api_base"] = config["api_base"]
+
+            # Add api_version (required for Azure)
+            if config.get("api_version"):
+                litellm_params["api_version"] = config["api_version"]
+
+            # Add any additional litellm parameters
+            if config.get("litellm_params"):
+                litellm_params.update(config["litellm_params"])
+
+            # All configs use same alias "auto" for unified routing
+            deployment: dict[str, Any] = {
+                "model_name": "auto",
+                "litellm_params": litellm_params,
+            }
+
+            # Add rate limits from config if available
+            if config.get("rpm"):
+                deployment["rpm"] = config["rpm"]
+            if config.get("tpm"):
+                deployment["tpm"] = config["tpm"]
+
+            return deployment
+
+        except Exception as e:
+            logger.warning(
+                f"Failed to convert image gen config to deployment: {e}"
+            )
+            return None
+
+    @classmethod
+    def get_router(cls) -> Router | None:
+        """Get the initialized router instance."""
+        instance = cls.get_instance()
+        return instance._router
+
+    @classmethod
+    def is_initialized(cls) -> bool:
+        """Check if the router has been initialized."""
+        instance = cls.get_instance()
+        return instance._initialized and instance._router is not None
+
+    @classmethod
+    def get_model_count(cls) -> int:
+        """Get the number of models in the router."""
+        instance = cls.get_instance()
+        return len(instance._model_list)
+
+    @classmethod
+    async def aimage_generation(
+        cls,
+        prompt: str,
+        model: str = "auto",
+        n: int | None = None,
+        quality: str | None = None,
+        size: str | None = None,
+        style: str | None = None,
+        response_format: str | None = None,
+        timeout: int = 600,
+        **kwargs,
+    ) -> ImageResponse:
+        """
+        Generate images using the router for load balancing.
+
+        Uses Router.aimage_generation() which distributes requests
+        across configured image generation deployments.
+
+        Args:
+            prompt: Text description of the desired image(s)
+            model: Model alias (default "auto" for router routing)
+            n: Number of images to generate
+            quality: Image quality setting
+            size: Image size
+            style: Style parameter
+            response_format: "url" or "b64_json"
+            timeout: Request timeout in seconds
+            **kwargs: Additional litellm params
+
+        Returns:
+            ImageResponse from litellm
+
+        Raises:
+            ValueError: If router is not initialized
+        """
+        instance = cls.get_instance()
+        if not instance._router:
+            raise ValueError(
+                "Image Generation Router not initialized. "
+                "Ensure global_llm_config.yaml has global_image_generation_configs."
+            )
+
+        # Build kwargs for aimage_generation
+        gen_kwargs: dict[str, Any] = {
+            "model": model,
+            "prompt": prompt,
+            "timeout": timeout,
+        }
+        if n is not None:
+            gen_kwargs["n"] = n
+        if quality is not None:
+            gen_kwargs["quality"] = quality
+        if size is not None:
+            gen_kwargs["size"] = size
+        if style is not None:
+            gen_kwargs["style"] = style
+        if response_format is not None:
+            gen_kwargs["response_format"] = response_format
+        gen_kwargs.update(kwargs)
+
+        return await instance._router.aimage_generation(**gen_kwargs)
+
+
+def is_image_gen_auto_mode(config_id: int | None) -> bool:
+    """
+    Check if the given config ID represents Image Generation Auto mode.
+
+    Args:
+        config_id: The config ID to check
+
+    Returns:
+        True if this is Auto mode, False otherwise
+    """
+    return config_id == IMAGE_GEN_AUTO_MODE_ID
diff --git a/surfsense_backend/app/utils/signed_image_urls.py b/surfsense_backend/app/utils/signed_image_urls.py
new file mode 100644
index 000000000..d8d0bb57e
--- /dev/null
+++ b/surfsense_backend/app/utils/signed_image_urls.py
@@ -0,0 +1,44 @@
+"""
+Access token utilities for generated images.
+
+Provides token generation and verification so that generated images can be
+served via <img> tags (which cannot pass auth headers) while still
+restricting access to authorised users.
+
+Each image generation record stores its own random access token.  The token
+is verified by comparing the incoming query-parameter value against the
+stored value in the database.  This approach:
+
+* Survives SECRET_KEY rotation — tokens are random, not derived from a key.
+* Allows explicit revocation — just clear the column.
+* Is immune to timing attacks — uses ``hmac.compare_digest``.
+"""
+
+import hmac
+import secrets
+
+
+def generate_image_token() -> str:
+    """
+    Generate a cryptographically random access token for an image.
+
+    Returns:
+        A 64-character URL-safe hex string.
+    """
+    return secrets.token_hex(32)
+
+
+def verify_image_token(stored_token: str | None, provided_token: str) -> bool:
+    """
+    Constant-time comparison of a stored token against a user-provided one.
+
+    Args:
+        stored_token: The token persisted on the ImageGeneration record.
+        provided_token: The token from the URL query parameter.
+
+    Returns:
+        True if the tokens match, False otherwise.
+    """
+    if not stored_token or not provided_token:
+        return False
+    return hmac.compare_digest(stored_token, provided_token)
diff --git a/surfsense_web/app/dashboard/[search_space_id]/settings/page.tsx b/surfsense_web/app/dashboard/[search_space_id]/settings/page.tsx
index 1a727f1b6..e6c973ac6 100644
--- a/surfsense_web/app/dashboard/[search_space_id]/settings/page.tsx
+++ b/surfsense_web/app/dashboard/[search_space_id]/settings/page.tsx
@@ -7,6 +7,7 @@ import {
 	ChevronRight,
 	FileText,
 	Globe,
+	ImageIcon,
 	type LucideIcon,
 	Menu,
 	MessageSquare,
@@ -19,6 +20,7 @@ import { useTranslations } from "next-intl";
 import { useCallback, useEffect, useState } from "react";
 import { PublicChatSnapshotsManager } from "@/components/public-chat-snapshots/public-chat-snapshots-manager";
 import { GeneralSettingsManager } from "@/components/settings/general-settings-manager";
+import { ImageModelManager } from "@/components/settings/image-model-manager";
 import { LLMRoleManager } from "@/components/settings/llm-role-manager";
 import { ModelConfigManager } from "@/components/settings/model-config-manager";
 import { PromptConfigManager } from "@/components/settings/prompt-config-manager";
@@ -52,6 +54,12 @@ const settingsNavItems: SettingsNavItem[] = [
 		descriptionKey: "nav_role_assignments_desc",
 		icon: Brain,
 	},
+	{
+		id: "image-models",
+		labelKey: "nav_image_models",
+		descriptionKey: "nav_image_models_desc",
+		icon: ImageIcon,
+	},
 	{
 		id: "prompts",
 		labelKey: "nav_system_instructions",
@@ -282,8 +290,11 @@ function SettingsContent({
 								<GeneralSettingsManager searchSpaceId={searchSpaceId} />
 							)}
 							{activeSection === "models" && <ModelConfigManager searchSpaceId={searchSpaceId} />}
-							{activeSection === "roles" && <LLMRoleManager searchSpaceId={searchSpaceId} />}
-							{activeSection === "prompts" && <PromptConfigManager searchSpaceId={searchSpaceId} />}
+						{activeSection === "roles" && <LLMRoleManager searchSpaceId={searchSpaceId} />}
+						{activeSection === "image-models" && (
+							<ImageModelManager searchSpaceId={searchSpaceId} />
+						)}
+						{activeSection === "prompts" && <PromptConfigManager searchSpaceId={searchSpaceId} />}
 							{activeSection === "public-links" && (
 								<PublicChatSnapshotsManager searchSpaceId={searchSpaceId} />
 							)}
diff --git a/surfsense_web/atoms/image-gen-config/image-gen-config-mutation.atoms.ts b/surfsense_web/atoms/image-gen-config/image-gen-config-mutation.atoms.ts
new file mode 100644
index 000000000..dbaf441d0
--- /dev/null
+++ b/surfsense_web/atoms/image-gen-config/image-gen-config-mutation.atoms.ts
@@ -0,0 +1,91 @@
+import { atomWithMutation } from "jotai-tanstack-query";
+import { toast } from "sonner";
+import type {
+	CreateImageGenConfigRequest,
+	GetImageGenConfigsResponse,
+	UpdateImageGenConfigRequest,
+	UpdateImageGenConfigResponse,
+} from "@/contracts/types/new-llm-config.types";
+import { imageGenConfigApiService } from "@/lib/apis/image-gen-config-api.service";
+import { cacheKeys } from "@/lib/query-client/cache-keys";
+import { queryClient } from "@/lib/query-client/client";
+import { activeSearchSpaceIdAtom } from "../search-spaces/search-space-query.atoms";
+
+/**
+ * Mutation atom for creating a new ImageGenerationConfig
+ */
+export const createImageGenConfigMutationAtom = atomWithMutation((get) => {
+	const searchSpaceId = get(activeSearchSpaceIdAtom);
+
+	return {
+		mutationKey: ["image-gen-configs", "create"],
+		enabled: !!searchSpaceId,
+		mutationFn: async (request: CreateImageGenConfigRequest) => {
+			return imageGenConfigApiService.createConfig(request);
+		},
+		onSuccess: () => {
+			toast.success("Image model configuration created");
+			queryClient.invalidateQueries({
+				queryKey: cacheKeys.imageGenConfigs.all(Number(searchSpaceId)),
+			});
+		},
+		onError: (error: Error) => {
+			toast.error(error.message || "Failed to create image model configuration");
+		},
+	};
+});
+
+/**
+ * Mutation atom for updating an existing ImageGenerationConfig
+ */
+export const updateImageGenConfigMutationAtom = atomWithMutation((get) => {
+	const searchSpaceId = get(activeSearchSpaceIdAtom);
+
+	return {
+		mutationKey: ["image-gen-configs", "update"],
+		enabled: !!searchSpaceId,
+		mutationFn: async (request: UpdateImageGenConfigRequest) => {
+			return imageGenConfigApiService.updateConfig(request);
+		},
+		onSuccess: (_: UpdateImageGenConfigResponse, request: UpdateImageGenConfigRequest) => {
+			toast.success("Image model configuration updated");
+			queryClient.invalidateQueries({
+				queryKey: cacheKeys.imageGenConfigs.all(Number(searchSpaceId)),
+			});
+			queryClient.invalidateQueries({
+				queryKey: cacheKeys.imageGenConfigs.byId(request.id),
+			});
+		},
+		onError: (error: Error) => {
+			toast.error(error.message || "Failed to update image model configuration");
+		},
+	};
+});
+
+/**
+ * Mutation atom for deleting an ImageGenerationConfig
+ */
+export const deleteImageGenConfigMutationAtom = atomWithMutation((get) => {
+	const searchSpaceId = get(activeSearchSpaceIdAtom);
+
+	return {
+		mutationKey: ["image-gen-configs", "delete"],
+		enabled: !!searchSpaceId,
+		mutationFn: async (id: number) => {
+			return imageGenConfigApiService.deleteConfig(id);
+		},
+		onSuccess: (_, id: number) => {
+			toast.success("Image model configuration deleted");
+			queryClient.setQueryData(
+				cacheKeys.imageGenConfigs.all(Number(searchSpaceId)),
+				(oldData: GetImageGenConfigsResponse | undefined) => {
+					if (!oldData) return oldData;
+					return oldData.filter((config) => config.id !== id);
+				}
+			);
+		},
+		onError: (error: Error) => {
+			toast.error(error.message || "Failed to delete image model configuration");
+		},
+	};
+});
diff --git a/surfsense_web/atoms/image-gen-config/image-gen-config-query.atoms.ts b/surfsense_web/atoms/image-gen-config/image-gen-config-query.atoms.ts
new file mode 100644
index 000000000..a45e69a03
--- /dev/null
+++ b/surfsense_web/atoms/image-gen-config/image-gen-config-query.atoms.ts
@@ -0,0 +1,33 @@
+import { atomWithQuery } from "jotai-tanstack-query";
+import { imageGenConfigApiService } from "@/lib/apis/image-gen-config-api.service";
+import { cacheKeys } from "@/lib/query-client/cache-keys";
+import { activeSearchSpaceIdAtom } from "../search-spaces/search-space-query.atoms";
+
+/**
+ * Query atom for fetching user-created image gen configs for the active search space
+ */
+export const imageGenConfigsAtom = atomWithQuery((get) => {
+	const searchSpaceId = get(activeSearchSpaceIdAtom);
+
+	return {
+		queryKey: cacheKeys.imageGenConfigs.all(Number(searchSpaceId)),
+		enabled: !!searchSpaceId,
+		staleTime: 5 * 60 * 1000, // 5 minutes
+		queryFn: async () => {
+			return imageGenConfigApiService.getConfigs(Number(searchSpaceId));
+		},
+	};
+});
+
+/**
+ * Query atom for fetching global image gen configs (from YAML, negative IDs)
+ */
+export const globalImageGenConfigsAtom = atomWithQuery(() => {
+	return {
+		queryKey: cacheKeys.imageGenConfigs.global(),
+		staleTime: 10 * 60 * 1000, // 10 minutes - global configs rarely change
+		queryFn: async () => {
+			return imageGenConfigApiService.getGlobalConfigs();
+		},
+	};
+});
diff --git a/surfsense_web/components/Logo.tsx b/surfsense_web/components/Logo.tsx
index 58f8d1c9f..9f5915777 100644
--- a/surfsense_web/components/Logo.tsx
+++ b/surfsense_web/components/Logo.tsx
@@ -4,16 +4,20 @@ import Image from "next/image";
 import Link from "next/link";
 import { cn } from "@/lib/utils";
 
-export const Logo = ({ className }: { className?: string }) => {
-	return (
-		<Link href="/">
-			<Image
-				src="/icon-128.svg"
-				className={cn("dark:invert", className)}
-				alt="logo"
-				width={128}
-				height={128}
-			/>
-		</Link>
+export const Logo = ({ className, disableLink = false }: { className?: string; disableLink?: boolean }) => {
+	const image = (
+		<Image
+			src="/icon-128.svg"
+			className={cn("dark:invert", className)}
+			alt="logo"
+			width={128}
+			height={128}
+		/>
 	);
+
+	if (disableLink) {
+		return image;
+	}
+
+	return <Link href="/">{image}</Link>;
 };
diff --git a/surfsense_web/components/homepage/navbar.tsx b/surfsense_web/components/homepage/navbar.tsx
index 4abd4031b..670e3c810 100644
--- a/surfsense_web/components/homepage/navbar.tsx
+++ b/surfsense_web/components/homepage/navbar.tsx
@@ -64,7 +64,7 @@ const DesktopNav = ({ navItems, isScrolled }: any) => {
 				href="/"
 				className="flex flex-1 flex-row items-center gap-0.5 hover:opacity-80 transition-opacity"
 			>
-				<Logo className="h-8 w-8 rounded-md" />
+				<Logo className="h-8 w-8 rounded-md" disableLink />
 				<span className="dark:text-white/90 text-gray-800 text-lg font-bold">SurfSense</span>
 			</Link>
 			<div className="hidden flex-1 flex-row items-center justify-center space-x-2 text-sm font-medium text-zinc-600 transition duration-200 hover:text-zinc-800 lg:flex lg:space-x-2">
@@ -145,7 +145,7 @@ const MobileNav = ({ navItems, isScrolled }: any) => {
 					href="/"
 					className="flex flex-row items-center gap-2 hover:opacity-80 transition-opacity"
 				>
-					<Logo className="h-8 w-8 rounded-md" />
+					<Logo className="h-8 w-8 rounded-md" disableLink />
 					<span className="dark:text-white/90 text-gray-800 text-lg font-bold">SurfSense</span>
 				</Link>
 				<button
diff --git a/surfsense_web/components/new-chat/chat-header.tsx b/surfsense_web/components/new-chat/chat-header.tsx
index a6cf8df3a..e19f97945 100644
--- a/surfsense_web/components/new-chat/chat-header.tsx
+++ b/surfsense_web/components/new-chat/chat-header.tsx
@@ -1,10 +1,12 @@
 "use client";
 
 import { useCallback, useState } from "react";
+import { useParams, useRouter } from "next/navigation";
 import type {
 	GlobalNewLLMConfig,
 	NewLLMConfigPublic,
 } from "@/contracts/types/new-llm-config.types";
+import { ImageModelSelector } from "./image-model-selector";
 import { ModelConfigSidebar } from "./model-config-sidebar";
 import { ModelSelector } from "./model-selector";
 
@@ -13,6 +15,7 @@ interface ChatHeaderProps {
 }
 
 export function ChatHeader({ searchSpaceId }: ChatHeaderProps) {
+	const router = useRouter();
 	const [sidebarOpen, setSidebarOpen] = useState(false);
 	const [selectedConfig, setSelectedConfig] = useState<
 		NewLLMConfigPublic | GlobalNewLLMConfig | null
@@ -40,14 +43,19 @@ export function ChatHeader({ searchSpaceId }: ChatHeaderProps) {
 	const handleSidebarClose = useCallback((open: boolean) => {
 		setSidebarOpen(open);
 		if (!open) {
-			// Reset state when closing
 			setSelectedConfig(null);
 		}
 	}, []);
 
+	const handleAddImageModel = useCallback(() => {
+		// Navigate to settings image-models tab
+		router.push(`/dashboard/${searchSpaceId}/settings?tab=image-models`);
+	}, [router, searchSpaceId]);
+
 	return (
 		<div className="flex items-center gap-2">
 			<ModelSelector onEdit={handleEditConfig} onAddNew={handleAddNew} />
+			<ImageModelSelector onAddNew={handleAddImageModel} />
 			<ModelConfigSidebar
 				open={sidebarOpen}
 				onOpenChange={handleSidebarClose}
diff --git a/surfsense_web/components/new-chat/image-model-selector.tsx b/surfsense_web/components/new-chat/image-model-selector.tsx
new file mode 100644
index 000000000..8cae10345
--- /dev/null
+++ b/surfsense_web/components/new-chat/image-model-selector.tsx
@@ -0,0 +1,335 @@
+"use client";
+
+import { useAtomValue } from "jotai";
+import {
+	Check,
+	ChevronDown,
+	ChevronRight,
+	Edit3,
+	Globe,
+	ImageIcon,
+	Plus,
+	Shuffle,
+	User,
+} from "lucide-react";
+import { useCallback, useMemo, useState } from "react";
+import { toast } from "sonner";
+import {
+	createImageGenConfigMutationAtom,
+	updateImageGenConfigMutationAtom,
+} from "@/atoms/image-gen-config/image-gen-config-mutation.atoms";
+import {
+	globalImageGenConfigsAtom,
+	imageGenConfigsAtom,
+} from "@/atoms/image-gen-config/image-gen-config-query.atoms";
+import { updateLLMPreferencesMutationAtom } from "@/atoms/new-llm-config/new-llm-config-mutation.atoms";
+import { llmPreferencesAtom } from "@/atoms/new-llm-config/new-llm-config-query.atoms";
+import { activeSearchSpaceIdAtom } from "@/atoms/search-spaces/search-space-query.atoms";
+import { Badge } from "@/components/ui/badge";
+import { Button } from "@/components/ui/button";
+import {
+	Command,
+	CommandEmpty,
+	CommandGroup,
+	CommandInput,
+	CommandItem,
+	CommandList,
+	CommandSeparator,
+} from "@/components/ui/command";
+import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover";
+import { Spinner } from "@/components/ui/spinner";
+import { cn } from "@/lib/utils";
+
+interface ImageModelSelectorProps {
+	className?: string;
+	onAddNew?: () => void;
+}
+
+export function ImageModelSelector({ className, onAddNew }: ImageModelSelectorProps) {
+	const [open, setOpen] = useState(false);
+	const [searchQuery, setSearchQuery] = useState("");
+
+	const { data: globalConfigs, isLoading: globalLoading } =
+		useAtomValue(globalImageGenConfigsAtom);
+	const { data: userConfigs, isLoading: userLoading } = useAtomValue(imageGenConfigsAtom);
+	const { data: preferences, isLoading: prefsLoading } = useAtomValue(llmPreferencesAtom);
+	const searchSpaceId = useAtomValue(activeSearchSpaceIdAtom);
+	const { mutateAsync: updatePreferences } = useAtomValue(updateLLMPreferencesMutationAtom);
+
+	const isLoading = globalLoading || userLoading || prefsLoading;
+
+	const currentConfig = useMemo(() => {
+		if (!preferences) return null;
+		const id = preferences.image_generation_config_id;
+		if (id === null || id === undefined) return null;
+		const globalMatch = globalConfigs?.find((c) => c.id === id);
+		if (globalMatch) return globalMatch;
+		return userConfigs?.find((c) => c.id === id) ?? null;
+	}, [preferences, globalConfigs, userConfigs]);
+
+	const isCurrentAutoMode = useMemo(() => {
+		return currentConfig && "is_auto_mode" in currentConfig && currentConfig.is_auto_mode;
+	}, [currentConfig]);
+
+	const filteredGlobal = useMemo(() => {
+		if (!globalConfigs) return [];
+		if (!searchQuery) return globalConfigs;
+		const q = searchQuery.toLowerCase();
+		return globalConfigs.filter(
+			(c) =>
+				c.name.toLowerCase().includes(q) ||
+				c.model_name.toLowerCase().includes(q) ||
+				c.provider.toLowerCase().includes(q)
+		);
+	}, [globalConfigs, searchQuery]);
+
+	const filteredUser = useMemo(() => {
+		if (!userConfigs) return [];
+		if (!searchQuery) return userConfigs;
+		const q = searchQuery.toLowerCase();
+		return userConfigs.filter(
+			(c) =>
+				c.name.toLowerCase().includes(q) ||
+				c.model_name.toLowerCase().includes(q) ||
+				c.provider.toLowerCase().includes(q)
+		);
+	}, [userConfigs, searchQuery]);
+
+	const totalModels = (globalConfigs?.length ?? 0) + (userConfigs?.length ?? 0);
+
+	const handleSelect = useCallback(
+		async (configId: number) => {
+			if (currentConfig?.id === configId) {
+				setOpen(false);
+				return;
+			}
+			if (!searchSpaceId) {
+				toast.error("No search space selected");
+				return;
+			}
+			try {
+				await updatePreferences({
+					search_space_id: Number(searchSpaceId),
+					data: { image_generation_config_id: configId },
+				});
+				toast.success("Image model updated");
+				setOpen(false);
+			} catch {
+				toast.error("Failed to switch image model");
+			}
+		},
+		[currentConfig, searchSpaceId, updatePreferences]
+	);
+
+	// Don't render if no configs at all
+	if (!isLoading && totalModels === 0) {
+		return (
+			<Button
+				variant="outline"
+				size="sm"
+				onClick={onAddNew}
+				className={cn("h-8 gap-2 px-3 text-sm border-border/60", className)}
+			>
+				<Plus className="size-4 text-teal-600" />
+				<span className="hidden md:inline">Add Image Model</span>
+			</Button>
+		);
+	}
+
+	return (
+		<Popover open={open} onOpenChange={setOpen}>
+			<PopoverTrigger asChild>
+				<Button
+					variant="outline"
+					size="sm"
+					role="combobox"
+					aria-expanded={open}
+					className={cn("h-8 gap-2 px-3 text-sm border-border/60", className)}
+				>
+					{isLoading ? (
+						<Spinner size="sm" className="text-muted-foreground" />
+					) : currentConfig ? (
+						<>
+							{isCurrentAutoMode ? (
+								<Shuffle className="size-4 text-violet-500" />
+							) : (
+								<ImageIcon className="size-4 text-teal-500" />
+							)}
+							<span className="max-w-[100px] md:max-w-[120px] truncate hidden md:inline">
+								{currentConfig.name}
+							</span>
+							{isCurrentAutoMode ? (
+								<Badge
+									variant="secondary"
+									className="ml-1 text-[10px] px-1.5 py-0 h-4 bg-violet-100 text-violet-700 dark:bg-violet-900/30 dark:text-violet-300"
+								>
+									Auto
+								</Badge>
+							) : (
+								<Badge
+									variant="secondary"
+									className="ml-1 text-[10px] px-1.5 py-0 h-4 bg-teal-50 text-teal-700 dark:bg-teal-900/30 dark:text-teal-300"
+								>
+									Image
+								</Badge>
+							)}
+						</>
+					) : (
+						<>
+							<ImageIcon className="h-4 w-4 text-muted-foreground" />
+							<span className="text-muted-foreground hidden md:inline">Image Model</span>
+						</>
+					)}
+					<ChevronDown
+						className={cn(
+							"h-3.5 w-3.5 text-muted-foreground ml-1 shrink-0 transition-transform duration-200",
+							open && "rotate-180"
+						)}
+					/>
+				</Button>
+			</PopoverTrigger>
+
+			<PopoverContent
+				className="w-[280px] md:w-[360px] p-0 rounded-lg shadow-lg border-border/60"
+				align="start"
+				sideOffset={8}
+			>
+				<Command shouldFilter={false} className="rounded-lg">
+					{totalModels > 3 && (
+						<div className="flex items-center gap-1 md:gap-2 px-2 md:px-3 py-1.5 md:py-2">
+							<CommandInput
+								placeholder="Search image models..."
+								value={searchQuery}
+								onValueChange={setSearchQuery}
+								className="h-7 md:h-8 text-xs md:text-sm border-0 bg-transparent focus:ring-0 placeholder:text-muted-foreground/60"
+							/>
+						</div>
+					)}
+					<CommandList className="max-h-[300px] md:max-h-[400px] overflow-y-auto">
+						<CommandEmpty className="py-8 text-center">
+							<div className="flex flex-col items-center gap-2">
+								<ImageIcon className="size-8 text-muted-foreground" />
+								<p className="text-sm text-muted-foreground">No image models found</p>
+							</div>
+						</CommandEmpty>
+
+						{/* Global Image Gen Configs */}
+						{filteredGlobal.length > 0 && (
+							<CommandGroup>
+								<div className="flex items-center gap-2 px-3 py-2 text-xs font-semibold text-muted-foreground tracking-wider">
+									<Globe className="size-3.5" />
+									Global Image Models
+								</div>
+								{filteredGlobal.map((config) => {
+									const isSelected = currentConfig?.id === config.id;
+									const isAuto = "is_auto_mode" in config && config.is_auto_mode;
+									return (
+										<CommandItem
+											key={`g-${config.id}`}
+											value={`g-${config.id}`}
+											onSelect={() => handleSelect(config.id)}
+											className={cn(
+												"mx-2 rounded-lg mb-1 cursor-pointer group transition-all hover:bg-accent/50",
+												isSelected && "bg-accent/80",
+												isAuto && "border border-violet-200 dark:border-violet-800/50"
+											)}
+										>
+											<div className="flex items-center gap-3 min-w-0 flex-1">
+												<div className="shrink-0">
+													{isAuto ? (
+														<Shuffle className="size-4 text-violet-500" />
+													) : (
+														<ImageIcon className="size-4 text-teal-500" />
+													)}
+												</div>
+												<div className="min-w-0 flex-1">
+													<div className="flex items-center gap-2">
+														<span className="font-medium truncate">{config.name}</span>
+														{isAuto && (
+															<Badge
+																variant="secondary"
+																className="text-[9px] px-1 py-0 h-3.5 bg-violet-100 text-violet-700 dark:bg-violet-900/30 dark:text-violet-300 border-0"
+															>
+																Recommended
+															</Badge>
+														)}
+														{isSelected && <Check className="size-3.5 text-primary shrink-0" />}
+													</div>
+													<span className="text-xs text-muted-foreground truncate block">
+														{isAuto ? "Auto load balancing" : config.model_name}
+													</span>
+												</div>
+											</div>
+										</CommandItem>
+									);
+								})}
+							</CommandGroup>
+						)}
+
+						{/* User Image Gen Configs */}
+						{filteredUser.length > 0 && (
+							<>
+								{filteredGlobal.length > 0 && <CommandSeparator className="my-1 bg-border/30" />}
+								<CommandGroup>
+									<div className="flex items-center gap-2 px-3 py-2 text-xs font-semibold text-muted-foreground tracking-wider">
+										<User className="size-3.5" />
+										Your Image Models
+									</div>
+									{filteredUser.map((config) => {
+										const isSelected = currentConfig?.id === config.id;
+										return (
+											<CommandItem
+												key={`u-${config.id}`}
+												value={`u-${config.id}`}
+												onSelect={() => handleSelect(config.id)}
+												className={cn(
+													"mx-2 rounded-lg mb-1 cursor-pointer group transition-all hover:bg-accent/50",
+													isSelected && "bg-accent/80"
+												)}
+											>
+												<div className="flex items-center gap-3 min-w-0 flex-1">
+													<div className="shrink-0">
+														<ImageIcon className="size-4 text-teal-500" />
+													</div>
+													<div className="min-w-0 flex-1">
+														<div className="flex items-center gap-2">
+															<span className="font-medium truncate">{config.name}</span>
+															{isSelected && (
+																<Check className="size-3.5 text-primary shrink-0" />
+															)}
+														</div>
+														<span className="text-xs text-muted-foreground truncate block">
+															{config.model_name}
+														</span>
+													</div>
+												</div>
+											</CommandItem>
+										);
+									})}
+								</CommandGroup>
+							</>
+						)}
+
+						{/* Add New */}
+						{onAddNew && (
+							<div className="p-2 bg-muted/20">
+								<Button
+									variant="ghost"
+									size="sm"
+									className="w-full justify-start gap-2 h-9 rounded-lg hover:bg-accent/50"
+									onClick={() => {
+										setOpen(false);
+										onAddNew();
+									}}
+								>
+									<Plus className="size-4 text-teal-600" />
+									<span className="text-sm font-medium">Add Image Model</span>
+								</Button>
+							</div>
+						)}
+					</CommandList>
+				</Command>
+			</PopoverContent>
+		</Popover>
+	);
+}
diff --git a/surfsense_web/components/new-chat/model-selector.tsx b/surfsense_web/components/new-chat/model-selector.tsx
index ec1143e04..148028df2 100644
--- a/surfsense_web/components/new-chat/model-selector.tsx
+++ b/surfsense_web/components/new-chat/model-selector.tsx
@@ -392,8 +392,8 @@ export function ModelSelector({ onEdit, onAddNew, className }: ModelSelectorProp
 							</CommandGroup>
 						)}
 
-						{/* Add New Config Button */}
-						<div className="p-2 bg-muted/20">
+					{/* Add New Config Button */}
+					<div className="p-2 bg-muted/20">
 							<Button
 								variant="ghost"
 								size="sm"
diff --git a/surfsense_web/components/settings/image-model-manager.tsx b/surfsense_web/components/settings/image-model-manager.tsx
new file mode 100644
index 000000000..4e3f1840a
--- /dev/null
+++ b/surfsense_web/components/settings/image-model-manager.tsx
@@ -0,0 +1,692 @@
+"use client";
+
+import { useAtomValue } from "jotai";
+import {
+	AlertCircle,
+	Check,
+	ChevronsUpDown,
+	Clock,
+	Edit3,
+	ImageIcon,
+	Key,
+	Plus,
+	RefreshCw,
+	Shuffle,
+	Sparkles,
+	Trash2,
+	Wand2,
+} from "lucide-react";
+import { AnimatePresence, motion } from "motion/react";
+import { useCallback, useEffect, useState } from "react";
+import { toast } from "sonner";
+import {
+	createImageGenConfigMutationAtom,
+	deleteImageGenConfigMutationAtom,
+	updateImageGenConfigMutationAtom,
+} from "@/atoms/image-gen-config/image-gen-config-mutation.atoms";
+import {
+	globalImageGenConfigsAtom,
+	imageGenConfigsAtom,
+} from "@/atoms/image-gen-config/image-gen-config-query.atoms";
+import { updateLLMPreferencesMutationAtom } from "@/atoms/new-llm-config/new-llm-config-mutation.atoms";
+import { llmPreferencesAtom } from "@/atoms/new-llm-config/new-llm-config-query.atoms";
+import { Alert, AlertDescription } from "@/components/ui/alert";
+import {
+	AlertDialog,
+	AlertDialogAction,
+	AlertDialogCancel,
+	AlertDialogContent,
+	AlertDialogDescription,
+	AlertDialogFooter,
+	AlertDialogHeader,
+	AlertDialogTitle,
+} from "@/components/ui/alert-dialog";
+import { Badge } from "@/components/ui/badge";
+import { Button } from "@/components/ui/button";
+import { Card, CardContent, CardDescription, CardHeader, CardTitle } from "@/components/ui/card";
+import {
+	Command,
+	CommandEmpty,
+	CommandGroup,
+	CommandInput,
+	CommandItem,
+	CommandList,
+} from "@/components/ui/command";
+import {
+	Dialog,
+	DialogContent,
+	DialogDescription,
+	DialogHeader,
+	DialogTitle,
+} from "@/components/ui/dialog";
+import { Input } from "@/components/ui/input";
+import { Label } from "@/components/ui/label";
+import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover";
+import {
+	Select,
+	SelectContent,
+	SelectItem,
+	SelectTrigger,
+	SelectValue,
+} from "@/components/ui/select";
+import { Separator } from "@/components/ui/separator";
+import { Spinner } from "@/components/ui/spinner";
+import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip";
+import {
+	IMAGE_GEN_PROVIDERS,
+	getImageGenModelsByProvider,
+} from "@/contracts/enums/image-gen-providers";
+import type { ImageGenerationConfig } from "@/contracts/types/new-llm-config.types";
+import { cn } from "@/lib/utils";
+
+interface ImageModelManagerProps {
+	searchSpaceId: number;
+}
+
+const container = {
+	hidden: { opacity: 0 },
+	show: { opacity: 1, transition: { staggerChildren: 0.05 } },
+};
+
+const item = {
+	hidden: { opacity: 0, y: 20 },
+	show: { opacity: 1, y: 0 },
+};
+
+export function ImageModelManager({ searchSpaceId }: ImageModelManagerProps) {
+	// Image gen config atoms
+	const { mutateAsync: createConfig, isPending: isCreating, error: createError } =
+		useAtomValue(createImageGenConfigMutationAtom);
+	const { mutateAsync: updateConfig, isPending: isUpdating, error: updateError } =
+		useAtomValue(updateImageGenConfigMutationAtom);
+	const { mutateAsync: deleteConfig, isPending: isDeleting, error: deleteError } =
+		useAtomValue(deleteImageGenConfigMutationAtom);
+	const { mutateAsync: updatePreferences } = useAtomValue(updateLLMPreferencesMutationAtom);
+
+	const { data: userConfigs, isFetching: configsLoading, error: fetchError, refetch: refreshConfigs } =
+		useAtomValue(imageGenConfigsAtom);
+	const { data: globalConfigs = [], isFetching: globalLoading } =
+		useAtomValue(globalImageGenConfigsAtom);
+	const { data: preferences = {}, isFetching: prefsLoading } = useAtomValue(llmPreferencesAtom);
+
+	// Local state
+	const [isDialogOpen, setIsDialogOpen] = useState(false);
+	const [editingConfig, setEditingConfig] = useState<ImageGenerationConfig | null>(null);
+	const [configToDelete, setConfigToDelete] = useState<ImageGenerationConfig | null>(null);
+
+	// Preference state
+	const [selectedPrefId, setSelectedPrefId] = useState<string | number>(
+		preferences.image_generation_config_id ?? ""
+	);
+	const [hasPrefChanges, setHasPrefChanges] = useState(false);
+	const [isSavingPref, setIsSavingPref] = useState(false);
+
+	useEffect(() => {
+		setSelectedPrefId(preferences.image_generation_config_id ?? "");
+		setHasPrefChanges(false);
+	}, [preferences]);
+
+	const isSubmitting = isCreating || isUpdating;
+	const isLoading = configsLoading || globalLoading || prefsLoading;
+	const errors = [createError, updateError, deleteError, fetchError].filter(Boolean) as Error[];
+
+	// Form state for create/edit dialog
+	const [formData, setFormData] = useState({
+		name: "",
+		description: "",
+		provider: "",
+		custom_provider: "",
+		model_name: "",
+		api_key: "",
+		api_base: "",
+		api_version: "",
+	});
+	const [modelComboboxOpen, setModelComboboxOpen] = useState(false);
+
+	const resetForm = () => {
+		setFormData({
+			name: "",
+			description: "",
+			provider: "",
+			custom_provider: "",
+			model_name: "",
+			api_key: "",
+			api_base: "",
+			api_version: "",
+		});
+	};
+
+	const handleFormSubmit = useCallback(async () => {
+		if (!formData.name || !formData.provider || !formData.model_name || !formData.api_key) {
+			toast.error("Please fill in all required fields");
+			return;
+		}
+		try {
+			if (editingConfig) {
+				await updateConfig({
+					id: editingConfig.id,
+					data: {
+						name: formData.name,
+						description: formData.description || undefined,
+						provider: formData.provider as any,
+						custom_provider: formData.custom_provider || undefined,
+						model_name: formData.model_name,
+						api_key: formData.api_key,
+						api_base: formData.api_base || undefined,
+						api_version: formData.api_version || undefined,
+					},
+				});
+			} else {
+				const result = await createConfig({
+					name: formData.name,
+					description: formData.description || undefined,
+					provider: formData.provider as any,
+					custom_provider: formData.custom_provider || undefined,
+					model_name: formData.model_name,
+					api_key: formData.api_key,
+					api_base: formData.api_base || undefined,
+					api_version: formData.api_version || undefined,
+					search_space_id: searchSpaceId,
+				});
+				// Auto-assign newly created config
+				if (result?.id) {
+					await updatePreferences({
+						search_space_id: searchSpaceId,
+						data: { image_generation_config_id: result.id },
+					});
+				}
+			}
+			setIsDialogOpen(false);
+			setEditingConfig(null);
+			resetForm();
+		} catch {
+			// Error handled by mutation
+		}
+	}, [editingConfig, formData, searchSpaceId, createConfig, updateConfig, updatePreferences]);
+
+	const handleDelete = async () => {
+		if (!configToDelete) return;
+		try {
+			await deleteConfig(configToDelete.id);
+			setConfigToDelete(null);
+		} catch {
+			// Error handled by mutation
+		}
+	};
+
+	const openEditDialog = (config: ImageGenerationConfig) => {
+		setEditingConfig(config);
+		setFormData({
+			name: config.name,
+			description: config.description || "",
+			provider: config.provider,
+			custom_provider: config.custom_provider || "",
+			model_name: config.model_name,
+			api_key: config.api_key,
+			api_base: config.api_base || "",
+			api_version: config.api_version || "",
+		});
+		setIsDialogOpen(true);
+	};
+
+	const openNewDialog = () => {
+		setEditingConfig(null);
+		resetForm();
+		setIsDialogOpen(true);
+	};
+
+	const handlePrefChange = (value: string) => {
+		const newVal = value === "unassigned" ? "" : parseInt(value);
+		setSelectedPrefId(newVal);
+		setHasPrefChanges(newVal !== (preferences.image_generation_config_id ?? ""));
+	};
+
+	const handleSavePref = async () => {
+		setIsSavingPref(true);
+		try {
+			await updatePreferences({
+				search_space_id: searchSpaceId,
+				data: {
+					image_generation_config_id:
+						typeof selectedPrefId === "string"
+							? selectedPrefId ? parseInt(selectedPrefId) : undefined
+							: selectedPrefId,
+				},
+			});
+			setHasPrefChanges(false);
+			toast.success("Image generation model preference saved!");
+		} catch {
+			toast.error("Failed to save preference");
+		} finally {
+			setIsSavingPref(false);
+		}
+	};
+
+	const allConfigs = [
+		...globalConfigs.map((c) => ({ ...c, _source: "global" as const })),
+		...(userConfigs ?? []).map((c) => ({ ...c, _source: "user" as const })),
+	];
+
+	const selectedProvider = IMAGE_GEN_PROVIDERS.find((p) => p.value === formData.provider);
+	const suggestedModels = getImageGenModelsByProvider(formData.provider);
+
+	return (
+		<div className="space-y-4 md:space-y-6">
+			{/* Header */}
+			<div className="flex flex-col space-y-4 sm:flex-row sm:items-center sm:justify-between sm:space-y-0">
+				<Button
+					variant="outline"
+					size="sm"
+					onClick={() => refreshConfigs()}
+					disabled={isLoading}
+					className="flex items-center gap-2 text-xs md:text-sm h-8 md:h-9"
+				>
+					<RefreshCw className={cn("h-3 w-3 md:h-4 md:w-4", configsLoading && "animate-spin")} />
+					Refresh
+				</Button>
+			</div>
+
+			{/* Errors */}
+			<AnimatePresence>
+				{errors.map((err) => (
+					<motion.div key={err?.message} initial={{ opacity: 0, y: -10 }} animate={{ opacity: 1, y: 0 }} exit={{ opacity: 0, y: -10 }}>
+						<Alert variant="destructive" className="py-3">
+							<AlertCircle className="h-3 w-3 md:h-4 md:w-4 shrink-0" />
+							<AlertDescription className="text-xs md:text-sm">{err?.message}</AlertDescription>
+						</Alert>
+					</motion.div>
+				))}
+			</AnimatePresence>
+
+			{/* Global info */}
+			{globalConfigs.filter((g) => !("is_auto_mode" in g && g.is_auto_mode)).length > 0 && (
+				<Alert className="border-teal-500/30 bg-teal-500/5 py-3">
+					<Sparkles className="h-3 w-3 md:h-4 md:w-4 text-teal-600 dark:text-teal-400 shrink-0" />
+					<AlertDescription className="text-teal-800 dark:text-teal-200 text-xs md:text-sm">
+						<span className="font-medium">
+							{globalConfigs.filter((g) => !("is_auto_mode" in g && g.is_auto_mode)).length} global image model(s)
+						</span>{" "}
+						available from your administrator.
+					</AlertDescription>
+				</Alert>
+			)}
+
+			{/* Active Preference Card */}
+			{!isLoading && allConfigs.length > 0 && (
+				<motion.div initial={{ opacity: 0, y: 10 }} animate={{ opacity: 1, y: 0 }}>
+					<Card className="border-l-4 border-l-teal-500">
+						<CardHeader className="pb-2 px-3 md:px-6 pt-3 md:pt-6">
+							<div className="flex items-center gap-2 md:gap-3">
+								<div className="p-1.5 md:p-2 rounded-lg bg-teal-100 text-teal-800">
+									<ImageIcon className="w-4 h-4 md:w-5 md:h-5" />
+								</div>
+								<div>
+									<CardTitle className="text-base md:text-lg">Active Image Model</CardTitle>
+									<CardDescription className="text-xs md:text-sm">
+										Select which model to use for image generation
+									</CardDescription>
+								</div>
+							</div>
+						</CardHeader>
+						<CardContent className="space-y-3 px-3 md:px-6 pb-3 md:pb-6">
+							<Select
+								value={selectedPrefId?.toString() || "unassigned"}
+								onValueChange={handlePrefChange}
+							>
+								<SelectTrigger className="h-9 md:h-10 text-xs md:text-sm">
+									<SelectValue placeholder="Select an image model" />
+								</SelectTrigger>
+								<SelectContent>
+									<SelectItem value="unassigned">
+										<span className="text-muted-foreground">Unassigned</span>
+									</SelectItem>
+									{globalConfigs.length > 0 && (
+										<>
+											<div className="px-2 py-1.5 text-xs font-semibold text-muted-foreground">Global</div>
+											{globalConfigs.map((c) => {
+												const isAuto = "is_auto_mode" in c && c.is_auto_mode;
+												return (
+													<SelectItem key={`g-${c.id}`} value={c.id.toString()}>
+														<div className="flex items-center gap-2">
+															{isAuto ? (
+																<Badge variant="outline" className="text-xs bg-violet-100 text-violet-700 dark:bg-violet-900/30 dark:text-violet-300 border-violet-200">
+																	<Shuffle className="size-3 mr-1" />AUTO
+																</Badge>
+															) : (
+																<Badge variant="outline" className="text-xs bg-teal-50 text-teal-700 dark:bg-teal-900/30 dark:text-teal-300 border-teal-200">
+																	{c.provider}
+																</Badge>
+															)}
+															<span>{c.name}</span>
+														</div>
+													</SelectItem>
+												);
+											})}
+										</>
+									)}
+									{(userConfigs?.length ?? 0) > 0 && (
+										<>
+											<div className="px-2 py-1.5 text-xs font-semibold text-muted-foreground">Your Models</div>
+											{userConfigs?.map((c) => (
+												<SelectItem key={`u-${c.id}`} value={c.id.toString()}>
+													<div className="flex items-center gap-2">
+														<Badge variant="outline" className="text-xs">{c.provider}</Badge>
+														<span>{c.name}</span>
+														<span className="text-muted-foreground">({c.model_name})</span>
+													</div>
+												</SelectItem>
+											))}
+										</>
+									)}
+								</SelectContent>
+							</Select>
+							{hasPrefChanges && (
+								<div className="flex gap-2 pt-1">
+									<Button size="sm" onClick={handleSavePref} disabled={isSavingPref} className="text-xs h-8">
+										{isSavingPref ? "Saving..." : "Save"}
+									</Button>
+									<Button size="sm" variant="outline" onClick={() => { setSelectedPrefId(preferences.image_generation_config_id ?? ""); setHasPrefChanges(false); }} className="text-xs h-8">
+										Reset
+									</Button>
+								</div>
+							)}
+						</CardContent>
+					</Card>
+				</motion.div>
+			)}
+
+			{/* Loading */}
+			{isLoading && (
+				<Card>
+					<CardContent className="flex items-center justify-center py-10">
+						<Spinner size="md" className="text-muted-foreground" />
+					</CardContent>
+				</Card>
+			)}
+
+			{/* User Configs */}
+			{!isLoading && (
+				<div className="space-y-4 md:space-y-6">
+					<div className="flex flex-col space-y-4 sm:flex-row sm:items-center sm:justify-between sm:space-y-0">
+						<h3 className="text-lg md:text-xl font-semibold tracking-tight">Your Image Models</h3>
+						<Button onClick={openNewDialog} className="flex items-center gap-2 text-xs md:text-sm h-8 md:h-9">
+							<Plus className="h-3 w-3 md:h-4 md:w-4" />
+							Add Image Model
+						</Button>
+					</div>
+
+					{(userConfigs?.length ?? 0) === 0 ? (
+						<Card className="border-dashed border-2 border-muted-foreground/25">
+							<CardContent className="flex flex-col items-center justify-center py-10 md:py-16 text-center">
+								<div className="rounded-full bg-gradient-to-br from-teal-500/10 to-cyan-500/10 p-4 md:p-6 mb-4">
+									<Wand2 className="h-8 w-8 md:h-12 md:w-12 text-teal-600 dark:text-teal-400" />
+								</div>
+								<h3 className="text-lg font-semibold mb-2">No Image Models Yet</h3>
+								<p className="text-xs md:text-sm text-muted-foreground max-w-sm mb-4">
+									Add your own image generation model (DALL-E 3, GPT Image 1, etc.)
+								</p>
+								<Button onClick={openNewDialog} size="lg" className="gap-2 text-xs md:text-sm">
+									<Plus className="h-3 w-3 md:h-4 md:w-4" />
+									Add First Image Model
+								</Button>
+							</CardContent>
+						</Card>
+					) : (
+						<motion.div variants={container} initial="hidden" animate="show" className="grid gap-4">
+							<AnimatePresence mode="popLayout">
+								{userConfigs?.map((config) => (
+									<motion.div key={config.id} variants={item} layout exit={{ opacity: 0, scale: 0.95 }}>
+										<Card className="group overflow-hidden hover:shadow-lg transition-all duration-300 border-muted-foreground/10 hover:border-teal-500/30">
+											<CardContent className="p-0">
+												<div className="flex">
+													<div className="w-1 md:w-1.5 bg-gradient-to-b from-teal-500/50 to-cyan-500/50 group-hover:from-teal-500 group-hover:to-cyan-500 transition-colors" />
+													<div className="flex-1 p-3 md:p-5">
+														<div className="flex items-start justify-between gap-2">
+															<div className="flex items-start gap-2 md:gap-4 flex-1 min-w-0">
+																<div className="flex h-10 w-10 md:h-12 md:w-12 items-center justify-center rounded-lg md:rounded-xl bg-gradient-to-br from-teal-500/10 to-cyan-500/10 shrink-0">
+																	<ImageIcon className="h-5 w-5 md:h-6 md:w-6 text-teal-600 dark:text-teal-400" />
+																</div>
+																<div className="flex-1 min-w-0 space-y-2">
+																	<div className="flex items-center gap-1.5 flex-wrap">
+																		<h4 className="text-sm md:text-base font-semibold truncate">{config.name}</h4>
+																		<Badge variant="secondary" className="text-[9px] md:text-[10px] px-1.5 py-0.5 bg-teal-500/10 text-teal-700 dark:text-teal-300 border-teal-500/20">
+																			{config.provider}
+																		</Badge>
+																	</div>
+																	<code className="text-[10px] md:text-xs font-mono text-muted-foreground bg-muted/50 px-1.5 py-0.5 rounded-md inline-block">
+																		{config.model_name}
+																	</code>
+																	{config.description && (
+																		<p className="text-[10px] md:text-xs text-muted-foreground line-clamp-1">{config.description}</p>
+																	)}
+																	<div className="flex items-center gap-1 text-[10px] md:text-xs text-muted-foreground pt-1">
+																		<Clock className="h-2.5 w-2.5 md:h-3 md:w-3" />
+																		{new Date(config.created_at).toLocaleDateString()}
+																	</div>
+																</div>
+															</div>
+															<div className="flex items-center gap-0.5 shrink-0 opacity-0 group-hover:opacity-100 transition-opacity">
+																<TooltipProvider>
+																	<Tooltip>
+																		<TooltipTrigger asChild>
+																			<Button variant="ghost" size="sm" onClick={() => openEditDialog(config)} className="h-7 w-7 p-0 text-muted-foreground hover:text-foreground">
+																				<Edit3 className="h-3.5 w-3.5" />
+																			</Button>
+																		</TooltipTrigger>
+																		<TooltipContent>Edit</TooltipContent>
+																	</Tooltip>
+																</TooltipProvider>
+																<TooltipProvider>
+																	<Tooltip>
+																		<TooltipTrigger asChild>
+																			<Button variant="ghost" size="sm" onClick={() => setConfigToDelete(config)} className="h-7 w-7 p-0 text-muted-foreground hover:text-destructive">
+																				<Trash2 className="h-3.5 w-3.5" />
+																			</Button>
+																		</TooltipTrigger>
+																		<TooltipContent>Delete</TooltipContent>
+																	</Tooltip>
+																</TooltipProvider>
+															</div>
+														</div>
+													</div>
+												</div>
+											</CardContent>
+										</Card>
+									</motion.div>
+								))}
+							</AnimatePresence>
+						</motion.div>
+					)}
+				</div>
+			)}
+
+			{/* Create/Edit Dialog */}
+			<Dialog open={isDialogOpen} onOpenChange={(open) => { if (!open) { setIsDialogOpen(false); setEditingConfig(null); resetForm(); } }}>
+				<DialogContent className="max-w-lg max-h-[90vh] overflow-y-auto">
+					<DialogHeader>
+						<DialogTitle className="flex items-center gap-2">
+							{editingConfig ? <Edit3 className="w-5 h-5 text-teal-600" /> : <Plus className="w-5 h-5 text-teal-600" />}
+							{editingConfig ? "Edit Image Model" : "Add Image Model"}
+						</DialogTitle>
+						<DialogDescription>
+							{editingConfig ? "Update your image generation model" : "Configure a new image generation model (DALL-E 3, GPT Image 1, etc.)"}
+						</DialogDescription>
+					</DialogHeader>
+
+					<div className="space-y-4 pt-2">
+						{/* Name */}
+						<div className="space-y-2">
+							<Label className="text-sm font-medium">Name *</Label>
+							<Input
+								placeholder="e.g., My DALL-E 3"
+								value={formData.name}
+								onChange={(e) => setFormData((p) => ({ ...p, name: e.target.value }))}
+							/>
+						</div>
+
+						{/* Description */}
+						<div className="space-y-2">
+							<Label className="text-sm font-medium">Description</Label>
+							<Input
+								placeholder="Optional description"
+								value={formData.description}
+								onChange={(e) => setFormData((p) => ({ ...p, description: e.target.value }))}
+							/>
+						</div>
+
+						<Separator />
+
+						{/* Provider */}
+						<div className="space-y-2">
+							<Label className="text-sm font-medium">Provider *</Label>
+							<Select
+								value={formData.provider}
+								onValueChange={(val) => setFormData((p) => ({ ...p, provider: val, model_name: "" }))}
+							>
+								<SelectTrigger>
+									<SelectValue placeholder="Select a provider" />
+								</SelectTrigger>
+								<SelectContent>
+									{IMAGE_GEN_PROVIDERS.map((p) => (
+										<SelectItem key={p.value} value={p.value}>
+											<div className="flex flex-col">
+												<span className="font-medium">{p.label}</span>
+												<span className="text-xs text-muted-foreground">{p.example}</span>
+											</div>
+										</SelectItem>
+									))}
+								</SelectContent>
+							</Select>
+						</div>
+
+						{/* Model Name */}
+						<div className="space-y-2">
+							<Label className="text-sm font-medium">Model Name *</Label>
+							{suggestedModels.length > 0 ? (
+								<Popover open={modelComboboxOpen} onOpenChange={setModelComboboxOpen}>
+									<PopoverTrigger asChild>
+										<Button variant="outline" role="combobox" className="w-full justify-between font-normal">
+											{formData.model_name || "Select or type a model..."}
+											<ChevronsUpDown className="ml-2 h-4 w-4 shrink-0 opacity-50" />
+										</Button>
+									</PopoverTrigger>
+									<PopoverContent className="w-full p-0" align="start">
+										<Command>
+											<CommandInput
+												placeholder="Search or type model name..."
+												value={formData.model_name}
+												onValueChange={(val) => setFormData((p) => ({ ...p, model_name: val }))}
+											/>
+											<CommandList>
+												<CommandEmpty>
+													<span className="text-xs text-muted-foreground">Type a custom model name</span>
+												</CommandEmpty>
+												<CommandGroup>
+													{suggestedModels.map((m) => (
+														<CommandItem
+															key={m.value}
+															value={m.value}
+															onSelect={() => {
+																setFormData((p) => ({ ...p, model_name: m.value }));
+																setModelComboboxOpen(false);
+															}}
+														>
+															<Check className={cn("mr-2 h-4 w-4", formData.model_name === m.value ? "opacity-100" : "opacity-0")} />
+															<span className="font-mono text-sm">{m.value}</span>
+															<span className="ml-2 text-xs text-muted-foreground">{m.label}</span>
+														</CommandItem>
+													))}
+												</CommandGroup>
+											</CommandList>
+										</Command>
+									</PopoverContent>
+								</Popover>
+							) : (
+								<Input
+									placeholder="e.g., dall-e-3"
+									value={formData.model_name}
+									onChange={(e) => setFormData((p) => ({ ...p, model_name: e.target.value }))}
+								/>
+							)}
+						</div>
+
+						{/* API Key */}
+						<div className="space-y-2">
+							<Label className="text-sm font-medium flex items-center gap-1.5">
+								<Key className="h-3.5 w-3.5" /> API Key *
+							</Label>
+							<Input
+								type="password"
+								placeholder="sk-..."
+								value={formData.api_key}
+								onChange={(e) => setFormData((p) => ({ ...p, api_key: e.target.value }))}
+							/>
+						</div>
+
+						{/* API Base (optional) */}
+						<div className="space-y-2">
+							<Label className="text-sm font-medium">API Base URL</Label>
+							<Input
+								placeholder={selectedProvider?.apiBase || "Optional"}
+								value={formData.api_base}
+								onChange={(e) => setFormData((p) => ({ ...p, api_base: e.target.value }))}
+							/>
+						</div>
+
+						{/* API Version (Azure) */}
+						{formData.provider === "AZURE_OPENAI" && (
+							<div className="space-y-2">
+								<Label className="text-sm font-medium">API Version (Azure)</Label>
+								<Input
+									placeholder="2024-02-15-preview"
+									value={formData.api_version}
+									onChange={(e) => setFormData((p) => ({ ...p, api_version: e.target.value }))}
+								/>
+							</div>
+						)}
+
+						{/* Actions */}
+						<div className="flex gap-3 pt-4 border-t">
+							<Button
+								variant="outline"
+								className="flex-1"
+								onClick={() => { setIsDialogOpen(false); setEditingConfig(null); resetForm(); }}
+							>
+								Cancel
+							</Button>
+							<Button
+								className="flex-1"
+								onClick={handleFormSubmit}
+								disabled={isSubmitting || !formData.name || !formData.provider || !formData.model_name || !formData.api_key}
+							>
+								{isSubmitting ? <Spinner size="sm" className="mr-2" /> : null}
+								{editingConfig ? "Save Changes" : "Create & Use"}
+							</Button>
+						</div>
+					</div>
+				</DialogContent>
+			</Dialog>
+
+			{/* Delete Confirmation */}
+			<AlertDialog open={!!configToDelete} onOpenChange={(open) => !open && setConfigToDelete(null)}>
+				<AlertDialogContent>
+					<AlertDialogHeader>
+						<AlertDialogTitle className="flex items-center gap-2">
+							<Trash2 className="h-5 w-5 text-destructive" />
+							Delete Image Model
+						</AlertDialogTitle>
+						<AlertDialogDescription>
+							Are you sure you want to delete <span className="font-semibold text-foreground">{configToDelete?.name}</span>?
+						</AlertDialogDescription>
+					</AlertDialogHeader>
+					<AlertDialogFooter>
+						<AlertDialogCancel disabled={isDeleting}>Cancel</AlertDialogCancel>
+						<AlertDialogAction onClick={handleDelete} disabled={isDeleting} className="bg-destructive text-destructive-foreground hover:bg-destructive/90">
+							{isDeleting ? <><Spinner size="sm" className="mr-2" />Deleting</> : <><Trash2 className="mr-2 h-4 w-4" />Delete</>}
+						</AlertDialogAction>
+					</AlertDialogFooter>
+				</AlertDialogContent>
+			</AlertDialog>
+		</div>
+	);
+}
diff --git a/surfsense_web/components/settings/llm-role-manager.tsx b/surfsense_web/components/settings/llm-role-manager.tsx
index dac68a358..22e3d8e08 100644
--- a/surfsense_web/components/settings/llm-role-manager.tsx
+++ b/surfsense_web/components/settings/llm-role-manager.tsx
@@ -255,15 +255,15 @@ export function LLMRoleManager({ searchSpaceId }: LLMRoleManagerProps) {
 						</Alert>
 					)}
 
-					{/* Role Assignment Cards */}
-					{availableConfigs.length > 0 && (
-						<div className="grid gap-4 md:gap-6">
-							{Object.entries(ROLE_DESCRIPTIONS).map(([key, role]) => {
-								const IconComponent = role.icon;
-								const currentAssignment = assignments[`${key}_llm_id` as keyof typeof assignments];
-								const assignedConfig = availableConfigs.find(
-									(config) => config.id === currentAssignment
-								);
+				{/* Role Assignment Cards */}
+				{availableConfigs.length > 0 && (
+					<div className="grid gap-4 md:gap-6">
+						{Object.entries(ROLE_DESCRIPTIONS).map(([key, role]) => {
+							const IconComponent = role.icon;
+							const currentAssignment = assignments[`${key}_llm_id` as keyof typeof assignments];
+							const assignedConfig = availableConfigs.find(
+								(config) => config.id === currentAssignment
+							);
 
 								return (
 									<motion.div
@@ -294,100 +294,100 @@ export function LLMRoleManager({ searchSpaceId }: LLMRoleManagerProps) {
 												</div>
 											</CardHeader>
 											<CardContent className="space-y-3 md:space-y-4 px-3 md:px-6 pb-3 md:pb-6">
-												<div className="space-y-1.5 md:space-y-2">
-													<Label className="text-xs md:text-sm font-medium">
-														Assign LLM Configuration:
-													</Label>
-													<Select
-														value={currentAssignment?.toString() || "unassigned"}
-														onValueChange={(value) => handleRoleAssignment(`${key}_llm_id`, value)}
-													>
-														<SelectTrigger className="h-9 md:h-10 text-xs md:text-sm">
-															<SelectValue placeholder="Select an LLM configuration" />
-														</SelectTrigger>
-														<SelectContent>
-															<SelectItem value="unassigned">
-																<span className="text-muted-foreground">Unassigned</span>
-															</SelectItem>
+											<div className="space-y-1.5 md:space-y-2">
+												<Label className="text-xs md:text-sm font-medium">
+													Assign LLM Configuration:
+												</Label>
+												<Select
+													value={currentAssignment?.toString() || "unassigned"}
+													onValueChange={(value) => handleRoleAssignment(`${key}_llm_id`, value)}
+												>
+													<SelectTrigger className="h-9 md:h-10 text-xs md:text-sm">
+														<SelectValue placeholder="Select an LLM configuration" />
+													</SelectTrigger>
+													<SelectContent>
+														<SelectItem value="unassigned">
+															<span className="text-muted-foreground">Unassigned</span>
+														</SelectItem>
 
-															{/* Global Configurations */}
-															{globalConfigs.length > 0 && (
-																<>
-																	<div className="px-2 py-1.5 text-xs font-semibold text-muted-foreground">
-																		Global Configurations
-																	</div>
-																	{globalConfigs.map((config) => {
-																		const isAutoMode =
-																			"is_auto_mode" in config && config.is_auto_mode;
-																		return (
-																			<SelectItem key={config.id} value={config.id.toString()}>
-																				<div className="flex items-center gap-2">
-																					{isAutoMode ? (
-																						<Badge
-																							variant="outline"
-																							className="text-xs bg-violet-100 text-violet-700 dark:bg-violet-900/30 dark:text-violet-300 border-violet-200 dark:border-violet-700"
-																						>
-																							<Shuffle className="size-3 mr-1" />
-																							AUTO
-																						</Badge>
-																					) : (
-																						<Badge variant="outline" className="text-xs">
-																							{config.provider}
-																						</Badge>
-																					)}
-																					<span>{config.name}</span>
-																					{!isAutoMode && (
-																						<span className="text-muted-foreground">
-																							({config.model_name})
-																						</span>
-																					)}
-																					{isAutoMode ? (
-																						<Badge
-																							variant="secondary"
-																							className="text-xs bg-violet-100 text-violet-700 dark:bg-violet-900/30 dark:text-violet-300"
-																						>
-																							Recommended
-																						</Badge>
-																					) : (
-																						<Badge variant="secondary" className="text-xs">
-																							🌐 Global
-																						</Badge>
-																					)}
-																				</div>
-																			</SelectItem>
-																		);
-																	})}
-																</>
-															)}
+														{/* Global Configurations */}
+														{globalConfigs.length > 0 && (
+																	<>
+																		<div className="px-2 py-1.5 text-xs font-semibold text-muted-foreground">
+																			Global Configurations
+																		</div>
+																		{globalConfigs.map((config) => {
+																			const isAutoMode =
+																				"is_auto_mode" in config && config.is_auto_mode;
+																			return (
+																				<SelectItem key={config.id} value={config.id.toString()}>
+																					<div className="flex items-center gap-2">
+																						{isAutoMode ? (
+																							<Badge
+																								variant="outline"
+																								className="text-xs bg-violet-100 text-violet-700 dark:bg-violet-900/30 dark:text-violet-300 border-violet-200 dark:border-violet-700"
+																							>
+																								<Shuffle className="size-3 mr-1" />
+																								AUTO
+																							</Badge>
+																						) : (
+																							<Badge variant="outline" className="text-xs">
+																								{config.provider}
+																							</Badge>
+																						)}
+																						<span>{config.name}</span>
+																						{!isAutoMode && (
+																							<span className="text-muted-foreground">
+																								({config.model_name})
+																							</span>
+																						)}
+																						{isAutoMode ? (
+																							<Badge
+																								variant="secondary"
+																								className="text-xs bg-violet-100 text-violet-700 dark:bg-violet-900/30 dark:text-violet-300"
+																							>
+																								Recommended
+																							</Badge>
+																						) : (
+																							<Badge variant="secondary" className="text-xs">
+																								🌐 Global
+																							</Badge>
+																						)}
+																					</div>
+																				</SelectItem>
+																			);
+																		})}
+																	</>
+																)}
 
-															{/* Custom Configurations */}
-															{newLLMConfigs.length > 0 && (
-																<>
-																	<div className="px-2 py-1.5 text-xs font-semibold text-muted-foreground">
-																		Your Configurations
-																	</div>
-																	{newLLMConfigs
-																		.filter(
-																			(config) => config.id && config.id.toString().trim() !== ""
-																		)
-																		.map((config) => (
-																			<SelectItem key={config.id} value={config.id.toString()}>
-																				<div className="flex items-center gap-2">
-																					<Badge variant="outline" className="text-xs">
-																						{config.provider}
-																					</Badge>
-																					<span>{config.name}</span>
-																					<span className="text-muted-foreground">
-																						({config.model_name})
-																					</span>
-																				</div>
-																			</SelectItem>
-																		))}
-																</>
-															)}
-														</SelectContent>
-													</Select>
-												</div>
+														{/* Custom Configurations */}
+														{newLLMConfigs.length > 0 && (
+															<>
+																<div className="px-2 py-1.5 text-xs font-semibold text-muted-foreground">
+																	Your Configurations
+																</div>
+																{newLLMConfigs
+																	.filter(
+																		(config) => config.id && config.id.toString().trim() !== ""
+																	)
+																	.map((config) => (
+																		<SelectItem key={config.id} value={config.id.toString()}>
+																			<div className="flex items-center gap-2">
+																				<Badge variant="outline" className="text-xs">
+																					{config.provider}
+																				</Badge>
+																				<span>{config.name}</span>
+																				<span className="text-muted-foreground">
+																					({config.model_name})
+																				</span>
+																			</div>
+																		</SelectItem>
+																	))}
+															</>
+														)}
+													</SelectContent>
+												</Select>
+											</div>
 
 												{assignedConfig && (
 													<div
diff --git a/surfsense_web/components/tool-ui/display-image.tsx b/surfsense_web/components/tool-ui/display-image.tsx
index 660e95bca..b5fccbc78 100644
--- a/surfsense_web/components/tool-ui/display-image.tsx
+++ b/surfsense_web/components/tool-ui/display-image.tsx
@@ -88,7 +88,7 @@ function ImageCancelledState({ src }: { src: string }) {
 function ParsedImage({ result }: { result: unknown }) {
 	const image = parseSerializableImage(result);
 
-	return <Image {...image} maxWidth="420px" />;
+	return <Image {...image} maxWidth="512px" />;
 }
 
 /**
diff --git a/surfsense_web/components/tool-ui/image/index.tsx b/surfsense_web/components/tool-ui/image/index.tsx
index 42725d258..9ecfe4cfa 100644
--- a/surfsense_web/components/tool-ui/image/index.tsx
+++ b/surfsense_web/components/tool-ui/image/index.tsx
@@ -1,6 +1,6 @@
 "use client";
 
-import { ExternalLinkIcon, ImageIcon } from "lucide-react";
+import { ExternalLinkIcon, ImageIcon, SparklesIcon } from "lucide-react";
 import NextImage from "next/image";
 import { Component, type ReactNode, useState } from "react";
 import { z } from "zod";
@@ -25,7 +25,7 @@ const SerializableImageSchema = z.object({
 	id: z.string(),
 	assetId: z.string(),
 	src: z.string(),
-	alt: z.string().nullish(), // Made optional - will use fallback if missing
+	alt: z.string().nullish(),
 	title: z.string().nullish(),
 	description: z.string().nullish(),
 	href: z.string().nullish(),
@@ -49,7 +49,7 @@ export interface ImageProps {
 	id: string;
 	assetId: string;
 	src: string;
-	alt?: string; // Optional with default fallback
+	alt?: string;
 	title?: string;
 	description?: string;
 	href?: string;
@@ -71,10 +71,8 @@ export function parseSerializableImage(result: unknown): SerializableImage & { a
 	if (!parsed.success) {
 		console.warn("Invalid image data:", parsed.error.issues);
 
-		// Try to extract basic info and return a fallback object
 		const obj = (result && typeof result === "object" ? result : {}) as Record<string, unknown>;
 
-		// If we have at least id, assetId, and src, we can still render the image
 		if (
 			typeof obj.id === "string" &&
 			typeof obj.assetId === "string" &&
@@ -89,7 +87,7 @@ export function parseSerializableImage(result: unknown): SerializableImage & { a
 				description: typeof obj.description === "string" ? obj.description : undefined,
 				href: typeof obj.href === "string" ? obj.href : undefined,
 				domain: typeof obj.domain === "string" ? obj.domain : undefined,
-				ratio: undefined, // Use default ratio
+				ratio: undefined,
 				source: undefined,
 			};
 		}
@@ -97,7 +95,6 @@ export function parseSerializableImage(result: unknown): SerializableImage & { a
 		throw new Error(`Invalid image: ${parsed.error.issues.map((i) => i.message).join(", ")}`);
 	}
 
-	// Provide fallback for alt if it's null/undefined
 	return {
 		...parsed.data,
 		alt: parsed.data.alt ?? "Image",
@@ -105,7 +102,7 @@ export function parseSerializableImage(result: unknown): SerializableImage & { a
 }
 
 /**
- * Get aspect ratio class based on ratio prop
+ * Get aspect ratio class based on ratio prop (used for fixed-ratio images only)
  */
 function getAspectRatioClass(ratio?: AspectRatio): string {
 	switch (ratio) {
@@ -119,7 +116,6 @@ function getAspectRatioClass(ratio?: AspectRatio): string {
 			return "aspect-[9/16]";
 		case "21:9":
 			return "aspect-[21/9]";
-		case "auto":
 		default:
 			return "aspect-[4/3]";
 	}
@@ -150,7 +146,7 @@ export class ImageErrorBoundary extends Component<
 		if (this.state.hasError) {
 			return (
 				<Card className="w-full max-w-md overflow-hidden">
-					<div className="aspect-[4/3] bg-muted flex items-center justify-center">
+					<div className="aspect-square bg-muted flex items-center justify-center">
 						<div className="flex flex-col items-center gap-2 text-muted-foreground">
 							<ImageIcon className="size-8" />
 							<p className="text-sm">Failed to load image</p>
@@ -167,10 +163,10 @@ export class ImageErrorBoundary extends Component<
 /**
  * Loading skeleton for Image
  */
-export function ImageSkeleton({ maxWidth = "420px" }: { maxWidth?: string }) {
+export function ImageSkeleton({ maxWidth = "512px" }: { maxWidth?: string }) {
 	return (
 		<Card className="w-full overflow-hidden animate-pulse" style={{ maxWidth }}>
-			<div className="aspect-[4/3] bg-muted flex items-center justify-center">
+			<div className="aspect-square bg-muted flex items-center justify-center">
 				<ImageIcon className="size-12 text-muted-foreground/30" />
 			</div>
 		</Card>
@@ -183,7 +179,7 @@ export function ImageSkeleton({ maxWidth = "420px" }: { maxWidth?: string }) {
 export function ImageLoading({ title = "Loading image..." }: { title?: string }) {
 	return (
 		<Card className="w-full max-w-md overflow-hidden">
-			<div className="aspect-[4/3] bg-muted flex items-center justify-center">
+			<div className="aspect-square bg-muted flex items-center justify-center">
 				<div className="flex flex-col items-center gap-3">
 					<Spinner size="lg" className="text-muted-foreground" />
 					<p className="text-muted-foreground text-sm">{title}</p>
@@ -197,7 +193,9 @@ export function ImageLoading({ title = "Loading image..." }: { title?: string })
  * Image Component
  *
  * Display images with metadata and attribution.
- * Features hover overlay with title and source attribution.
+ * - For "auto" ratio: renders the image at natural dimensions (no cropping)
+ * - For fixed ratios: uses a fixed aspect container with object-cover
+ * - Features hover overlay with title, description, and source attribution.
  */
 export function Image({
 	id,
@@ -207,16 +205,18 @@ export function Image({
 	description,
 	href,
 	domain,
-	ratio = "4:3",
+	ratio = "auto",
 	fit = "cover",
 	source,
-	maxWidth = "420px",
+	maxWidth = "512px",
 	className,
 }: ImageProps) {
 	const [isHovered, setIsHovered] = useState(false);
 	const [imageError, setImageError] = useState(false);
-	const aspectRatioClass = getAspectRatioClass(ratio);
+	const [imageLoaded, setImageLoaded] = useState(false);
 	const displayDomain = domain || source?.label;
+	const isGenerated = domain === "ai-generated";
+	const isAutoRatio = !ratio || ratio === "auto";
 
 	const handleClick = () => {
 		const targetUrl = href || source?.url || src;
@@ -228,7 +228,7 @@ export function Image({
 	if (imageError) {
 		return (
 			<Card id={id} className={cn("w-full overflow-hidden", className)} style={{ maxWidth }}>
-				<div className={cn("bg-muted flex items-center justify-center", aspectRatioClass)}>
+				<div className="aspect-square bg-muted flex items-center justify-center">
 					<div className="flex flex-col items-center gap-2 text-muted-foreground">
 						<ImageIcon className="size-8" />
 						<p className="text-sm">Image not available</p>
@@ -243,6 +243,7 @@ export function Image({
 			id={id}
 			className={cn(
 				"group w-full overflow-hidden cursor-pointer transition-shadow duration-200 hover:shadow-lg",
+				isGenerated && "ring-1 ring-primary/10",
 				className
 			)}
 			style={{ maxWidth }}
@@ -258,71 +259,98 @@ export function Image({
 			role="button"
 			tabIndex={0}
 		>
-			<div className={cn("relative w-full overflow-hidden bg-muted", aspectRatioClass)}>
-				{/* Image */}
-				<NextImage
-					src={src}
-					alt={alt}
-					fill
-					className={cn(
-						"transition-transform duration-300",
-						fit === "cover" ? "object-cover" : "object-contain",
-						isHovered && "scale-105"
-					)}
-					unoptimized
-					onError={() => setImageError(true)}
-				/>
+			<div className="relative w-full overflow-hidden bg-muted">
+				{isAutoRatio ? (
+					/* Auto ratio: image renders at natural dimensions, no cropping */
+					<>
+						{!imageLoaded && (
+							<div className="aspect-square flex items-center justify-center">
+								<Spinner size="lg" className="text-muted-foreground" />
+							</div>
+						)}
+						{/* eslint-disable-next-line @next/next/no-img-element */}
+						<img
+							src={src}
+							alt={alt}
+							className={cn(
+								"w-full h-auto transition-transform duration-300",
+								isHovered && "scale-[1.02]",
+								!imageLoaded && "hidden"
+							)}
+							onLoad={() => setImageLoaded(true)}
+							onError={() => setImageError(true)}
+						/>
+					</>
+				) : (
+					/* Fixed ratio: constrained aspect container with fill */
+					<div className={getAspectRatioClass(ratio)}>
+						<NextImage
+							src={src}
+							alt={alt}
+							fill
+							className={cn(
+								"transition-transform duration-300",
+								fit === "cover" ? "object-cover" : "object-contain",
+								isHovered && "scale-105"
+							)}
+							unoptimized
+							onError={() => setImageError(true)}
+						/>
+					</div>
+				)}
 
-				{/* Hover overlay - appears on hover */}
+				{/* Hover overlay */}
 				<div
 					className={cn(
-						"absolute inset-0 bg-gradient-to-t from-black/80 via-black/20 to-transparent",
+						"absolute inset-0 bg-gradient-to-t from-black/70 via-transparent to-transparent",
 						"transition-opacity duration-200",
 						isHovered ? "opacity-100" : "opacity-0"
 					)}
 				>
-					{/* Content at bottom */}
-					<div className="absolute bottom-0 left-0 right-0 p-4">
-						{/* Title */}
+					<div className="absolute bottom-0 left-0 right-0 p-3">
 						{title && (
-							<h3 className="font-semibold text-white text-base leading-tight line-clamp-2 mb-1">
+							<h3 className="font-semibold text-white text-sm leading-tight line-clamp-2 mb-0.5">
 								{title}
 							</h3>
 						)}
-
-						{/* Description */}
 						{description && (
-							<p className="text-white/80 text-sm line-clamp-2 mb-2">{description}</p>
+							<p className="text-white/80 text-xs line-clamp-2 mb-1.5">{description}</p>
 						)}
-
-						{/* Source attribution */}
 						{displayDomain && (
 							<div className="flex items-center gap-1.5">
-								{source?.iconUrl ? (
+								{isGenerated ? (
+									<SparklesIcon className="size-3.5 text-white/70" />
+								) : source?.iconUrl ? (
 									<NextImage
 										src={source.iconUrl}
 										alt={source.label}
-										width={16}
-										height={16}
+										width={14}
+										height={14}
 										className="rounded"
 										unoptimized
 									/>
 								) : (
-									<ExternalLinkIcon className="size-4 text-white/70" />
+									<ExternalLinkIcon className="size-3.5 text-white/70" />
 								)}
-								<span className="text-white/70 text-sm">{displayDomain}</span>
+								<span className="text-white/70 text-xs">{displayDomain}</span>
 							</div>
 						)}
 					</div>
 				</div>
 
-				{/* Always visible domain badge (bottom right, shown when NOT hovered) */}
+				{/* Badge when not hovered */}
 				{displayDomain && !isHovered && (
 					<div className="absolute bottom-2 right-2">
 						<Badge
 							variant="secondary"
-							className="bg-black/60 text-white border-0 text-xs backdrop-blur-sm"
+							className={cn(
+								"border-0 text-xs backdrop-blur-sm",
+								isGenerated
+									? "bg-primary/80 text-primary-foreground"
+									: "bg-black/60 text-white"
+							)}
 						>
+							{isGenerated && <SparklesIcon className="size-3 mr-1" />}
 							{displayDomain}
 						</Badge>
 					</div>
diff --git a/surfsense_web/contracts/enums/image-gen-providers.ts b/surfsense_web/contracts/enums/image-gen-providers.ts
new file mode 100644
index 000000000..8410aeb4b
--- /dev/null
+++ b/surfsense_web/contracts/enums/image-gen-providers.ts
@@ -0,0 +1,105 @@
+export interface ImageGenProvider {
+	value: string;
+	label: string;
+	example: string;
+	description: string;
+	apiBase?: string;
+}
+
+/**
+ * Image generation providers supported by LiteLLM.
+ * See: https://docs.litellm.ai/docs/image_generation#supported-providers
+ */
+export const IMAGE_GEN_PROVIDERS: ImageGenProvider[] = [
+	{
+		value: "OPENAI",
+		label: "OpenAI",
+		example: "dall-e-3, gpt-image-1, dall-e-2",
+		description: "DALL-E and GPT Image models",
+	},
+	{
+		value: "AZURE_OPENAI",
+		label: "Azure OpenAI",
+		example: "azure/dall-e-3, azure/gpt-image-1",
+		description: "OpenAI image models on Azure",
+	},
+	{
+		value: "GOOGLE",
+		label: "Google AI Studio",
+		example: "gemini/imagen-3.0-generate-002",
+		description: "Google AI Studio image generation",
+	},
+	{
+		value: "VERTEX_AI",
+		label: "Google Vertex AI",
+		example: "vertex_ai/imagegeneration@006",
+		description: "Vertex AI image generation models",
+	},
+	{
+		value: "BEDROCK",
+		label: "AWS Bedrock",
+		example: "bedrock/stability.stable-diffusion-xl-v0",
+		description: "Stable Diffusion on AWS Bedrock",
+	},
+	{
+		value: "RECRAFT",
+		label: "Recraft",
+		example: "recraft/recraftv3",
+		description: "AI-powered design and image generation",
+	},
+	{
+		value: "OPENROUTER",
+		label: "OpenRouter",
+		example: "openrouter/google/gemini-2.5-flash-image",
+		description: "Image generation via OpenRouter",
+	},
+	{
+		value: "XINFERENCE",
+		label: "Xinference",
+		example: "xinference/stable-diffusion-xl",
+		description: "Self-hosted Stable Diffusion models",
+	},
+	{
+		value: "NSCALE",
+		label: "Nscale",
+		example: "nscale/flux.1-schnell",
+		description: "Nscale image generation",
+	},
+];
+
+/**
+ * Image generation models organized by provider.
+ */
+export interface ImageGenModel {
+	value: string;
+	label: string;
+	provider: string;
+}
+
+export const IMAGE_GEN_MODELS: ImageGenModel[] = [
+	// OpenAI
+	{ value: "gpt-image-1", label: "GPT Image 1", provider: "OPENAI" },
+	{ value: "dall-e-3", label: "DALL-E 3", provider: "OPENAI" },
+	{ value: "dall-e-2", label: "DALL-E 2", provider: "OPENAI" },
+	// Azure OpenAI
+	{ value: "azure/dall-e-3", label: "DALL-E 3 (Azure)", provider: "AZURE_OPENAI" },
+	{ value: "azure/gpt-image-1", label: "GPT Image 1 (Azure)", provider: "AZURE_OPENAI" },
+	// Recraft
+	{ value: "recraft/recraftv3", label: "Recraft V3", provider: "RECRAFT" },
+	// Bedrock
+	{
+		value: "bedrock/stability.stable-diffusion-xl-v0",
+		label: "Stable Diffusion XL",
+		provider: "BEDROCK",
+	},
+	// Vertex AI
+	{
+		value: "vertex_ai/imagegeneration@006",
+		label: "Imagen 3",
+		provider: "VERTEX_AI",
+	},
+];
+
+export function getImageGenModelsByProvider(provider: string): ImageGenModel[] {
+	return IMAGE_GEN_MODELS.filter((m) => m.provider === provider);
+}
diff --git a/surfsense_web/contracts/types/new-llm-config.types.ts b/surfsense_web/contracts/types/new-llm-config.types.ts
index f397d4f08..3f0d39e5a 100644
--- a/surfsense_web/contracts/types/new-llm-config.types.ts
+++ b/surfsense_web/contracts/types/new-llm-config.types.ts
@@ -161,19 +161,105 @@ export const globalNewLLMConfig = z.object({
 
 export const getGlobalNewLLMConfigsResponse = z.array(globalNewLLMConfig);
 
+// =============================================================================
+// Image Generation Config (separate table from NewLLMConfig)
+// =============================================================================
+
+/**
+ * ImageGenProvider enum - only providers that support image generation
+ * See: https://docs.litellm.ai/docs/image_generation#supported-providers
+ */
+export const imageGenProviderEnum = z.enum([
+	"OPENAI",
+	"AZURE_OPENAI",
+	"GOOGLE",
+	"VERTEX_AI",
+	"BEDROCK",
+	"RECRAFT",
+	"OPENROUTER",
+	"XINFERENCE",
+	"NSCALE",
+]);
+
+export type ImageGenProvider = z.infer<typeof imageGenProviderEnum>;
+
+/**
+ * ImageGenerationConfig - user-created image gen model configs
+ * Separate from NewLLMConfig: no system_instructions, no citations_enabled.
+ */
+export const imageGenerationConfig = z.object({
+	id: z.number(),
+	name: z.string().max(100),
+	description: z.string().max(500).nullable().optional(),
+	provider: imageGenProviderEnum,
+	custom_provider: z.string().max(100).nullable().optional(),
+	model_name: z.string().max(100),
+	api_key: z.string(),
+	api_base: z.string().max(500).nullable().optional(),
+	api_version: z.string().max(50).nullable().optional(),
+	litellm_params: z.record(z.string(), z.any()).nullable().optional(),
+	created_at: z.string(),
+	search_space_id: z.number(),
+});
+
+export const createImageGenConfigRequest = imageGenerationConfig.omit({
+	id: true,
+	created_at: true,
+});
+
+export const createImageGenConfigResponse = imageGenerationConfig;
+
+export const getImageGenConfigsResponse = z.array(imageGenerationConfig);
+
+export const updateImageGenConfigRequest = z.object({
+	id: z.number(),
+	data: imageGenerationConfig
+		.omit({ id: true, created_at: true, search_space_id: true })
+		.partial(),
+});
+
+export const updateImageGenConfigResponse = imageGenerationConfig;
+
+export const deleteImageGenConfigResponse = z.object({
+	message: z.string(),
+	id: z.number(),
+});
+
+/**
+ * Global Image Generation Config - from YAML, has negative IDs
+ * ID 0 is reserved for "Auto" mode (LiteLLM Router load balancing)
+ */
+export const globalImageGenConfig = z.object({
+	id: z.number(),
+	name: z.string(),
+	description: z.string().nullable().optional(),
+	provider: z.string(),
+	custom_provider: z.string().nullable().optional(),
+	model_name: z.string(),
+	api_base: z.string().nullable().optional(),
+	api_version: z.string().nullable().optional(),
+	litellm_params: z.record(z.string(), z.any()).nullable().optional(),
+	is_global: z.literal(true),
+	is_auto_mode: z.boolean().optional().default(false),
+});
+
+export const getGlobalImageGenConfigsResponse = z.array(globalImageGenConfig);
+
 // =============================================================================
 // LLM Preferences (Role Assignments)
 // =============================================================================
 
 /**
  * LLM Preferences schemas - for role assignments
- * The agent_llm and document_summary_llm fields contain the full NewLLMConfig objects
+ * image_generation uses image_generation_config_id (not llm_id)
  */
 export const llmPreferences = z.object({
 	agent_llm_id: z.union([z.number(), z.null()]).optional(),
 	document_summary_llm_id: z.union([z.number(), z.null()]).optional(),
+	image_generation_config_id: z.union([z.number(), z.null()]).optional(),
 	agent_llm: z.union([z.record(z.string(), z.unknown()), z.null()]).optional(),
 	document_summary_llm: z.union([z.record(z.string(), z.unknown()), z.null()]).optional(),
+	image_generation_config: z.union([z.record(z.string(), z.unknown()), z.null()]).optional(),
 });
 
 /**
@@ -193,6 +279,7 @@ export const updateLLMPreferencesRequest = z.object({
 	data: llmPreferences.pick({
 		agent_llm_id: true,
 		document_summary_llm_id: true,
+		image_generation_config_id: true,
 	}),
 });
 
@@ -219,6 +306,15 @@ export type GetDefaultSystemInstructionsResponse = z.infer<
 >;
 export type GlobalNewLLMConfig = z.infer<typeof globalNewLLMConfig>;
 export type GetGlobalNewLLMConfigsResponse = z.infer<typeof getGlobalNewLLMConfigsResponse>;
+export type ImageGenerationConfig = z.infer<typeof imageGenerationConfig>;
+export type CreateImageGenConfigRequest = z.infer<typeof createImageGenConfigRequest>;
+export type CreateImageGenConfigResponse = z.infer<typeof createImageGenConfigResponse>;
+export type GetImageGenConfigsResponse = z.infer<typeof getImageGenConfigsResponse>;
+export type UpdateImageGenConfigRequest = z.infer<typeof updateImageGenConfigRequest>;
+export type UpdateImageGenConfigResponse = z.infer<typeof updateImageGenConfigResponse>;
+export type DeleteImageGenConfigResponse = z.infer<typeof deleteImageGenConfigResponse>;
+export type GlobalImageGenConfig = z.infer<typeof globalImageGenConfig>;
+export type GetGlobalImageGenConfigsResponse = z.infer<typeof getGlobalImageGenConfigsResponse>;
 export type LLMPreferences = z.infer<typeof llmPreferences>;
 export type GetLLMPreferencesRequest = z.infer<typeof getLLMPreferencesRequest>;
 export type GetLLMPreferencesResponse = z.infer<typeof getLLMPreferencesResponse>;
diff --git a/surfsense_web/lib/apis/image-gen-config-api.service.ts b/surfsense_web/lib/apis/image-gen-config-api.service.ts
new file mode 100644
index 000000000..84aeed3d8
--- /dev/null
+++ b/surfsense_web/lib/apis/image-gen-config-api.service.ts
@@ -0,0 +1,83 @@
+import {
+	type CreateImageGenConfigRequest,
+	createImageGenConfigRequest,
+	createImageGenConfigResponse,
+	type UpdateImageGenConfigRequest,
+	updateImageGenConfigRequest,
+	updateImageGenConfigResponse,
+	deleteImageGenConfigResponse,
+	getImageGenConfigsResponse,
+	getGlobalImageGenConfigsResponse,
+} from "@/contracts/types/new-llm-config.types";
+import { ValidationError } from "../error";
+import { baseApiService } from "./base-api.service";
+
+class ImageGenConfigApiService {
+	/**
+	 * Get all global image generation configs (from YAML, negative IDs)
+	 */
+	getGlobalConfigs = async () => {
+		return baseApiService.get(
+			`/api/v1/global-image-generation-configs`,
+			getGlobalImageGenConfigsResponse
+		);
+	};
+
+	/**
+	 * Create a new image generation config for a search space
+	 */
+	createConfig = async (request: CreateImageGenConfigRequest) => {
+		const parsed = createImageGenConfigRequest.safeParse(request);
+		if (!parsed.success) {
+			const msg = parsed.error.issues.map((i) => i.message).join(", ");
+			throw new ValidationError(`Invalid request: ${msg}`);
+		}
+		return baseApiService.post(
+			`/api/v1/image-generation-configs`,
+			createImageGenConfigResponse,
+			{ body: parsed.data }
+		);
+	};
+
+	/**
+	 * Get image generation configs for a search space
+	 */
+	getConfigs = async (searchSpaceId: number) => {
+		const params = new URLSearchParams({
+			search_space_id: String(searchSpaceId),
+		}).toString();
+		return baseApiService.get(
+			`/api/v1/image-generation-configs?${params}`,
+			getImageGenConfigsResponse
+		);
+	};
+
+	/**
+	 * Update an existing image generation config
+	 */
+	updateConfig = async (request: UpdateImageGenConfigRequest) => {
+		const parsed = updateImageGenConfigRequest.safeParse(request);
+		if (!parsed.success) {
+			const msg = parsed.error.issues.map((i) => i.message).join(", ");
+			throw new ValidationError(`Invalid request: ${msg}`);
+		}
+		const { id, data } = parsed.data;
+		return baseApiService.put(
+			`/api/v1/image-generation-configs/${id}`,
+			updateImageGenConfigResponse,
+			{ body: data }
+		);
+	};
+
+	/**
+	 * Delete an image generation config
+	 */
+	deleteConfig = async (id: number) => {
+		return baseApiService.delete(
+			`/api/v1/image-generation-configs/${id}`,
+			deleteImageGenConfigResponse
+		);
+	};
+}
+
+export const imageGenConfigApiService = new ImageGenConfigApiService();
diff --git a/surfsense_web/lib/query-client/cache-keys.ts b/surfsense_web/lib/query-client/cache-keys.ts
index 4d220a62a..c6981b28a 100644
--- a/surfsense_web/lib/query-client/cache-keys.ts
+++ b/surfsense_web/lib/query-client/cache-keys.ts
@@ -34,6 +34,11 @@ export const cacheKeys = {
 		defaultInstructions: () => ["new-llm-configs", "default-instructions"] as const,
 		global: () => ["new-llm-configs", "global"] as const,
 	},
+	imageGenConfigs: {
+		all: (searchSpaceId: number) => ["image-gen-configs", searchSpaceId] as const,
+		byId: (configId: number) => ["image-gen-configs", "detail", configId] as const,
+		global: () => ["image-gen-configs", "global"] as const,
+	},
 	auth: {
 		user: ["auth", "user"] as const,
 	},
diff --git a/surfsense_web/messages/en.json b/surfsense_web/messages/en.json
index a1ef1f248..5a18f80c3 100644
--- a/surfsense_web/messages/en.json
+++ b/surfsense_web/messages/en.json
@@ -740,6 +740,8 @@
 		"nav_agent_configs_desc": "LLM models with prompts & citations",
 		"nav_role_assignments": "Role Assignments",
 		"nav_role_assignments_desc": "Assign configs to agent roles",
+		"nav_image_models": "Image Models",
+		"nav_image_models_desc": "Configure image generation models",
 		"nav_system_instructions": "System Instructions",
 		"nav_system_instructions_desc": "SearchSpace-wide AI instructions",
 		"nav_public_links": "Public Chat Links",
diff --git a/surfsense_web/messages/zh.json b/surfsense_web/messages/zh.json
index 60a0d279f..1046b7296 100644
--- a/surfsense_web/messages/zh.json
+++ b/surfsense_web/messages/zh.json
@@ -725,6 +725,8 @@
 		"nav_agent_configs_desc": "LLM 模型配置提示词和引用",
 		"nav_role_assignments": "角色分配",
 		"nav_role_assignments_desc": "为代理角色分配配置",
+		"nav_image_models": "图像模型",
+		"nav_image_models_desc": "配置图像生成模型",
 		"nav_system_instructions": "系统指令",
 		"nav_system_instructions_desc": "搜索空间级别的 AI 指令",
 		"nav_public_links": "公开聊天链接",

From f85adefe5effc91680526cf0ce9ecb2e50a8385f Mon Sep 17 00:00:00 2001
From: "DESKTOP-RTLN3BA\\$punk" <vermarohanfinal@gmail.com>
Date: Thu, 5 Feb 2026 17:18:27 -0800
Subject: [PATCH 36/40] chore: made generate_image more agnostic

---
 .../93_add_image_generations_table.py         |  11 +-
 .../app/agents/new_chat/system_prompt.py      |  10 +-
 .../agents/new_chat/tools/generate_image.py   |  24 +--
 .../app/connectors/airtable_history.py        |   4 +-
 .../app/connectors/confluence_history.py      |   4 +-
 .../app/connectors/jira_history.py            |   4 +-
 .../app/connectors/linear_connector.py        |   4 +-
 surfsense_backend/app/db.py                   |   9 +-
 surfsense_backend/app/routes/__init__.py      |   2 +-
 .../app/routes/image_generation_routes.py     | 176 ++++++++++++------
 surfsense_backend/app/schemas/__init__.py     |   6 +-
 .../app/schemas/image_generation.py           |   1 -
 .../app/services/image_gen_router_service.py  |  12 +-
 .../app/tasks/chat/stream_new_chat.py         |  17 +-
 surfsense_backend/app/users.py                |   4 +-
 15 files changed, 176 insertions(+), 112 deletions(-)

diff --git a/surfsense_backend/alembic/versions/93_add_image_generations_table.py b/surfsense_backend/alembic/versions/93_add_image_generations_table.py
index 151208229..f24adc68f 100644
--- a/surfsense_backend/alembic/versions/93_add_image_generations_table.py
+++ b/surfsense_backend/alembic/versions/93_add_image_generations_table.py
@@ -60,8 +60,15 @@ def upgrade() -> None:
             sa.Column(
                 "provider",
                 sa.Enum(
-                    "OPENAI", "AZURE_OPENAI", "GOOGLE", "VERTEX_AI", "BEDROCK",
-                    "RECRAFT", "OPENROUTER", "XINFERENCE", "NSCALE",
+                    "OPENAI",
+                    "AZURE_OPENAI",
+                    "GOOGLE",
+                    "VERTEX_AI",
+                    "BEDROCK",
+                    "RECRAFT",
+                    "OPENROUTER",
+                    "XINFERENCE",
+                    "NSCALE",
                     name="imagegenprovider",
                     create_type=False,
                 ),
diff --git a/surfsense_backend/app/agents/new_chat/system_prompt.py b/surfsense_backend/app/agents/new_chat/system_prompt.py
index 38bae230d..01c762197 100644
--- a/surfsense_backend/app/agents/new_chat/system_prompt.py
+++ b/surfsense_backend/app/agents/new_chat/system_prompt.py
@@ -106,8 +106,6 @@ You have access to the following tools:
   - Trigger phrases: "generate an image of", "create a picture of", "draw me", "make an image", "design a logo", "create artwork"
   - Args:
     - prompt: A detailed text description of the image to generate. Be specific about subject, style, colors, composition, and mood.
-    - size: Image size. Options: "1024x1024" (square, default), "1536x1024" (landscape), "1024x1536" (portrait), "1792x1024" (wide)
-    - quality: Image quality. Options: "auto" (default), "high", "medium", "low"
     - n: Number of images to generate (1-4, default: 1)
   - Returns: A dictionary with the generated image URL in the "src" field, along with metadata.
   - CRITICAL: After calling generate_image, you MUST call `display_image` with the returned "src" URL
@@ -300,19 +298,19 @@ You have access to the following tools:
   - Then provide your explanation, referencing the displayed image
 
 - User: "Generate an image of a cat"
-  - Step 1: `generate_image(prompt="A fluffy orange tabby cat sitting on a windowsill, bathed in warm golden sunlight, soft bokeh background with green houseplants, photorealistic style, cozy atmosphere", size="1024x1024", quality="auto")`
+  - Step 1: `generate_image(prompt="A fluffy orange tabby cat sitting on a windowsill, bathed in warm golden sunlight, soft bokeh background with green houseplants, photorealistic style, cozy atmosphere")`
   - Step 2: Use the returned "src" URL to display it: `display_image(src="<returned_url>", alt="A fluffy orange tabby cat on a windowsill", title="Generated Image")`
 
 - User: "Create a landscape painting of mountains"
-  - Step 1: `generate_image(prompt="Majestic snow-capped mountain range at sunset, dramatic orange and purple sky, alpine meadow with wildflowers in the foreground, oil painting style with visible brushstrokes, inspired by the Hudson River School art movement", size="1536x1024", quality="high")`
+  - Step 1: `generate_image(prompt="Majestic snow-capped mountain range at sunset, dramatic orange and purple sky, alpine meadow with wildflowers in the foreground, oil painting style with visible brushstrokes, inspired by the Hudson River School art movement")`
   - Step 2: `display_image(src="<returned_url>", alt="Mountain landscape painting", title="Generated Image")`
 
 - User: "Draw me a logo for a coffee shop called Bean Dream"
-  - Step 1: `generate_image(prompt="Minimalist modern logo design for a coffee shop called 'Bean Dream', featuring a stylized coffee bean with dream-like swirls of steam, clean vector style, warm brown and cream color palette, white background, professional branding", size="1024x1024", quality="high")`
+  - Step 1: `generate_image(prompt="Minimalist modern logo design for a coffee shop called 'Bean Dream', featuring a stylized coffee bean with dream-like swirls of steam, clean vector style, warm brown and cream color palette, white background, professional branding")`
   - Step 2: `display_image(src="<returned_url>", alt="Bean Dream coffee shop logo", title="Generated Image")`
 
 - User: "Make a wide banner image for my blog about AI"
-  - Step 1: `generate_image(prompt="Wide banner illustration for an AI technology blog, featuring abstract neural network patterns, glowing blue and purple connections, modern futuristic aesthetic, digital art style, clean and professional", size="1792x1024", quality="high")`
+  - Step 1: `generate_image(prompt="Wide banner illustration for an AI technology blog, featuring abstract neural network patterns, glowing blue and purple connections, modern futuristic aesthetic, digital art style, clean and professional")`
   - Step 2: `display_image(src="<returned_url>", alt="AI blog banner", title="Generated Image")`
 </tool_call_examples>
 """
diff --git a/surfsense_backend/app/agents/new_chat/tools/generate_image.py b/surfsense_backend/app/agents/new_chat/tools/generate_image.py
index 091fb122f..8ffa4ecde 100644
--- a/surfsense_backend/app/agents/new_chat/tools/generate_image.py
+++ b/surfsense_backend/app/agents/new_chat/tools/generate_image.py
@@ -21,12 +21,12 @@ from sqlalchemy.ext.asyncio import AsyncSession
 
 from app.config import config
 from app.db import ImageGeneration, ImageGenerationConfig, SearchSpace
-from app.utils.signed_image_urls import generate_image_token
 from app.services.image_gen_router_service import (
     IMAGE_GEN_AUTO_MODE_ID,
     ImageGenRouterService,
     is_image_gen_auto_mode,
 )
+from app.utils.signed_image_urls import generate_image_token
 
 logger = logging.getLogger(__name__)
 
@@ -76,8 +76,6 @@ def create_generate_image_tool(
     @tool
     async def generate_image(
         prompt: str,
-        size: str = "1024x1024",
-        quality: str = "auto",
         n: int = 1,
     ) -> dict[str, Any]:
         """
@@ -89,10 +87,6 @@ def create_generate_image_tool(
         Args:
             prompt: A detailed text description of the image to generate.
                     Be specific about subject, style, colors, composition, and mood.
-            size: Image size. Options: "1024x1024" (square), "1536x1024" (landscape),
-                  "1024x1536" (portrait), "1792x1024" (wide). Default: "1024x1024"
-            quality: Image quality. Options: "auto" (default), "high", "medium", "low".
-                     Default: "auto"
             n: Number of images to generate (1-4). Default: 1
 
         Returns:
@@ -112,18 +106,14 @@ def create_generate_image_tool(
             )
 
             # Build generation kwargs
-            # NOTE: 'style' is intentionally excluded from gen_kwargs because
-            # it is only supported by DALL-E 3 and causes errors with other
-            # models (e.g. gpt-image-1 rejects it as an unknown parameter).
-            # Since we can't predict which model auto-mode will route to,
-            # it's safest to omit it.
+            # NOTE: size, quality, and style are intentionally NOT passed.
+            # Different models support different values for these params
+            # (e.g. DALL-E 3 wants "hd"/"standard" for quality while
+            # gpt-image-1 wants "high"/"medium"/"low"; size options also
+            # differ). Letting the model use its own defaults avoids errors.
             gen_kwargs: dict[str, Any] = {}
             if n is not None and n > 1:
                 gen_kwargs["n"] = n
-            if quality:
-                gen_kwargs["quality"] = quality
-            if size:
-                gen_kwargs["size"] = size
 
             # Call litellm based on config type
             if is_image_gen_auto_mode(config_id):
@@ -199,8 +189,6 @@ def create_generate_image_tool(
                 prompt=prompt,
                 model=getattr(response, "_hidden_params", {}).get("model"),
                 n=n,
-                quality=quality,
-                size=size,
                 image_generation_config_id=config_id,
                 response_data=response_dict,
                 search_space_id=search_space_id,
diff --git a/surfsense_backend/app/connectors/airtable_history.py b/surfsense_backend/app/connectors/airtable_history.py
index 092485f77..49c2fcbdd 100644
--- a/surfsense_backend/app/connectors/airtable_history.py
+++ b/surfsense_backend/app/connectors/airtable_history.py
@@ -108,7 +108,9 @@ class AirtableHistoryConnector:
 
             # Final validation after decryption
             final_token = config_data.get("access_token")
-            if not final_token or (isinstance(final_token, str) and not final_token.strip()):
+            if not final_token or (
+                isinstance(final_token, str) and not final_token.strip()
+            ):
                 raise ValueError(
                     "Airtable access token is invalid or empty. "
                     "Please reconnect your Airtable account."
diff --git a/surfsense_backend/app/connectors/confluence_history.py b/surfsense_backend/app/connectors/confluence_history.py
index 908f532db..5d19edc54 100644
--- a/surfsense_backend/app/connectors/confluence_history.py
+++ b/surfsense_backend/app/connectors/confluence_history.py
@@ -128,7 +128,9 @@ class ConfluenceHistoryConnector:
 
                 # Final validation after decryption
                 final_token = config_data.get("access_token")
-                if not final_token or (isinstance(final_token, str) and not final_token.strip()):
+                if not final_token or (
+                    isinstance(final_token, str) and not final_token.strip()
+                ):
                     raise ValueError(
                         "Confluence access token is invalid or empty. "
                         "Please reconnect your Confluence account."
diff --git a/surfsense_backend/app/connectors/jira_history.py b/surfsense_backend/app/connectors/jira_history.py
index 46a28324d..e9f28a2c4 100644
--- a/surfsense_backend/app/connectors/jira_history.py
+++ b/surfsense_backend/app/connectors/jira_history.py
@@ -129,7 +129,9 @@ class JiraHistoryConnector:
 
                 # Final validation after decryption
                 final_token = config_data.get("access_token")
-                if not final_token or (isinstance(final_token, str) and not final_token.strip()):
+                if not final_token or (
+                    isinstance(final_token, str) and not final_token.strip()
+                ):
                     raise ValueError(
                         "Jira access token is invalid or empty. "
                         "Please reconnect your Jira account."
diff --git a/surfsense_backend/app/connectors/linear_connector.py b/surfsense_backend/app/connectors/linear_connector.py
index 6500b9027..534d70b89 100644
--- a/surfsense_backend/app/connectors/linear_connector.py
+++ b/surfsense_backend/app/connectors/linear_connector.py
@@ -153,7 +153,9 @@ class LinearConnector:
 
             # Final validation after decryption
             final_token = config_data.get("access_token")
-            if not final_token or (isinstance(final_token, str) and not final_token.strip()):
+            if not final_token or (
+                isinstance(final_token, str) and not final_token.strip()
+            ):
                 raise ValueError(
                     "Linear access token is invalid or empty. "
                     "Please reconnect your Linear account."
diff --git a/surfsense_backend/app/db.py b/surfsense_backend/app/db.py
index 320ff6d8d..a82c18470 100644
--- a/surfsense_backend/app/db.py
+++ b/surfsense_backend/app/db.py
@@ -100,7 +100,6 @@ class PodcastStatus(str, Enum):
     FAILED = "failed"
 
 
-
 class LiteLLMProvider(str, Enum):
     """
     Enum for LLM providers supported by LiteLLM.
@@ -941,7 +940,9 @@ class ImageGenerationConfig(BaseModel, TimestampMixin):
     search_space_id = Column(
         Integer, ForeignKey("searchspaces.id", ondelete="CASCADE"), nullable=False
     )
-    search_space = relationship("SearchSpace", back_populates="image_generation_configs")
+    search_space = relationship(
+        "SearchSpace", back_populates="image_generation_configs"
+    )
 
 
 class ImageGeneration(BaseModel, TimestampMixin):
@@ -973,9 +974,7 @@ class ImageGeneration(BaseModel, TimestampMixin):
         String(50), nullable=True
     )  # "1024x1024", "1536x1024", "1024x1536", etc.
     style = Column(String(50), nullable=True)  # Model-specific style parameter
-    response_format = Column(
-        String(50), nullable=True
-    )  # "url" or "b64_json"
+    response_format = Column(String(50), nullable=True)  # "url" or "b64_json"
 
     # Image generation config reference
     # 0 = Auto mode (router), negative IDs = global configs from YAML,
diff --git a/surfsense_backend/app/routes/__init__.py b/surfsense_backend/app/routes/__init__.py
index 683f3548b..d9353284c 100644
--- a/surfsense_backend/app/routes/__init__.py
+++ b/surfsense_backend/app/routes/__init__.py
@@ -11,7 +11,6 @@ from .confluence_add_connector_route import router as confluence_add_connector_r
 from .discord_add_connector_route import router as discord_add_connector_router
 from .documents_routes import router as documents_router
 from .editor_routes import router as editor_router
-from .image_generation_routes import router as image_generation_router
 from .google_calendar_add_connector_route import (
     router as google_calendar_add_connector_router,
 )
@@ -21,6 +20,7 @@ from .google_drive_add_connector_route import (
 from .google_gmail_add_connector_route import (
     router as google_gmail_add_connector_router,
 )
+from .image_generation_routes import router as image_generation_router
 from .incentive_tasks_routes import router as incentive_tasks_router
 from .jira_add_connector_route import router as jira_add_connector_router
 from .linear_add_connector_route import router as linear_add_connector_router
diff --git a/surfsense_backend/app/routes/image_generation_routes.py b/surfsense_backend/app/routes/image_generation_routes.py
index 9b79771eb..9406867c6 100644
--- a/surfsense_backend/app/routes/image_generation_routes.py
+++ b/surfsense_backend/app/routes/image_generation_routes.py
@@ -54,9 +54,9 @@ logger = logging.getLogger(__name__)
 _PROVIDER_MAP = {
     "OPENAI": "openai",
     "AZURE_OPENAI": "azure",
-    "GOOGLE": "gemini",       # Google AI Studio
+    "GOOGLE": "gemini",  # Google AI Studio
     "VERTEX_AI": "vertex_ai",
-    "BEDROCK": "bedrock",     # AWS Bedrock
+    "BEDROCK": "bedrock",  # AWS Bedrock
     "RECRAFT": "recraft",
     "OPENROUTER": "openrouter",
     "XINFERENCE": "xinference",
@@ -82,7 +82,9 @@ def _get_global_image_gen_config(config_id: int) -> dict | None:
     return None
 
 
-def _build_model_string(provider: str, model_name: str, custom_provider: str | None) -> str:
+def _build_model_string(
+    provider: str, model_name: str, custom_provider: str | None
+) -> str:
     """Build a litellm model string from provider + model_name."""
     if custom_provider:
         return f"{custom_provider}/{model_name}"
@@ -210,38 +212,44 @@ async def get_global_image_gen_configs(
         safe_configs = []
 
         if global_configs and len(global_configs) > 0:
-            safe_configs.append({
-                "id": 0,
-                "name": "Auto (Load Balanced)",
-                "description": "Automatically routes across available image generation providers.",
-                "provider": "AUTO",
-                "custom_provider": None,
-                "model_name": "auto",
-                "api_base": None,
-                "api_version": None,
-                "litellm_params": {},
-                "is_global": True,
-                "is_auto_mode": True,
-            })
+            safe_configs.append(
+                {
+                    "id": 0,
+                    "name": "Auto (Load Balanced)",
+                    "description": "Automatically routes across available image generation providers.",
+                    "provider": "AUTO",
+                    "custom_provider": None,
+                    "model_name": "auto",
+                    "api_base": None,
+                    "api_version": None,
+                    "litellm_params": {},
+                    "is_global": True,
+                    "is_auto_mode": True,
+                }
+            )
 
         for cfg in global_configs:
-            safe_configs.append({
-                "id": cfg.get("id"),
-                "name": cfg.get("name"),
-                "description": cfg.get("description"),
-                "provider": cfg.get("provider"),
-                "custom_provider": cfg.get("custom_provider"),
-                "model_name": cfg.get("model_name"),
-                "api_base": cfg.get("api_base") or None,
-                "api_version": cfg.get("api_version") or None,
-                "litellm_params": cfg.get("litellm_params", {}),
-                "is_global": True,
-            })
+            safe_configs.append(
+                {
+                    "id": cfg.get("id"),
+                    "name": cfg.get("name"),
+                    "description": cfg.get("description"),
+                    "provider": cfg.get("provider"),
+                    "custom_provider": cfg.get("custom_provider"),
+                    "model_name": cfg.get("model_name"),
+                    "api_base": cfg.get("api_base") or None,
+                    "api_version": cfg.get("api_version") or None,
+                    "litellm_params": cfg.get("litellm_params", {}),
+                    "is_global": True,
+                }
+            )
 
         return safe_configs
     except Exception as e:
         logger.exception("Failed to fetch global image generation configs")
-        raise HTTPException(status_code=500, detail=f"Failed to fetch configs: {e!s}") from e
+        raise HTTPException(
+            status_code=500, detail=f"Failed to fetch configs: {e!s}"
+        ) from e
 
 
 # =============================================================================
@@ -258,7 +266,9 @@ async def create_image_gen_config(
     """Create a new image generation config for a search space."""
     try:
         await check_permission(
-            session, user, config_data.search_space_id,
+            session,
+            user,
+            config_data.search_space_id,
             Permission.IMAGE_GENERATIONS_CREATE.value,
             "You don't have permission to create image generation configs in this search space",
         )
@@ -274,7 +284,9 @@ async def create_image_gen_config(
     except Exception as e:
         await session.rollback()
         logger.exception("Failed to create ImageGenerationConfig")
-        raise HTTPException(status_code=500, detail=f"Failed to create config: {e!s}") from e
+        raise HTTPException(
+            status_code=500, detail=f"Failed to create config: {e!s}"
+        ) from e
 
 
 @router.get("/image-generation-configs", response_model=list[ImageGenerationConfigRead])
@@ -288,7 +300,9 @@ async def list_image_gen_configs(
     """List image generation configs for a search space."""
     try:
         await check_permission(
-            session, user, search_space_id,
+            session,
+            user,
+            search_space_id,
             Permission.IMAGE_GENERATIONS_READ.value,
             "You don't have permission to view image generation configs in this search space",
         )
@@ -306,10 +320,14 @@ async def list_image_gen_configs(
         raise
     except Exception as e:
         logger.exception("Failed to list ImageGenerationConfigs")
-        raise HTTPException(status_code=500, detail=f"Failed to fetch configs: {e!s}") from e
+        raise HTTPException(
+            status_code=500, detail=f"Failed to fetch configs: {e!s}"
+        ) from e
 
 
-@router.get("/image-generation-configs/{config_id}", response_model=ImageGenerationConfigRead)
+@router.get(
+    "/image-generation-configs/{config_id}", response_model=ImageGenerationConfigRead
+)
 async def get_image_gen_config(
     config_id: int,
     session: AsyncSession = Depends(get_async_session),
@@ -325,7 +343,9 @@ async def get_image_gen_config(
             raise HTTPException(status_code=404, detail="Config not found")
 
         await check_permission(
-            session, user, db_config.search_space_id,
+            session,
+            user,
+            db_config.search_space_id,
             Permission.IMAGE_GENERATIONS_READ.value,
             "You don't have permission to view image generation configs in this search space",
         )
@@ -335,10 +355,14 @@ async def get_image_gen_config(
         raise
     except Exception as e:
         logger.exception("Failed to get ImageGenerationConfig")
-        raise HTTPException(status_code=500, detail=f"Failed to fetch config: {e!s}") from e
+        raise HTTPException(
+            status_code=500, detail=f"Failed to fetch config: {e!s}"
+        ) from e
 
 
-@router.put("/image-generation-configs/{config_id}", response_model=ImageGenerationConfigRead)
+@router.put(
+    "/image-generation-configs/{config_id}", response_model=ImageGenerationConfigRead
+)
 async def update_image_gen_config(
     config_id: int,
     update_data: ImageGenerationConfigUpdate,
@@ -355,7 +379,9 @@ async def update_image_gen_config(
             raise HTTPException(status_code=404, detail="Config not found")
 
         await check_permission(
-            session, user, db_config.search_space_id,
+            session,
+            user,
+            db_config.search_space_id,
             Permission.IMAGE_GENERATIONS_CREATE.value,
             "You don't have permission to update image generation configs in this search space",
         )
@@ -372,7 +398,9 @@ async def update_image_gen_config(
     except Exception as e:
         await session.rollback()
         logger.exception("Failed to update ImageGenerationConfig")
-        raise HTTPException(status_code=500, detail=f"Failed to update config: {e!s}") from e
+        raise HTTPException(
+            status_code=500, detail=f"Failed to update config: {e!s}"
+        ) from e
 
 
 @router.delete("/image-generation-configs/{config_id}", response_model=dict)
@@ -391,21 +419,28 @@ async def delete_image_gen_config(
             raise HTTPException(status_code=404, detail="Config not found")
 
         await check_permission(
-            session, user, db_config.search_space_id,
+            session,
+            user,
+            db_config.search_space_id,
             Permission.IMAGE_GENERATIONS_DELETE.value,
             "You don't have permission to delete image generation configs in this search space",
         )
 
         await session.delete(db_config)
         await session.commit()
-        return {"message": "Image generation config deleted successfully", "id": config_id}
+        return {
+            "message": "Image generation config deleted successfully",
+            "id": config_id,
+        }
 
     except HTTPException:
         raise
     except Exception as e:
         await session.rollback()
         logger.exception("Failed to delete ImageGenerationConfig")
-        raise HTTPException(status_code=500, detail=f"Failed to delete config: {e!s}") from e
+        raise HTTPException(
+            status_code=500, detail=f"Failed to delete config: {e!s}"
+        ) from e
 
 
 # =============================================================================
@@ -422,7 +457,9 @@ async def create_image_generation(
     """Create and execute an image generation request."""
     try:
         await check_permission(
-            session, user, data.search_space_id,
+            session,
+            user,
+            data.search_space_id,
             Permission.IMAGE_GENERATIONS_CREATE.value,
             "You don't have permission to create image generations in this search space",
         )
@@ -463,11 +500,15 @@ async def create_image_generation(
         raise
     except SQLAlchemyError:
         await session.rollback()
-        raise HTTPException(status_code=500, detail="Database error during image generation") from None
+        raise HTTPException(
+            status_code=500, detail="Database error during image generation"
+        ) from None
     except Exception as e:
         await session.rollback()
         logger.exception("Failed to create image generation")
-        raise HTTPException(status_code=500, detail=f"Image generation failed: {e!s}") from e
+        raise HTTPException(
+            status_code=500, detail=f"Image generation failed: {e!s}"
+        ) from e
 
 
 @router.get("/image-generations", response_model=list[ImageGenerationListRead])
@@ -487,7 +528,9 @@ async def list_image_generations(
     try:
         if search_space_id is not None:
             await check_permission(
-                session, user, search_space_id,
+                session,
+                user,
+                search_space_id,
                 Permission.IMAGE_GENERATIONS_READ.value,
                 "You don't have permission to read image generations in this search space",
             )
@@ -495,7 +538,8 @@ async def list_image_generations(
                 select(ImageGeneration)
                 .filter(ImageGeneration.search_space_id == search_space_id)
                 .order_by(ImageGeneration.created_at.desc())
-                .offset(skip).limit(limit)
+                .offset(skip)
+                .limit(limit)
             )
         else:
             result = await session.execute(
@@ -504,15 +548,21 @@ async def list_image_generations(
                 .join(SearchSpaceMembership)
                 .filter(SearchSpaceMembership.user_id == user.id)
                 .order_by(ImageGeneration.created_at.desc())
-                .offset(skip).limit(limit)
+                .offset(skip)
+                .limit(limit)
             )
 
-        return [ImageGenerationListRead.from_orm_with_count(img) for img in result.scalars().all()]
+        return [
+            ImageGenerationListRead.from_orm_with_count(img)
+            for img in result.scalars().all()
+        ]
 
     except HTTPException:
         raise
     except SQLAlchemyError:
-        raise HTTPException(status_code=500, detail="Database error fetching image generations") from None
+        raise HTTPException(
+            status_code=500, detail="Database error fetching image generations"
+        ) from None
 
 
 @router.get("/image-generations/{image_gen_id}", response_model=ImageGenerationRead)
@@ -531,7 +581,9 @@ async def get_image_generation(
             raise HTTPException(status_code=404, detail="Image generation not found")
 
         await check_permission(
-            session, user, image_gen.search_space_id,
+            session,
+            user,
+            image_gen.search_space_id,
             Permission.IMAGE_GENERATIONS_READ.value,
             "You don't have permission to read image generations in this search space",
         )
@@ -540,7 +592,9 @@ async def get_image_generation(
     except HTTPException:
         raise
     except SQLAlchemyError:
-        raise HTTPException(status_code=500, detail="Database error fetching image generation") from None
+        raise HTTPException(
+            status_code=500, detail="Database error fetching image generation"
+        ) from None
 
 
 @router.delete("/image-generations/{image_gen_id}", response_model=dict)
@@ -559,7 +613,9 @@ async def delete_image_generation(
             raise HTTPException(status_code=404, detail="Image generation not found")
 
         await check_permission(
-            session, user, db_image_gen.search_space_id,
+            session,
+            user,
+            db_image_gen.search_space_id,
             Permission.IMAGE_GENERATIONS_DELETE.value,
             "You don't have permission to delete image generations in this search space",
         )
@@ -572,13 +628,16 @@ async def delete_image_generation(
         raise
     except SQLAlchemyError:
         await session.rollback()
-        raise HTTPException(status_code=500, detail="Database error deleting image generation") from None
+        raise HTTPException(
+            status_code=500, detail="Database error deleting image generation"
+        ) from None
 
 
 # =============================================================================
 # Image Serving (serves generated images from DB, protected by signed tokens)
 # =============================================================================
 
+
 @router.get("/image-generations/{image_gen_id}/image")
 async def serve_generated_image(
     image_gen_id: int,
@@ -616,13 +675,16 @@ async def serve_generated_image(
 
         images = image_gen.response_data.get("data", [])
         if not images or index >= len(images):
-            raise HTTPException(status_code=404, detail="Image not found at the specified index")
+            raise HTTPException(
+                status_code=404, detail="Image not found at the specified index"
+            )
 
         image_entry = images[index]
 
         # If there's a URL, redirect to it
         if image_entry.get("url"):
             from fastapi.responses import RedirectResponse
+
             return RedirectResponse(url=image_entry["url"])
 
         # If there's b64_json data, decode and serve it
@@ -643,4 +705,6 @@ async def serve_generated_image(
         raise
     except Exception as e:
         logger.exception("Failed to serve generated image")
-        raise HTTPException(status_code=500, detail=f"Failed to serve image: {e!s}") from e
+        raise HTTPException(
+            status_code=500, detail=f"Failed to serve image: {e!s}"
+        ) from e
diff --git a/surfsense_backend/app/schemas/__init__.py b/surfsense_backend/app/schemas/__init__.py
index 332af55fd..ad5abf777 100644
--- a/surfsense_backend/app/schemas/__init__.py
+++ b/surfsense_backend/app/schemas/__init__.py
@@ -113,10 +113,12 @@ __all__ = [
     "DriveItem",
     "ExtensionDocumentContent",
     "ExtensionDocumentMetadata",
+    "GlobalImageGenConfigRead",
     "GlobalNewLLMConfigRead",
     "GoogleDriveIndexRequest",
     "GoogleDriveIndexingOptions",
-    "GlobalImageGenConfigRead",
+    # Base schemas
+    "IDModel",
     # Image Generation Config schemas
     "ImageGenerationConfigCreate",
     "ImageGenerationConfigPublic",
@@ -126,8 +128,6 @@ __all__ = [
     "ImageGenerationCreate",
     "ImageGenerationListRead",
     "ImageGenerationRead",
-    # Base schemas
-    "IDModel",
     # RBAC schemas
     "InviteAcceptRequest",
     "InviteAcceptResponse",
diff --git a/surfsense_backend/app/schemas/image_generation.py b/surfsense_backend/app/schemas/image_generation.py
index 367a35a77..6ef4feff8 100644
--- a/surfsense_backend/app/schemas/image_generation.py
+++ b/surfsense_backend/app/schemas/image_generation.py
@@ -13,7 +13,6 @@ from pydantic import BaseModel, ConfigDict, Field
 
 from app.db import ImageGenProvider
 
-
 # =============================================================================
 # ImageGenerationConfig CRUD Schemas
 # =============================================================================
diff --git a/surfsense_backend/app/services/image_gen_router_service.py b/surfsense_backend/app/services/image_gen_router_service.py
index 3b8a15d2a..eb6936efd 100644
--- a/surfsense_backend/app/services/image_gen_router_service.py
+++ b/surfsense_backend/app/services/image_gen_router_service.py
@@ -31,9 +31,9 @@ IMAGE_GEN_AUTO_MODE_ID = 0
 IMAGE_GEN_PROVIDER_MAP = {
     "OPENAI": "openai",
     "AZURE_OPENAI": "azure",
-    "GOOGLE": "gemini",       # Google AI Studio
+    "GOOGLE": "gemini",  # Google AI Studio
     "VERTEX_AI": "vertex_ai",
-    "BEDROCK": "bedrock",     # AWS Bedrock
+    "BEDROCK": "bedrock",  # AWS Bedrock
     "RECRAFT": "recraft",
     "OPENROUTER": "openrouter",
     "XINFERENCE": "xinference",
@@ -156,9 +156,7 @@ class ImageGenRouterService:
                 model_string = f"{config['custom_provider']}/{config['model_name']}"
             else:
                 provider = config.get("provider", "").upper()
-                provider_prefix = IMAGE_GEN_PROVIDER_MAP.get(
-                    provider, provider.lower()
-                )
+                provider_prefix = IMAGE_GEN_PROVIDER_MAP.get(provider, provider.lower())
                 model_string = f"{provider_prefix}/{config['model_name']}"
 
             # Build litellm params
@@ -194,9 +192,7 @@ class ImageGenRouterService:
             return deployment
 
         except Exception as e:
-            logger.warning(
-                f"Failed to convert image gen config to deployment: {e}"
-            )
+            logger.warning(f"Failed to convert image gen config to deployment: {e}")
             return None
 
     @classmethod
diff --git a/surfsense_backend/app/tasks/chat/stream_new_chat.py b/surfsense_backend/app/tasks/chat/stream_new_chat.py
index a9751e5d1..685f77e39 100644
--- a/surfsense_backend/app/tasks/chat/stream_new_chat.py
+++ b/surfsense_backend/app/tasks/chat/stream_new_chat.py
@@ -27,12 +27,12 @@ from app.agents.new_chat.llm_config import (
     load_llm_config_from_yaml,
 )
 from app.db import Document, SurfsenseDocsDocument
+from app.prompts import TITLE_GENERATION_PROMPT_TEMPLATE
 from app.schemas.new_chat import ChatAttachment
 from app.services.chat_session_state_service import (
     clear_ai_responding,
     set_ai_responding,
 )
-from app.prompts import TITLE_GENERATION_PROMPT_TEMPLATE
 from app.services.connector_service import ConnectorService
 from app.services.new_streaming_service import VercelStreamingService
 from app.utils.content_utils import bootstrap_history_from_db
@@ -1211,9 +1211,10 @@ async def stream_new_chat(
 
         # Generate LLM title for new chats after first response
         # Check if this is the first assistant response by counting existing assistant messages
-        from app.db import NewChatMessage, NewChatThread
         from sqlalchemy import func
 
+        from app.db import NewChatMessage, NewChatThread
+
         assistant_count_result = await session.execute(
             select(func.count(NewChatMessage.id)).filter(
                 NewChatMessage.thread_id == chat_id,
@@ -1231,10 +1232,12 @@ async def stream_new_chat(
                 # Truncate inputs to avoid context length issues
                 truncated_query = user_query[:500]
                 truncated_response = accumulated_text[:1000]
-                title_result = await title_chain.ainvoke({
-                    "user_query": truncated_query,
-                    "assistant_response": truncated_response,
-                })
+                title_result = await title_chain.ainvoke(
+                    {
+                        "user_query": truncated_query,
+                        "assistant_response": truncated_response,
+                    }
+                )
 
                 # Extract and clean the title
                 if title_result and hasattr(title_result, "content"):
@@ -1242,7 +1245,7 @@ async def stream_new_chat(
                     # Validate the title (reasonable length)
                     if raw_title and len(raw_title) <= 100:
                         # Remove any quotes or extra formatting
-                        generated_title = raw_title.strip('"\'')
+                        generated_title = raw_title.strip("\"'")
             except Exception:
                 generated_title = None
 
diff --git a/surfsense_backend/app/users.py b/surfsense_backend/app/users.py
index 696cdf25e..ee07ba88f 100644
--- a/surfsense_backend/app/users.py
+++ b/surfsense_backend/app/users.py
@@ -219,7 +219,9 @@ class CustomBearerTransport(BearerTransport):
 
         # Decode JWT to get user_id for refresh token creation
         try:
-            payload = jwt.decode(token, SECRET, algorithms=["HS256"], options={"verify_aud": False})
+            payload = jwt.decode(
+                token, SECRET, algorithms=["HS256"], options={"verify_aud": False}
+            )
             user_id = uuid.UUID(payload.get("sub"))
             refresh_token = await create_refresh_token(user_id)
         except Exception as e:

From af16b6656c83b162caa2ddcfafe43a41db7f8dbc Mon Sep 17 00:00:00 2001
From: "DESKTOP-RTLN3BA\\$punk" <vermarohanfinal@gmail.com>
Date: Thu, 5 Feb 2026 17:32:43 -0800
Subject: [PATCH 37/40] feat: implement connection pooling for
 AsyncPostgresSaver in checkpointer

---
 .../app/agents/new_chat/checkpointer.py       | 84 +++++++++++++++----
 1 file changed, 67 insertions(+), 17 deletions(-)

diff --git a/surfsense_backend/app/agents/new_chat/checkpointer.py b/surfsense_backend/app/agents/new_chat/checkpointer.py
index 637b2926f..04ecfbdea 100644
--- a/surfsense_backend/app/agents/new_chat/checkpointer.py
+++ b/surfsense_backend/app/agents/new_chat/checkpointer.py
@@ -3,15 +3,25 @@ PostgreSQL-based checkpointer for LangGraph agents.
 
 This module provides a persistent checkpointer using AsyncPostgresSaver
 that stores conversation state in the PostgreSQL database.
+
+Uses a connection pool (psycopg_pool.AsyncConnectionPool) to handle
+connection lifecycle, health checks, and automatic reconnection,
+preventing 'the connection is closed' errors in long-running deployments.
 """
 
+import logging
+
 from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
+from psycopg.rows import dict_row
+from psycopg_pool import AsyncConnectionPool
 
 from app.config import config
 
+logger = logging.getLogger(__name__)
+
 # Global checkpointer instance (initialized lazily)
 _checkpointer: AsyncPostgresSaver | None = None
-_checkpointer_context = None  # Store the context manager for cleanup
+_connection_pool: AsyncConnectionPool | None = None
 _checkpointer_initialized: bool = False
 
 
@@ -38,26 +48,65 @@ def get_postgres_connection_string() -> str:
     return db_url
 
 
+async def _create_checkpointer() -> AsyncPostgresSaver:
+    """
+    Create a new AsyncPostgresSaver backed by a connection pool.
+
+    The connection pool automatically handles:
+    - Connection health checks before use
+    - Reconnection when connections die (idle timeout, DB restart, etc.)
+    - Connection lifecycle management (max_lifetime, max_idle)
+    """
+    global _connection_pool
+
+    conn_string = get_postgres_connection_string()
+
+    _connection_pool = AsyncConnectionPool(
+        conninfo=conn_string,
+        min_size=2,
+        max_size=10,
+        # Connections are recycled after 30 minutes to avoid stale connections
+        max_lifetime=1800,
+        # Idle connections are closed after 5 minutes
+        max_idle=300,
+        open=False,
+        # Connection kwargs required by AsyncPostgresSaver:
+        # - autocommit: required for .setup() to commit checkpoint tables
+        # - prepare_threshold: disable prepared statements for compatibility
+        # - row_factory: checkpointer accesses rows as dicts (row["column"])
+        kwargs={
+            "autocommit": True,
+            "prepare_threshold": 0,
+            "row_factory": dict_row,
+        },
+    )
+    await _connection_pool.open(wait=True)
+
+    checkpointer = AsyncPostgresSaver(conn=_connection_pool)
+    logger.info("[Checkpointer] Created AsyncPostgresSaver with connection pool")
+    return checkpointer
+
+
 async def get_checkpointer() -> AsyncPostgresSaver:
     """
     Get or create the global AsyncPostgresSaver instance.
 
     This function:
-    1. Creates the checkpointer if it doesn't exist
+    1. Creates the checkpointer with a connection pool if it doesn't exist
     2. Sets up the required database tables on first call
     3. Returns the cached instance on subsequent calls
 
+    The underlying connection pool handles reconnection automatically,
+    so a stale/closed connection will not cause OperationalError.
+
     Returns:
         AsyncPostgresSaver: The configured checkpointer instance
     """
-    global _checkpointer, _checkpointer_context, _checkpointer_initialized
+    global _checkpointer, _checkpointer_initialized
 
     if _checkpointer is None:
-        conn_string = get_postgres_connection_string()
-        # from_conn_string returns an async context manager
-        # We need to enter the context to get the actual checkpointer
-        _checkpointer_context = AsyncPostgresSaver.from_conn_string(conn_string)
-        _checkpointer = await _checkpointer_context.__aenter__()
+        _checkpointer = await _create_checkpointer()
+        _checkpointer_initialized = False
 
     # Setup tables on first call (idempotent)
     if not _checkpointer_initialized:
@@ -75,20 +124,21 @@ async def setup_checkpointer_tables() -> None:
     tables exist before any agent calls.
     """
     await get_checkpointer()
-    print("[Checkpointer] PostgreSQL checkpoint tables ready")
+    logger.info("[Checkpointer] PostgreSQL checkpoint tables ready")
 
 
 async def close_checkpointer() -> None:
     """
-    Close the checkpointer connection.
+    Close the checkpointer connection pool.
 
     This should be called during application shutdown.
     """
-    global _checkpointer, _checkpointer_context, _checkpointer_initialized
+    global _checkpointer, _connection_pool, _checkpointer_initialized
 
-    if _checkpointer_context is not None:
-        await _checkpointer_context.__aexit__(None, None, None)
-        _checkpointer = None
-        _checkpointer_context = None
-        _checkpointer_initialized = False
-        print("[Checkpointer] PostgreSQL connection closed")
+    if _connection_pool is not None:
+        await _connection_pool.close()
+        logger.info("[Checkpointer] PostgreSQL connection pool closed")
+
+    _checkpointer = None
+    _connection_pool = None
+    _checkpointer_initialized = False

From 91fe7222b1e5d1020c2add815066a48e0c5c5026 Mon Sep 17 00:00:00 2001
From: "DESKTOP-RTLN3BA\\$punk" <vermarohanfinal@gmail.com>
Date: Thu, 5 Feb 2026 17:41:28 -0800
Subject: [PATCH 38/40] feat: enhance chat header with image generation
 configuration support and sidebar integration

---
 .../components/new-chat/chat-header.tsx       |  53 +-
 .../new-chat/image-config-sidebar.tsx         | 522 ++++++++++++++++++
 .../new-chat/image-model-selector.tsx         | 175 +++---
 3 files changed, 668 insertions(+), 82 deletions(-)
 create mode 100644 surfsense_web/components/new-chat/image-config-sidebar.tsx

diff --git a/surfsense_web/components/new-chat/chat-header.tsx b/surfsense_web/components/new-chat/chat-header.tsx
index e19f97945..8a8fa11a0 100644
--- a/surfsense_web/components/new-chat/chat-header.tsx
+++ b/surfsense_web/components/new-chat/chat-header.tsx
@@ -1,11 +1,13 @@
 "use client";
 
 import { useCallback, useState } from "react";
-import { useParams, useRouter } from "next/navigation";
 import type {
+	GlobalImageGenConfig,
 	GlobalNewLLMConfig,
+	ImageGenerationConfig,
 	NewLLMConfigPublic,
 } from "@/contracts/types/new-llm-config.types";
+import { ImageConfigSidebar } from "./image-config-sidebar";
 import { ImageModelSelector } from "./image-model-selector";
 import { ModelConfigSidebar } from "./model-config-sidebar";
 import { ModelSelector } from "./model-selector";
@@ -15,7 +17,7 @@ interface ChatHeaderProps {
 }
 
 export function ChatHeader({ searchSpaceId }: ChatHeaderProps) {
-	const router = useRouter();
+	// LLM config sidebar state
 	const [sidebarOpen, setSidebarOpen] = useState(false);
 	const [selectedConfig, setSelectedConfig] = useState<
 		NewLLMConfigPublic | GlobalNewLLMConfig | null
@@ -23,6 +25,15 @@ export function ChatHeader({ searchSpaceId }: ChatHeaderProps) {
 	const [isGlobal, setIsGlobal] = useState(false);
 	const [sidebarMode, setSidebarMode] = useState<"create" | "edit" | "view">("view");
 
+	// Image config sidebar state
+	const [imageSidebarOpen, setImageSidebarOpen] = useState(false);
+	const [selectedImageConfig, setSelectedImageConfig] = useState<
+		ImageGenerationConfig | GlobalImageGenConfig | null
+	>(null);
+	const [isImageGlobal, setIsImageGlobal] = useState(false);
+	const [imageSidebarMode, setImageSidebarMode] = useState<"create" | "edit" | "view">("view");
+
+	// LLM handlers
 	const handleEditConfig = useCallback(
 		(config: NewLLMConfigPublic | GlobalNewLLMConfig, global: boolean) => {
 			setSelectedConfig(config);
@@ -42,20 +53,36 @@ export function ChatHeader({ searchSpaceId }: ChatHeaderProps) {
 
 	const handleSidebarClose = useCallback((open: boolean) => {
 		setSidebarOpen(open);
-		if (!open) {
-			setSelectedConfig(null);
-		}
+		if (!open) setSelectedConfig(null);
 	}, []);
 
+	// Image model handlers
 	const handleAddImageModel = useCallback(() => {
-		// Navigate to settings image-models tab
-		router.push(`/dashboard/${searchSpaceId}/settings?tab=image-models`);
-	}, [router, searchSpaceId]);
+		setSelectedImageConfig(null);
+		setIsImageGlobal(false);
+		setImageSidebarMode("create");
+		setImageSidebarOpen(true);
+	}, []);
+
+	const handleEditImageConfig = useCallback(
+		(config: ImageGenerationConfig | GlobalImageGenConfig, global: boolean) => {
+			setSelectedImageConfig(config);
+			setIsImageGlobal(global);
+			setImageSidebarMode(global ? "view" : "edit");
+			setImageSidebarOpen(true);
+		},
+		[]
+	);
+
+	const handleImageSidebarClose = useCallback((open: boolean) => {
+		setImageSidebarOpen(open);
+		if (!open) setSelectedImageConfig(null);
+	}, []);
 
 	return (
 		<div className="flex items-center gap-2">
 			<ModelSelector onEdit={handleEditConfig} onAddNew={handleAddNew} />
-			<ImageModelSelector onAddNew={handleAddImageModel} />
+			<ImageModelSelector onEdit={handleEditImageConfig} onAddNew={handleAddImageModel} />
 			<ModelConfigSidebar
 				open={sidebarOpen}
 				onOpenChange={handleSidebarClose}
@@ -64,6 +91,14 @@ export function ChatHeader({ searchSpaceId }: ChatHeaderProps) {
 				searchSpaceId={searchSpaceId}
 				mode={sidebarMode}
 			/>
+			<ImageConfigSidebar
+				open={imageSidebarOpen}
+				onOpenChange={handleImageSidebarClose}
+				config={selectedImageConfig}
+				isGlobal={isImageGlobal}
+				searchSpaceId={searchSpaceId}
+				mode={imageSidebarMode}
+			/>
 		</div>
 	);
 }
diff --git a/surfsense_web/components/new-chat/image-config-sidebar.tsx b/surfsense_web/components/new-chat/image-config-sidebar.tsx
new file mode 100644
index 000000000..18f98acb7
--- /dev/null
+++ b/surfsense_web/components/new-chat/image-config-sidebar.tsx
@@ -0,0 +1,522 @@
+"use client";
+
+import { useAtomValue } from "jotai";
+import {
+	AlertCircle,
+	Check,
+	ChevronsUpDown,
+	Globe,
+	ImageIcon,
+	Key,
+	Shuffle,
+	X,
+	Zap,
+} from "lucide-react";
+import { AnimatePresence, motion } from "motion/react";
+import { useCallback, useEffect, useMemo, useState } from "react";
+import { createPortal } from "react-dom";
+import { toast } from "sonner";
+import {
+	createImageGenConfigMutationAtom,
+	updateImageGenConfigMutationAtom,
+} from "@/atoms/image-gen-config/image-gen-config-mutation.atoms";
+import { updateLLMPreferencesMutationAtom } from "@/atoms/new-llm-config/new-llm-config-mutation.atoms";
+import { Alert, AlertDescription } from "@/components/ui/alert";
+import { Badge } from "@/components/ui/badge";
+import { Button } from "@/components/ui/button";
+import {
+	Command,
+	CommandEmpty,
+	CommandGroup,
+	CommandInput,
+	CommandItem,
+	CommandList,
+} from "@/components/ui/command";
+import { Input } from "@/components/ui/input";
+import { Label } from "@/components/ui/label";
+import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover";
+import {
+	Select,
+	SelectContent,
+	SelectItem,
+	SelectTrigger,
+	SelectValue,
+} from "@/components/ui/select";
+import { Separator } from "@/components/ui/separator";
+import { Spinner } from "@/components/ui/spinner";
+import { IMAGE_GEN_MODELS, IMAGE_GEN_PROVIDERS } from "@/contracts/enums/image-gen-providers";
+import type {
+	GlobalImageGenConfig,
+	ImageGenerationConfig,
+} from "@/contracts/types/new-llm-config.types";
+import { cn } from "@/lib/utils";
+
+interface ImageConfigSidebarProps {
+	open: boolean;
+	onOpenChange: (open: boolean) => void;
+	config: ImageGenerationConfig | GlobalImageGenConfig | null;
+	isGlobal: boolean;
+	searchSpaceId: number;
+	mode: "create" | "edit" | "view";
+}
+
+const INITIAL_FORM = {
+	name: "",
+	description: "",
+	provider: "",
+	model_name: "",
+	api_key: "",
+	api_base: "",
+	api_version: "",
+};
+
+export function ImageConfigSidebar({
+	open,
+	onOpenChange,
+	config,
+	isGlobal,
+	searchSpaceId,
+	mode,
+}: ImageConfigSidebarProps) {
+	const [isSubmitting, setIsSubmitting] = useState(false);
+	const [mounted, setMounted] = useState(false);
+	const [formData, setFormData] = useState(INITIAL_FORM);
+	const [modelComboboxOpen, setModelComboboxOpen] = useState(false);
+
+	useEffect(() => {
+		setMounted(true);
+	}, []);
+
+	// Reset form when opening
+	useEffect(() => {
+		if (open) {
+			if (mode === "edit" && config && !isGlobal) {
+				setFormData({
+					name: config.name || "",
+					description: config.description || "",
+					provider: config.provider || "",
+					model_name: config.model_name || "",
+					api_key: (config as ImageGenerationConfig).api_key || "",
+					api_base: config.api_base || "",
+					api_version: config.api_version || "",
+				});
+			} else if (mode === "create") {
+				setFormData(INITIAL_FORM);
+			}
+		}
+	}, [open, mode, config, isGlobal]);
+
+	// Mutations
+	const { mutateAsync: createConfig } = useAtomValue(createImageGenConfigMutationAtom);
+	const { mutateAsync: updateConfig } = useAtomValue(updateImageGenConfigMutationAtom);
+	const { mutateAsync: updatePreferences } = useAtomValue(updateLLMPreferencesMutationAtom);
+
+	// Escape key
+	useEffect(() => {
+		const handleEscape = (e: KeyboardEvent) => {
+			if (e.key === "Escape" && open) onOpenChange(false);
+		};
+		window.addEventListener("keydown", handleEscape);
+		return () => window.removeEventListener("keydown", handleEscape);
+	}, [open, onOpenChange]);
+
+	const isAutoMode = config && "is_auto_mode" in config && config.is_auto_mode;
+
+	const suggestedModels = useMemo(() => {
+		if (!formData.provider) return [];
+		return IMAGE_GEN_MODELS.filter((m) => m.provider === formData.provider);
+	}, [formData.provider]);
+
+	const getTitle = () => {
+		if (mode === "create") return "Add Image Model";
+		if (isAutoMode) return "Auto Mode (Load Balanced)";
+		if (isGlobal) return "View Global Image Model";
+		return "Edit Image Model";
+	};
+
+	const handleSubmit = useCallback(async () => {
+		setIsSubmitting(true);
+		try {
+			if (mode === "create") {
+				const result = await createConfig({
+					name: formData.name,
+					provider: formData.provider,
+					model_name: formData.model_name,
+					api_key: formData.api_key,
+					api_base: formData.api_base || undefined,
+					api_version: formData.api_version || undefined,
+					description: formData.description || undefined,
+					search_space_id: searchSpaceId,
+				});
+				// Set as active image model
+				if (result?.id) {
+					await updatePreferences({
+						search_space_id: searchSpaceId,
+						data: { image_generation_config_id: result.id },
+					});
+				}
+				toast.success("Image model created and assigned!");
+				onOpenChange(false);
+			} else if (!isGlobal && config) {
+				await updateConfig({
+					id: config.id,
+					data: {
+						name: formData.name,
+						description: formData.description || undefined,
+						provider: formData.provider,
+						model_name: formData.model_name,
+						api_key: formData.api_key,
+						api_base: formData.api_base || undefined,
+						api_version: formData.api_version || undefined,
+					},
+				});
+				toast.success("Image model updated!");
+				onOpenChange(false);
+			}
+		} catch (error) {
+			console.error("Failed to save image config:", error);
+			toast.error("Failed to save image model");
+		} finally {
+			setIsSubmitting(false);
+		}
+	}, [mode, isGlobal, config, formData, searchSpaceId, createConfig, updateConfig, updatePreferences, onOpenChange]);
+
+	const handleUseGlobalConfig = useCallback(async () => {
+		if (!config || !isGlobal) return;
+		setIsSubmitting(true);
+		try {
+			await updatePreferences({
+				search_space_id: searchSpaceId,
+				data: { image_generation_config_id: config.id },
+			});
+			toast.success(`Now using ${config.name}`);
+			onOpenChange(false);
+		} catch (error) {
+			console.error("Failed to set image model:", error);
+			toast.error("Failed to set image model");
+		} finally {
+			setIsSubmitting(false);
+		}
+	}, [config, isGlobal, searchSpaceId, updatePreferences, onOpenChange]);
+
+	const isFormValid = formData.name && formData.provider && formData.model_name && formData.api_key;
+	const selectedProvider = IMAGE_GEN_PROVIDERS.find((p) => p.value === formData.provider);
+
+	if (!mounted) return null;
+
+	const sidebarContent = (
+		<AnimatePresence>
+			{open && (
+				<>
+					{/* Backdrop */}
+					<motion.div
+						initial={{ opacity: 0 }}
+						animate={{ opacity: 1 }}
+						exit={{ opacity: 0 }}
+						transition={{ duration: 0.2 }}
+						className="fixed inset-0 z-50 bg-black/20 backdrop-blur-sm"
+						onClick={() => onOpenChange(false)}
+					/>
+
+					{/* Sidebar */}
+					<motion.div
+						initial={{ x: "100%", opacity: 0 }}
+						animate={{ x: 0, opacity: 1 }}
+						exit={{ x: "100%", opacity: 0 }}
+						transition={{ type: "spring", damping: 30, stiffness: 300 }}
+						className={cn(
+							"fixed right-0 top-0 z-50 h-full w-full sm:w-[480px] lg:w-[540px]",
+							"bg-background border-l border-border/50 shadow-2xl",
+							"flex flex-col"
+						)}
+					>
+						{/* Header */}
+						<div
+							className={cn(
+								"flex items-center justify-between px-6 py-4 border-b border-border/50",
+								isAutoMode
+									? "bg-gradient-to-r from-violet-500/10 to-purple-500/10"
+									: "bg-gradient-to-r from-teal-500/10 to-cyan-500/10"
+							)}
+						>
+							<div className="flex items-center gap-3">
+								<div
+									className={cn(
+										"flex items-center justify-center size-10 rounded-xl",
+										isAutoMode
+											? "bg-gradient-to-br from-violet-500 to-purple-600"
+											: "bg-gradient-to-br from-teal-500 to-cyan-600"
+									)}
+								>
+									{isAutoMode ? (
+										<Shuffle className="size-5 text-white" />
+									) : (
+										<ImageIcon className="size-5 text-white" />
+									)}
+								</div>
+								<div>
+									<h2 className="text-base sm:text-lg font-semibold">{getTitle()}</h2>
+									<div className="flex items-center gap-2 mt-0.5">
+										{isAutoMode ? (
+											<Badge
+												variant="secondary"
+												className="gap-1 text-xs bg-violet-100 text-violet-700 dark:bg-violet-900/30 dark:text-violet-300"
+											>
+												<Zap className="size-3" />
+												Recommended
+											</Badge>
+										) : isGlobal ? (
+											<Badge variant="secondary" className="gap-1 text-xs">
+												<Globe className="size-3" />
+												Global
+											</Badge>
+										) : null}
+										{config && !isAutoMode && (
+											<span className="text-xs text-muted-foreground">{config.model_name}</span>
+										)}
+									</div>
+								</div>
+							</div>
+							<Button
+								variant="ghost"
+								size="icon"
+								onClick={() => onOpenChange(false)}
+								className="h-8 w-8 rounded-full"
+							>
+								<X className="h-4 w-4" />
+								<span className="sr-only">Close</span>
+							</Button>
+						</div>
+
+						{/* Content */}
+						<div className="flex-1 overflow-y-auto">
+							<div className="p-6">
+								{/* Auto mode */}
+								{isAutoMode && (
+									<>
+										<Alert className="mb-6 border-violet-500/30 bg-violet-500/5">
+											<Shuffle className="size-4 text-violet-500" />
+											<AlertDescription className="text-sm text-violet-700 dark:text-violet-400">
+												Auto mode distributes image generation requests across all configured providers for optimal performance and rate limit protection.
+											</AlertDescription>
+										</Alert>
+										<div className="flex gap-3 pt-4 border-t border-border/50">
+											<Button variant="outline" className="flex-1" onClick={() => onOpenChange(false)}>
+												Close
+											</Button>
+											<Button
+												className="flex-1 gap-2 bg-gradient-to-r from-violet-500 to-purple-600 hover:from-violet-600 hover:to-purple-700"
+												onClick={handleUseGlobalConfig}
+												disabled={isSubmitting}
+											>
+												{isSubmitting ? "Loading..." : "Use Auto Mode"}
+											</Button>
+										</div>
+									</>
+								)}
+
+								{/* Global config (read-only) */}
+								{isGlobal && !isAutoMode && config && (
+									<>
+										<Alert className="mb-6 border-amber-500/30 bg-amber-500/5">
+											<AlertCircle className="size-4 text-amber-500" />
+											<AlertDescription className="text-sm text-amber-700 dark:text-amber-400">
+												Global configurations are read-only. To customize, create a new model.
+											</AlertDescription>
+										</Alert>
+										<div className="space-y-4">
+											<div className="grid gap-4 sm:grid-cols-2">
+												<div className="space-y-1.5">
+													<div className="text-xs font-medium text-muted-foreground uppercase tracking-wider">Name</div>
+													<p className="text-sm font-medium">{config.name}</p>
+												</div>
+												{config.description && (
+													<div className="space-y-1.5">
+														<div className="text-xs font-medium text-muted-foreground uppercase tracking-wider">Description</div>
+														<p className="text-sm text-muted-foreground">{config.description}</p>
+													</div>
+												)}
+											</div>
+											<Separator />
+											<div className="grid gap-4 sm:grid-cols-2">
+												<div className="space-y-1.5">
+													<div className="text-xs font-medium text-muted-foreground uppercase tracking-wider">Provider</div>
+													<p className="text-sm font-medium">{config.provider}</p>
+												</div>
+												<div className="space-y-1.5">
+													<div className="text-xs font-medium text-muted-foreground uppercase tracking-wider">Model</div>
+													<p className="text-sm font-medium font-mono">{config.model_name}</p>
+												</div>
+											</div>
+										</div>
+										<div className="flex gap-3 pt-6 border-t border-border/50 mt-6">
+											<Button variant="outline" className="flex-1" onClick={() => onOpenChange(false)}>
+												Close
+											</Button>
+											<Button className="flex-1 gap-2" onClick={handleUseGlobalConfig} disabled={isSubmitting}>
+												{isSubmitting ? "Loading..." : "Use This Model"}
+											</Button>
+										</div>
+									</>
+								)}
+
+								{/* Create / Edit form */}
+								{(mode === "create" || (mode === "edit" && !isGlobal)) && (
+									<div className="space-y-4">
+										{/* Name */}
+										<div className="space-y-2">
+											<Label className="text-sm font-medium">Name *</Label>
+											<Input
+												placeholder="e.g., My DALL-E 3"
+												value={formData.name}
+												onChange={(e) => setFormData((p) => ({ ...p, name: e.target.value }))}
+											/>
+										</div>
+
+										{/* Description */}
+										<div className="space-y-2">
+											<Label className="text-sm font-medium">Description</Label>
+											<Input
+												placeholder="Optional description"
+												value={formData.description}
+												onChange={(e) => setFormData((p) => ({ ...p, description: e.target.value }))}
+											/>
+										</div>
+
+										<Separator />
+
+										{/* Provider */}
+										<div className="space-y-2">
+											<Label className="text-sm font-medium">Provider *</Label>
+											<Select
+												value={formData.provider}
+												onValueChange={(val) => setFormData((p) => ({ ...p, provider: val, model_name: "" }))}
+											>
+												<SelectTrigger>
+													<SelectValue placeholder="Select a provider" />
+												</SelectTrigger>
+												<SelectContent>
+													{IMAGE_GEN_PROVIDERS.map((p) => (
+														<SelectItem key={p.value} value={p.value}>
+															<div className="flex flex-col">
+																<span className="font-medium">{p.label}</span>
+																<span className="text-xs text-muted-foreground">{p.example}</span>
+															</div>
+														</SelectItem>
+													))}
+												</SelectContent>
+											</Select>
+										</div>
+
+										{/* Model Name */}
+										<div className="space-y-2">
+											<Label className="text-sm font-medium">Model Name *</Label>
+											{suggestedModels.length > 0 ? (
+												<Popover open={modelComboboxOpen} onOpenChange={setModelComboboxOpen}>
+													<PopoverTrigger asChild>
+														<Button variant="outline" role="combobox" className="w-full justify-between font-normal">
+															{formData.model_name || "Select or type a model..."}
+															<ChevronsUpDown className="ml-2 h-4 w-4 shrink-0 opacity-50" />
+														</Button>
+													</PopoverTrigger>
+													<PopoverContent className="w-full p-0" align="start">
+														<Command>
+															<CommandInput
+																placeholder="Search or type model..."
+																value={formData.model_name}
+																onValueChange={(val) => setFormData((p) => ({ ...p, model_name: val }))}
+															/>
+															<CommandList>
+																<CommandEmpty>
+																	<span className="text-xs text-muted-foreground">Type a custom model name</span>
+																</CommandEmpty>
+																<CommandGroup>
+																	{suggestedModels.map((m) => (
+																		<CommandItem
+																			key={m.value}
+																			value={m.value}
+																			onSelect={() => {
+																				setFormData((p) => ({ ...p, model_name: m.value }));
+																				setModelComboboxOpen(false);
+																			}}
+																		>
+																			<Check className={cn("mr-2 h-4 w-4", formData.model_name === m.value ? "opacity-100" : "opacity-0")} />
+																			<span className="font-mono text-sm">{m.value}</span>
+																			<span className="ml-2 text-xs text-muted-foreground">{m.label}</span>
+																		</CommandItem>
+																	))}
+																</CommandGroup>
+															</CommandList>
+														</Command>
+													</PopoverContent>
+												</Popover>
+											) : (
+												<Input
+													placeholder="e.g., dall-e-3"
+													value={formData.model_name}
+													onChange={(e) => setFormData((p) => ({ ...p, model_name: e.target.value }))}
+												/>
+											)}
+										</div>
+
+										{/* API Key */}
+										<div className="space-y-2">
+											<Label className="text-sm font-medium flex items-center gap-1.5">
+												<Key className="h-3.5 w-3.5" /> API Key *
+											</Label>
+											<Input
+												type="password"
+												placeholder="sk-..."
+												value={formData.api_key}
+												onChange={(e) => setFormData((p) => ({ ...p, api_key: e.target.value }))}
+											/>
+										</div>
+
+										{/* API Base */}
+										<div className="space-y-2">
+											<Label className="text-sm font-medium">API Base URL</Label>
+											<Input
+												placeholder={selectedProvider?.apiBase || "Optional"}
+												value={formData.api_base}
+												onChange={(e) => setFormData((p) => ({ ...p, api_base: e.target.value }))}
+											/>
+										</div>
+
+										{/* Azure API Version */}
+										{formData.provider === "AZURE_OPENAI" && (
+											<div className="space-y-2">
+												<Label className="text-sm font-medium">API Version (Azure)</Label>
+												<Input
+													placeholder="2024-02-15-preview"
+													value={formData.api_version}
+													onChange={(e) => setFormData((p) => ({ ...p, api_version: e.target.value }))}
+												/>
+											</div>
+										)}
+
+										{/* Actions */}
+										<div className="flex gap-3 pt-4 border-t">
+											<Button variant="outline" className="flex-1" onClick={() => onOpenChange(false)}>
+												Cancel
+											</Button>
+											<Button
+												className="flex-1 gap-2 bg-gradient-to-r from-teal-500 to-cyan-600 hover:from-teal-600 hover:to-cyan-700"
+												onClick={handleSubmit}
+												disabled={isSubmitting || !isFormValid}
+											>
+												{isSubmitting ? <Spinner size="sm" className="mr-2" /> : null}
+												{mode === "edit" ? "Save Changes" : "Create & Use"}
+											</Button>
+										</div>
+									</div>
+								)}
+							</div>
+						</div>
+					</motion.div>
+				</>
+			)}
+		</AnimatePresence>
+	);
+
+	return typeof document !== "undefined" ? createPortal(sidebarContent, document.body) : null;
+}
diff --git a/surfsense_web/components/new-chat/image-model-selector.tsx b/surfsense_web/components/new-chat/image-model-selector.tsx
index 8cae10345..b3422b264 100644
--- a/surfsense_web/components/new-chat/image-model-selector.tsx
+++ b/surfsense_web/components/new-chat/image-model-selector.tsx
@@ -38,14 +38,19 @@ import {
 } from "@/components/ui/command";
 import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover";
 import { Spinner } from "@/components/ui/spinner";
+import type {
+	GlobalImageGenConfig,
+	ImageGenerationConfig,
+} from "@/contracts/types/new-llm-config.types";
 import { cn } from "@/lib/utils";
 
 interface ImageModelSelectorProps {
 	className?: string;
 	onAddNew?: () => void;
+	onEdit?: (config: ImageGenerationConfig | GlobalImageGenConfig, isGlobal: boolean) => void;
 }
 
-export function ImageModelSelector({ className, onAddNew }: ImageModelSelectorProps) {
+export function ImageModelSelector({ className, onAddNew, onEdit }: ImageModelSelectorProps) {
 	const [open, setOpen] = useState(false);
 	const [searchQuery, setSearchQuery] = useState("");
 
@@ -220,49 +225,59 @@ export function ImageModelSelector({ className, onAddNew }: ImageModelSelectorPr
 									<Globe className="size-3.5" />
 									Global Image Models
 								</div>
-								{filteredGlobal.map((config) => {
-									const isSelected = currentConfig?.id === config.id;
-									const isAuto = "is_auto_mode" in config && config.is_auto_mode;
-									return (
-										<CommandItem
-											key={`g-${config.id}`}
-											value={`g-${config.id}`}
-											onSelect={() => handleSelect(config.id)}
-											className={cn(
-												"mx-2 rounded-lg mb-1 cursor-pointer group transition-all hover:bg-accent/50",
-												isSelected && "bg-accent/80",
-												isAuto && "border border-violet-200 dark:border-violet-800/50"
-											)}
-										>
-											<div className="flex items-center gap-3 min-w-0 flex-1">
-												<div className="shrink-0">
-													{isAuto ? (
-														<Shuffle className="size-4 text-violet-500" />
-													) : (
-														<ImageIcon className="size-4 text-teal-500" />
-													)}
-												</div>
-												<div className="min-w-0 flex-1">
-													<div className="flex items-center gap-2">
-														<span className="font-medium truncate">{config.name}</span>
-														{isAuto && (
-															<Badge
-																variant="secondary"
-																className="text-[9px] px-1 py-0 h-3.5 bg-violet-100 text-violet-700 dark:bg-violet-900/30 dark:text-violet-300 border-0"
-															>
-																Recommended
-															</Badge>
-														)}
-														{isSelected && <Check className="size-3.5 text-primary shrink-0" />}
-													</div>
-													<span className="text-xs text-muted-foreground truncate block">
-														{isAuto ? "Auto load balancing" : config.model_name}
-													</span>
-												</div>
+							{filteredGlobal.map((config) => {
+								const isSelected = currentConfig?.id === config.id;
+								const isAuto = "is_auto_mode" in config && config.is_auto_mode;
+								return (
+									<CommandItem
+										key={`g-${config.id}`}
+										value={`g-${config.id}`}
+										onSelect={() => handleSelect(config.id)}
+										className={cn(
+											"mx-2 rounded-lg mb-1 cursor-pointer group transition-all hover:bg-accent/50",
+											isSelected && "bg-accent/80",
+											isAuto && "border border-violet-200 dark:border-violet-800/50"
+										)}
+									>
+										<div className="flex items-center gap-3 min-w-0 flex-1">
+											<div className="shrink-0">
+												{isAuto ? (
+													<Shuffle className="size-4 text-violet-500" />
+												) : (
+													<ImageIcon className="size-4 text-teal-500" />
+												)}
 											</div>
-										</CommandItem>
-									);
-								})}
+											<div className="min-w-0 flex-1">
+												<div className="flex items-center gap-2">
+													<span className="font-medium truncate">{config.name}</span>
+													{isAuto && (
+														<Badge
+															variant="secondary"
+															className="text-[9px] px-1 py-0 h-3.5 bg-violet-100 text-violet-700 dark:bg-violet-900/30 dark:text-violet-300 border-0"
+														>
+															Recommended
+														</Badge>
+													)}
+													{isSelected && <Check className="size-3.5 text-primary shrink-0" />}
+												</div>
+												<span className="text-xs text-muted-foreground truncate block">
+													{isAuto ? "Auto load balancing" : config.model_name}
+												</span>
+											</div>
+											{onEdit && (
+												<ChevronRight
+													className="size-3.5 text-muted-foreground shrink-0 opacity-0 group-hover:opacity-100 transition-opacity cursor-pointer"
+													onClick={(e) => {
+														e.stopPropagation();
+														setOpen(false);
+														onEdit(config, true);
+													}}
+												/>
+											)}
+										</div>
+									</CommandItem>
+								);
+							})}
 							</CommandGroup>
 						)}
 
@@ -275,37 +290,51 @@ export function ImageModelSelector({ className, onAddNew }: ImageModelSelectorPr
 										<User className="size-3.5" />
 										Your Image Models
 									</div>
-									{filteredUser.map((config) => {
-										const isSelected = currentConfig?.id === config.id;
-										return (
-											<CommandItem
-												key={`u-${config.id}`}
-												value={`u-${config.id}`}
-												onSelect={() => handleSelect(config.id)}
-												className={cn(
-													"mx-2 rounded-lg mb-1 cursor-pointer group transition-all hover:bg-accent/50",
-													isSelected && "bg-accent/80"
-												)}
-											>
-												<div className="flex items-center gap-3 min-w-0 flex-1">
-													<div className="shrink-0">
-														<ImageIcon className="size-4 text-teal-500" />
-													</div>
-													<div className="min-w-0 flex-1">
-														<div className="flex items-center gap-2">
-															<span className="font-medium truncate">{config.name}</span>
-															{isSelected && (
-																<Check className="size-3.5 text-primary shrink-0" />
-															)}
-														</div>
-														<span className="text-xs text-muted-foreground truncate block">
-															{config.model_name}
-														</span>
-													</div>
+								{filteredUser.map((config) => {
+									const isSelected = currentConfig?.id === config.id;
+									return (
+										<CommandItem
+											key={`u-${config.id}`}
+											value={`u-${config.id}`}
+											onSelect={() => handleSelect(config.id)}
+											className={cn(
+												"mx-2 rounded-lg mb-1 cursor-pointer group transition-all hover:bg-accent/50",
+												isSelected && "bg-accent/80"
+											)}
+										>
+											<div className="flex items-center gap-3 min-w-0 flex-1">
+												<div className="shrink-0">
+													<ImageIcon className="size-4 text-teal-500" />
 												</div>
-											</CommandItem>
-										);
-									})}
+												<div className="min-w-0 flex-1">
+													<div className="flex items-center gap-2">
+														<span className="font-medium truncate">{config.name}</span>
+														{isSelected && (
+															<Check className="size-3.5 text-primary shrink-0" />
+														)}
+													</div>
+													<span className="text-xs text-muted-foreground truncate block">
+														{config.model_name}
+													</span>
+												</div>
+												{onEdit && (
+													<Button
+														variant="ghost"
+														size="icon"
+														className="h-7 w-7 shrink-0 opacity-0 group-hover:opacity-100 transition-opacity"
+														onClick={(e) => {
+															e.stopPropagation();
+															setOpen(false);
+															onEdit(config, false);
+														}}
+													>
+														<Edit3 className="size-3.5 text-muted-foreground" />
+													</Button>
+												)}
+											</div>
+										</CommandItem>
+									);
+								})}
 								</CommandGroup>
 							</>
 						)}

From 0d031cb2c235d3b0147b8e3aa7a1bb1b89610939 Mon Sep 17 00:00:00 2001
From: "DESKTOP-RTLN3BA\\$punk" <vermarohanfinal@gmail.com>
Date: Thu, 5 Feb 2026 18:07:27 -0800
Subject: [PATCH 39/40] refactor: update image generation configuration to
 remove TPM references and clarify RPM usage in comments

---
 .../app/config/global_llm_config.example.yaml | 16 ++---------
 .../app/services/image_gen_router_service.py  | 28 ++++++-------------
 2 files changed, 11 insertions(+), 33 deletions(-)

diff --git a/surfsense_backend/app/config/global_llm_config.example.yaml b/surfsense_backend/app/config/global_llm_config.example.yaml
index e727b8d56..0bb00c398 100644
--- a/surfsense_backend/app/config/global_llm_config.example.yaml
+++ b/surfsense_backend/app/config/global_llm_config.example.yaml
@@ -208,8 +208,7 @@ global_image_generation_configs:
     model_name: "dall-e-3"
     api_key: "sk-your-openai-api-key-here"
     api_base: ""
-    rpm: 50
-    tpm: 100000
+    rpm: 50  # Requests per minute (image gen is rate-limited by RPM, not tokens)
     litellm_params: {}
 
   # Example: OpenAI GPT Image 1
@@ -221,7 +220,6 @@ global_image_generation_configs:
     api_key: "sk-your-openai-api-key-here"
     api_base: ""
     rpm: 50
-    tpm: 100000
     litellm_params: {}
 
   # Example: Azure OpenAI DALL-E 3
@@ -234,7 +232,6 @@ global_image_generation_configs:
     api_base: "https://your-resource.openai.azure.com"
     api_version: "2024-02-15-preview"
     rpm: 50
-    tpm: 100000
     litellm_params:
       base_model: "dall-e-3"
 
@@ -247,7 +244,6 @@ global_image_generation_configs:
   #   api_key: "your-openrouter-api-key-here"
   #   api_base: ""
   #   rpm: 30
-  #   tpm: 50000
   #   litellm_params: {}
 
 # Notes:
@@ -262,17 +258,11 @@ global_image_generation_configs:
 # - rpm/tpm: Optional rate limits for load balancing (requests/tokens per minute)
 #   These help the router distribute load evenly and avoid rate limit errors
 #
-# AZURE-SPECIFIC NOTES:
-# - Always add 'base_model' in litellm_params for Azure deployments
-# - This fixes "Could not identify azure model 'X'" warnings
-# - base_model should match the underlying OpenAI model (e.g., gpt-4o, gpt-4-turbo, gpt-3.5-turbo)
-# - model_name format: "azure/<your-deployment-name>"
-# - api_version: Use a recent Azure API version (e.g., "2024-02-15-preview")
-# - See: https://docs.litellm.ai/docs/proxy/cost_tracking#spend-tracking-for-azure-openai-models
 #
 # IMAGE GENERATION NOTES:
 # - Image generation configs use the same ID scheme as LLM configs (negative for global)
 # - Supported models: dall-e-2, dall-e-3, gpt-image-1 (OpenAI), azure/* (Azure), 
 #   bedrock/* (AWS), vertex_ai/* (Google), recraft/* (Recraft), openrouter/* (OpenRouter)
 # - The router uses litellm.aimage_generation() for async image generation
-# - api_version is required for Azure image generation deployments
+# - Only RPM (requests per minute) is relevant for image generation rate limiting.
+#   TPM (tokens per minute) does not apply since image APIs are billed/rate-limited per request, not per token.
diff --git a/surfsense_backend/app/services/image_gen_router_service.py b/surfsense_backend/app/services/image_gen_router_service.py
index eb6936efd..f45a6ab63 100644
--- a/surfsense_backend/app/services/image_gen_router_service.py
+++ b/surfsense_backend/app/services/image_gen_router_service.py
@@ -183,11 +183,11 @@ class ImageGenRouterService:
                 "litellm_params": litellm_params,
             }
 
-            # Add rate limits from config if available
+            # Add RPM rate limit from config if available
+            # Note: TPM (tokens per minute) is not applicable for image generation
+            # since image APIs are rate-limited by requests, not tokens.
             if config.get("rpm"):
                 deployment["rpm"] = config["rpm"]
-            if config.get("tpm"):
-                deployment["tpm"] = config["tpm"]
 
             return deployment
 
@@ -219,10 +219,6 @@ class ImageGenRouterService:
         prompt: str,
         model: str = "auto",
         n: int | None = None,
-        quality: str | None = None,
-        size: str | None = None,
-        style: str | None = None,
-        response_format: str | None = None,
         timeout: int = 600,
         **kwargs,
     ) -> ImageResponse:
@@ -232,16 +228,16 @@ class ImageGenRouterService:
         Uses Router.aimage_generation() which distributes requests
         across configured image generation deployments.
 
+        Parameters like size, quality, style, and response_format are intentionally
+        omitted to keep the interface model-agnostic. Providers use their own
+        sensible defaults. If needed, pass them via **kwargs.
+
         Args:
             prompt: Text description of the desired image(s)
             model: Model alias (default "auto" for router routing)
             n: Number of images to generate
-            quality: Image quality setting
-            size: Image size
-            style: Style parameter
-            response_format: "url" or "b64_json"
             timeout: Request timeout in seconds
-            **kwargs: Additional litellm params
+            **kwargs: Additional provider-specific params (size, quality, etc.)
 
         Returns:
             ImageResponse from litellm
@@ -264,14 +260,6 @@ class ImageGenRouterService:
         }
         if n is not None:
             gen_kwargs["n"] = n
-        if quality is not None:
-            gen_kwargs["quality"] = quality
-        if size is not None:
-            gen_kwargs["size"] = size
-        if style is not None:
-            gen_kwargs["style"] = style
-        if response_format is not None:
-            gen_kwargs["response_format"] = response_format
         gen_kwargs.update(kwargs)
 
         return await instance._router.aimage_generation(**gen_kwargs)

From 7ce6493caa79a5cb0038c0a9a1caa4fac717bbf9 Mon Sep 17 00:00:00 2001
From: "DESKTOP-RTLN3BA\\$punk" <vermarohanfinal@gmail.com>
Date: Thu, 5 Feb 2026 18:16:21 -0800
Subject: [PATCH 40/40] fix: update ENUM import for PostgreSQL compatibility in
 image generations table migration

---
 .../alembic/versions/93_add_image_generations_table.py         | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/surfsense_backend/alembic/versions/93_add_image_generations_table.py b/surfsense_backend/alembic/versions/93_add_image_generations_table.py
index f24adc68f..eba9d7c86 100644
--- a/surfsense_backend/alembic/versions/93_add_image_generations_table.py
+++ b/surfsense_backend/alembic/versions/93_add_image_generations_table.py
@@ -13,6 +13,7 @@ Changes:
 from collections.abc import Sequence
 
 import sqlalchemy as sa
+from sqlalchemy.dialects.postgresql import ENUM as PG_ENUM
 from sqlalchemy.dialects.postgresql import JSONB, UUID
 
 from alembic import op
@@ -59,7 +60,7 @@ def upgrade() -> None:
             sa.Column("description", sa.String(500), nullable=True),
             sa.Column(
                 "provider",
-                sa.Enum(
+                PG_ENUM(
                     "OPENAI",
                     "AZURE_OPENAI",
                     "GOOGLE",