clawdie-ai/.pi/extensions/clawdie-harness/index.ts

import fs from "node:fs";
import path from "node:path";

import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";

import { Type } from "@sinclair/typebox";
import { jailStatusTool, systemHealthTool, zfsSnapshotsTool } from "./system-tools.js";
import { skillsSearchTool, skillsContentTool } from "./skill-tools.js";
import { taskCreateTool, taskStatusTool, taskDecomposeTool } from "./controlplane-tools.js";
import { callHostd } from "./hostd-bridge.js";

type SafetyAction = "allow" | "ask" | "deny";

interface SafetyRule {
  id: string;
  tool?: string;
  action: SafetyAction;
  reason?: string;
  patterns?: string[];
  path_globs?: string[];
}

interface BashPatternRule {
  id: string;
  pattern: string;
  action?: Exclude<SafetyAction, "allow">;
  reason?: string;
}

interface SafetyConfig {
  version?: number;
  default_action?: SafetyAction;
  rules?: SafetyRule[];
  bash_patterns?: BashPatternRule[];
  bashToolPatterns?: BashPatternRule[];
  zero_access_paths?: string[];
  zeroAccessPaths?: string[];
  read_only_paths?: string[];
  readOnlyPaths?: string[];
  no_delete_paths?: string[];
  noDeletePaths?: string[];
}

const PURPOSE_ENTRY = "clawdie-purpose";
const DEFAULT_RULES: SafetyRule[] = [
  {
    id: "block-rm-rf",
    tool: "bash",
    action: "deny",
    reason: "Blocked dangerous rm -rf usage",
    patterns: ["rm -rf", "rm -fr"],
  },
  {
    id: "confirm-sudo",
    tool: "bash",
    action: "ask",
    reason: "Confirm sudo commands",
    patterns: ["sudo "],
  },
  {
    id: "block-env-write",
    tool: "write",
    action: "deny",
    reason: "Blocked writing .env files",
    path_globs: ["**/.env", "**/.env.*"],
  },
  {
    id: "block-ssh-read",
    tool: "read",
    action: "deny",
    reason: "Blocked reading SSH keys",
    path_globs: ["**/.ssh/**"],
  },
  {
    id: "block-ssh-write",
    tool: "write",
    action: "deny",
    reason: "Blocked writing SSH keys",
    path_globs: ["**/.ssh/**"],
  },
];

interface HarnessState {
  purpose: string | null;
  toolCounts: Map<string, number>;
  toolTotal: number;
  safetyRules: SafetyRule[];
  bashPatterns: BashPatternRule[];
  zeroAccessPaths: string[];
  readOnlyPaths: string[];
  noDeletePaths: string[];
  defaultAction: SafetyAction;
}

const state: HarnessState = {
  purpose: null,
  toolCounts: new Map(),
  toolTotal: 0,
  safetyRules: DEFAULT_RULES,
  bashPatterns: [],
  zeroAccessPaths: [],
  readOnlyPaths: [],
  noDeletePaths: [],
  defaultAction: "allow",
};

function globToRegExp(glob: string): RegExp {
  const escaped = glob
    .replace(/[.+^${}()|[\]\\]/g, "\\$&")
    .replace(/\*\*/g, "§§");
  const withSingle = escaped.replace(/\*/g, "[^/]*").replace(/\?/g, ".");
  const withDouble = withSingle.replace(/§§/g, ".*");
  return new RegExp(`^${withDouble}$`);
}

function matchesGlobList(target: string, globs: string[]): boolean {
  return globs.some((pattern) => globToRegExp(pattern).test(target));
}

function formatCounts(counts: Map<string, number>): string {
  const entries = Array.from(counts.entries()).sort(([a], [b]) =>
    a.localeCompare(b),
  );
  if (entries.length === 0) return "none";
  return entries.map(([tool, count]) => `${tool}=${count}`).join(" ");
}

function updateHarnessUI(ctx: ExtensionContext): void {
  if (!ctx.hasUI) return;
  const purpose = state.purpose ? state.purpose : "not set";
  const toolCounts = formatCounts(state.toolCounts);
  const safetyCount =
    state.safetyRules.length +
    state.bashPatterns.length +
    state.zeroAccessPaths.length +
    state.readOnlyPaths.length +
    state.noDeletePaths.length;
  ctx.ui.setStatus(
    "clawdie-harness",
    `Purpose: ${purpose} | Tools: ${toolCounts} | Safety: ${safetyCount}`,
  );
  ctx.ui.setWidget("clawdie-harness", [
    `Purpose: ${purpose}`,
    `Tools: ${toolCounts} (total=${state.toolTotal})`,
    `Safety rules: ${safetyCount}`,
  ]);
}

function normalizeList(value?: string[] | null): string[] {
  if (!value) return [];
  return value.map((entry) => entry.trim()).filter(Boolean);
}

function normalizeSafetyConfig(parsed: SafetyConfig): void {
  state.safetyRules = parsed.rules?.length ? parsed.rules : DEFAULT_RULES;
  state.defaultAction = parsed.default_action ?? "allow";
  state.bashPatterns = parsed.bash_patterns?.length
    ? parsed.bash_patterns
    : parsed.bashToolPatterns?.length
      ? parsed.bashToolPatterns
      : [];
  state.zeroAccessPaths = normalizeList(parsed.zero_access_paths ?? parsed.zeroAccessPaths);
  state.readOnlyPaths = normalizeList(parsed.read_only_paths ?? parsed.readOnlyPaths);
  state.noDeletePaths = normalizeList(parsed.no_delete_paths ?? parsed.noDeletePaths);
}

async function loadSafetyConfig(cwd: string): Promise<void> {
  const filePath = path.join(cwd, ".agent", "harness", "safety.yaml");
  if (!fs.existsSync(filePath)) {
    state.safetyRules = DEFAULT_RULES;
    state.bashPatterns = [];
    state.zeroAccessPaths = [];
    state.readOnlyPaths = [];
    state.noDeletePaths = [];
    state.defaultAction = "allow";
    return;
  }

  try {
    const yamlModule = await import("yaml");
    const parsed = yamlModule.parse(
      fs.readFileSync(filePath, "utf8"),
    ) as SafetyConfig;
    normalizeSafetyConfig(parsed);
  } catch {
    state.safetyRules = DEFAULT_RULES;
    state.bashPatterns = [];
    state.zeroAccessPaths = [];
    state.readOnlyPaths = [];
    state.noDeletePaths = [];
    state.defaultAction = "allow";
  }
}

function extractPath(input: Record<string, unknown>): string | null {
  const candidates = ["path", "filePath", "filepath", "target"];
  for (const key of candidates) {
    const value = input[key];
    if (typeof value === "string" && value.length > 0) return value;
  }
  return null;
}

function resolvePath(cwd: string, target: string): string {
  if (target.startsWith("~")) {
    return path.join(process.env.HOME || "", target.slice(1));
  }
  return path.isAbsolute(target) ? target : path.resolve(cwd, target);
}

function matchesPathGlobs(cwd: string, targetPath: string, globs: string[]): boolean {
  const resolved = resolvePath(cwd, targetPath);
  const relative = path.relative(cwd, resolved) || targetPath;
  return (
    matchesGlobList(targetPath, globs) ||
    matchesGlobList(resolved, globs) ||
    matchesGlobList(relative, globs)
  );
}

function isGlobPattern(value: string): boolean {
  return /[*?]/.test(value);
}

function shouldMatchRule(
  rule: SafetyRule,
  toolName: string,
  input: Record<string, unknown>,
  cwd: string,
): boolean {
  if (rule.tool && rule.tool !== toolName) return false;
  if (rule.patterns && rule.patterns.length > 0) {
    const command = typeof input.command === "string" ? input.command : "";
    if (!rule.patterns.some((pattern) => command.includes(pattern))) return false;
  }
  if (rule.path_globs && rule.path_globs.length > 0) {
    const targetPath = extractPath(input);
    if (!targetPath) return false;
    if (!matchesPathGlobs(cwd, targetPath, rule.path_globs)) return false;
  }
  return true;
}

function restorePurpose(ctx: ExtensionContext): string | null {
  const entries = ctx.sessionManager.getEntries();
  for (let i = entries.length - 1; i >= 0; i -= 1) {
    const entry = entries[i];
    if (entry.type === "custom" && entry.customType === PURPOSE_ENTRY) {
      const data = entry.data as Record<string, unknown> | undefined;
      const value = data?.value;
      if (typeof value === "string" && value.trim().length > 0) return value.trim();
    }
  }
  return null;
}

function setPurpose(pi: ExtensionAPI, ctx: ExtensionContext, purpose: string): void {
  const trimmed = purpose.trim();
  if (!trimmed) return;
  state.purpose = trimmed;
  pi.appendEntry(PURPOSE_ENTRY, { value: trimmed });
  if (ctx.hasUI) ctx.ui.notify(`Purpose set: ${trimmed}`, "info");
  updateHarnessUI(ctx);
}

export default function registerHarness(pi: ExtensionAPI): void {
  pi.on("session_start", async (_event, ctx) => {
    await loadSafetyConfig(ctx.cwd);
    const restored = restorePurpose(ctx);
    if (restored) state.purpose = restored;
    if (!state.purpose && ctx.hasUI) {
      const purpose = await ctx.ui.input(
        "Session purpose",
        "Describe the goal for this session",
      );
      if (purpose) setPurpose(pi, ctx, purpose);
    }
    updateHarnessUI(ctx);
  });

  pi.on("tool_execution_end", (event, ctx) => {
    const count = state.toolCounts.get(event.toolName) ?? 0;
    state.toolCounts.set(event.toolName, count + 1);
    state.toolTotal += 1;
    updateHarnessUI(ctx);
  });

  pi.on("before_agent_start", (event) => {
    if (!state.purpose) return undefined;
    return {
      systemPrompt:
        `${event.systemPrompt}\n\n<purpose>\nYour singular purpose this session: ${state.purpose}\nStay focused on this goal. If a request drifts from this purpose, gently remind the user.\n</purpose>`,
    };
  });

  pi.on("input", async (_event, ctx) => {
    if (!ctx.hasUI) return { action: "continue" as const };
    if (!state.purpose) {
      ctx.ui.notify("Set a purpose first.", "warning");
      return { action: "handled" as const };
    }
    return { action: "continue" as const };
  });

  pi.on("tool_call", async (event, ctx) => {
    const input = (event.input ?? {}) as Record<string, unknown>;
    for (const rule of state.safetyRules) {
      if (!shouldMatchRule(rule, event.toolName, input, ctx.cwd)) continue;
      const action = rule.action ?? state.defaultAction;
      const reason = rule.reason ?? `Blocked by safety rule ${rule.id}`;
      if (action === "deny") {
        return { block: true, reason };
      }
      if (action === "ask") {
        if (!ctx.hasUI) return { block: true, reason };
        const ok = await ctx.ui.confirm("Safety check", reason);
        if (!ok) return { block: true, reason };
      }
    }

    const targetPath = extractPath(input);
    const resolvedTarget = targetPath ? resolvePath(ctx.cwd, targetPath) : null;
    const command = typeof input.command === "string" ? input.command : "";
    const isWriteTool = event.toolName === "write" || event.toolName === "edit";

    if (resolvedTarget && state.zeroAccessPaths.length > 0) {
      if (matchesPathGlobs(ctx.cwd, resolvedTarget, state.zeroAccessPaths)) {
        return { block: true, reason: "Blocked access to protected path" };
      }
    }

    if (resolvedTarget && isWriteTool && state.readOnlyPaths.length > 0) {
      if (matchesPathGlobs(ctx.cwd, resolvedTarget, state.readOnlyPaths)) {
        return { block: true, reason: "Blocked write to read-only path" };
      }
    }

    if (event.toolName === "bash") {
      for (const rule of state.bashPatterns) {
        const regex = new RegExp(rule.pattern);
        if (!regex.test(command)) continue;
        const action = rule.action ?? "deny";
        const reason = rule.reason ?? "Blocked by bash safety rule";
        if (action === "ask") {
          if (!ctx.hasUI) return { block: true, reason };
          const ok = await ctx.ui.confirm("Safety check", reason);
          if (!ok) return { block: true, reason };
        } else {
          return { block: true, reason };
        }
      }

      for (const pathPattern of state.zeroAccessPaths) {
        if (!isGlobPattern(pathPattern) && command.includes(pathPattern)) {
          return { block: true, reason: `Blocked bash access to ${pathPattern}` };
        }
      }

      for (const pathPattern of state.readOnlyPaths) {
        if (
          !isGlobPattern(pathPattern) &&
          command.includes(pathPattern) &&
          /(^|\\s)(rm|mv|sed|tee|>)/.test(command)
        ) {
          return { block: true, reason: `Blocked bash write to ${pathPattern}` };
        }
      }

      for (const pathPattern of state.noDeletePaths) {
        if (!isGlobPattern(pathPattern) && command.includes(pathPattern) && /(^|\\s)(rm|mv)/.test(command)) {
          return { block: true, reason: `Blocked bash delete of ${pathPattern}` };
        }
      }
    }
    return undefined;
  });

  pi.registerCommand("purpose", {
    description: "Set or update the session purpose",
    handler: async (args, ctx) => {
      const raw = args?.trim();
      const purpose = raw
        ? raw
        : await ctx.ui.input("Session purpose", "Describe the goal for this session");
      if (purpose) setPurpose(pi, ctx, purpose);
    },
  });

  pi.registerCommand("harness-reload", {
    description: "Reload harness safety rules from .agent/harness/safety.yaml",
    handler: async (_args, ctx) => {
      await loadSafetyConfig(ctx.cwd);
      updateHarnessUI(ctx);
      if (ctx.hasUI) ctx.ui.notify("Harness rules reloaded", "info");
    },
  });

  // ── System tools ────────────────────────────────────────────────────────
  pi.registerTool(jailStatusTool);
  pi.registerTool(systemHealthTool);
  pi.registerTool(zfsSnapshotsTool);

  // ── Skill library tools ─────────────────────────────────────────────────
  pi.registerTool(skillsSearchTool);
  pi.registerTool(skillsContentTool);

  // ── Controlplane task tools ─────────────────────────────────────────────
  pi.registerTool(taskDecomposeTool);
  pi.registerTool(taskCreateTool);
  pi.registerTool(taskStatusTool);

  // ── Hostd passthrough ───────────────────────────────────────────────────
  pi.registerTool({
    name: "hostd",
    label: "Hostd",
    description:
      "Execute a privileged hostd operation on the FreeBSD host. " +
      "Available ops: bastille-start, bastille-stop, bastille-restart, bastille-list, " +
      "zfs-snapshot, zfs-list, zfs-create, zfs-rollback, pf-reload, pf-enable, " +
      "browser-clone-create, browser-clone-destroy, browser-clone-reap, browser-clone-force-unmount, " +
      "service-start, service-stop, service-restart, service-status, " +
      "pkg-version, pkg-audit, freebsd-update-status, freebsd-version, " +
      "bastille-pkg-install, bastille-mount-pkg-cache.",
    parameters: Type.Object({
      op: Type.String({ description: "The hostd operation name" }),
      params: Type.Optional(
        Type.Record(Type.String(), Type.String(), {
          description: "Operation parameters as key-value pairs",
        }),
      ),
    }),
    async execute(_toolCallId: string, params: { op: string; params?: Record<string, string> }) {
      const result = await callHostd(params.op, params.params ?? {});
      return {
        content: [{ type: "text" as const, text: result.output }],
        details: result,
      };
    },
  });
}