diff --git a/apps/desktop/src/app/chat/composer/model-pill.test.tsx b/apps/desktop/src/app/chat/composer/model-pill.test.tsx index c7c1d08bc8..bf7acb2ac9 100644 --- a/apps/desktop/src/app/chat/composer/model-pill.test.tsx +++ b/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'] => ({ canSwitch: true, @@ -80,6 +84,49 @@ describe('ModelPill pinned-override badge', () => { }) }) +function MenuChoice() { + const close = useContext(ModelMenuCloseContext) + + return +} + +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( })} />) + + 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') diff --git a/apps/desktop/src/app/chat/composer/model-pill.tsx b/apps/desktop/src/app/chat/composer/model-pill.tsx index 0581754353..5d8cef2bfd 100644 --- a/apps/desktop/src/app/chat/composer/model-pill.tsx +++ b/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({ - + { + if (restoreSelection.current) { + event.preventDefault() + restoreSelection.current() + restoreSelection.current = null + } + }} + onInteractOutside={() => { + restoreSelection.current = null + }} + side="top" + sideOffset={8} + > setMenuOpen(false)}> {model.modelMenuContent}