From 8ea042e88c59030b9c21fa085cfd079a37c33c13 Mon Sep 17 00:00:00 2001
From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com>
Date: Tue, 12 May 2026 20:57:15 +0530
Subject: [PATCH] refactor(chat): improve user query handling and mention chip
 functionality

---
 .../app/tasks/chat/stream_new_chat.py         |  22 ++--
 .../assistant-ui/inline-mention-editor.tsx    | 117 +++++++++++++-----
 .../components/assistant-ui/mention-chip.tsx  |  10 +-
 3 files changed, 107 insertions(+), 42 deletions(-)

diff --git a/surfsense_backend/app/tasks/chat/stream_new_chat.py b/surfsense_backend/app/tasks/chat/stream_new_chat.py
index 818282996..aed62ece9 100644
--- a/surfsense_backend/app/tasks/chat/stream_new_chat.py
+++ b/surfsense_backend/app/tasks/chat/stream_new_chat.py
@@ -1507,14 +1507,20 @@ async def stream_new_chat(
 
         # Resolve @-mention chips to canonical virtual paths and rewrite
         # the user-typed text so the LLM sees ``\`/documents/...\``` instead
-        # of bare ``@title``. The persisted user-message text keeps
-        # ``@title`` so chip rendering on reload is unchanged — see
-        # ``persistence._build_user_content``.
+        # of bare ``@title``. The substitution lands in ``agent_user_query``
+        # ONLY — the original ``user_query`` (with ``@title`` tokens) flows
+        # untouched into ``persist_user_turn`` below so chip rendering on
+        # reload still works (``UserTextPart`` → ``parseMentionSegments``
+        # matches ``@title``, not ``\`/documents/...\```). It also feeds
+        # the human-readable surfaces — SSE "Processing X" status, auto
+        # thread title, memory seed — which all want what the user typed.
+        # See ``persistence._build_user_content``.
         #
         # Cloud mode only: local-folder mode keeps the legacy
         # ``@title`` text path; mention support there is a follow-up
         # task because the path scheme (mount-rooted) and the picker
         # UI both need separate work.
+        agent_user_query = user_query
         accepted_folder_ids: list[int] = []
         if fs_mode == FilesystemMode.CLOUD.value and (
             mentioned_document_ids
@@ -1549,11 +1555,13 @@ async def stream_new_chat(
                 mentioned_surfsense_doc_ids=mentioned_surfsense_doc_ids,
                 mentioned_folder_ids=mentioned_folder_ids,
             )
-            user_query = substitute_in_text(user_query, resolved.token_to_path)
+            agent_user_query = substitute_in_text(user_query, resolved.token_to_path)
             accepted_folder_ids = resolved.mentioned_folder_ids
 
-        # Format the user query with context (SurfSense docs + reports only)
-        final_query = user_query
+        # Format the user query with context (SurfSense docs + reports only).
+        # Uses ``agent_user_query`` so the LLM sees backtick-wrapped paths
+        # instead of bare ``@title`` tokens.
+        final_query = agent_user_query
         context_parts = []
 
         if mentioned_surfsense_docs:
@@ -1584,7 +1592,7 @@ async def stream_new_chat(
 
         if context_parts:
             context = "\n\n".join(context_parts)
-            final_query = f"{context}\n\n<user_query>{user_query}</user_query>"
+            final_query = f"{context}\n\n<user_query>{agent_user_query}</user_query>"
 
         if visibility == ChatVisibility.SEARCH_SPACE and current_user_display_name:
             final_query = f"**[{current_user_display_name}]:** {final_query}"
diff --git a/surfsense_web/components/assistant-ui/inline-mention-editor.tsx b/surfsense_web/components/assistant-ui/inline-mention-editor.tsx
index e12556486..c2b896794 100644
--- a/surfsense_web/components/assistant-ui/inline-mention-editor.tsx
+++ b/surfsense_web/components/assistant-ui/inline-mention-editor.tsx
@@ -1,6 +1,6 @@
 "use client";
 
-import { Folder as FolderIcon } from "lucide-react";
+import { Folder as FolderIcon, X as XIcon } from "lucide-react";
 import type { PlateElementProps } from "platejs/react";
 import {
 	createPlatePlugin,
@@ -9,7 +9,16 @@ import {
 	PlateContent,
 	usePlateEditor,
 } from "platejs/react";
-import { type FC, forwardRef, useCallback, useImperativeHandle, useMemo, useRef } from "react";
+import {
+	createContext,
+	type FC,
+	forwardRef,
+	useCallback,
+	useContext,
+	useImperativeHandle,
+	useMemo,
+	useRef,
+} from "react";
 import { FOLDER_MENTION_DOCUMENT_TYPE } from "@/atoms/chat/mentioned-documents.atom";
 import { getConnectorIcon } from "@/contracts/enums/connectorIcons";
 import type { Document } from "@/contracts/types/document.types";
@@ -107,13 +116,25 @@ type ComposerValue = ComposerParagraph[];
 
 const MENTION_TYPE = "mention";
 const MENTION_CHIP_CLASSNAME =
-	"inline-flex h-5 items-center gap-1 mx-0.5 rounded bg-primary/10 px-1 text-xs font-bold text-primary/60 select-none align-middle leading-none";
+	"group inline-flex h-5 items-center gap-1 mx-0.5 rounded bg-primary/10 px-1 text-xs font-bold text-primary/60 select-none align-middle leading-none";
 const MENTION_CHIP_ICON_CLASSNAME = "flex items-center text-muted-foreground leading-none";
 const MENTION_CHIP_TITLE_CLASSNAME = "max-w-[120px] truncate leading-none";
 const COMPOSER_TEXT_METRICS_CLASSNAME = "text-sm leading-6";
 
 const EMPTY_VALUE: ComposerValue = [{ type: "p", children: [{ text: "" }] }];
 
+/**
+ * Internal seam that lets ``MentionElement`` (a Plate render component
+ * with no React props beyond ``element``) reach the editor's chip-removal
+ * function. Mirrors the Backspace path in ``handleKeyDown`` so the X
+ * button delegates to the exact same combined call site — no extra
+ * state, no atom coupling leaking into the chip.
+ */
+type MentionEditorContextValue = {
+	removeChip: (docId: number, docType: string | undefined) => void;
+};
+const MentionEditorContext = createContext<MentionEditorContextValue | null>(null);
+
 const MentionElement: FC<PlateElementProps<MentionElementNode>> = ({
 	attributes,
 	children,
@@ -127,16 +148,36 @@ const MentionElement: FC<PlateElementProps<MentionElementNode>> = ({
 				: "text-amber-700";
 
 	const isFolder = element.kind === "folder";
+	const ctx = useContext(MentionEditorContext);
 
 	return (
 		<span {...attributes} className="inline-flex align-middle">
 			<span contentEditable={false} className={`${MENTION_CHIP_CLASSNAME} cursor-default`}>
 				<span className={MENTION_CHIP_ICON_CLASSNAME}>
-					{isFolder ? (
-						<FolderIcon className="h-3 w-3" />
-					) : (
-						getConnectorIcon(element.document_type ?? "UNKNOWN", "h-3 w-3")
-					)}
+					<span className="relative flex h-3 w-3 items-center justify-center">
+						<span className="flex items-center justify-center transition-opacity group-hover:opacity-0">
+							{isFolder ? (
+								<FolderIcon className="h-3 w-3" />
+							) : (
+								getConnectorIcon(element.document_type ?? "UNKNOWN", "h-3 w-3")
+							)}
+						</span>
+						{ctx ? (
+							<button
+								type="button"
+								aria-label={`Remove mention ${element.title}`}
+								title={`Remove ${element.title}`}
+								onMouseDown={(e) => e.preventDefault()}
+								onClick={(e) => {
+									e.stopPropagation();
+									ctx.removeChip(element.id, element.document_type);
+								}}
+								className="absolute inset-0 flex items-center justify-center rounded-sm opacity-0 transition-opacity hover:text-primary focus-visible:opacity-100 focus-visible:outline-none group-hover:opacity-100"
+							>
+								<XIcon className="h-3 w-3" />
+							</button>
+						) : null}
+					</span>
 				</span>
 				<span className={MENTION_CHIP_TITLE_CLASSNAME} title={element.title}>
 					{element.title}
@@ -464,6 +505,18 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
 			[getCurrentValue, setValue]
 		);
 
+		// Combined "remove chip end-to-end" used by both the Backspace
+		// keybinding and the in-chip X button. Keeping these two surfaces
+		// pinned to a single helper guarantees they can never diverge —
+		// e.g. one path forgetting to notify the parent atom.
+		const removeChip = useCallback(
+			(docId: number, docType: string | undefined) => {
+				removeDocumentChip(docId, docType);
+				onDocumentRemove?.(docId, docType);
+			},
+			[onDocumentRemove, removeDocumentChip]
+		);
+
 		const setDocumentChipStatus = useCallback(
 			(
 				docId: number,
@@ -568,10 +621,9 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
 				if (!isMentionNode(prev)) return;
 
 				e.preventDefault();
-				removeDocumentChip(prev.id, prev.document_type);
-				onDocumentRemove?.(prev.id, prev.document_type);
+				removeChip(prev.id, prev.document_type);
 			},
-			[editor.selection, getCurrentValue, onDocumentRemove, onKeyDown, onSubmit, removeDocumentChip]
+			[editor.selection, getCurrentValue, onKeyDown, onSubmit, removeChip]
 		);
 
 		const editableProps = useMemo(
@@ -588,26 +640,33 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
 			[editor, handleKeyDown, placeholder]
 		);
 
+		const mentionEditorContextValue = useMemo<MentionEditorContextValue>(
+			() => ({ removeChip }),
+			[removeChip]
+		);
+
 		return (
 			<div className="relative w-full">
-				<Plate
-					editor={editor}
-					onChange={({ value }) => {
-						emitState(value as ComposerValue);
-					}}
-				>
-					<PlateContent
-						ref={editableRef}
-						readOnly={disabled}
-						{...editableProps}
-						className={cn(
-							"min-h-[24px] max-h-32 overflow-y-auto outline-none whitespace-pre-wrap wrap-break-word",
-							COMPOSER_TEXT_METRICS_CLASSNAME,
-							disabled && "opacity-50 cursor-not-allowed",
-							className
-						)}
-					/>
-				</Plate>
+				<MentionEditorContext.Provider value={mentionEditorContextValue}>
+					<Plate
+						editor={editor}
+						onChange={({ value }) => {
+							emitState(value as ComposerValue);
+						}}
+					>
+						<PlateContent
+							ref={editableRef}
+							readOnly={disabled}
+							{...editableProps}
+							className={cn(
+								"min-h-[24px] max-h-32 overflow-y-auto outline-none whitespace-pre-wrap wrap-break-word",
+								COMPOSER_TEXT_METRICS_CLASSNAME,
+								disabled && "opacity-50 cursor-not-allowed",
+								className
+							)}
+						/>
+					</Plate>
+				</MentionEditorContext.Provider>
 			</div>
 		);
 	}
diff --git a/surfsense_web/components/assistant-ui/mention-chip.tsx b/surfsense_web/components/assistant-ui/mention-chip.tsx
index 9f9c9b177..7197fccdc 100644
--- a/surfsense_web/components/assistant-ui/mention-chip.tsx
+++ b/surfsense_web/components/assistant-ui/mention-chip.tsx
@@ -66,23 +66,21 @@ export function MentionChip({
 			disabled={disabled}
 			aria-label={ariaLabel ?? label}
 			className={cn(
-				"inline-flex max-w-[220px] items-center gap-1.5 rounded-md border bg-background px-2 py-0.5 align-middle text-xs font-medium text-foreground leading-5 transition-colors focus-visible:outline-none focus-visible:ring-1 focus-visible:ring-ring",
-				isInteractive
-					? "cursor-pointer hover:bg-accent hover:text-accent-foreground"
-					: "cursor-default",
+				"inline-flex h-5 items-center gap-1 rounded bg-primary/10 px-1 align-middle text-xs font-bold text-primary/60 leading-none focus-visible:outline-none focus-visible:ring-1 focus-visible:ring-ring",
+				isInteractive ? "cursor-pointer" : "cursor-default",
 				disabled && "opacity-60",
 				className
 			)}
 		>
 			<span className="inline-flex shrink-0 text-muted-foreground">{icon}</span>
-			<span className="truncate">{label}</span>
+			<span className="max-w-[120px] truncate leading-none">{label}</span>
 		</button>
 	);
 
 	if (!tooltip) return chip;
 
 	return (
-		<Tooltip>
+		<Tooltip delayDuration={600}>
 			<TooltipTrigger asChild>{chip}</TooltipTrigger>
 			<TooltipContent side="top" className="max-w-xs break-all">
 				{tooltip}