From 36bc50b1c0dcd9c3ce7f5a7431e2d8f09b9a8da0 Mon Sep 17 00:00:00 2001
From: Stephen Zhou <38493346+hyoban@users.noreply.github.com>
Date: Tue, 21 Oct 2025 12:27:39 +0800
Subject: [PATCH] feat: separation for mentions

---
 .../editor/plugins/mention/MentionPlugin.tsx  |  2 -
 .../mention/components/MentionDropdown.tsx    | 66 ++++++++++++-
 .../mention/hooks/dateMentionConfig.ts        |  2 +-
 .../plugins/mention/hooks/useMentionSearch.ts | 16 +---
 .../mention/hooks/useMentionSearchService.ts  | 13 ++-
 .../shared/components/TypeaheadDropdown.tsx   | 95 +++++++++++++++----
 locales/ai/en.json                            |  5 +
 locales/ai/ja.json                            |  5 +
 locales/ai/zh-CN.json                         |  5 +
 9 files changed, 166 insertions(+), 43 deletions(-)

diff --git a/apps/desktop/layer/renderer/src/modules/ai-chat/editor/plugins/mention/MentionPlugin.tsx b/apps/desktop/layer/renderer/src/modules/ai-chat/editor/plugins/mention/MentionPlugin.tsx
index 22b898c75..57f62b22d 100644
--- a/apps/desktop/layer/renderer/src/modules/ai-chat/editor/plugins/mention/MentionPlugin.tsx
+++ b/apps/desktop/layer/renderer/src/modules/ai-chat/editor/plugins/mention/MentionPlugin.tsx
@@ -2,7 +2,6 @@ import * as React from "react"
 import { Suspense, useMemo } from "react"
 
 import { MentionDropdown } from "./components/MentionDropdown"
-import { DEFAULT_MAX_SUGGESTIONS } from "./constants"
 import { useMentionKeyboard } from "./hooks/useMentionKeyboard"
 import { useMentionSearch } from "./hooks/useMentionSearch"
 import { useMentionSearchService } from "./hooks/useMentionSearchService"
@@ -31,7 +30,6 @@ export function MentionPlugin() {
     hasResults,
   } = useMentionSearch({
     onSearch: searchMentions,
-    maxSuggestions: DEFAULT_MAX_SUGGESTIONS,
   })
 
   // Hook for handling mention selection
diff --git a/apps/desktop/layer/renderer/src/modules/ai-chat/editor/plugins/mention/components/MentionDropdown.tsx b/apps/desktop/layer/renderer/src/modules/ai-chat/editor/plugins/mention/components/MentionDropdown.tsx
index 491aa87e9..ef37ca9e0 100644
--- a/apps/desktop/layer/renderer/src/modules/ai-chat/editor/plugins/mention/components/MentionDropdown.tsx
+++ b/apps/desktop/layer/renderer/src/modules/ai-chat/editor/plugins/mention/components/MentionDropdown.tsx
@@ -1,8 +1,9 @@
 import { cn, thenable } from "@follow/utils"
 import * as React from "react"
-import { useCallback } from "react"
+import { useCallback, useMemo } from "react"
 import { useTranslation } from "react-i18next"
 
+import type { TypeaheadGroup } from "../../shared/components/TypeaheadDropdown"
 import { TypeaheadDropdown } from "../../shared/components/TypeaheadDropdown"
 import { MENTION_TRIGGER_PATTERN } from "../constants"
 import { RANGE_WITH_LABEL_KEY } from "../hooks/dateMentionConfig"
@@ -110,6 +111,41 @@ const MentionSuggestionItem = React.memo(
 
 MentionSuggestionItem.displayName = "MentionSuggestionItem"
 
+const MentionGroupHeader = React.memo(({ type }: { type: MentionData["type"] }) => {
+  const { t } = useTranslation("ai")
+
+  const label = useMemo(() => {
+    switch (type) {
+      case "date": {
+        return t("mentions.section.date")
+      }
+      case "entry": {
+        return t("mentions.section.entry")
+      }
+      case "feed": {
+        return t("mentions.section.feed")
+      }
+      case "category": {
+        return t("mentions.section.category")
+      }
+      case "view": {
+        return t("mentions.section.view")
+      }
+      default: {
+        return ""
+      }
+    }
+  }, [type, t])
+
+  return (
+    <div className="text-text-tertiary mb-1 mt-2 px-2.5 text-xs font-medium first:mt-0">
+      {label}
+    </div>
+  )
+})
+
+MentionGroupHeader.displayName = "MentionGroupHeader"
+
 export const MentionDropdown: React.FC<MentionDropdownProps> = ({
   isVisible,
   suggestions,
@@ -125,10 +161,30 @@ export const MentionDropdown: React.FC<MentionDropdownProps> = ({
 }) => {
   if (!isVisible) throw thenable
 
+  // Group suggestions by type
+  const groupedSuggestions = useMemo<TypeaheadGroup<MentionData, MentionData["type"]>[]>(() => {
+    const groups: TypeaheadGroup<MentionData, MentionData["type"]>[] = []
+    let currentType: MentionData["type"] | null = null
+
+    suggestions.forEach((mention) => {
+      if (mention.type !== currentType) {
+        currentType = mention.type
+        groups.push({ key: currentType, items: [mention] })
+      } else {
+        const lastGroup = groups.at(-1)
+        if (lastGroup) {
+          lastGroup.items.push(mention)
+        }
+      }
+    })
+
+    return groups
+  }, [suggestions])
+
   return (
-    <TypeaheadDropdown
+    <TypeaheadDropdown<MentionData, MentionData["type"]>
       isVisible={isVisible}
-      items={suggestions}
+      items={groupedSuggestions}
       selectedIndex={selectedIndex}
       isLoading={isLoading}
       onSelect={onSelect}
@@ -136,10 +192,9 @@ export const MentionDropdown: React.FC<MentionDropdownProps> = ({
       onClose={onClose}
       query={query}
       ariaLabel="Mention suggestions"
-      getKey={(m) => `${m.type}-${m.id}`}
+      getKey={(mention) => `${mention.type}-${mention.id}`}
       renderItem={(mention, _index, isSelected, handlers) => (
         <MentionSuggestionItem
-          key={`${mention.type}-${mention.id}`}
           mention={mention}
           isSelected={isSelected}
           onMouseMove={handlers.onMouseMove}
@@ -147,6 +202,7 @@ export const MentionDropdown: React.FC<MentionDropdownProps> = ({
           query={query}
         />
       )}
+      renderGroupHeader={(groupKey) => <MentionGroupHeader type={groupKey} />}
       anchor={anchor}
       showSearchInput={showSearchInput}
       onQueryChange={onQueryChange}
diff --git a/apps/desktop/layer/renderer/src/modules/ai-chat/editor/plugins/mention/hooks/dateMentionConfig.ts b/apps/desktop/layer/renderer/src/modules/ai-chat/editor/plugins/mention/hooks/dateMentionConfig.ts
index b4ada16ba..a858bbfe8 100644
--- a/apps/desktop/layer/renderer/src/modules/ai-chat/editor/plugins/mention/hooks/dateMentionConfig.ts
+++ b/apps/desktop/layer/renderer/src/modules/ai-chat/editor/plugins/mention/hooks/dateMentionConfig.ts
@@ -2,7 +2,7 @@ import type { Dayjs } from "dayjs"
 
 import { clampRangeToPastMonth } from "./dateMentionUtils"
 
-export const MAX_INLINE_DATE_SUGGESTIONS = 3
+export const MAX_INLINE_DATE_SUGGESTIONS = 2
 
 export type DateRangeFactory = (today: Dayjs) => ReturnType<typeof clampRangeToPastMonth>
 
diff --git a/apps/desktop/layer/renderer/src/modules/ai-chat/editor/plugins/mention/hooks/useMentionSearch.ts b/apps/desktop/layer/renderer/src/modules/ai-chat/editor/plugins/mention/hooks/useMentionSearch.ts
index 685bc3125..7472d63b3 100644
--- a/apps/desktop/layer/renderer/src/modules/ai-chat/editor/plugins/mention/hooks/useMentionSearch.ts
+++ b/apps/desktop/layer/renderer/src/modules/ai-chat/editor/plugins/mention/hooks/useMentionSearch.ts
@@ -1,6 +1,5 @@
 import { useCallback, useRef, useState, useTransition } from "react"
 
-import { DEFAULT_MAX_SUGGESTIONS } from "../constants"
 import type { MentionData, MentionSearchState, MentionType } from "../types"
 import { getMentionType, shouldTriggerMention } from "../utils/triggerDetection"
 
@@ -8,17 +7,14 @@ interface UseMentionSearchOptions {
   onSearch?: (
     query: string,
     type: MentionType | undefined,
+    maxSuggestions?: number,
   ) => Promise<MentionData[]> | MentionData[]
-  maxSuggestions?: number
 }
 
 // Default search function
 const defaultSearchFn = async (): Promise<MentionData[]> => []
 
-export const useMentionSearch = ({
-  onSearch = defaultSearchFn,
-  maxSuggestions = DEFAULT_MAX_SUGGESTIONS,
-}: UseMentionSearchOptions = {}) => {
+export const useMentionSearch = ({ onSearch = defaultSearchFn }: UseMentionSearchOptions = {}) => {
   const [searchState, setSearchState] = useState<MentionSearchState>({
     suggestions: [],
     selectedIndex: -1,
@@ -27,15 +23,13 @@ export const useMentionSearch = ({
 
   const [isPending, startTransition] = useTransition()
   const onSearchRef = useRef(onSearch)
-  const maxSuggestionsRef = useRef(maxSuggestions)
   const abortControllerRef = useRef<AbortController | null>(null)
 
   // Update refs when props change to avoid stale closures
   onSearchRef.current = onSearch
-  maxSuggestionsRef.current = maxSuggestions
 
   const searchMentions = useCallback(
-    async (query: string) => {
+    async (query: string, maxSuggestions?: number) => {
       // Cancel any pending search
       if (abortControllerRef.current) {
         abortControllerRef.current.abort()
@@ -63,7 +57,7 @@ export const useMentionSearch = ({
       try {
         const [mentionType, cleanQuery] = getMentionType(query)
 
-        const results = await onSearchRef.current(cleanQuery, mentionType)
+        const results = await onSearchRef.current(cleanQuery, mentionType, maxSuggestions)
 
         // Check if this search was aborted
         if (abortController.signal.aborted) {
@@ -72,7 +66,7 @@ export const useMentionSearch = ({
 
         startTransition(() => {
           setSearchState({
-            suggestions: results.slice(0, maxSuggestionsRef.current),
+            suggestions: results,
             selectedIndex: results.length > 0 ? 0 : -1,
             isLoading: false,
           })
diff --git a/apps/desktop/layer/renderer/src/modules/ai-chat/editor/plugins/mention/hooks/useMentionSearchService.ts b/apps/desktop/layer/renderer/src/modules/ai-chat/editor/plugins/mention/hooks/useMentionSearchService.ts
index 085dabd67..a91c25c9f 100644
--- a/apps/desktop/layer/renderer/src/modules/ai-chat/editor/plugins/mention/hooks/useMentionSearchService.ts
+++ b/apps/desktop/layer/renderer/src/modules/ai-chat/editor/plugins/mention/hooks/useMentionSearchService.ts
@@ -19,7 +19,11 @@ export const useMentionSearchService = () => {
   const buildDateMentions = useMemo(() => createDateMentionBuilder({ t, language }), [t, language])
 
   const searchMentions = useMemo(() => {
-    return async (query: string, type?: MentionType): Promise<MentionData[]> => {
+    return async (
+      query: string,
+      type?: MentionType,
+      maxSuggestions = 10,
+    ): Promise<MentionData[]> => {
       const trimmedQuery = query.trim()
       const results: MentionData[] = []
       const seen = new Set<string>()
@@ -38,7 +42,7 @@ export const useMentionSearchService = () => {
       }
 
       if (type === "feed" || type === "entry" || type === "category") {
-        const searchResults = search(trimmedQuery, type, 10)
+        const searchResults = search(trimmedQuery, type, maxSuggestions)
         searchResults.forEach((item) =>
           pushResult({
             id: item.id,
@@ -71,7 +75,10 @@ export const useMentionSearchService = () => {
       const dateSuggestions = buildDateMentions(trimmedQuery)
       dateSuggestions.slice(0, MAX_INLINE_DATE_SUGGESTIONS).forEach(pushResult)
 
-      const searchResults = search(trimmedQuery, undefined, 10)
+      // Calculate remaining slots for search results
+      const remainingSlots = Math.max(0, maxSuggestions - results.length)
+
+      const searchResults = search(trimmedQuery, undefined, remainingSlots)
       searchResults.forEach((item) =>
         pushResult({
           id: item.id,
diff --git a/apps/desktop/layer/renderer/src/modules/ai-chat/editor/plugins/shared/components/TypeaheadDropdown.tsx b/apps/desktop/layer/renderer/src/modules/ai-chat/editor/plugins/shared/components/TypeaheadDropdown.tsx
index 0a01a65fa..22aa89d57 100644
--- a/apps/desktop/layer/renderer/src/modules/ai-chat/editor/plugins/shared/components/TypeaheadDropdown.tsx
+++ b/apps/desktop/layer/renderer/src/modules/ai-chat/editor/plugins/shared/components/TypeaheadDropdown.tsx
@@ -16,9 +16,14 @@ import { useEffect, useMemo, useRef, useState } from "react"
 
 import { calculateDropdownPosition } from "../utils/positioning"
 
-export interface TypeaheadDropdownProps<TItem> {
-  isVisible: boolean
+export interface TypeaheadGroup<TItem, TGroupKey = string> {
+  key: TGroupKey
   items: TItem[]
+}
+
+export interface TypeaheadDropdownProps<TItem, TGroupKey = string> {
+  isVisible: boolean
+  items: TItem[] | TypeaheadGroup<TItem, TGroupKey>[]
   selectedIndex: number
   isLoading: boolean
   onSelect: (item: TItem) => void
@@ -39,6 +44,8 @@ export interface TypeaheadDropdownProps<TItem> {
   anchor?: HTMLElement | null
   showSearchInput?: boolean
   onQueryChange?: (query: string) => void
+  // Group support
+  renderGroupHeader?: (groupKey: TGroupKey) => React.ReactNode
 }
 
 function useOptionalLexicalEditor() {
@@ -50,7 +57,7 @@ function useOptionalLexicalEditor() {
   }
 }
 
-export function TypeaheadDropdown<TItem>({
+export function TypeaheadDropdown<TItem, TGroupKey = string>({
   isVisible,
   items,
   selectedIndex,
@@ -68,13 +75,24 @@ export function TypeaheadDropdown<TItem>({
   anchor,
   showSearchInput = false,
   onQueryChange,
-}: TypeaheadDropdownProps<TItem>) {
+  renderGroupHeader,
+}: TypeaheadDropdownProps<TItem, TGroupKey>) {
   if (!isVisible) throw thenable
 
   const editor = useOptionalLexicalEditor()
   const dropdownRef = useRef<HTMLDivElement>(null)
   const [referenceWidth, setReferenceWidth] = useState<number>(320)
 
+  // Check if items are grouped
+  const isGrouped =
+    items.length > 0 && typeof items[0] === "object" && items[0] !== null && "key" in items[0]
+
+  // Flatten grouped items for selection logic
+  const flatItems = useMemo(() => {
+    if (!isGrouped) return items as TItem[]
+    return (items as TypeaheadGroup<TItem, TGroupKey>[]).flatMap((group) => group.items)
+  }, [items, isGrouped])
+
   const virtualReference = useRef({
     getBoundingClientRect: () => {
       // If anchor is provided, use it
@@ -192,7 +210,8 @@ export function TypeaheadDropdown<TItem>({
       )
     }
 
-    if (items.length === 0) {
+    const totalItems = isGrouped ? flatItems.length : items.length
+    if (totalItems === 0) {
       return (
         <div className="text-text-tertiary px-2.5 py-1.5 text-center">
           <span className="text-sm">{emptyMessage}</span>
@@ -201,20 +220,49 @@ export function TypeaheadDropdown<TItem>({
       )
     }
 
+    if (!isGrouped) {
+      // Render flat list
+      return (
+        <div role="listbox" aria-label={ariaLabel}>
+          {(items as TItem[]).map((item, index) => {
+            const isSelected = index === selectedIndex
+            const handlers = {
+              onMouseMove: () => onSetSelectIndex(index),
+              onClick: () => onSelect(item),
+            }
+            return (
+              <React.Fragment key={getKey(item)}>
+                {renderItem(item, index, isSelected, handlers)}
+              </React.Fragment>
+            )
+          })}
+        </div>
+      )
+    }
+
+    // Render grouped list
+    let itemIndex = 0
     return (
       <div role="listbox" aria-label={ariaLabel}>
-        {items.map((item, index) => {
-          const isSelected = index === selectedIndex
-          const handlers = {
-            onMouseMove: () => onSetSelectIndex(index),
-            onClick: () => onSelect(item),
-          }
-          return (
-            <React.Fragment key={getKey(item)}>
-              {renderItem(item, index, isSelected, handlers)}
-            </React.Fragment>
-          )
-        })}
+        {(items as TypeaheadGroup<TItem, TGroupKey>[]).map((group) => (
+          <React.Fragment key={String(group.key)}>
+            {renderGroupHeader && renderGroupHeader(group.key)}
+            {group.items.map((item) => {
+              const currentIndex = itemIndex
+              itemIndex++
+              const isSelected = currentIndex === selectedIndex
+              const handlers = {
+                onMouseMove: () => onSetSelectIndex(currentIndex),
+                onClick: () => onSelect(item),
+              }
+              return (
+                <React.Fragment key={getKey(item)}>
+                  {renderItem(item, currentIndex, isSelected, handlers)}
+                </React.Fragment>
+              )
+            })}
+          </React.Fragment>
+        ))}
       </div>
     )
   }, [
@@ -230,6 +278,9 @@ export function TypeaheadDropdown<TItem>({
     emptyHint,
     onSetSelectIndex,
     onSelect,
+    isGrouped,
+    flatItems,
+    renderGroupHeader,
   ])
 
   return (
@@ -260,10 +311,12 @@ export function TypeaheadDropdown<TItem>({
                   value={query}
                   onChange={(e) => onQueryChange(e.target.value)}
                   onKeyDown={(e) => {
-                    const suggestion = items[selectedIndex] || items[0]
-                    if (e.key === "Enter" && suggestion) {
-                      e.preventDefault()
-                      onSelect(suggestion)
+                    if (e.key === "Enter") {
+                      const suggestion = flatItems[selectedIndex] || flatItems[0]
+                      if (suggestion) {
+                        e.preventDefault()
+                        onSelect(suggestion)
+                      }
                     }
                   }}
                   placeholder="Search for context..."
diff --git a/locales/ai/en.json b/locales/ai/en.json
index 4024b3b63..5bce59634 100644
--- a/locales/ai/en.json
+++ b/locales/ai/en.json
@@ -121,6 +121,11 @@
   "mentions.date.weekday.prefix.last.search": "last|last week",
   "mentions.date.weekday.prefix.this.search": "this|this week",
   "mentions.date.weekday.this.label": "This {{weekday}}",
+  "mentions.section.category": "Categories",
+  "mentions.section.date": "Dates",
+  "mentions.section.entry": "Entries",
+  "mentions.section.feed": "Feeds",
+  "mentions.section.view": "Views",
   "personalize.description": "Tell me about yourself to get personalized AI responses.",
   "personalize.prompt.help": "This helps AI provide personalized responses based on your preferences.",
   "personalize.prompt.placeholder": "Tell me about yourself and how you prefer to read content...",
diff --git a/locales/ai/ja.json b/locales/ai/ja.json
index 56df9744f..a24b8bf00 100644
--- a/locales/ai/ja.json
+++ b/locales/ai/ja.json
@@ -115,6 +115,11 @@
   "mentions.date.weekday.prefix.last.search": "先週|せんしゅう",
   "mentions.date.weekday.prefix.this.search": "今週|こんしゅう",
   "mentions.date.weekday.this.label": "今週{{weekday}}",
+  "mentions.section.category": "カテゴリー",
+  "mentions.section.date": "日付",
+  "mentions.section.entry": "エントリー",
+  "mentions.section.feed": "フィード",
+  "mentions.section.view": "ビュー",
   "personalize.description": "あなた自身について教えてください。そうすれば、パーソナライズされたAIの応答が得られます。",
   "personalize.prompt.help": "これにより、AIはあなたの好みに基づいてパーソナライズされた応答を提供できます。",
   "personalize.prompt.placeholder": "あなた自身について、どのようにコンテンツを読みたいかを教えてください...",
diff --git a/locales/ai/zh-CN.json b/locales/ai/zh-CN.json
index a8c7afa8d..de9569e7a 100644
--- a/locales/ai/zh-CN.json
+++ b/locales/ai/zh-CN.json
@@ -117,6 +117,11 @@
   "mentions.date.weekday.prefix.last.search": "上周|上一周|上星期|上个星期|上礼拜|上个礼拜",
   "mentions.date.weekday.prefix.this.search": "这周|本周|这星期|本星期|这礼拜|本礼拜",
   "mentions.date.weekday.this.label": "本周{{weekday}}",
+  "mentions.section.category": "分类",
+  "mentions.section.date": "日期",
+  "mentions.section.entry": "条目",
+  "mentions.section.feed": "订阅源",
+  "mentions.section.view": "视图",
   "personalize.description": "介绍一下您自己，以获得个性化的AI响应。",
   "personalize.prompt.help": "这有助于AI根据您的偏好提供个性化响应。",
   "personalize.prompt.placeholder": "介绍一下您自己以及您希望如何阅读内容...",