62 lines
2.3 KiB
TypeScript
62 lines
2.3 KiB
TypeScript
import * as fs from "node:fs";
|
|
import * as os from "node:os";
|
|
import * as path from "node:path";
|
|
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
|
|
|
|
const PATH_TOOLS = new Set(["read", "write", "edit", "grep", "find", "ls"]);
|
|
|
|
function isWithin(root: string, candidate: string): boolean {
|
|
const relative = path.relative(root, candidate);
|
|
return relative === "" || (!relative.startsWith(`..${path.sep}`) && relative !== ".." && !path.isAbsolute(relative));
|
|
}
|
|
|
|
function canonicalizePotentialPath(inputPath: string, cwd: string): string {
|
|
let normalized = inputPath.startsWith("@") ? inputPath.slice(1) : inputPath;
|
|
if (normalized === "~") normalized = os.homedir();
|
|
else if (normalized.startsWith(`~${path.sep}`)) normalized = path.join(os.homedir(), normalized.slice(2));
|
|
|
|
const resolved = path.resolve(cwd, normalized);
|
|
const missingSegments: string[] = [];
|
|
let existing = resolved;
|
|
|
|
while (!fs.existsSync(existing)) {
|
|
const parent = path.dirname(existing);
|
|
if (parent === existing) break;
|
|
missingSegments.unshift(path.basename(existing));
|
|
existing = parent;
|
|
}
|
|
|
|
const canonicalBase = fs.existsSync(existing) ? fs.realpathSync(existing) : existing;
|
|
return path.join(canonicalBase, ...missingSegments);
|
|
}
|
|
|
|
export function registerWorkerGuard(pi: ExtensionAPI, configuredRoot: string): void {
|
|
const root = fs.realpathSync(configuredRoot);
|
|
|
|
pi.on("tool_call", async (event) => {
|
|
if (!PATH_TOOLS.has(event.toolName)) return;
|
|
|
|
const input = event.input as { path?: unknown };
|
|
if (input.path === undefined && (event.toolName === "grep" || event.toolName === "find" || event.toolName === "ls")) {
|
|
return;
|
|
}
|
|
if (typeof input.path !== "string") {
|
|
return { block: true, reason: `${event.toolName} requires a path inside the assigned worktree.` };
|
|
}
|
|
|
|
const target = canonicalizePotentialPath(input.path, root);
|
|
if (!isWithin(root, target)) {
|
|
return {
|
|
block: true,
|
|
reason: `Worker file tools are restricted to ${root}. Attempted path: ${input.path}`,
|
|
};
|
|
}
|
|
});
|
|
|
|
pi.on("before_agent_start", async (event) => ({
|
|
systemPrompt:
|
|
event.systemPrompt +
|
|
`\n\nYou are running as a worktree worker. Your assigned repository root is ${root}. ` +
|
|
"Keep all file mutations and shell work inside this worktree. Do not access or modify sibling worktrees.",
|
|
}));
|
|
}
|