Merge remote-tracking branch 'upstream/dev' into feat/ui-revamp

surfsense_web/components/assistant-ui/thread.tsx

@@ -12,6 +12,7 @@ import {
 	AlertCircle,
 	ArrowDownIcon,
 	ArrowUpIcon,
+	Camera,
 	ChevronDown,
 	ChevronUp,
 	Clipboard,
@@ -39,6 +40,7 @@ import { chatSessionStateAtom } from "@/atoms/chat/chat-session-state.atom";
 import {
 	mentionedDocumentsAtom,
 } from "@/atoms/chat/mentioned-documents.atom";
+import { pendingUserImageDataUrlsAtom } from "@/atoms/chat/pending-user-images.atom";
 import { connectorDialogOpenAtom } from "@/atoms/connector-dialog/connector-dialog.atoms";
 import { connectorsAtom } from "@/atoms/connectors/connector-query.atoms";
 import { membersAtom } from "@/atoms/members/members-query.atoms";
@@ -87,6 +89,8 @@ import { useBatchCommentsPreload } from "@/hooks/use-comments";
 import { useCommentsSync } from "@/hooks/use-comments-sync";
 import { useMediaQuery } from "@/hooks/use-media-query";
 import { useElectronAPI } from "@/hooks/use-platform";
+import { getMentionDocKey } from "@/lib/chat/mention-doc-key";
+import { captureDisplayToPngDataUrl } from "@/lib/chat/display-media-capture";
 import { SLIDEOUT_PANEL_OPENED_EVENT } from "@/lib/layout-events";
 import { cn } from "@/lib/utils";
 
@@ -293,6 +297,32 @@ const ConnectToolsBanner: FC<{ isThreadEmpty: boolean }> = ({ isThreadEmpty }) =
 	);
 };
 
+const PendingScreenImageStrip: FC = () => {
+	const [urls, setUrls] = useAtom(pendingUserImageDataUrlsAtom);
+	if (urls.length === 0) return null;
+	return (
+		<div className="mx-3 mt-2 flex flex-wrap gap-2">
+			{urls.map((url, index) => (
+				<div
+					key={url}
+					className="group relative h-14 w-14 shrink-0 overflow-hidden rounded-md border border-border/50 bg-muted"
+				>
+					{/* biome-ignore lint/performance/noImgElement: data URL thumbnails from capture */}
+					<img src={url} alt="" className="size-full object-cover" draggable={false} />
+					<button
+						type="button"
+						onClick={() => setUrls((prev) => prev.filter((_, i) => i !== index))}
+						className="absolute right-0.5 top-0.5 flex size-5 items-center justify-center rounded-full bg-background/90 text-muted-foreground shadow-sm transition-opacity hover:text-foreground sm:opacity-0 sm:group-hover:opacity-100"
+						aria-label="Remove screenshot"
+					>
+						<X className="size-3" />
+					</button>
+				</div>
+			))}
+		</div>
+	);
+};
+
 const ClipboardChip: FC<{ text: string; onDismiss: () => void }> = ({ text, onDismiss }) => {
 	const [expanded, setExpanded] = useState(false);
 	const isLong = text.length > 120;
@@ -338,6 +368,9 @@ const Composer: FC = () => {
 	const [mentionQuery, setMentionQuery] = useState("");
 	const [actionQuery, setActionQuery] = useState("");
 	const editorRef = useRef<InlineMentionEditorRef>(null);
+	const prevMentionedDocsRef = useRef<
+		Map<string, Pick<Document, "id" | "title" | "document_type">>
+	>(new Map());
 	const documentPickerRef = useRef<DocumentMentionPickerRef>(null);
 	const promptPickerRef = useRef<PromptPickerRef>(null);
 	const viewportRef = useRef<Element | null>(null);
@@ -624,60 +657,64 @@ const Composer: FC = () => {
 
 	const handleDocumentRemove = useCallback(
 		(docId: number, docType?: string) => {
-			setMentionedDocuments((prev) =>
-				prev.filter((doc) => !(doc.id === docId && doc.document_type === docType))
-			);
+			setMentionedDocuments((prev) => {
+				if (!docType) {
+					// Defensive fallback: keep UI in sync even when chip type is unavailable.
+					return prev.filter((doc) => doc.id !== docId);
+				}
+				const removedKey = getMentionDocKey({ id: docId, document_type: docType });
+				return prev.filter((doc) => getMentionDocKey(doc) !== removedKey);
+			});
 		},
 		[setMentionedDocuments]
 	);
 
 	const handleDocumentsMention = useCallback(
 		(documents: Pick<Document, "id" | "title" | "document_type">[]) => {
-			const existingKeys = new Set(mentionedDocuments.map((d) => `${d.document_type}:${d.id}`));
-			const newDocs = documents.filter(
-				(doc) => !existingKeys.has(`${doc.document_type}:${doc.id}`)
-			);
+			const editorMentionedDocs = editorRef.current?.getMentionedDocuments() ?? [];
+			const editorDocKeys = new Set(editorMentionedDocs.map((doc) => getMentionDocKey(doc)));
 
-			for (const doc of newDocs) {
+			for (const doc of documents) {
+				const key = getMentionDocKey(doc);
+				if (editorDocKeys.has(key)) continue;
 				editorRef.current?.insertDocumentChip(doc);
 			}
 
 			setMentionedDocuments((prev) => {
-				const existingKeySet = new Set(prev.map((d) => `${d.document_type}:${d.id}`));
-				const uniqueNewDocs = documents.filter(
-					(doc) => !existingKeySet.has(`${doc.document_type}:${doc.id}`)
-				);
+				const existingKeySet = new Set(prev.map((d) => getMentionDocKey(d)));
+				const uniqueNewDocs = documents.filter((doc) => !existingKeySet.has(getMentionDocKey(doc)));
 				return [...prev, ...uniqueNewDocs];
 			});
 
 			setMentionQuery("");
 		},
-		[mentionedDocuments, setMentionedDocuments]
+		[setMentionedDocuments]
 	);
 
 	useEffect(() => {
 		const editor = editorRef.current;
-		if (!editor) return;
+		const nextDocsMap = new Map(mentionedDocuments.map((doc) => [getMentionDocKey(doc), doc]));
+		const prevDocsMap = prevMentionedDocsRef.current;
 
-		const toKey = (doc: { id: number; document_type?: string }) =>
-			`${doc.document_type ?? "UNKNOWN"}:${doc.id}`;
-
-		const atomDocs = mentionedDocuments;
-		const editorDocs = editor.getMentionedDocuments();
-		const atomKeys = new Set(atomDocs.map(toKey));
-		const editorKeys = new Set(editorDocs.map(toKey));
-
-		for (const doc of atomDocs) {
-			if (!editorKeys.has(toKey(doc))) {
-				editor.insertDocumentChip(doc, { removeTriggerText: false });
-			}
+		if (!editor) {
+			prevMentionedDocsRef.current = nextDocsMap;
+			return;
 		}
 
-		for (const doc of editorDocs) {
-			if (!atomKeys.has(toKey(doc))) {
+		const editorKeys = new Set(editor.getMentionedDocuments().map(getMentionDocKey));
+
+		for (const [key, doc] of nextDocsMap) {
+			if (prevDocsMap.has(key) || editorKeys.has(key)) continue;
+			editor.insertDocumentChip(doc, { removeTriggerText: false });
+		}
+
+		for (const [key, doc] of prevDocsMap) {
+			if (!nextDocsMap.has(key)) {
 				editor.removeDocumentChip(doc.id, doc.document_type);
 			}
 		}
+
+		prevMentionedDocsRef.current = nextDocsMap;
 	}, [mentionedDocuments]);
 
 	return (
@@ -722,6 +759,7 @@ const Composer: FC = () => {
 				</div>
 			)}
 			<div className="aui-composer-attachment-dropzone flex w-full flex-col overflow-hidden rounded-2xl border-input bg-muted pt-2 outline-none transition-shadow">
+				<PendingScreenImageStrip />
 				{clipboardInitialText && (
 					<ClipboardChip
 						text={clipboardInitialText}
@@ -779,11 +817,23 @@ const ComposerAction: FC<ComposerActionProps> = ({ isBlockedByOtherUser = false
 		},
 		[]
 	);
+	const pendingScreenImages = useAtomValue(pendingUserImageDataUrlsAtom);
+	const setPendingScreenImages = useSetAtom(pendingUserImageDataUrlsAtom);
+	const electronAPI = useElectronAPI();
+
 	const isComposerTextEmpty = useAuiState(({ composer }) => {
 		const text = composer.text?.trim() || "";
 		return text.length === 0;
 	});
-	const isComposerEmpty = isComposerTextEmpty && mentionedDocuments.length === 0;
+	const isComposerEmpty =
+		isComposerTextEmpty && mentionedDocuments.length === 0 && pendingScreenImages.length === 0;
+
+	const handleScreenCapture = useCallback(async () => {
+		const url = electronAPI?.captureFullScreen
+			? await electronAPI.captureFullScreen()
+			: await captureDisplayToPngDataUrl();
+		if (url) setPendingScreenImages((prev) => [...prev, url]);
+	}, [electronAPI, setPendingScreenImages]);
 
 	const { data: userConfigs } = useAtomValue(newLLMConfigsAtom);
 	const { data: globalConfigs } = useAtomValue(globalNewLLMConfigsAtom);
@@ -1210,6 +1260,17 @@ const ComposerAction: FC<ComposerActionProps> = ({ isBlockedByOtherUser = false
 				</div>
 			)}
 			<div className="flex items-center gap-2">
+				<TooltipIconButton
+					tooltip="Capture screen"
+					type="button"
+					variant="ghost"
+					size="icon"
+					className="size-8 rounded-full"
+					aria-label="Capture screen"
+					onClick={() => void handleScreenCapture()}
+				>
+					<Camera className="size-4" />
+				</TooltipIconButton>
 				<AuiIf condition={({ thread }) => !thread.isRunning}>
 					<ComposerPrimitive.Send asChild disabled={isSendDisabled}>
 						<TooltipIconButton
@@ -1219,7 +1280,7 @@ const ComposerAction: FC<ComposerActionProps> = ({ isBlockedByOtherUser = false
 									: !hasModelConfigured
 										? "Please select a model from the header to start chatting"
 										: isComposerEmpty
-											? "Enter a message to send"
+											? "Enter a message or add a screenshot to send"
 											: "Send message"
 							}
 							side="bottom"