Skip to content

fix(agent): optimize file mention search #1283

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ErlichLiu wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(agent): optimize file mention search #1283

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
188 changes: 4 additions & 184 deletions apps/electron/src/main/ipc.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,7 +64,6 @@ import type {
WorkspaceCapabilities,
WorkspaceMemorySummary,
FileEntry,
FileSearchResult,
EnvironmentCheckResult,
InstallerManifest,
InstallerDownloadRequest,
Expand Down Expand Up @@ -259,6 +258,7 @@ import {
getReleaseByTag,
} from './lib/github-release-service'
import { watchAttachedDirectory, unwatchAttachedDirectory } from './lib/workspace-watcher'
import { searchWorkspaceFiles } from './lib/workspace-file-search'
import {
getFeishuConfig,
saveFeishuConfig,
Expand Down Expand Up @@ -3390,191 +3390,11 @@ export function registerIpcHandlers(): void {
}
)

// 搜索工作区文件（用于 @ 引用，递归扫描，支持附加目录）
// 搜索工作区文件（用于 @ 引用，使用缓存索引避免输入期间重复扫描）
ipcMain.handle(
AGENT_IPC_CHANNELS.SEARCH_WORKSPACE_FILES,
async (_, rootPath: string, query: string, limit = 20, additionalPaths?: string[], sessionPaths?: string[]): Promise<FileSearchResult> => {
const { readdirSync, statSync } = await import('node:fs')
const { resolve, relative, basename } = await import('node:path')

const safeRoot = resolve(rootPath)
const ignoreDirs = new Set(['node_modules', '.git', 'dist', '.next', '__pycache__', '.venv', 'build', '.cache'])
const ignoreFiles = new Set(['.DS_Store', '.Spotlight-V100', '.Trashes', 'Thumbs.db', 'desktop.ini'])
const BROWSE_LIMIT_PER_GROUP = 2000
const BROWSE_TOTAL_CAP = 3000

// 按来源分组收集文件
type Entry = { name: string; path: string; type: 'file' | 'dir'; source: 'session' | 'workspace' }
const rootEntries: Entry[] = []
const workspaceEntries: Entry[] = []

function scan(
dir: string,
depth: number,
baseRoot: string,
target: Entry[],
useAbsPath: boolean,
source: 'session' | 'workspace',
): void {
if (depth > 10) return
try {
const items = readdirSync(dir, { withFileTypes: true })
for (const item of items) {
if (ignoreFiles.has(item.name)) continue
if (item.isDirectory() && ignoreDirs.has(item.name)) continue

const fullPath = resolve(dir, item.name)
const entryPath = useAbsPath ? fullPath : relative(baseRoot, fullPath)
target.push({
name: item.name,
path: entryPath,
type: item.isDirectory() ? 'dir' : 'file',
source,
})

if (item.isDirectory()) {
scan(fullPath, depth + 1, baseRoot, target, useAbsPath, source)
}
}
} catch {
// 忽略无权限的目录
}
}

function addAttachedPath(pathValue: string, target: Entry[], source: 'session' | 'workspace'): void {
try {
const attachedPath = resolve(pathValue)
const name = basename(attachedPath)
if (ignoreFiles.has(name)) return

const stats = statSync(attachedPath)
if (stats.isFile()) {
target.push({
name,
path: attachedPath,
type: 'file',
source,
})
return
}

if (!stats.isDirectory()) return
if (ignoreDirs.has(name)) return

target.push({
name: name === 'workspace-files' ? '工作文件' : name,
path: attachedPath,
type: 'dir',
source,
})
scan(attachedPath, 0, attachedPath, target, true, source)
} catch {
// 忽略不存在或无权限的附加路径
}
}

// session 目录：相对路径
scan(safeRoot, 0, safeRoot, rootEntries, false, 'session')

// 会话级附加路径：绝对路径，标记为 session（归入会话文件分组）
if (sessionPaths && sessionPaths.length > 0) {
for (const sp of sessionPaths) {
addAttachedPath(sp, rootEntries, 'session')
}
}

// 工作区文件 + 工作区级附加路径：绝对路径，标记为 workspace
if (additionalPaths && additionalPaths.length > 0) {
for (const addPath of additionalPaths) {
addAttachedPath(addPath, workspaceEntries, 'workspace')
}
}

// 组内排序：目录优先，前缀匹配优先，路径短优先
function sortGroup(entries: Entry[], q: string): void {
entries.sort((a, b) => {
const aStartsWith = a.name.toLowerCase().startsWith(q) ? 0 : 1
const bStartsWith = b.name.toLowerCase().startsWith(q) ? 0 : 1
if (aStartsWith !== bStartsWith) return aStartsWith - bStartsWith
if (a.type === 'dir' && b.type !== 'dir') return -1
if (a.type !== 'dir' && b.type === 'dir') return 1
return a.path.length - b.path.length
})
}

function matchEntries(entries: Entry[], q: string): Entry[] {
return entries.filter((entry) => {
const nameLower = entry.name.toLowerCase()
const pathLower = entry.path.toLowerCase()
if (nameLower.startsWith(q)) return true
if (nameLower.includes(q) || pathLower.includes(q)) return true
let qi = 0
for (let i = 0; i < nameLower.length && qi < q.length; i++) {
if (nameLower[i] === q[qi]) qi++
}
return qi === q.length
})
}

// 目录优先排序：确保截断前所有目录（特别是顶层目录）排在前面
function sortDirsFirst(entries: Entry[]): void {
entries.sort((a, b) => {
if (a.type === 'dir' && b.type !== 'dir') return -1
if (a.type !== 'dir' && b.type === 'dir') return 1
return a.path.length - b.path.length || a.name.localeCompare(b.name)
})
}

const q = query.toLowerCase()

if (!q) {
// 空 query：目录优先排序后再截断，保证文件夹结构完整可见
sortDirsFirst(rootEntries)
sortDirsFirst(workspaceEntries)
const maxPerGroup = Math.max(limit, BROWSE_LIMIT_PER_GROUP)
const sessionSlice = rootEntries.slice(0, maxPerGroup)
const workspaceSlice = workspaceEntries.slice(0, maxPerGroup)
const combined = [...sessionSlice, ...workspaceSlice]
const capped = combined.length > BROWSE_TOTAL_CAP ? combined.slice(0, BROWSE_TOTAL_CAP) : combined
return {
entries: capped,
total: rootEntries.length + workspaceEntries.length,
sessionEntries: sessionSlice,
workspaceEntries: workspaceSlice,
}
}

const sessionMatched = matchEntries(rootEntries, q)
const workspaceMatched = matchEntries(workspaceEntries, q)
sortGroup(sessionMatched, q)
sortGroup(workspaceMatched, q)

const totalMatched = sessionMatched.length + workspaceMatched.length
let sessionSlice: Entry[]
let workspaceSlice: Entry[]
if (totalMatched <= limit) {
sessionSlice = sessionMatched
workspaceSlice = workspaceMatched
} else {
const sessionQuota = Math.max(
sessionMatched.length > 0 ? 1 : 0,
Math.round(limit * sessionMatched.length / totalMatched),
)
const workspaceQuota = Math.max(
workspaceMatched.length > 0 ? 1 : 0,
limit - sessionQuota,
)
sessionSlice = sessionMatched.slice(0, sessionQuota)
workspaceSlice = workspaceMatched.slice(0, workspaceQuota)
}

return {
entries: [...sessionSlice, ...workspaceSlice],
total: sessionMatched.length + workspaceMatched.length,
sessionEntries: sessionSlice,
workspaceEntries: workspaceSlice,
}
}
(_, rootPath: string, query: string, limit = 20, additionalPaths?: string[], sessionPaths?: string[]) =>
searchWorkspaceFiles(rootPath, query, limit, additionalPaths, sessionPaths),
)

// ===== 系统提示词管理 =====
Expand Down
85 changes: 85 additions & 0 deletions apps/electron/src/main/lib/workspace-file-search.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,85 @@
import { afterEach, describe, expect, test } from 'bun:test'
import { mkdirSync, mkdtempSync, rmSync, writeFileSync } from 'node:fs'
import { tmpdir } from 'node:os'
import { join } from 'node:path'
import {
clearWorkspaceFileSearchCacheForTest,
invalidateWorkspaceFileSearchCache,
searchWorkspaceFiles,
} from './workspace-file-search'

const testRoots: string[] = []

function createRoot(): string {
const root = mkdtempSync(join(tmpdir(), 'proma-file-mention-search-'))
testRoots.push(root)
return root
}

afterEach(() => {
clearWorkspaceFileSearchCacheForTest()
for (const root of testRoots.splice(0)) {
rmSync(root, { recursive: true, force: true })
}
})

describe('workspace file search', () => {
test('Given a nested matching file When searching Then returns its ancestor directories for tree rendering', async () => {
const root = createRoot()
mkdirSync(join(root, 'src', 'components'), { recursive: true })
writeFileSync(join(root, 'src', 'components', 'MentionPicker.tsx'), '')

const result = await searchWorkspaceFiles(root, 'picker')

expect(result.entries.map((entry) => entry.path)).toEqual(['src/components/MentionPicker.tsx'])
expect(result.sessionEntries.map((entry) => entry.path)).toEqual([
'src',
'src/components',
'src/components/MentionPicker.tsx',
])
})

test('Given multiple matches When searching Then flat results retain relevance order without tree ancestors', async () => {
const root = createRoot()
mkdirSync(join(root, 'alpha'), { recursive: true })
writeFileSync(join(root, 'z-target.md'), '')
writeFileSync(join(root, 'target-top.md'), '')
writeFileSync(join(root, 'alpha', 'target.md'), '')

const result = await searchWorkspaceFiles(root, 'target')

expect(result.entries.map((entry) => entry.path)).toEqual([
'target-top.md',
'alpha/target.md',
'z-target.md',
])
expect(result.sessionEntries.map((entry) => entry.path)).toContain('alpha')
})

test('Given an attached directory When a query only matches its absolute parent path Then it returns no unrelated files', async () => {
const root = createRoot()
const attachedDirectory = join(root, 'attached-project')
mkdirSync(attachedDirectory, { recursive: true })
writeFileSync(join(attachedDirectory, 'unrelated.sql'), '')

const result = await searchWorkspaceFiles(root, 'proma-file-mention-search', 80, [attachedDirectory])

expect(result.workspaceEntries).toEqual([])
})

test('Given a cached index When files change and the watcher invalidates it Then the next search includes the new file', async () => {
const root = createRoot()
writeFileSync(join(root, 'existing.md'), '')

await searchWorkspaceFiles(root, 'existing')
writeFileSync(join(root, 'new-file.md'), '')

expect((await searchWorkspaceFiles(root, 'new-file')).sessionEntries).toEqual([])

invalidateWorkspaceFileSearchCache()

expect((await searchWorkspaceFiles(root, 'new-file')).sessionEntries.map((entry) => entry.name)).toEqual([
'new-file.md',
])
})
})
Loading