pi-better-btw-plus 1.0.0

package/srcs/tool-wrapper.ts

@@ -0,0 +1,225 @@
+import type { AgentTool } from "@earendil-works/pi-agent-core";
+import type { FileActivityTracker } from "./file-activity-tracker.ts";
+
+export function wrapToolsWithOverlapDetection(
+  tools: AgentTool[],
+  tracker: FileActivityTracker,
+  cwd: string,
+  confirmOverlap: (path: string) => Promise<boolean>,
+): AgentTool[] {
+  const writingTools = ["write", "edit", "bash"];
+  return tools.map((tool) =>
+    writingTools.includes(tool.name) ? wrapTool(tool, tracker, cwd, confirmOverlap) : tool
+  );
+}
+
+function wrapTool(
+  tool: AgentTool,
+  tracker: FileActivityTracker,
+  cwd: string,
+  confirmOverlap: (path: string) => Promise<boolean>,
+): AgentTool {
+  return {
+    ...tool,
+    execute: async (toolCallId, args, signal, onUpdate) => {
+      const paths = extractWritePaths(tool.name, args);
+
+      for (const path of paths) {
+        if (tracker.hasWritten(path, cwd)) {
+          const proceed = await confirmOverlap(path);
+          if (!proceed) {
+            return {
+              content: [{ type: "text", text: `Skipped: ${path} (main agent has modified it)` }],
+              details: undefined,
+            };
+          }
+        }
+      }
+
+      return tool.execute(toolCallId, args, signal, onUpdate);
+    },
+  };
+}
+
+export function extractWritePaths(toolName: string, args: unknown): string[] {
+  const record = args && typeof args === "object" ? (args as Record<string, unknown>) : {};
+
+  switch (toolName) {
+    case "write":
+    case "edit":
+      return typeof record.path === "string" ? [record.path] : [];
+    case "bash":
+      return typeof record.command === "string" ? parseBashWritePaths(record.command) : [];
+    default:
+      return [];
+  }
+}
+
+function parseBashWritePaths(command: string): string[] {
+  const tokens = tokenizeShell(command);
+  const paths: string[] = [];
+  let segment: ShellToken[] = [];
+
+  for (const token of tokens) {
+    if (token.type === "op" && isCommandSeparator(token.value)) {
+      collectSegmentWritePaths(segment, paths);
+      segment = [];
+      continue;
+    }
+    segment.push(token);
+  }
+
+  collectSegmentWritePaths(segment, paths);
+  return [...new Set(paths)];
+}
+
+type ShellToken =
+  | { type: "word"; value: string }
+  | { type: "op"; value: ">" | ">>" | "|" | "||" | "&" | "&&" | ";" };
+
+function tokenizeShell(command: string): ShellToken[] {
+  const tokens: ShellToken[] = [];
+
+  for (let i = 0; i < command.length;) {
+    const char = command[i];
+
+    if (/\s/.test(char)) {
+      i++;
+      continue;
+    }
+
+    const twoCharOp = command.slice(i, i + 2);
+    if (twoCharOp === ">>" || twoCharOp === "||" || twoCharOp === "&&") {
+      tokens.push({ type: "op", value: twoCharOp });
+      i += 2;
+      continue;
+    }
+
+    if (char === ">" || char === "|" || char === "&" || char === ";") {
+      tokens.push({ type: "op", value: char });
+      i++;
+      continue;
+    }
+
+    let value = "";
+    while (i < command.length) {
+      const current = command[i];
+
+      if (/\s/.test(current) || current === ">" || current === "|" || current === "&" || current === ";") {
+        break;
+      }
+
+      if (current === "\\") {
+        if (i + 1 < command.length) {
+          value += command[i + 1];
+          i += 2;
+        } else {
+          i++;
+        }
+        continue;
+      }
+
+      if (current === "'") {
+        i++;
+        while (i < command.length && command[i] !== "'") {
+          value += command[i];
+          i++;
+        }
+        if (command[i] === "'") i++;
+        continue;
+      }
+
+      if (current === '"') {
+        i++;
+        while (i < command.length && command[i] !== '"') {
+          if (command[i] === "\\" && i + 1 < command.length && /["\\$`]/.test(command[i + 1])) {
+            value += command[i + 1];
+            i += 2;
+          } else {
+            value += command[i];
+            i++;
+          }
+        }
+        if (command[i] === '"') i++;
+        continue;
+      }
+
+      value += current;
+      i++;
+    }
+
+    if (value) tokens.push({ type: "word", value });
+  }
+
+  return tokens;
+}
+
+function collectSegmentWritePaths(segment: ShellToken[], paths: string[]): void {
+  for (let i = 0; i < segment.length; i++) {
+    const token = segment[i];
+    if (token.type === "op" && (token.value === ">" || token.value === ">>")) {
+      if (segment[i + 1]?.type === "word") {
+        pushPath(paths, segment[i + 1].value);
+        i++;
+      }
+    }
+  }
+
+  const commandIndex = segment.findIndex((token) => token.type === "word");
+  if (commandIndex === -1) return;
+
+  const command = segment[commandIndex];
+
+  const operands = collectCommandOperands(segment.slice(commandIndex + 1));
+  if (command.value === "tee" || command.value === "touch" || command.value === "rm") {
+    for (const operand of operands) pushPath(paths, operand);
+  }
+  if ((command.value === "cp" || command.value === "mv") && operands.length >= 2) {
+    pushPath(paths, operands[operands.length - 1]);
+  }
+}
+
+function collectCommandOperands(tokens: ShellToken[]): string[] {
+  const operands: string[] = [];
+  let parsingOptions = true;
+
+  for (let i = 0; i < tokens.length; i++) {
+    const token = tokens[i];
+
+    if (token.type === "op") {
+      if (token.value === ">" || token.value === ">>") {
+        if (tokens[i + 1]?.type === "word") i++;
+      }
+      continue;
+    }
+
+    if (parsingOptions) {
+      if (token.value === "--") {
+        parsingOptions = false;
+        continue;
+      }
+      if (token.value.startsWith("-")) {
+        continue;
+      }
+      parsingOptions = false;
+    }
+
+    operands.push(token.value);
+  }
+
+  return operands;
+}
+
+function isCommandSeparator(value: ShellToken["value"]): boolean {
+  return value === "|" || value === "||" || value === "&" || value === "&&" || value === ";";
+}
+
+function pushPath(paths: string[], path: string): void {
+  if (path && !path.startsWith("-") && !isIgnoredWritePath(path)) {
+    paths.push(path);
+  }
+}
+
+function isIgnoredWritePath(path: string): boolean {
+  return path === "/dev/null" || path === "/dev/stdout" || path === "/dev/stderr" || path.startsWith("/dev/fd/");
+}