Initial commit: omp-ast DSH plugin (oh-my-pi port) with docs and tests
This commit is contained in:
Agent committed
commit
408d0b3066
9 files changed
+488
No files matched your search
@@ -0,0 +1,266 @@
|
||||
/**
|
||||
* @local/omp-ast — oh-my-pi's `ast_grep` and `ast_edit` for DSH, backed by the ast-grep CLI.
|
||||
*
|
||||
* ast_grep structural search: `$NAME` captures one node, `$$$NAME` zero-or-more, `$_`/`$$$`
|
||||
* match without binding. Results are grouped per file with line numbers and captures.
|
||||
* ast_edit structural rewrite, preview first: returns a diff preview and a proposal id; calling
|
||||
* again with `apply: "<id>"` writes it (refused if any file changed since the preview).
|
||||
* `reject: "<id>"` drops it.
|
||||
*/
|
||||
import crypto from 'node:crypto'
|
||||
import path from 'node:path'
|
||||
import { findAstGrep, search, applyReplacements } from './lib/sg.js'
|
||||
|
||||
export const name = 'omp-ast'
|
||||
export const inject = ['tools', 'fs']
|
||||
|
||||
const PATTERN_HELP = 'Pattern grammar: real code for the target language with metavariables — `$NAME` captures one AST node, `$$$NAME` captures zero or more nodes, `$_` / `$$$` match without binding. Names are UPPERCASE and stand for whole nodes (not parts of identifiers or strings); reusing a name requires identical code. The pattern must parse as one valid node (C/C++/Java statements need their trailing `;`).'
|
||||
|
||||
export function apply(ctx, config = {}) {
|
||||
const cfg = { maxMatches: 200, maxFiles: 1000, proposalTtlMinutes: 30, ...config }
|
||||
const cwdOf = (exec) => exec?.agent?.session?.header?.cwd ?? process.cwd()
|
||||
const bin = () => {
|
||||
const b = findAstGrep(cfg.binary ? [path.dirname(cfg.binary)] : [])
|
||||
if (!b) throw new Error('ast-grep is not installed. Install it with `npm i -g @ast-grep/cli` (or `npm i -g --prefix ~/.local @ast-grep/cli`).')
|
||||
return b
|
||||
}
|
||||
const splitPaths = (p) => (Array.isArray(p) ? p : String(p ?? '').split(';')).map((s) => String(s).trim()).filter(Boolean)
|
||||
const rel = (f, cwd) => { const r = path.relative(cwd, path.resolve(cwd, f)); return r && !r.startsWith('..') ? r : f }
|
||||
|
||||
/** Sandbox fence for writes: same rule as DSH workspace-write (inside workspace root). */
|
||||
const policyFor = (exec) => { try { return ctx.get('sandboxPolicy')?.resolve({ ...(exec?.agent ? { session: exec.agent.session } : {}) }) } catch { return undefined } }
|
||||
|
||||
// ---------------------------------------------------------------- ast_grep
|
||||
ctx.tools.register({
|
||||
name: 'ast_grep',
|
||||
description: `Structural code search with ast-grep over tree-sitter grammars (js/ts/tsx, python, go, rust, java, c/cpp, c#, kotlin, swift, ruby, php, lua, bash, html, css, json, yaml, …). Use it instead of regex grep when the shape of the code matters (calls with any arguments, functions with a decorator, a specific nesting). ${PATTERN_HELP} Results are grouped per file with \`LINE:TEXT\` and the captured metavariables.`,
|
||||
parameters: {
|
||||
type: 'object',
|
||||
additionalProperties: false,
|
||||
required: ['pat'],
|
||||
properties: {
|
||||
pat: { type: 'string', description: 'One AST pattern, e.g. `console.log($$$ARGS)` or `def $F($$$): $$$BODY`.' },
|
||||
path: { type: 'string', description: 'File, directory or glob; several separated by `;`. Defaults to the session workspace.' },
|
||||
lang: { type: 'string', description: 'Language override (e.g. ts, tsx, python, cpp). Inferred from file extensions when omitted.' },
|
||||
globs: { type: 'array', items: { type: 'string' }, description: 'Include/exclude globs (prefix ! to exclude), e.g. ["src/**", "!**/*.test.ts"].' },
|
||||
selector: { type: 'string', description: 'AST kind to select inside the pattern (advanced; e.g. `pair` for a JSON key pattern).' },
|
||||
skip: { type: 'integer', description: 'Skip this many matches (paging).' },
|
||||
},
|
||||
},
|
||||
output: { schema: { type: 'object', additionalProperties: true, properties: { text: { type: 'string' } } }, render: (_a, v) => [{ type: 'text', text: v.text }] },
|
||||
timeoutMs: 150_000,
|
||||
isConcurrencySafe: () => true,
|
||||
async execute(args, exec) {
|
||||
const pat = String(args.pat ?? '').trim()
|
||||
if (!pat) throw new Error('pat must be a non-empty AST pattern')
|
||||
const cwd = cwdOf(exec)
|
||||
const { matches, warnings } = await search({ bin: bin(), pattern: pat, lang: args.lang, paths: splitPaths(args.path), globs: args.globs, selector: args.selector, cwd, signal: exec?.signal })
|
||||
const skip = Math.max(0, Math.floor(args.skip ?? 0))
|
||||
const page = matches.slice(skip, skip + cfg.maxMatches)
|
||||
const byFile = new Map()
|
||||
for (const m of page) { if (!byFile.has(m.file)) byFile.set(m.file, []); byFile.get(m.file).push(m) }
|
||||
const out = []
|
||||
for (const [file, ms] of byFile) {
|
||||
out.push(rel(file, cwd))
|
||||
for (const m of ms) {
|
||||
const start = m.range.start.line + 1
|
||||
const lines = String(m.lines ?? m.text).split('\n')
|
||||
lines.forEach((l, i) => out.push(`${start + i}:${l.length > 300 ? `${l.slice(0, 300)}…` : l}`))
|
||||
const caps = captures(m)
|
||||
if (caps) out.push(` ↳ ${caps}`)
|
||||
}
|
||||
}
|
||||
const files = new Set(matches.map((m) => m.file)).size
|
||||
let header = matches.length ? `${matches.length} match${matches.length === 1 ? '' : 'es'} in ${files} file${files === 1 ? '' : 's'}` : 'No matches.'
|
||||
if (matches.length > skip + cfg.maxMatches) header += ` (showing ${skip + 1}-${skip + page.length}; use skip=${skip + page.length} for more)`
|
||||
const warn = warnings.filter((w) => !/^$/.test(w))
|
||||
return { text: [header, ...out, ...(warn.length ? ['', 'ast-grep:', ...warn.map((w) => ` ${w}`)] : [])].join('\n'), count: matches.length }
|
||||
},
|
||||
})
|
||||
|
||||
// ---------------------------------------------------------------- ast_edit
|
||||
/** @type {Map<string, {at:number, files: Map<string,{before:Buffer, after:Buffer}>, summary:string, owner:any}>} */
|
||||
const proposals = new Map()
|
||||
const gc = () => { const ttl = cfg.proposalTtlMinutes * 60_000; for (const [id, p] of proposals) if (Date.now() - p.at > ttl) proposals.delete(id) }
|
||||
|
||||
ctx.tools.register({
|
||||
name: 'ast_edit',
|
||||
description: `Structural rewrite with ast-grep, previewed before it is applied. Pass \`ops\` (each {pat, out}: captures from \`pat\` are substituted into \`out\`; empty \`out\` deletes the node) and \`paths\`. The call returns a diff preview and a proposal id WITHOUT modifying files; call ast_edit again with \`apply: "<id>"\` to write it (refused if a file changed since the preview) or \`reject: "<id>"\` to drop it. Each rewrite is a 1:1 node substitution. ${PATTERN_HELP}`,
|
||||
parameters: {
|
||||
type: 'object',
|
||||
additionalProperties: false,
|
||||
properties: {
|
||||
ops: {
|
||||
type: 'array',
|
||||
description: 'Rewrite rules applied in order.',
|
||||
items: { type: 'object', additionalProperties: false, required: ['pat', 'out'], properties: { pat: { type: 'string' }, out: { type: 'string' } } },
|
||||
},
|
||||
paths: { type: 'array', items: { type: 'string' }, description: 'Files, directories or globs to rewrite.' },
|
||||
lang: { type: 'string', description: 'Language override.' },
|
||||
globs: { type: 'array', items: { type: 'string' }, description: 'Include/exclude globs.' },
|
||||
apply: { type: 'string', description: 'Proposal id from a previous preview: write it.' },
|
||||
reject: { type: 'string', description: 'Proposal id to discard.' },
|
||||
},
|
||||
},
|
||||
output: {
|
||||
schema: { type: 'object', additionalProperties: true, properties: { text: { type: 'string' } } },
|
||||
render: (_a, v) => [{ type: 'text', text: v.text }],
|
||||
presentationMeta: (_a, v) => ({ diffs: Array.isArray(v.diffs) ? v.diffs : [] }),
|
||||
},
|
||||
presentResult(args, result) {
|
||||
if (result.isError) return undefined
|
||||
const diffs = result.meta?.diffs
|
||||
if (!Array.isArray(diffs) || !diffs.length) return undefined
|
||||
return { card: 'diff', title: args.apply ? 'AST edit (applied)' : 'AST edit (preview)', diffs }
|
||||
},
|
||||
timeoutMs: 300_000,
|
||||
async execute(args, exec) {
|
||||
gc()
|
||||
const cwd = cwdOf(exec)
|
||||
if (args.reject) {
|
||||
const ok = proposals.delete(args.reject)
|
||||
return { text: ok ? `Proposal ${args.reject} discarded; no files were changed.` : `No pending proposal ${args.reject}.` }
|
||||
}
|
||||
if (args.apply) return applyProposal(args.apply, exec, cwd)
|
||||
const ops = Array.isArray(args.ops) ? args.ops : []
|
||||
if (!ops.length) throw new Error('ops must contain at least one {pat, out} rewrite (or pass apply/reject with a proposal id)')
|
||||
const seen = new Set()
|
||||
for (const [i, o] of ops.entries()) {
|
||||
if (!o || typeof o.pat !== 'string' || !o.pat.trim()) throw new Error(`ops[${i}].pat must be a non-empty pattern`)
|
||||
if (typeof o.out !== 'string') throw new Error(`ops[${i}].out must be a string (empty deletes)`)
|
||||
if (seen.has(o.pat)) throw new Error(`ops[${i}] duplicates pattern ${JSON.stringify(o.pat)}`)
|
||||
seen.add(o.pat)
|
||||
}
|
||||
const paths = splitPaths(args.paths)
|
||||
if (!paths.length) throw new Error('paths must list at least one file, directory or glob')
|
||||
const fsp = await import('node:fs/promises')
|
||||
// Apply ops sequentially in memory: each op's rewrite is computed on the result of the previous.
|
||||
const files = new Map() // abs -> {before, after}
|
||||
const counts = []
|
||||
const warnings = []
|
||||
for (const op of ops) {
|
||||
const { matches, warnings: w } = await search({ bin: bin(), pattern: op.pat, rewrite: op.out, lang: args.lang, paths, globs: args.globs, cwd, signal: exec?.signal })
|
||||
warnings.push(...w)
|
||||
// ast-grep reads from disk. Files an earlier op already rewrote are re-searched on a temp copy
|
||||
// of their in-memory version; the rest use this op's disk matches.
|
||||
const rep = (m) => ({ start: m.replacementOffsets?.start ?? m.range.byteOffset.start, end: m.replacementOffsets?.end ?? m.range.byteOffset.end, text: m.replacement ?? '' })
|
||||
const byFile = new Map()
|
||||
for (const m of matches) { const abs = path.resolve(cwd, m.file); if (!byFile.has(abs)) byFile.set(abs, []); byFile.get(abs).push(m) }
|
||||
let n = 0
|
||||
for (const [abs, entry] of files) {
|
||||
if (entry.after === entry.before) continue
|
||||
byFile.delete(abs)
|
||||
const tmp = path.join(path.dirname(abs), `.omp-ast-${crypto.randomBytes(4).toString('hex')}${path.extname(abs)}`)
|
||||
await fsp.writeFile(tmp, entry.after)
|
||||
try {
|
||||
const r = await search({ bin: bin(), pattern: op.pat, rewrite: op.out, lang: args.lang, paths: [tmp], cwd, signal: exec?.signal })
|
||||
if (r.matches.length) { entry.after = applyReplacements(entry.after, r.matches.map(rep)); n += r.matches.length }
|
||||
} finally { await fsp.rm(tmp, { force: true }) }
|
||||
}
|
||||
for (const [abs, ms] of byFile) {
|
||||
let entry = files.get(abs)
|
||||
if (!entry) { const b = await fsp.readFile(abs); entry = { before: b, after: b }; files.set(abs, entry) }
|
||||
entry.after = applyReplacements(entry.after, ms.map(rep))
|
||||
n += ms.length
|
||||
}
|
||||
counts.push(`${JSON.stringify(op.pat)} → ${n} replacement${n === 1 ? '' : 's'}`)
|
||||
}
|
||||
for (const [abs, e] of files) if (Buffer.compare(e.before, e.after) === 0) files.delete(abs)
|
||||
if (files.size > cfg.maxFiles) throw new Error(`rewrite touches ${files.size} files (max ${cfg.maxFiles}); narrow paths/globs`)
|
||||
if (!files.size) return { text: ['No changes: no pattern matched.', ...counts, ...(warnings.length ? ['ast-grep:', ...warnings] : [])].join('\n') }
|
||||
const id = crypto.randomBytes(3).toString('hex')
|
||||
proposals.set(id, { at: Date.now(), files, owner: exec?.agent?.session })
|
||||
const diffs = []
|
||||
const lines = []
|
||||
for (const [abs, e] of files) {
|
||||
const d = lineDiff(e.before.toString('utf8'), e.after.toString('utf8'))
|
||||
diffs.push(...d.map((h) => ({ path: rel(abs, cwd), ...h })))
|
||||
lines.push(`${rel(abs, cwd)}:`)
|
||||
for (const h of d.slice(0, 20)) {
|
||||
for (const l of (h.oldText ?? '').split('\n')) lines.push(`- ${l}`)
|
||||
for (const l of h.newText.split('\n')) lines.push(`+ ${l}`)
|
||||
}
|
||||
}
|
||||
const shown = lines.length > 400 ? [...lines.slice(0, 400), `… ${lines.length - 400} more preview lines`] : lines
|
||||
return {
|
||||
text: [
|
||||
`Staged as proposal ${id} — files NOT modified yet. ${files.size} file${files.size === 1 ? '' : 's'}:`,
|
||||
...counts, '', ...shown, '',
|
||||
`Apply: ast_edit {"apply": "${id}"} · Discard: ast_edit {"reject": "${id}"}`,
|
||||
...(warnings.length ? ['', 'ast-grep:', ...warnings.map((w) => ` ${w}`)] : []),
|
||||
].join('\n'),
|
||||
proposal: id,
|
||||
diffs,
|
||||
}
|
||||
},
|
||||
})
|
||||
|
||||
async function applyProposal(id, exec, cwd) {
|
||||
const p = proposals.get(id)
|
||||
if (!p) throw new Error(`No pending proposal ${id} (expired after ${cfg.proposalTtlMinutes} min, already applied, or never issued). Run the preview again.`)
|
||||
if (p.owner && exec?.agent?.session && p.owner !== exec.agent.session) throw new Error(`Proposal ${id} belongs to another session.`)
|
||||
const fsp = await import('node:fs/promises')
|
||||
// Verify nothing changed since the preview, then write all.
|
||||
for (const [abs, e] of p.files) {
|
||||
const now = await fsp.readFile(abs).catch(() => null)
|
||||
if (!now || Buffer.compare(now, e.before) !== 0) throw new Error(`${rel(abs, cwd)} changed since the preview; proposal ${id} not applied. Run ast_edit again.`)
|
||||
}
|
||||
const policy = policyFor(exec)
|
||||
const written = []
|
||||
const diffs = []
|
||||
for (const [abs, e] of p.files) {
|
||||
const target = await ctx.fs.resolve(abs, { signal: exec?.signal })
|
||||
const outcome = await ctx.fs.writeText(target, e.after.toString('utf8'), undefined, exec?.signal, policy)
|
||||
ctx.emit('fs/observed', target, { kind: 'present', version: outcome.version }, exec)
|
||||
written.push(rel(abs, cwd))
|
||||
diffs.push(...lineDiff(e.before.toString('utf8'), e.after.toString('utf8')).map((h) => ({ path: rel(abs, cwd), ...h })))
|
||||
}
|
||||
proposals.delete(id)
|
||||
return { text: `Applied proposal ${id} to ${written.length} file${written.length === 1 ? '' : 's'}:\n${written.join('\n')}`, results: written.map((f) => ({ path: f, operation: 'update' })), diffs }
|
||||
}
|
||||
|
||||
ctx.effect(() => () => proposals.clear(), 'omp-ast: proposals')
|
||||
}
|
||||
|
||||
function captures(m) {
|
||||
const mv = m.metaVariables
|
||||
if (!mv) return ''
|
||||
const parts = []
|
||||
for (const [k, v] of Object.entries(mv.single ?? {})) parts.push(`$${k}=${short(v.text)}`)
|
||||
for (const [k, v] of Object.entries(mv.multi ?? {})) parts.push(`$$$${k}=[${v.filter((x) => !/^[,;]$/.test(x.text)).map((x) => short(x.text)).join(', ')}]`)
|
||||
return parts.join(' ')
|
||||
}
|
||||
const short = (s) => { const t = String(s).replace(/\s+/g, ' '); return t.length > 60 ? `${t.slice(0, 60)}…` : t }
|
||||
|
||||
/** Minimal line diff → hunks with 2 lines of context (LCS on lines). */
|
||||
export function lineDiff(a, b, context = 2) {
|
||||
const A = a.split('\n'); const B = b.split('\n')
|
||||
let s = 0
|
||||
while (s < A.length && s < B.length && A[s] === B[s]) s++
|
||||
let ea = A.length; let eb = B.length
|
||||
while (ea > s && eb > s && A[ea - 1] === B[eb - 1]) { ea--; eb-- }
|
||||
if (s === ea && s === eb) return []
|
||||
// Split the middle into hunks by common lines (simple LCS for moderate sizes).
|
||||
const midA = A.slice(s, ea); const midB = B.slice(s, eb)
|
||||
const n = midA.length; const m = midB.length
|
||||
const hunks = []
|
||||
if (n * m <= 1_000_000) {
|
||||
const dp = Array.from({ length: n + 1 }, () => new Uint32Array(m + 1))
|
||||
for (let i = n - 1; i >= 0; i--) for (let j = m - 1; j >= 0; j--) dp[i][j] = midA[i] === midB[j] ? dp[i + 1][j + 1] + 1 : Math.max(dp[i + 1][j], dp[i][j + 1])
|
||||
let i = 0; let j = 0; let cur = null
|
||||
const flush = () => { if (cur) { hunks.push(cur); cur = null } }
|
||||
while (i < n || j < m) {
|
||||
if (i < n && j < m && midA[i] === midB[j]) { flush(); i++; j++; continue }
|
||||
cur ??= { a0: s + i, b0: s + j, old: [], neu: [] }
|
||||
if (j < m && (i >= n || dp[i][j + 1] >= dp[i + 1][j])) cur.neu.push(midB[j++])
|
||||
else cur.old.push(midA[i++])
|
||||
}
|
||||
flush()
|
||||
} else hunks.push({ a0: s, b0: s, old: midA, neu: midB })
|
||||
return hunks.map((h) => {
|
||||
const pre = A.slice(Math.max(0, h.a0 - context), h.a0)
|
||||
const postStart = h.a0 + h.old.length
|
||||
const post = A.slice(postStart, postStart + context)
|
||||
return { oldText: [...pre, ...h.old, ...post].join('\n') || null, newText: [...pre, ...h.neu, ...post].join('\n') }
|
||||
})
|
||||
}
|
||||
Reference in new issue
Block a user