From 381d17d9b36d943116b2dabb06f8c3774550e54b Mon Sep 17 00:00:00 2001
From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com>
Date: Sun, 14 Jun 2026 12:40:49 +0530
Subject: [PATCH] feat(chat-ui): integrate chat model selection handling in
 Composer and related components

---
 surfsense_web/components/assistant-ui/thread.tsx     |  9 +++++++++
 surfsense_web/components/new-chat/chat-header.tsx    |  9 +++++++--
 surfsense_web/components/new-chat/model-selector.tsx | 10 +++++++++-
 3 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/surfsense_web/components/assistant-ui/thread.tsx b/surfsense_web/components/assistant-ui/thread.tsx
index 583814acb..9f74895d1 100644
--- a/surfsense_web/components/assistant-ui/thread.tsx
+++ b/surfsense_web/components/assistant-ui/thread.tsx
@@ -522,6 +522,11 @@ const Composer: FC = () => {
 		editorRef.current?.focus();
 	}, [isDesktop, showDocumentPopover, showPromptPicker, threadId]);
 
+	const handleChatModelSelected = useCallback(() => {
+		if (!isDesktop) return;
+		editorRef.current?.focus();
+	}, [isDesktop]);
+
 	// Close document picker when a sidebar slide-out panel (inbox, etc.) opens.
 	// React only on changes to the tick — comparing against the previously-seen
 	// value preserves the one-shot semantics of the prior window-event approach
@@ -935,6 +940,7 @@ const Composer: FC = () => {
 					<ComposerAction
 						isBlockedByOtherUser={isBlockedByOtherUser}
 						searchSpaceId={Number(search_space_id)}
+						onChatModelSelected={handleChatModelSelected}
 					/>
 					<ConnectorIndicator showTrigger={false} />
 				</div>
@@ -955,11 +961,13 @@ const Composer: FC = () => {
 interface ComposerActionProps {
 	isBlockedByOtherUser?: boolean;
 	searchSpaceId: number;
+	onChatModelSelected?: () => void;
 }
 
 const ComposerAction: FC<ComposerActionProps> = ({
 	isBlockedByOtherUser = false,
 	searchSpaceId,
+	onChatModelSelected,
 }) => {
 	const mentionedDocuments = useAtomValue(mentionedDocumentsAtom);
 	const setConnectorDialogOpen = useSetAtom(connectorDialogOpenAtom);
@@ -1573,6 +1581,7 @@ const ComposerAction: FC<ComposerActionProps> = ({
 				<ChatHeader
 					searchSpaceId={searchSpaceId}
 					className="h-9 max-w-[44vw] px-2 sm:max-w-[220px] sm:px-3"
+					onChatModelSelected={onChatModelSelected}
 				/>
 				<AuiIf condition={({ thread }) => !thread.isRunning}>
 					<ComposerPrimitive.Send asChild disabled={isSendDisabled}>
diff --git a/surfsense_web/components/new-chat/chat-header.tsx b/surfsense_web/components/new-chat/chat-header.tsx
index d65dc93a7..9923d8ebb 100644
--- a/surfsense_web/components/new-chat/chat-header.tsx
+++ b/surfsense_web/components/new-chat/chat-header.tsx
@@ -5,12 +5,17 @@ import { ModelSelector } from "./model-selector";
 interface ChatHeaderProps {
 	searchSpaceId: number;
 	className?: string;
+	onChatModelSelected?: () => void;
 }
 
-export function ChatHeader({ searchSpaceId, className }: ChatHeaderProps) {
+export function ChatHeader({ searchSpaceId, className, onChatModelSelected }: ChatHeaderProps) {
 	return (
 		<div className="flex items-center gap-2">
-			<ModelSelector searchSpaceId={searchSpaceId} className={className} />
+			<ModelSelector
+				searchSpaceId={searchSpaceId}
+				className={className}
+				onChatModelSelected={onChatModelSelected}
+			/>
 		</div>
 	);
 }
diff --git a/surfsense_web/components/new-chat/model-selector.tsx b/surfsense_web/components/new-chat/model-selector.tsx
index a28f49e15..bad48d9ab 100644
--- a/surfsense_web/components/new-chat/model-selector.tsx
+++ b/surfsense_web/components/new-chat/model-selector.tsx
@@ -33,6 +33,7 @@ import { providerDisplay } from "../settings/model-connections/provider-metadata
 interface ModelSelectorProps {
 	searchSpaceId: number;
 	className?: string;
+	onChatModelSelected?: () => void;
 }
 
 type ChatModel = ModelRead & {
@@ -82,7 +83,11 @@ function groupedModels(models: ChatModel[]) {
 	}, {});
 }
 
-export function ModelSelector({ searchSpaceId, className }: ModelSelectorProps) {
+export function ModelSelector({
+	searchSpaceId,
+	className,
+	onChatModelSelected,
+}: ModelSelectorProps) {
 	const router = useRouter();
 	const isMobile = useIsMobile();
 	const [open, setOpen] = useState(false);
@@ -114,6 +119,9 @@ export function ModelSelector({ searchSpaceId, className }: ModelSelectorProps)
 	function selectModel(modelId: number) {
 		updateRoles.mutate({ chat_model_id: modelId });
 		setOpen(false);
+		requestAnimationFrame(() => {
+			onChatModelSelected?.();
+		});
 	}
 
 	function manageModelConnections() {