From a1d0a3a666a4e1146e9943a40a489c7579d6216d Mon Sep 17 00:00:00 2001
From: LUIS NOVO <lfnovo@gmail.com>
Date: Fri, 19 Dec 2025 16:55:26 -0300
Subject: [PATCH] fix: allow model override before chat session exists

- Add pendingModelOverride state to useNotebookChat hook
- Store model selection when no session exists yet
- Apply pending model override when session is auto-created on first message
- Simplify ChatColumn by using new setModelOverride function
---
 .../notebooks/components/ChatColumn.tsx       |  8 ++----
 frontend/src/lib/hooks/useNotebookChat.ts     | 25 ++++++++++++++++++-
 2 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/frontend/src/app/(dashboard)/notebooks/components/ChatColumn.tsx b/frontend/src/app/(dashboard)/notebooks/components/ChatColumn.tsx
index b65d4a3..be8ce32 100644
--- a/frontend/src/app/(dashboard)/notebooks/components/ChatColumn.tsx
+++ b/frontend/src/app/(dashboard)/notebooks/components/ChatColumn.tsx
@@ -95,12 +95,8 @@ export function ChatColumn({ notebookId, contextSelections }: ChatColumnProps) {
       isStreaming={chat.isSending}
       contextIndicators={null}
       onSendMessage={(message, modelOverride) => chat.sendMessage(message, modelOverride)}
-      modelOverride={chat.currentSession?.model_override ?? undefined}
-      onModelChange={(model) => {
-        if (chat.currentSessionId) {
-          chat.updateSession(chat.currentSessionId, { model_override: model ?? null })
-        }
-      }}
+      modelOverride={chat.currentSession?.model_override ?? chat.pendingModelOverride ?? undefined}
+      onModelChange={(model) => chat.setModelOverride(model ?? null)}
       sessions={chat.sessions}
       currentSessionId={chat.currentSessionId}
       onCreateSession={(title) => chat.createSession(title)}
diff --git a/frontend/src/lib/hooks/useNotebookChat.ts b/frontend/src/lib/hooks/useNotebookChat.ts
index c10249c..e97cd18 100644
--- a/frontend/src/lib/hooks/useNotebookChat.ts
+++ b/frontend/src/lib/hooks/useNotebookChat.ts
@@ -28,6 +28,8 @@ export function useNotebookChat({ notebookId, sources, notes, contextSelections
   const [isSending, setIsSending] = useState(false)
   const [tokenCount, setTokenCount] = useState<number>(0)
   const [charCount, setCharCount] = useState<number>(0)
+  // Pending model override for when user changes model before a session exists
+  const [pendingModelOverride, setPendingModelOverride] = useState<string | null>(null)
 
   // Fetch sessions for this notebook
   const {
@@ -176,10 +178,14 @@ export function useNotebookChat({ notebookId, sources, notes, contextSelections
           : message
         const newSession = await chatApi.createSession({
           notebook_id: notebookId,
-          title: defaultTitle
+          title: defaultTitle,
+          // Include pending model override when creating session
+          model_override: pendingModelOverride ?? undefined
         })
         sessionId = newSession.id
         setCurrentSessionId(sessionId)
+        // Clear pending model override now that it's applied to the session
+        setPendingModelOverride(null)
         queryClient.invalidateQueries({
           queryKey: QUERY_KEYS.notebookChatSessions(notebookId)
         })
@@ -226,6 +232,7 @@ export function useNotebookChat({ notebookId, sources, notes, contextSelections
     notebookId,
     currentSessionId,
     currentSession,
+    pendingModelOverride,
     buildContext,
     refetchCurrentSession,
     queryClient
@@ -257,6 +264,20 @@ export function useNotebookChat({ notebookId, sources, notes, contextSelections
     return deleteSessionMutation.mutate(sessionId)
   }, [deleteSessionMutation])
 
+  // Set model override - handles both existing sessions and pending state
+  const setModelOverride = useCallback((model: string | null) => {
+    if (currentSessionId) {
+      // Session exists - update it directly
+      updateSessionMutation.mutate({
+        sessionId: currentSessionId,
+        data: { model_override: model }
+      })
+    } else {
+      // No session yet - store as pending
+      setPendingModelOverride(model)
+    }
+  }, [currentSessionId, updateSessionMutation])
+
   // Update token/char counts when context selections change
   useEffect(() => {
     const updateContextCounts = async () => {
@@ -279,6 +300,7 @@ export function useNotebookChat({ notebookId, sources, notes, contextSelections
     loadingSessions,
     tokenCount,
     charCount,
+    pendingModelOverride,
 
     // Actions
     createSession,
@@ -286,6 +308,7 @@ export function useNotebookChat({ notebookId, sources, notes, contextSelections
     deleteSession,
     switchSession,
     sendMessage,
+    setModelOverride,
     refetchSessions
   }
 }