|
| 1 | +import { execFile } from "node:child_process"; |
| 2 | +import { realpathSync } from "node:fs"; |
| 3 | +import { delimiter, isAbsolute, relative, sep } from "node:path"; |
| 4 | +import type { ExtensionAPI } from "@oh-my-pi/pi-coding-agent"; |
| 5 | +import { |
| 6 | + expandDelimitedPathEntries, |
| 7 | + isInternalUrlPath, |
| 8 | + isReadableUrlPath, |
| 9 | + normalizePathLikeInput, |
| 10 | + parseSearchPath, |
| 11 | + resolveReadPath, |
| 12 | + splitPathAndSelPreferringLiteral, |
| 13 | +} from "@oh-my-pi/pi-coding-agent/tools/path-utils"; |
| 14 | + |
| 15 | +const INPUT_LIMIT = 256 * 1024; |
| 16 | +const OUTPUT_LIMIT = 64 * 1024; |
| 17 | +const TIMEOUT_MS = 2000; |
| 18 | +const CONTEXT_TYPE = "graphify-guard"; |
| 19 | +const TOOL_NAMES = { bash: "Bash", grep: "Grep", read: "Read", glob: "Glob" } as const; |
| 20 | + |
| 21 | +function isRemote(path: string): boolean { |
| 22 | + return isInternalUrlPath(path) || isReadableUrlPath(path) || path.includes("://"); |
| 23 | +} |
| 24 | + |
| 25 | +function installedCommand(cwd: string): string | undefined { |
| 26 | + // Never resolve a bare/relative PATH entry against the project or use a |
| 27 | + // project-supplied Python module as a fallback for an absent installation. |
| 28 | + const path = (process.env.PATH ?? "").split(delimiter).filter(isAbsolute).join(delimiter); |
| 29 | + const found = path ? Bun.which("graphify", { PATH: path }) : null; |
| 30 | + if (!found) return; |
| 31 | + const command = realpathSync(found); |
| 32 | + const within = relative(realpathSync(cwd), command); |
| 33 | + if (!within || (!within.startsWith(`..${sep}`) && !isAbsolute(within))) return; |
| 34 | + return command; |
| 35 | +} |
| 36 | + |
| 37 | +function runGuard(command: string, kind: string, payload: string, cwd: string, signal: AbortSignal, timeout: number, maxBuffer: number): Promise<string | undefined> { |
| 38 | + if (Buffer.byteLength(payload) > INPUT_LIMIT || signal.aborted) return Promise.resolve(undefined); |
| 39 | + const { promise, resolve } = Promise.withResolvers<string | undefined>(); |
| 40 | + const child = execFile(command, ["hook-guard", kind], { |
| 41 | + cwd, encoding: "utf8", timeout, killSignal: "SIGKILL", maxBuffer, signal, |
| 42 | + env: { ...process.env, CLAUDE_PROJECT_DIR: cwd }, |
| 43 | + }, (error, stdout) => resolve(error ? undefined : stdout.trim() || undefined)); |
| 44 | + child.stdin?.on("error", () => {}); |
| 45 | + child.stdin?.end(payload); |
| 46 | + return promise; |
| 47 | +} |
| 48 | + |
| 49 | +export default function graphify(api: ExtensionAPI): void { |
| 50 | + const guidance = new Set<string>(); |
| 51 | + let guidanceBytes = 0; |
| 52 | + let generation = 0; |
| 53 | + let controller = new AbortController(); |
| 54 | + const reset = () => { |
| 55 | + generation++; |
| 56 | + controller.abort(); |
| 57 | + controller = new AbortController(); |
| 58 | + guidance.clear(); |
| 59 | + guidanceBytes = 0; |
| 60 | + }; |
| 61 | + api.on("session_start", reset); |
| 62 | + api.on("session_switch", reset); |
| 63 | + api.on("session_branch", reset); |
| 64 | + api.on("session_tree", reset); |
| 65 | + api.on("session_shutdown", reset); |
| 66 | + api.on("before_agent_start", reset); |
| 67 | + |
| 68 | + api.on("tool_call", async (event, ctx) => { |
| 69 | + if (!ctx.isProjectTrusted()) { reset(); return; } |
| 70 | + if (!Object.hasOwn(TOOL_NAMES, event.toolName)) return; |
| 71 | + const current = generation; |
| 72 | + const signal = controller.signal; |
| 73 | + const deadline = performance.now() + TIMEOUT_MS; |
| 74 | + try { |
| 75 | + if (Buffer.byteLength(JSON.stringify(event.input)) > INPUT_LIMIT) return; |
| 76 | + const command = installedCommand(ctx.cwd); |
| 77 | + if (!command) return; |
| 78 | + const toolName = TOOL_NAMES[event.toolName as keyof typeof TOOL_NAMES]; |
| 79 | + const rawPath = typeof event.input.path === "string" ? normalizePathLikeInput(event.input.path) : "."; |
| 80 | + if (isRemote(rawPath)) return; |
| 81 | + const paths = toolName === "Read" || toolName === "Bash" ? [rawPath] : |
| 82 | + await expandDelimitedPathEntries([rawPath], ctx.cwd, { splitter: parseSearchPath }); |
| 83 | + let remainingOutput = OUTPUT_LIMIT; |
| 84 | + for (const path of paths) { |
| 85 | + if (isRemote(path)) continue; |
| 86 | + const input: Record<string, unknown> = { ...event.input }; |
| 87 | + if (toolName !== "Bash") { |
| 88 | + const target = toolName === "Glob" ? path : (await splitPathAndSelPreferringLiteral(path, ctx.cwd)).path; |
| 89 | + const resolved = resolveReadPath(target, ctx.cwd); |
| 90 | + delete input.path; |
| 91 | + if (toolName === "Read") input.file_path = resolved; |
| 92 | + else if (toolName === "Glob") input.pattern = resolved; |
| 93 | + else input.path = resolved; |
| 94 | + } |
| 95 | + const timeout = Math.floor(deadline - performance.now()); |
| 96 | + if (current !== generation || !ctx.isProjectTrusted() || timeout <= 0 || remainingOutput <= 0) return; |
| 97 | + const output = await runGuard(command, toolName === "Bash" || toolName === "Grep" ? "search" : "read", JSON.stringify({ |
| 98 | + session_id: ctx.sessionManager.getSessionId(), cwd: ctx.cwd, tool_name: toolName, tool_input: input, |
| 99 | + }), ctx.cwd, signal, timeout, remainingOutput); |
| 100 | + if (current !== generation || !ctx.isProjectTrusted()) return; |
| 101 | + if (!output) continue; |
| 102 | + remainingOutput -= Buffer.byteLength(output); |
| 103 | + const hook = JSON.parse(output)?.hookSpecificOutput; |
| 104 | + if (hook?.hookEventName !== "PreToolUse") continue; |
| 105 | + if (hook.permissionDecision === "deny" && typeof hook.permissionDecisionReason === "string" && hook.permissionDecisionReason.trim()) { |
| 106 | + return { block: true, reason: hook.permissionDecisionReason }; |
| 107 | + } |
| 108 | + if (typeof hook.additionalContext === "string" && hook.additionalContext.trim() && !guidance.has(hook.additionalContext)) { |
| 109 | + const bytes = Buffer.byteLength(hook.additionalContext) + 2; |
| 110 | + if (guidanceBytes + bytes <= OUTPUT_LIMIT) { |
| 111 | + guidance.add(hook.additionalContext); |
| 112 | + guidanceBytes += bytes; |
| 113 | + } |
| 114 | + } |
| 115 | + // The search CLI does not inspect individual targets; one call suffices. |
| 116 | + if (toolName === "Grep") break; |
| 117 | + } |
| 118 | + } catch { |
| 119 | + // Optional guidance must not break native tools on missing executables, |
| 120 | + // invalid paths, malformed hook output, timeout, or cancellation. |
| 121 | + } |
| 122 | + }); |
| 123 | + |
| 124 | + api.on("context", (event, ctx) => { |
| 125 | + if (!ctx.isProjectTrusted()) reset(); |
| 126 | + // Context transforms are not persisted. Keep one current-run message across |
| 127 | + // provider requests, and discard any prior generation's injected message. |
| 128 | + const messages = event.messages.filter(message => message.role !== "custom" || message.customType !== CONTEXT_TYPE); |
| 129 | + if (guidance.size) messages.push({ |
| 130 | + role: "custom", customType: CONTEXT_TYPE, content: [...guidance].join("\n\n"), display: false, timestamp: Date.now(), |
| 131 | + }); |
| 132 | + if (guidance.size || messages.length !== event.messages.length) return { messages }; |
| 133 | + }); |
| 134 | +} |
0 commit comments