feat: improve direct model workspace tools
Direct models now use bundled ripgrep for workspace discovery, read files through resumable pages, and edit multiple files with apply_patch. Ask can search and read while Execute retains write access, and fixed total input limits no longer constrain these tools. Release note: 直连模型现可使用随包 ripgrep 搜索项目、分页读取任意大小的文本，并通过多文件补丁修改代码；设置中的相关能力统一显示为工作区文件和开发工具。
This commit is contained in:
mesalogo
@@ -0,0 +1,338 @@
|
||||
import { randomUUID } from 'node:crypto'
|
||||
import {
|
||||
lstat,
|
||||
mkdir,
|
||||
open,
|
||||
readFile,
|
||||
realpath,
|
||||
rename,
|
||||
rm,
|
||||
stat
|
||||
} from 'node:fs/promises'
|
||||
import { dirname, isAbsolute, join, relative, sep, win32 } from 'node:path'
|
||||
import type { WorkspaceAccess } from '../workspace'
|
||||
|
||||
type PatchOperation =
|
||||
| { kind: 'add'; path: string; content: string }
|
||||
| { kind: 'update'; path: string; hunks: PatchHunk[] }
|
||||
| { kind: 'delete'; path: string }
|
||||
|
||||
type PatchHunk = {
|
||||
section?: string
|
||||
lines: Array<{ kind: 'context' | 'add' | 'delete'; text: string }>
|
||||
}
|
||||
|
||||
type PreparedOperation =
|
||||
| { kind: 'write'; path: string; target: string; content: string; mode: number }
|
||||
| { kind: 'delete'; path: string; target: string }
|
||||
|
||||
const MAXIMUM_RESULT_BYTES = 200 * 1024
|
||||
|
||||
function normalizedPatchPath(value: string): string {
|
||||
const path = value.trim().replaceAll('\\', '/')
|
||||
if (
|
||||
!path ||
|
||||
isAbsolute(path) ||
|
||||
win32.isAbsolute(path) ||
|
||||
path.includes('\0')
|
||||
) {
|
||||
throw new Error(`补丁路径无效：${value}`)
|
||||
}
|
||||
const segments = path.split('/')
|
||||
if (segments.some((segment) => !segment || segment === '.' || segment === '..')) {
|
||||
throw new Error(`补丁路径不能超出工作区：${value}`)
|
||||
}
|
||||
return segments.join('/')
|
||||
}
|
||||
|
||||
function parseUpdateHunks(lines: string[], path: string): PatchHunk[] {
|
||||
const hunks: PatchHunk[] = []
|
||||
let current: PatchHunk | undefined
|
||||
for (const line of lines) {
|
||||
if (line === '@@' || line.startsWith('@@ ')) {
|
||||
current = {
|
||||
...(line.length > 3 ? { section: line.slice(3) } : {}),
|
||||
lines: []
|
||||
}
|
||||
hunks.push(current)
|
||||
continue
|
||||
}
|
||||
if (!current) {
|
||||
if (!line) continue
|
||||
throw new Error(`更新文件 ${path} 的补丁缺少 @@ hunk`)
|
||||
}
|
||||
const prefix = line[0]
|
||||
if (prefix !== ' ' && prefix !== '+' && prefix !== '-') {
|
||||
throw new Error(`更新文件 ${path} 包含无效补丁行`)
|
||||
}
|
||||
current.lines.push({
|
||||
kind: prefix === '+' ? 'add' : prefix === '-' ? 'delete' : 'context',
|
||||
text: line.slice(1)
|
||||
})
|
||||
}
|
||||
if (hunks.length === 0 || hunks.some((hunk) => hunk.lines.length === 0)) {
|
||||
throw new Error(`更新文件 ${path} 的补丁 hunk 为空`)
|
||||
}
|
||||
return hunks
|
||||
}
|
||||
|
||||
export function parseWorkspacePatch(patch: string): PatchOperation[] {
|
||||
const lines = patch.replaceAll('\r\n', '\n').split('\n')
|
||||
if (lines.shift() !== '*** Begin Patch') {
|
||||
throw new Error('补丁必须以 *** Begin Patch 开始')
|
||||
}
|
||||
const operations: PatchOperation[] = []
|
||||
while (lines.length > 0) {
|
||||
const header = lines.shift()!
|
||||
if (header === '*** End Patch') {
|
||||
if (lines.some(Boolean)) {
|
||||
throw new Error('*** End Patch 后不能包含其他内容')
|
||||
}
|
||||
if (operations.length === 0) {
|
||||
throw new Error('补丁不包含文件操作')
|
||||
}
|
||||
return operations
|
||||
}
|
||||
const match = header.match(/^\*\*\* (Add|Update|Delete) File: (.+)$/u)
|
||||
if (!match) {
|
||||
throw new Error(`无效补丁文件头：${header}`)
|
||||
}
|
||||
const path = normalizedPatchPath(match[2]!)
|
||||
const body: string[] = []
|
||||
while (lines.length > 0 && !lines[0]!.startsWith('*** ')) {
|
||||
body.push(lines.shift()!)
|
||||
}
|
||||
if (match[1] === 'Add') {
|
||||
if (body.some((line) => !line.startsWith('+'))) {
|
||||
throw new Error(`新增文件 ${path} 的每一行都必须以 + 开始`)
|
||||
}
|
||||
operations.push({
|
||||
kind: 'add',
|
||||
path,
|
||||
content: body.map((line) => line.slice(1)).join('\n')
|
||||
})
|
||||
} else if (match[1] === 'Update') {
|
||||
operations.push({ kind: 'update', path, hunks: parseUpdateHunks(body, path) })
|
||||
} else {
|
||||
if (body.some(Boolean)) {
|
||||
throw new Error(`删除文件 ${path} 不能包含补丁正文`)
|
||||
}
|
||||
operations.push({ kind: 'delete', path })
|
||||
}
|
||||
}
|
||||
throw new Error('补丁缺少 *** End Patch')
|
||||
}
|
||||
|
||||
function sequenceMatches(
|
||||
source: string[],
|
||||
expected: string[],
|
||||
index: number,
|
||||
relaxed: boolean
|
||||
): boolean {
|
||||
return expected.every((line, offset) =>
|
||||
relaxed
|
||||
? source[index + offset]?.trimEnd() === line.trimEnd()
|
||||
: source[index + offset] === line
|
||||
)
|
||||
}
|
||||
|
||||
function findSequence(
|
||||
source: string[],
|
||||
expected: string[],
|
||||
start: number
|
||||
): number {
|
||||
if (expected.length === 0) return start
|
||||
for (const relaxed of [false, true]) {
|
||||
for (let index = start; index <= source.length - expected.length; index += 1) {
|
||||
if (sequenceMatches(source, expected, index, relaxed)) return index
|
||||
}
|
||||
}
|
||||
return -1
|
||||
}
|
||||
|
||||
export function applyUpdateHunks(
|
||||
content: string,
|
||||
hunks: PatchHunk[],
|
||||
path: string
|
||||
): string {
|
||||
const newline = content.includes('\r\n') ? '\r\n' : '\n'
|
||||
let source = content.replaceAll('\r\n', '\n').split('\n')
|
||||
let cursor = 0
|
||||
for (const hunk of hunks) {
|
||||
if (hunk.section) {
|
||||
const sectionIndex = source.findIndex(
|
||||
(line, index) => index >= cursor && line.includes(hunk.section!)
|
||||
)
|
||||
if (sectionIndex < 0) {
|
||||
throw new Error(`无法在 ${path} 中找到补丁区段：${hunk.section}`)
|
||||
}
|
||||
cursor = sectionIndex
|
||||
}
|
||||
const expected = hunk.lines
|
||||
.filter((line) => line.kind !== 'add')
|
||||
.map((line) => line.text)
|
||||
const replacement = hunk.lines
|
||||
.filter((line) => line.kind !== 'delete')
|
||||
.map((line) => line.text)
|
||||
const index = findSequence(source, expected, cursor)
|
||||
if (index < 0) {
|
||||
throw new Error(`无法在 ${path} 中匹配补丁上下文`)
|
||||
}
|
||||
source = [
|
||||
...source.slice(0, index),
|
||||
...replacement,
|
||||
...source.slice(index + expected.length)
|
||||
]
|
||||
cursor = index + replacement.length
|
||||
}
|
||||
return source.join(newline)
|
||||
}
|
||||
|
||||
function isInside(root: string, target: string): boolean {
|
||||
const value = relative(root, target)
|
||||
return value === '' || (!value.startsWith(`..${sep}`) && value !== '..' && !isAbsolute(value))
|
||||
}
|
||||
|
||||
async function resolveTarget(
|
||||
root: string,
|
||||
patchPath: string,
|
||||
createParents: boolean
|
||||
): Promise<string> {
|
||||
const segments = patchPath.split('/')
|
||||
let current = root
|
||||
for (const [index, segment] of segments.slice(0, -1).entries()) {
|
||||
current = join(current, segment)
|
||||
const metadata = await lstat(current).catch((error: NodeJS.ErrnoException) => {
|
||||
if (error.code === 'ENOENT') return undefined
|
||||
throw error
|
||||
})
|
||||
if (!metadata) {
|
||||
if (!createParents) throw new Error(`补丁父目录不存在：${patchPath}`)
|
||||
return join(current, ...segments.slice(index + 1))
|
||||
} else if (metadata.isSymbolicLink() || !metadata.isDirectory()) {
|
||||
throw new Error(`补丁路径包含非普通目录：${patchPath}`)
|
||||
}
|
||||
}
|
||||
const canonicalParent = await realpath(join(root, ...segments.slice(0, -1)))
|
||||
if (!isInside(root, canonicalParent)) {
|
||||
throw new Error(`补丁路径不能通过符号链接超出工作区：${patchPath}`)
|
||||
}
|
||||
return join(canonicalParent, segments.at(-1)!)
|
||||
}
|
||||
|
||||
async function readUtf8(path: string): Promise<string> {
|
||||
const bytes = await readFile(path)
|
||||
try {
|
||||
return new TextDecoder('utf-8', { fatal: true }).decode(bytes)
|
||||
} catch (error) {
|
||||
throw new Error('补丁只能修改有效 UTF-8 文本文件', { cause: error })
|
||||
}
|
||||
}
|
||||
|
||||
async function writeAtomic(target: string, content: string, mode: number): Promise<void> {
|
||||
const temporary = `${target}.${randomUUID()}.tmp`
|
||||
const handle = await open(temporary, 'wx', mode)
|
||||
try {
|
||||
try {
|
||||
await handle.writeFile(content, 'utf8')
|
||||
} finally {
|
||||
await handle.close()
|
||||
}
|
||||
await rename(temporary, target)
|
||||
} catch (error) {
|
||||
await rm(temporary, { force: true }).catch(() => undefined)
|
||||
throw error
|
||||
}
|
||||
}
|
||||
|
||||
export async function applyWorkspacePatch(
|
||||
patch: string,
|
||||
workspace: WorkspaceAccess,
|
||||
signal: AbortSignal
|
||||
): Promise<string> {
|
||||
signal.throwIfAborted()
|
||||
const identity = await workspace.getIdentity()
|
||||
if (identity.kind !== 'local') {
|
||||
throw new Error('工作区补丁当前只支持本机工作区')
|
||||
}
|
||||
const root = await realpath(identity.canonicalDisplayPath)
|
||||
const operations = parseWorkspacePatch(patch)
|
||||
const seen = new Set<string>()
|
||||
const prepared: PreparedOperation[] = []
|
||||
for (const operation of operations) {
|
||||
signal.throwIfAborted()
|
||||
const targetKey = process.platform === 'win32' ? operation.path.toLowerCase() : operation.path
|
||||
if (seen.has(targetKey)) {
|
||||
throw new Error(`同一补丁不能重复操作文件：${operation.path}`)
|
||||
}
|
||||
seen.add(targetKey)
|
||||
const target = await resolveTarget(root, operation.path, operation.kind === 'add')
|
||||
const metadata = await lstat(target).catch((error: NodeJS.ErrnoException) => {
|
||||
if (error.code === 'ENOENT') return undefined
|
||||
throw error
|
||||
})
|
||||
if (metadata?.isSymbolicLink() || (metadata && !metadata.isFile())) {
|
||||
throw new Error(`补丁目标不是普通文件：${operation.path}`)
|
||||
}
|
||||
if (operation.kind === 'add') {
|
||||
if (metadata) throw new Error(`新增文件已经存在：${operation.path}`)
|
||||
prepared.push({
|
||||
kind: 'write',
|
||||
path: operation.path,
|
||||
target,
|
||||
content: operation.content,
|
||||
mode: 0o600
|
||||
})
|
||||
continue
|
||||
}
|
||||
if (!metadata) throw new Error(`补丁目标不存在：${operation.path}`)
|
||||
if (operation.kind === 'delete') {
|
||||
prepared.push({ kind: 'delete', path: operation.path, target })
|
||||
continue
|
||||
}
|
||||
const content = applyUpdateHunks(
|
||||
await readUtf8(target),
|
||||
operation.hunks,
|
||||
operation.path
|
||||
)
|
||||
prepared.push({
|
||||
kind: 'write',
|
||||
path: operation.path,
|
||||
target,
|
||||
content,
|
||||
mode: (await stat(target)).mode & 0o777
|
||||
})
|
||||
}
|
||||
|
||||
const applied: string[] = []
|
||||
try {
|
||||
for (const operation of prepared) {
|
||||
signal.throwIfAborted()
|
||||
if (operation.kind === 'delete') {
|
||||
await rm(operation.target)
|
||||
} else {
|
||||
await mkdir(dirname(operation.target), { recursive: true, mode: 0o700 })
|
||||
await writeAtomic(operation.target, operation.content, operation.mode)
|
||||
}
|
||||
applied.push(operation.path)
|
||||
}
|
||||
} catch (error) {
|
||||
throw new Error(
|
||||
`补丁仅完成 ${applied.length}/${prepared.length} 个文件：${applied.join(', ') || '无'}`,
|
||||
{ cause: error }
|
||||
)
|
||||
}
|
||||
const header = `Applied patch to ${applied.length} file(s):`
|
||||
const result = [header]
|
||||
let resultBytes = Buffer.byteLength(header) + 1
|
||||
for (const path of applied) {
|
||||
const bytes = Buffer.byteLength(path) + 1
|
||||
if (resultBytes + bytes > MAXIMUM_RESULT_BYTES) {
|
||||
result.push(`...[${applied.length - result.length + 1} more files]`)
|
||||
break
|
||||
}
|
||||
result.push(path)
|
||||
resultBytes += bytes
|
||||
}
|
||||
return result.join('\n')
|
||||
}
|
||||
Reference in New Issue
Block a user