fix(desktop): preserve composer selection across model picking

apps/desktop/src/app/chat/composer/model-pill.test.tsx

@@ -1,12 +1,16 @@
-import { cleanup, render, screen } from '@testing-library/react'
+import { act, cleanup, fireEvent, render, screen, waitFor } from '@testing-library/react'
 import { atom } from 'nanostores'
+import { useContext } from 'react'
 import { afterEach, describe, expect, it } from 'vitest'
 
 import type { ChatBarState } from '@/app/chat/composer/types'
 import { type SessionView, SessionViewProvider } from '@/app/chat/session-view'
+import { ModelMenuCloseContext } from '@/app/shell/model-menu-panel'
 import { $activeSessionId, $currentModel, setCurrentModel, setCurrentModelSource } from '@/store/session'
 
+import { requestModelMenuToggle } from './focus'
 import { ModelPill } from './model-pill'
+import { RICH_INPUT_SLOT } from './rich-editor'
 
 const modelState = (over: Partial<ChatBarState['model']> = {}): ChatBarState['model'] => ({
   canSwitch: true,
@@ -80,6 +84,49 @@ describe('ModelPill pinned-override badge', () => {
   })
 })
 
+function MenuChoice() {
+  const close = useContext(ModelMenuCloseContext)
+
+  return <button onClick={() => close?.()}>Choose model</button>
+}
+
+it('returns to the exact caret or backward selection after the model menu closes', async () => {
+  const surface = document.createElement('div')
+  surface.dataset.composerTarget = 'main'
+  const editor = document.createElement('div')
+  editor.dataset.slot = RICH_INPUT_SLOT
+  editor.contentEditable = 'true'
+  editor.tabIndex = 0
+  editor.textContent = 'before and after'
+  surface.append(editor)
+  document.body.append(surface)
+
+  try {
+    render(<ModelPill disabled={false} model={modelState({ modelMenuContent: <MenuChoice /> })} />)
+
+    for (const [anchor, focus] of [
+      [3, 3],
+      [10, 4]
+    ]) {
+      editor.focus()
+      window.getSelection()!.setBaseAndExtent(editor.firstChild!, anchor, editor.firstChild!, focus)
+      await act(async () => {
+        requestModelMenuToggle()
+        await new Promise(resolve => setTimeout(resolve, 0))
+      })
+      const choice = await screen.findByText('Choose model')
+      choice.focus()
+      window.getSelection()!.removeAllRanges()
+      fireEvent.click(choice)
+      await waitFor(() => expect(document.activeElement).toBe(editor))
+      expect(window.getSelection()!.anchorOffset).toBe(anchor)
+      expect(window.getSelection()!.focusOffset).toBe(focus)
+    }
+  } finally {
+    surface.remove()
+  }
+})
+
 describe('ModelPill per-surface model label', () => {
   it('shows the chat-bar model even when the primary global differs', () => {
     setCurrentModel('primary/model')

apps/desktop/src/app/chat/composer/model-pill.tsx

@@ -1,9 +1,10 @@
 import { useStore } from '@nanostores/react'
-import { useEffect, useState } from 'react'
+import { useEffect, useRef, useState } from 'react'
 
 import { useSessionView } from '@/app/chat/session-view'
 import { useTourMarker } from '@/app/chat/tour-marker'
 import { ModelMenuCloseContext } from '@/app/shell/model-menu-panel'
+import { isElementInHiddenPane } from '@/components/pane-shell/pane-visibility'
 import { Button } from '@/components/ui/button'
 import { DropdownMenu, DropdownMenuContent, DropdownMenuTrigger } from '@/components/ui/dropdown-menu'
 import { GlyphSpinner } from '@/components/ui/glyph-spinner'
@@ -16,6 +17,7 @@ import { cn } from '@/lib/utils'
 import { $currentModelSource, $defaultReasoningEffort, setModelPickerOpen } from '@/store/session'
 
 import { onComposerModelMenuRequest } from './focus'
+import { RICH_INPUT_SLOT } from './rich-editor'
 import { useComposerScope } from './scope'
 import type { ChatBarState } from './types'
 
@@ -60,6 +62,7 @@ export function ModelPill({
   const defaultEffort = useStore($defaultReasoningEffort)
   const runtimeId = useStore(view.$runtimeId)
   const [open, setOpen] = useState(false)
+  const restoreSelection = useRef<(() => void) | null>(null)
   const scope = useComposerScope()
   const hasLiveMenu = Boolean(model.modelMenuContent)
 
@@ -75,6 +78,34 @@ export function ModelPill({
         }
 
         if (hasLiveMenu) {
+          const editor = document.activeElement
+          const selection = window.getSelection()
+
+          if (
+            editor instanceof HTMLElement &&
+            editor.dataset.slot === RICH_INPUT_SLOT &&
+            selection?.anchorNode &&
+            selection.focusNode &&
+            editor.contains(selection.anchorNode) &&
+            editor.contains(selection.focusNode)
+          ) {
+            const { anchorNode, anchorOffset, focusNode, focusOffset } = selection
+
+            restoreSelection.current = () => {
+              if (
+                !editor.isConnected ||
+                isElementInHiddenPane(editor) ||
+                !editor.contains(anchorNode) ||
+                !editor.contains(focusNode)
+              ) {
+                return
+              }
+
+              editor.focus({ preventScroll: true })
+              window.getSelection()?.setBaseAndExtent(anchorNode, anchorOffset, focusNode, focusOffset)
+            }
+          }
+
           setOpen(prev => !prev)
         } else {
           setModelPickerOpen(true)
@@ -178,7 +209,22 @@ export function ModelPill({
           </Button>
         </DropdownMenuTrigger>
       </Tip>
-      <DropdownMenuContent align="end" className="w-64 p-0" side="top" sideOffset={8}>
+      <DropdownMenuContent
+        align="end"
+        className="w-64 p-0"
+        onCloseAutoFocus={event => {
+          if (restoreSelection.current) {
+            event.preventDefault()
+            restoreSelection.current()
+            restoreSelection.current = null
+          }
+        }}
+        onInteractOutside={() => {
+          restoreSelection.current = null
+        }}
+        side="top"
+        sideOffset={8}
+      >
         <ModelMenuCloseContext.Provider value={() => setMenuOpen(false)}>
           {model.modelMenuContent}
         </ModelMenuCloseContext.Provider>