Merge pull request #5542 from menloresearch/fix/increase-context-window-could-not-popup

Author: louis-jan

web-app/src/containers/ChatInput.tsx

@@ -35,7 +35,6 @@ import { ModelLoader } from '@/containers/loaders/ModelLoader'
 import DropdownToolsAvailable from '@/containers/DropdownToolsAvailable'
 import { getConnectedServers } from '@/services/mcp'
 import { stopAllModels } from '@/services/models'
-import { useOutOfContextPromiseModal } from './dialogs/OutOfContextDialog'
 
 type ChatInputProps = {
   className?: string
@@ -55,8 +54,6 @@ const ChatInput = ({ model, className, initialMessage }: ChatInputProps) => {
   const { t } = useTranslation()
   const { spellCheckChatInput } = useGeneralSetting()
 
-  const { showModal, PromiseModal: OutOfContextModal } =
-    useOutOfContextPromiseModal()
   const maxRows = 10
 
   const { selectedModel } = useModelProvider()
@@ -107,7 +104,7 @@ const ChatInput = ({ model, className, initialMessage }: ChatInputProps) => {
       return
     }
     setMessage('')
-    sendMessage(prompt, showModal)
+    sendMessage(prompt)
   }
 
   useEffect(() => {
@@ -599,7 +596,6 @@ const ChatInput = ({ model, className, initialMessage }: ChatInputProps) => {
           </div>
         </div>
       )}
-      <OutOfContextModal />
     </div>
   )
 }

web-app/src/containers/ThreadContent.tsx

@@ -83,7 +83,6 @@ export const ThreadContent = memo(
       // eslint-disable-next-line @typescript-eslint/no-explicit-any
       streamTools?: any
       contextOverflowModal?: React.ReactNode | null
-      showContextOverflowModal?: () => Promise<unknown>
     }
   ) => {
     const [message, setMessage] = useState(item.content?.[0]?.text?.value || '')
@@ -134,10 +133,7 @@ export const ThreadContent = memo(
       }
       if (toSendMessage) {
         deleteMessage(toSendMessage.thread_id, toSendMessage.id ?? '')
-        sendMessage(
-          toSendMessage.content?.[0]?.text?.value || '',
-          item.showContextOverflowModal
-        )
+        sendMessage(toSendMessage.content?.[0]?.text?.value || '')
       }
     }, [deleteMessage, getMessages, item, sendMessage])
 
@@ -179,16 +175,9 @@ export const ThreadContent = memo(
           deleteMessage(threadMessages[i].thread_id, threadMessages[i].id)
         }
 
-        sendMessage(message, item.showContextOverflowModal)
+        sendMessage(message)
       },
-      [
-        deleteMessage,
-        getMessages,
-        item.thread_id,
-        message,
-        sendMessage,
-        item.showContextOverflowModal,
-      ]
+      [deleteMessage, getMessages, item.thread_id, message, sendMessage]
     )
 
     const isToolCalls =

web-app/src/containers/dialogs/OutOfContextDialog.tsx

@@ -8,108 +8,76 @@ import {
   DialogTitle,
 } from '@/components/ui/dialog'
 
-import { ReactNode, useCallback, useState } from 'react'
 import { Button } from '@/components/ui/button'
+import { useContextSizeApproval } from '@/hooks/useModelContextApproval'
 
-export function useOutOfContextPromiseModal() {
-  const [isOpen, setIsOpen] = useState(false)
-  const [modalProps, setModalProps] = useState<{
-    resolveRef:
-      | ((value: 'ctx_len' | 'context_shift' | undefined) => void)
-      | null
-  }>({
-    resolveRef: null,
-  })
-  // Function to open the modal and return a Promise
-  const showModal = useCallback(() => {
-    return new Promise((resolve) => {
-      setModalProps({
-        resolveRef: resolve,
-      })
-      setIsOpen(true)
-    })
-  }, [])
+export default function OutOfContextPromiseModal() {
+  const { isModalOpen, modalProps, setModalOpen } = useContextSizeApproval()
+  if (!modalProps) {
+    return null
+  }
+  const { onApprove, onDeny } = modalProps
 
-  const PromiseModal = useCallback((): ReactNode => {
-    if (!isOpen) {
-      return null
-    }
+  const handleContextLength = () => {
+    onApprove('ctx_len')
+  }
 
-    const handleContextLength = () => {
-      setIsOpen(false)
-      if (modalProps.resolveRef) {
-        modalProps.resolveRef('ctx_len')
-      }
-    }
+  const handleContextShift = () => {
+    onApprove('context_shift')
+  }
 
-    const handleContextShift = () => {
-      setIsOpen(false)
-      if (modalProps.resolveRef) {
-        modalProps.resolveRef('context_shift')
-      }
-    }
-    const handleCancel = () => {
-      setIsOpen(false)
-      if (modalProps.resolveRef) {
-        modalProps.resolveRef(undefined)
-      }
+  const handleDialogOpen = (open: boolean) => {
+    setModalOpen(open)
+    if (!open) {
+      onDeny()
     }
+  }
 
-    return (
-      <Dialog
-        open={isOpen}
-        onOpenChange={(open) => {
-          setIsOpen(open)
-          if (!open) handleCancel()
-        }}
-      >
-        <DialogContent>
-          <DialogHeader>
-            <DialogTitle>
-              {t('outOfContextError.title', 'Out of context error')}
-            </DialogTitle>
-          </DialogHeader>
-          <DialogDescription>
-            {t(
-              'outOfContextError.description',
-              'This chat is reaching the AI’s memory limit, like a whiteboard filling up. We can expand the memory window (called context size) so it remembers more, but it may use more of your computer’s memory. We can also truncate the input, which means it will forget some of the chat history to make room for new messages.'
-            )}
-            <br />
-            <br />
-            {t(
-              'outOfContextError.increaseContextSizeDescription',
-              'Do you want to increase the context size?'
-            )}
-          </DialogDescription>
-          <DialogFooter className="flex gap-2">
-            <Button
-              variant="default"
-              className="bg-transparent border border-main-view-fg/20 hover:bg-main-view-fg/4"
-              onClick={() => {
-                handleContextShift()
-                setIsOpen(false)
-              }}
-            >
-              {t('outOfContextError.truncateInput', 'Truncate Input')}
-            </Button>
-            <Button
-              asChild
-              onClick={() => {
-                handleContextLength()
-                setIsOpen(false)
-              }}
-            >
-              <span className="text-main-view-fg/70">
-                {t(
-                  'outOfContextError.increaseContextSize',
-                  'Increase Context Size'
-                )}
-              </span>
-            </Button>
-          </DialogFooter>
-        </DialogContent>
-      </Dialog>
-    )
-  }, [isOpen, modalProps])
-  return { showModal, PromiseModal }
+  return (
+    <Dialog open={isModalOpen} onOpenChange={handleDialogOpen}>
+      <DialogContent>
+        <DialogHeader>
+          <DialogTitle>
+            {t('outOfContextError.title', 'Out of context error')}
+          </DialogTitle>
+        </DialogHeader>
+        <DialogDescription>
+          {t(
+            'outOfContextError.description',
+            'This chat is reaching the AI’s memory limit, like a whiteboard filling up. We can expand the memory window (called context size) so it remembers more, but it may use more of your computer’s memory. We can also truncate the input, which means it will forget some of the chat history to make room for new messages.'
+          )}
+          <br />
+          <br />
+          {t(
+            'outOfContextError.increaseContextSizeDescription',
+            'Do you want to increase the context size?'
+          )}
+        </DialogDescription>
+        <DialogFooter className="flex gap-2">
+          <Button
+            variant="default"
+            className="bg-transparent border border-main-view-fg/20 hover:bg-main-view-fg/4"
+            onClick={() => {
+              handleContextShift()
+            }}
+          >
+            {t('outOfContextError.truncateInput', 'Truncate Input')}
+          </Button>
+          <Button
+            asChild
+            onClick={() => {
+              handleContextLength()
+            }}
+          >
+            <span className="text-main-view-fg/70">
+              {t(
+                'outOfContextError.increaseContextSize',
+                'Increase Context Size'
+              )}
+            </span>
+          </Button>
+        </DialogFooter>
+      </DialogContent>
+    </Dialog>
+  )
 }

web-app/src/hooks/useChat.ts

@@ -30,6 +30,7 @@ import { useToolApproval } from '@/hooks/useToolApproval'
 import { useToolAvailable } from '@/hooks/useToolAvailable'
 import { OUT_OF_CONTEXT_SIZE } from '@/utils/error'
 import { updateSettings } from '@/services/providers'
+import { useContextSizeApproval } from './useModelContextApproval'
 
 export const useChat = () => {
   const { prompt, setPrompt } = usePrompt()
@@ -47,6 +48,8 @@ export const useChat = () => {
 
   const { approvedTools, showApprovalModal, allowAllMCPPermissions } =
     useToolApproval()
+  const { showApprovalModal: showIncreaseContextSizeModal } =
+    useContextSizeApproval()
   const { getDisabledToolsForThread } = useToolAvailable()
 
   const { getProviderByName, selectedModel, selectedProvider } =
@@ -223,11 +226,7 @@ export const useChat = () => {
   )
 
   const sendMessage = useCallback(
-    async (
-      message: string,
-      showModal?: () => Promise<unknown>,
-      troubleshooting = true
-    ) => {
+    async (message: string, troubleshooting = true) => {
       const activeThread = await getCurrentThread()
 
       resetTokenSpeed()
@@ -361,7 +360,7 @@ export const useChat = () => {
               selectedModel &&
               troubleshooting
             ) {
-              const method = await showModal?.()
+              const method = await showIncreaseContextSizeModal()
               if (method === 'ctx_len') {
                 /// Increase context size
                 activeProvider = await increaseModelContextSize(
@@ -447,15 +446,15 @@ export const useChat = () => {
       updateThreadTimestamp,
       setPrompt,
       selectedModel,
-      currentAssistant?.instructions,
-      currentAssistant.parameters,
+      currentAssistant,
       tools,
       updateLoadingModel,
       getDisabledToolsForThread,
       approvedTools,
       allowAllMCPPermissions,
       showApprovalModal,
       updateTokenSpeed,
+      showIncreaseContextSizeModal,
       increaseModelContextSize,
       toggleOnContextShifting,
     ]

web-app/src/hooks/useModelContextApproval.ts

@@ -0,0 +1,53 @@
+import { create } from 'zustand'
+
+export type ApprovalModalProps = {
+  onApprove: (method: 'ctx_len' | 'context_shift') => void
+  onDeny: () => void
+}
+
+type ApprovalState = {
+  // Modal state
+  isModalOpen: boolean
+  modalProps: ApprovalModalProps | null
+
+  showApprovalModal: () => Promise<'ctx_len' | 'context_shift' | undefined>
+  closeModal: () => void
+  setModalOpen: (open: boolean) => void
+}
+
+export const useContextSizeApproval = create<ApprovalState>()((set, get) => ({
+  isModalOpen: false,
+  modalProps: null,
+
+  showApprovalModal: async () => {
+    return new Promise<'ctx_len' | 'context_shift' | undefined>((resolve) => {
+      set({
+        isModalOpen: true,
+        modalProps: {
+          onApprove: (method) => {
+            get().closeModal()
+            resolve(method)
+          },
+          onDeny: () => {
+            get().closeModal()
+            resolve(undefined)
+          },
+        },
+      })
+    })
+  },
+
+  closeModal: () => {
+    set({
+      isModalOpen: false,
+      modalProps: null,
+    })
+  },
+
+  setModalOpen: (open: boolean) => {
+    set({ isModalOpen: open })
+    if (!open) {
+      get().closeModal()
+    }
+  },
+}))

web-app/src/routes/__root.tsx

@@ -18,6 +18,7 @@ import { AnalyticProvider } from '@/providers/AnalyticProvider'
 import { useLeftPanel } from '@/hooks/useLeftPanel'
 import { cn } from '@/lib/utils'
 import ToolApproval from '@/containers/dialogs/ToolApproval'
+import OutOfContextPromiseModal from '@/containers/dialogs/OutOfContextDialog'
 
 export const Route = createRootRoute({
   component: RootLayout,
@@ -94,6 +95,7 @@ function RootLayout() {
       {/* <TanStackRouterDevtools position="bottom-right" /> */}
       <CortexFailureDialog />
       <ToolApproval />
+      <OutOfContextPromiseModal />
     </Fragment>
   )
 }