From 3d9200dd3b7e1b035df4396a3cbfb6e7b666eba3 Mon Sep 17 00:00:00 2001
From: Lifei Zhou <lifei@squareup.com>
Date: Mon, 24 Aug 2026 07:58:22 +0000
Subject: [PATCH] fix (desktop): fix chat input freeze after opening model
 picker (#11525)

---
 .../bottom_bar/ModelsBottomBar.test.tsx       | 85 ++++++++++++++++++-
 .../models/bottom_bar/ModelsBottomBar.tsx     | 47 ++++++++--
 2 files changed, 121 insertions(+), 11 deletions(-)

diff --git a/ui/desktop/src/components/settings/models/bottom_bar/ModelsBottomBar.test.tsx b/ui/desktop/src/components/settings/models/bottom_bar/ModelsBottomBar.test.tsx
index 877703609..bcf744903 100644
--- a/ui/desktop/src/components/settings/models/bottom_bar/ModelsBottomBar.test.tsx
+++ b/ui/desktop/src/components/settings/models/bottom_bar/ModelsBottomBar.test.tsx
@@ -1,5 +1,5 @@
 import { describe, it, expect, vi, beforeEach } from 'vitest';
-import { render, type RenderOptions, screen } from '@testing-library/react';
+import { fireEvent, render, type RenderOptions, screen } from '@testing-library/react';
 import ModelsBottomBar from './ModelsBottomBar';
 import { IntlTestWrapper } from '../../../../i18n/test-utils';
 
@@ -13,6 +13,7 @@ let mockCurrentModel: string | null = 'config-model';
 let mockCurrentProvider: string | null = 'config-provider';
 const mockGetProviders = vi.fn();
 const mockOnModelChanged = vi.fn();
+const mockPreventCloseAutoFocus = vi.fn();
 
 vi.mock('../../../ModelAndProviderContext', () => ({
   useModelAndProvider: () => ({
@@ -41,10 +42,47 @@ vi.mock('../../../bottom_menu/BottomMenuAlertPopover', () => ({
 }));
 
 vi.mock('../../../ui/dropdown-menu', () => ({
-  DropdownMenu: ({ children }: { children: React.ReactNode }) => <div>{children}</div>,
+  DropdownMenu: ({
+    children,
+    open,
+    onOpenChange,
+  }: {
+    children: React.ReactNode;
+    open: boolean;
+    onOpenChange: (open: boolean) => void;
+  }) => (
+    <div data-testid="model-menu" data-open={open}>
+      <button onClick={() => onOpenChange(true)}>Open model menu</button>
+      {children}
+    </div>
+  ),
   DropdownMenuTrigger: ({ children }: { children: React.ReactNode }) => <div>{children}</div>,
-  DropdownMenuContent: ({ children }: { children: React.ReactNode }) => <div>{children}</div>,
-  DropdownMenuItem: ({ children }: { children: React.ReactNode }) => <div>{children}</div>,
+  DropdownMenuContent: ({
+    children,
+    onCloseAutoFocus,
+  }: {
+    children: React.ReactNode;
+    onCloseAutoFocus?: (event: Pick<Event, 'preventDefault'>) => void;
+  }) => (
+    <div>
+      <button onClick={() => onCloseAutoFocus?.({ preventDefault: mockPreventCloseAutoFocus })}>
+        Complete model menu close
+      </button>
+      {children}
+    </div>
+  ),
+  DropdownMenuItem: ({
+    children,
+    onSelect,
+  }: {
+    children: React.ReactNode;
+    onSelect?: () => void;
+  }) => <button onClick={onSelect}>{children}</button>,
+  DropdownMenuSeparator: () => null,
+}));
+
+vi.mock('../subcomponents/SwitchModelModal', () => ({
+  SwitchModelModal: () => <div data-testid="switch-model-modal" />,
 }));
 
 vi.mock('../../localInference/ModelSettingsPanel', () => ({
@@ -107,4 +145,43 @@ describe('ModelsBottomBar', () => {
     expect(screen.getByText('config-model')).toBeInTheDocument();
     expect(screen.queryByTestId('model-loading-state')).not.toBeInTheDocument();
   });
+
+  it('opens model overlays after the menu closes with the appropriate focus behavior', () => {
+    renderWithIntl(
+      <ModelsBottomBar
+        sessionId="session-123"
+        dropdownRef={createDropdownRef()}
+        setView={vi.fn()}
+        sessionModel="local-model"
+        sessionProvider="local"
+        onModelChanged={mockOnModelChanged}
+        sessionLoaded={true}
+      />
+    );
+
+    fireEvent.click(screen.getByRole('button', { name: 'Open model menu' }));
+    expect(screen.getByTestId('model-menu')).toHaveAttribute('data-open', 'true');
+
+    fireEvent.click(screen.getByRole('button', { name: 'Local Model Settings' }));
+    expect(screen.getByTestId('model-menu')).toHaveAttribute('data-open', 'false');
+    expect(
+      screen.queryByRole('heading', { name: 'Local Model Settings — Display local-model' })
+    ).not.toBeInTheDocument();
+
+    fireEvent.click(screen.getByRole('button', { name: 'Complete model menu close' }));
+    expect(
+      screen.getByRole('heading', { name: 'Local Model Settings — Display local-model' })
+    ).toBeInTheDocument();
+    expect(mockPreventCloseAutoFocus).not.toHaveBeenCalled();
+
+    fireEvent.click(screen.getByRole('button', { name: '×' }));
+    fireEvent.click(screen.getByRole('button', { name: 'Open model menu' }));
+    fireEvent.click(screen.getByRole('button', { name: 'Change Model' }));
+    expect(screen.getByTestId('model-menu')).toHaveAttribute('data-open', 'false');
+    expect(screen.queryByTestId('switch-model-modal')).not.toBeInTheDocument();
+
+    fireEvent.click(screen.getByRole('button', { name: 'Complete model menu close' }));
+    expect(screen.getByTestId('switch-model-modal')).toBeInTheDocument();
+    expect(mockPreventCloseAutoFocus).toHaveBeenCalledOnce();
+  });
 });
diff --git a/ui/desktop/src/components/settings/models/bottom_bar/ModelsBottomBar.tsx b/ui/desktop/src/components/settings/models/bottom_bar/ModelsBottomBar.tsx
index 11e88bc24..edc26bd82 100644
--- a/ui/desktop/src/components/settings/models/bottom_bar/ModelsBottomBar.tsx
+++ b/ui/desktop/src/components/settings/models/bottom_bar/ModelsBottomBar.tsx
@@ -1,5 +1,5 @@
 import { Sliders, Bot, LoaderCircle, Settings, History } from 'lucide-react';
-import React, { useCallback, useEffect, useMemo, useState } from 'react';
+import React, { useCallback, useEffect, useMemo, useRef, useState } from 'react';
 import { useModelAndProvider } from '../../../ModelAndProviderContext';
 import { SwitchModelModal } from '../subcomponents/SwitchModelModal';
 import { View } from '../../../../utils/navigationUtils';
@@ -68,6 +68,8 @@ interface ModelsBottomBarProps {
   sessionLoaded?: boolean;
 }
 
+type ModelMenuModal = 'switch-model' | 'local-model-settings';
+
 export default function ModelsBottomBar({
   sessionId,
   dropdownRef,
@@ -80,7 +82,11 @@ export default function ModelsBottomBar({
 }: ModelsBottomBarProps) {
   // ChatInput owns the override state and passes effective model/provider as sessionModel/sessionProvider.
   // Fall back to config defaults when no session-specific model is available.
-  const { currentModel: configModel, currentProvider: configProvider, changeModel } = useModelAndProvider();
+  const {
+    currentModel: configModel,
+    currentProvider: configProvider,
+    changeModel,
+  } = useModelAndProvider();
   const currentModel = sessionModel ?? configModel;
   const currentProvider = sessionProvider ?? configProvider;
 
@@ -89,6 +95,8 @@ export default function ModelsBottomBar({
   const [displayModelName, setDisplayModelName] = useState<string>(
     intl.formatMessage(i18n.selectModel)
   );
+  const [isModelMenuOpen, setIsModelMenuOpen] = useState(false);
+  const pendingModalRef = useRef<ModelMenuModal | null>(null);
   const [isAddModelModalOpen, setIsAddModelModalOpen] = useState(false);
   const [isLocalModelSettingsOpen, setIsLocalModelSettingsOpen] = useState(false);
   const [providerDefaultModel, setProviderDefaultModel] = useState<string | null>(null);
@@ -162,6 +170,24 @@ export default function ModelsBottomBar({
     onModelChanged({ model, provider });
   };
 
+  const openModalAfterMenuCloses = (modal: ModelMenuModal) => {
+    pendingModalRef.current = modal;
+    setIsModelMenuOpen(false);
+  };
+
+  const handleModelMenuCloseAutoFocus = (event: Event) => {
+    const pendingModal = pendingModalRef.current;
+    if (!pendingModal) return;
+
+    pendingModalRef.current = null;
+    if (pendingModal === 'switch-model') {
+      event.preventDefault();
+      setIsAddModelModalOpen(true);
+    } else {
+      setIsLocalModelSettingsOpen(true);
+    }
+  };
+
   const handleRecentModelClick = async (recent: RecentModel) => {
     const previousModel = currentModel;
     const previousProvider = currentProvider;
@@ -195,7 +221,7 @@ export default function ModelsBottomBar({
 
   return (
     <div className="relative flex items-center" ref={dropdownRef}>
-      <DropdownMenu>
+      <DropdownMenu open={isModelMenuOpen} onOpenChange={setIsModelMenuOpen}>
         <DropdownMenuTrigger className="flex items-center hover:cursor-pointer max-w-[180px] md:max-w-[200px] lg:max-w-[380px] min-w-0 text-text-primary/70 hover:text-text-primary transition-colors">
           <div className="flex items-center truncate max-w-[130px] md:max-w-[200px] lg:max-w-[360px] min-w-0">
             <Bot className="mr-1 h-4 w-4 flex-shrink-0" />
@@ -212,7 +238,12 @@ export default function ModelsBottomBar({
             )}
           </div>
         </DropdownMenuTrigger>
-        <DropdownMenuContent side="top" align="center" className="w-64 text-sm">
+        <DropdownMenuContent
+          side="top"
+          align="center"
+          className="w-64 text-sm"
+          onCloseAutoFocus={handleModelMenuCloseAutoFocus}
+        >
           <h6 className="text-xs text-text-primary mt-2 ml-2">
             {intl.formatMessage(i18n.currentModel)}
           </h6>
@@ -241,18 +272,20 @@ export default function ModelsBottomBar({
                   onClick={() => void handleRecentModelClick(recent)}
                 >
                   <History className="mr-2 h-3.5 w-3.5 flex-shrink-0 text-text-secondary" />
-                  <span className="truncate">{getModelDisplayName(recent.model)} — {recent.provider}</span>
+                  <span className="truncate">
+                    {getModelDisplayName(recent.model)} — {recent.provider}
+                  </span>
                 </DropdownMenuItem>
               ))}
               <DropdownMenuSeparator />
             </>
           )}
-          <DropdownMenuItem onClick={() => setIsAddModelModalOpen(true)}>
+          <DropdownMenuItem onSelect={() => openModalAfterMenuCloses('switch-model')}>
             <span>{intl.formatMessage(i18n.changeModel)}</span>
             <Sliders className="ml-auto h-4 w-4 rotate-90" />
           </DropdownMenuItem>
           {currentProvider === 'local' && currentModel && (
-            <DropdownMenuItem onClick={() => setIsLocalModelSettingsOpen(true)}>
+            <DropdownMenuItem onSelect={() => openModalAfterMenuCloses('local-model-settings')}>
               <span>{intl.formatMessage(i18n.localModelSettings)}</span>
               <Settings className="ml-auto h-4 w-4" />
             </DropdownMenuItem>