@ssweens/pi-leash 0.12.0

package/src/hooks/policies.ts

@@ -0,0 +1,315 @@
+import { stat } from "node:fs/promises";
+import { isAbsolute, relative, resolve } from "node:path";
+import { parse } from "../vendor/aliou-sh/index.js";
+import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
+import type { PolicyRule, Protection, ResolvedConfig } from "../config";
+import { emitBlocked } from "../utils/events";
+import { expandGlob, hasGlobChars } from "../utils/glob-expander";
+import {
+  type CompiledPattern,
+  compileFilePatterns,
+  normalizeFilePath,
+} from "../utils/matching";
+import { walkCommands, wordToString } from "../utils/shell-utils";
+
+const DEFAULT_BLOCK_MESSAGES: Record<Protection, string> = {
+  noAccess:
+    "Accessing {file} is not allowed. This file is protected. Ask the user if changes are needed.",
+  readOnly:
+    "Writing to {file} is not allowed. This file is read-only. Use the read tool to inspect it instead of bash commands like cat or ls.",
+  none: "",
+};
+
+const BLOCKED_TOOLS: Record<Protection, Set<string>> = {
+  noAccess: new Set(["read", "write", "edit", "bash", "grep", "find", "ls"]),
+  readOnly: new Set(["write", "edit", "bash"]),
+  none: new Set(),
+};
+
+interface CompiledRule {
+  id: string;
+  protection: Protection;
+  patterns: CompiledPattern[];
+  allowedPatterns: CompiledPattern[];
+  onlyIfExists: boolean;
+  blockMessage: string;
+  enabled: boolean;
+}
+
+async function fileExists(cwd: string, filePath: string): Promise<boolean> {
+  try {
+    await stat(resolve(cwd, filePath));
+    return true;
+  } catch {
+    return false;
+  }
+}
+
+function protectionRank(protection: Protection): number {
+  switch (protection) {
+    case "none":
+      return 0;
+    case "readOnly":
+      return 1;
+    case "noAccess":
+      return 2;
+  }
+}
+
+function compileRules(rules: PolicyRule[]): CompiledRule[] {
+  const compiled: CompiledRule[] = [];
+
+  for (const rule of rules) {
+    const id = rule.id?.trim();
+    if (!id) {
+      console.warn("[pi-leash] Skipping policy rule without id.");
+      continue;
+    }
+
+    if (
+      rule.protection !== "none" &&
+      rule.protection !== "readOnly" &&
+      rule.protection !== "noAccess"
+    ) {
+      console.warn(
+        `[pi-leash] Skipping policy rule "${id}": invalid protection.`,
+      );
+      continue;
+    }
+
+    const normalizedPatterns = (rule.patterns ?? []).filter(
+      (pattern) => pattern.pattern.trim().length > 0,
+    );
+    if (normalizedPatterns.length === 0) {
+      console.warn(
+        `[pi-leash] Skipping policy rule "${id}": missing non-empty patterns.`,
+      );
+      continue;
+    }
+
+    const normalizedAllowedPatterns = (rule.allowedPatterns ?? []).filter(
+      (pattern) => pattern.pattern.trim().length > 0,
+    );
+
+    compiled.push({
+      id,
+      protection: rule.protection,
+      patterns: compileFilePatterns(normalizedPatterns),
+      allowedPatterns: compileFilePatterns(normalizedAllowedPatterns),
+      onlyIfExists: rule.onlyIfExists ?? true,
+      blockMessage:
+        rule.blockMessage ?? DEFAULT_BLOCK_MESSAGES[rule.protection] ?? "",
+      enabled: rule.enabled ?? true,
+    });
+  }
+
+  return compiled;
+}
+
+function maybePathLike(token: string): boolean {
+  return (
+    token.includes("/") ||
+    token.includes(".") ||
+    token.startsWith("~") ||
+    token.startsWith("./") ||
+    token.startsWith("../")
+  );
+}
+
+function normalizeTargetForPolicy(filePath: string, cwd: string): string {
+  const absolute = resolve(cwd, filePath);
+  const rel = relative(cwd, absolute);
+
+  const candidate =
+    rel && !rel.startsWith("..") && !isAbsolute(rel) ? rel : absolute;
+
+  return normalizeFilePath(candidate);
+}
+
+function matchesAnyPolicyPattern(
+  filePath: string,
+  rules: CompiledRule[],
+): boolean {
+  return rules.some(
+    (rule) =>
+      rule.enabled && rule.patterns.some((pattern) => pattern.test(filePath)),
+  );
+}
+
+async function expandCandidate(candidate: string): Promise<string[]> {
+  if (!hasGlobChars(candidate)) return [candidate];
+
+  const matches = await expandGlob(candidate);
+  if (matches.length > 0) return matches;
+
+  return [candidate];
+}
+
+async function extractBashFileTargets(
+  command: string,
+  rules: CompiledRule[],
+  cwd: string,
+): Promise<string[]> {
+  const targets = new Set<string>();
+
+  const maybeAddTarget = async (candidate: string): Promise<void> => {
+    if (!candidate || candidate.startsWith("-")) return;
+
+    const expanded = await expandCandidate(candidate);
+    for (const file of expanded) {
+      const normalized = normalizeTargetForPolicy(file, cwd);
+      if (matchesAnyPolicyPattern(normalized, rules)) {
+        targets.add(normalized);
+      }
+    }
+  };
+
+  try {
+    const { ast } = parse(command);
+    const pending: Promise<void>[] = [];
+
+    walkCommands(ast, (cmd) => {
+      const words = (cmd.words ?? []).map(wordToString);
+      for (let i = 1; i < words.length; i++) {
+        const arg = words[i] as string;
+        pending.push(maybeAddTarget(arg));
+      }
+
+      for (const redir of cmd.redirects ?? []) {
+        const target = wordToString(redir.target);
+        pending.push(maybeAddTarget(target));
+      }
+
+      return false;
+    });
+
+    await Promise.all(pending);
+
+    return [...targets];
+  } catch {
+    const tokenRegex = /"([^"]+)"|'([^']+)'|`([^`]+)`|([^\s"'`<|;&]+)/g;
+
+    for (const match of command.matchAll(tokenRegex)) {
+      const token = match[1] ?? match[2] ?? match[3] ?? match[4] ?? "";
+      if (!token || token.startsWith("-") || !maybePathLike(token)) {
+        continue;
+      }
+
+      const expanded = await expandCandidate(token);
+      for (const file of expanded) {
+        const normalized = normalizeTargetForPolicy(file, cwd);
+        if (matchesAnyPolicyPattern(normalized, rules)) {
+          targets.add(normalized);
+        }
+      }
+    }
+
+    return [...targets];
+  }
+}
+
+async function getEffectiveProtection(
+  filePath: string,
+  compiledRules: CompiledRule[],
+  cwd: string,
+): Promise<{
+  protection: Protection;
+  blockMessage: string;
+  ruleId: string;
+} | null> {
+  let bestMatch: {
+    protection: Protection;
+    blockMessage: string;
+    ruleId: string;
+    rank: number;
+  } | null = null;
+
+  for (const rule of compiledRules) {
+    if (!rule.enabled) continue;
+
+    const matched = rule.patterns.some((pattern) => pattern.test(filePath));
+    if (!matched) continue;
+
+    const allowed = rule.allowedPatterns.some((pattern) =>
+      pattern.test(filePath),
+    );
+    if (allowed) continue;
+
+    if (rule.onlyIfExists && !(await fileExists(cwd, filePath))) continue;
+
+    const rank = protectionRank(rule.protection);
+    if (!bestMatch || rank > bestMatch.rank) {
+      bestMatch = {
+        protection: rule.protection,
+        blockMessage: rule.blockMessage,
+        ruleId: rule.id,
+        rank,
+      };
+    }
+  }
+
+  if (!bestMatch || bestMatch.protection === "none") return null;
+
+  return {
+    protection: bestMatch.protection,
+    blockMessage: bestMatch.blockMessage,
+    ruleId: bestMatch.ruleId,
+  };
+}
+
+function extractPathTarget(input: Record<string, unknown>): string[] {
+  const target = String(input.file_path ?? input.path ?? "").trim();
+  return target ? [target] : [];
+}
+
+export function setupPoliciesHook(pi: ExtensionAPI, config: ResolvedConfig) {
+  if (!config.features.policies) return;
+
+  const compiledRules = compileRules(config.policies.rules);
+
+  pi.on("tool_call", async (event, ctx) => {
+    const toolName = event.toolName;
+    let targets: string[] = [];
+
+    if (["read", "write", "edit", "grep", "find", "ls"].includes(toolName)) {
+      targets = extractPathTarget(event.input);
+    } else if (toolName === "bash") {
+      const command = String(event.input.command ?? "");
+      targets = await extractBashFileTargets(command, compiledRules, ctx.cwd);
+    } else {
+      return;
+    }
+
+    for (const target of targets) {
+      const normalizedTarget = normalizeTargetForPolicy(target, ctx.cwd);
+
+      const effective = await getEffectiveProtection(
+        normalizedTarget,
+        compiledRules,
+        ctx.cwd,
+      );
+      if (!effective) continue;
+
+      const blockedTools = BLOCKED_TOOLS[effective.protection];
+      if (!blockedTools.has(toolName)) continue;
+
+      ctx.ui.notify(
+        `Blocked ${toolName} on protected file: ${normalizedTarget} (${effective.ruleId})`,
+        "warning",
+      );
+
+      const reason = effective.blockMessage.replace("{file}", normalizedTarget);
+
+      emitBlocked(pi, {
+        feature: "policies",
+        toolName,
+        input: event.input,
+        reason,
+      });
+
+      return { block: true, reason };
+    }
+
+    return;
+  });
+}

package/src/index.ts

@@ -0,0 +1,38 @@
+import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
+import { loadConfig } from "./config";
+import { setupLeashHooks } from "./hooks";
+
+/**
+ * Pi Leash Extension
+ *
+ * Security hooks to prevent potentially dangerous operations:
+ * - policies: File access policies with per-rule protection levels
+ * - permission-gate: Prompts for confirmation on dangerous commands
+ * - sudo-mode: Secure password handling for sudo commands
+ *
+ * Configuration:
+ * - Global: ~/.pi/agent/settings/pi-leash.json
+ *
+ * Example config:
+ * {
+ *   "enabled": true,
+ *   "features": {
+ *     "policies": true,
+ *     "permissionGate": true
+ *   },
+ *   "permissionGate": {
+ *     "sudoMode": {
+ *       "enabled": true,
+ *       "timeout": 30000,
+ *       "preserveEnv": false
+ *     }
+ *   }
+ * }
+ */
+export default async function (pi: ExtensionAPI) {
+  const config = loadConfig();
+
+  if (!config.enabled) return;
+
+  setupLeashHooks(pi, config);
+}

package/src/lib/executor.ts

@@ -0,0 +1,280 @@
+/**
+ * Core subagent executor.
+ *
+ * Uses createAgentSession from the SDK for all subagent patterns.
+ * Supports streaming text updates, tool execution tracking, and usage tracking.
+ */
+
+import type { AssistantMessage } from "@mariozechner/pi-ai";
+import type { ExtensionContext } from "@mariozechner/pi-coding-agent";
+import {
+  createAgentSession,
+  DefaultResourceLoader,
+  getAgentDir,
+  SessionManager,
+  SettingsManager,
+} from "@mariozechner/pi-coding-agent";
+import {
+  createExecutionTimer,
+  markExecutionEnd,
+  markExecutionStart,
+} from "./timing";
+import type {
+  OnTextUpdate,
+  OnToolUpdate,
+  SubagentConfig,
+  SubagentResult,
+  SubagentToolCall,
+  SubagentUsage,
+} from "./types";
+
+function generateRunId(name: string): string {
+  const slug =
+    name
+      .trim()
+      .toLowerCase()
+      .replace(/[^a-z0-9]+/g, "-") || "subagent";
+  const randomPart =
+    typeof globalThis.crypto?.randomUUID === "function"
+      ? globalThis.crypto.randomUUID().slice(0, 8)
+      : Date.now().toString(36);
+  return `${slug}-${randomPart}`;
+}
+
+/**
+ * Execute a subagent with the given configuration.
+ *
+ * @param config - Subagent configuration
+ * @param userMessage - The user's prompt
+ * @param ctx - Extension context
+ * @param onTextUpdate - Callback for streaming text
+ * @param signal - Abort signal
+ * @param onToolUpdate - Callback for tool execution updates
+ */
+export async function executeSubagent(
+  config: SubagentConfig,
+  userMessage: string,
+  ctx: ExtensionContext,
+  onTextUpdate?: OnTextUpdate,
+  signal?: AbortSignal,
+  onToolUpdate?: OnToolUpdate,
+): Promise<SubagentResult> {
+  const runId = generateRunId(config.name);
+  const executionTimer = createExecutionTimer();
+
+  const agentDir = getAgentDir();
+  const settingsManager = SettingsManager.create(ctx.cwd, agentDir);
+  const resourceLoader = new DefaultResourceLoader({
+    cwd: ctx.cwd,
+    agentDir,
+    settingsManager,
+    noExtensions: true,
+    noPromptTemplates: true,
+    noThemes: true,
+    noSkills: true,
+    systemPromptOverride: () => config.systemPrompt,
+    appendSystemPromptOverride: () => [],
+    agentsFilesOverride: () => ({ agentsFiles: [] }),
+    skillsOverride: () => ({
+      skills: config.skills ?? [],
+      diagnostics: [],
+    }),
+  });
+  await resourceLoader.reload();
+
+  const { session } = await createAgentSession({
+    model: config.model,
+    tools: config.tools ?? [],
+    customTools: config.customTools ?? [],
+    sessionManager: SessionManager.inMemory(),
+    thinkingLevel: config.thinkingLevel ?? "low",
+    modelRegistry: ctx.modelRegistry,
+    resourceLoader,
+  });
+
+  let accumulated = "";
+  let finalResponse = "";
+  let aborted = false;
+  const toolCalls = new Map<string, SubagentToolCall>();
+
+  let toolsHaveStarted = false;
+  let toolsHaveCompleted = false;
+
+  const usage: SubagentUsage = {
+    inputTokens: 0,
+    outputTokens: 0,
+    cacheReadTokens: 0,
+    cacheWriteTokens: 0,
+    estimatedTokens: 0,
+    llmCost: 0,
+    toolCost: 0,
+    totalCost: 0,
+  };
+
+  const unsubscribe = session.subscribe((event) => {
+    if (event.type === "message_update") {
+      if (event.assistantMessageEvent.type === "text_delta") {
+        const delta = event.assistantMessageEvent.delta;
+        accumulated += delta;
+
+        if (toolsHaveCompleted) {
+          finalResponse += delta;
+        }
+
+        onTextUpdate?.(delta, accumulated);
+      }
+    }
+
+    if (event.type === "tool_execution_start") {
+      toolsHaveStarted = true;
+      toolsHaveCompleted = false;
+      finalResponse = "";
+      const toolCall: SubagentToolCall = {
+        toolCallId: event.toolCallId,
+        toolName: event.toolName,
+        args: event.args ?? {},
+        status: "running",
+      };
+      markExecutionStart(toolCall);
+      toolCalls.set(event.toolCallId, toolCall);
+      onToolUpdate?.([...toolCalls.values()]);
+    }
+
+    if (event.type === "tool_execution_update") {
+      const existing = toolCalls.get(event.toolCallId);
+      if (existing) {
+        existing.args = event.args ?? existing.args;
+        if (event.partialResult) {
+          existing.partialResult = event.partialResult as {
+            content: Array<{ type: string; text?: string }>;
+            details?: unknown;
+          };
+        }
+        onToolUpdate?.([...toolCalls.values()]);
+      }
+    }
+
+    if (event.type === "tool_execution_end") {
+      const existing = toolCalls.get(event.toolCallId);
+      if (existing) {
+        existing.status = event.isError ? "error" : "done";
+        existing.result = event.result;
+        markExecutionEnd(existing);
+        if (event.isError && event.result) {
+          existing.error =
+            typeof event.result === "string"
+              ? event.result
+              : JSON.stringify(event.result);
+        }
+        onToolUpdate?.([...toolCalls.values()]);
+
+        const resultDetails = event.result?.details as
+          | { cost?: number }
+          | undefined;
+        if (resultDetails?.cost !== undefined) {
+          usage.toolCost = (usage.toolCost ?? 0) + resultDetails.cost;
+        }
+      }
+
+      const allDone = [...toolCalls.values()].every(
+        (tc) => tc.status === "done" || tc.status === "error",
+      );
+      if (allDone) {
+        toolsHaveCompleted = true;
+      }
+    }
+
+    if (event.type === "turn_end") {
+      const msg = event.message;
+      if (msg.role === "assistant") {
+        const assistantMsg = msg as AssistantMessage;
+        const msgUsage = assistantMsg.usage;
+        if (msgUsage) {
+          usage.inputTokens = (usage.inputTokens ?? 0) + msgUsage.input;
+          usage.outputTokens = (usage.outputTokens ?? 0) + msgUsage.output;
+          usage.cacheReadTokens =
+            (usage.cacheReadTokens ?? 0) + msgUsage.cacheRead;
+          usage.cacheWriteTokens =
+            (usage.cacheWriteTokens ?? 0) + msgUsage.cacheWrite;
+          usage.llmCost = (usage.llmCost ?? 0) + msgUsage.cost.total;
+        }
+      }
+    }
+  });
+
+  if (signal) {
+    if (signal.aborted) {
+      unsubscribe();
+      session.dispose();
+      return {
+        content: "",
+        aborted: true,
+        toolCalls: [],
+        totalDurationMs: executionTimer.getDurationMs(),
+        runId,
+        usage,
+      };
+    }
+
+    signal.addEventListener(
+      "abort",
+      () => {
+        session.abort();
+        aborted = true;
+      },
+      { once: true },
+    );
+  }
+
+  let error: string | undefined;
+
+  try {
+    await session.prompt(userMessage);
+  } catch (err) {
+    if (signal?.aborted) {
+      aborted = true;
+    } else {
+      error =
+        err instanceof Error
+          ? err.message
+          : typeof err === "string"
+            ? err
+            : JSON.stringify(err);
+    }
+  } finally {
+    unsubscribe();
+    session.dispose();
+  }
+
+  const responseText = toolsHaveStarted ? finalResponse : accumulated;
+  const cleanedContent = filterThinkingTags(responseText);
+
+  const totalRealTokens =
+    (usage.inputTokens ?? 0) +
+    (usage.outputTokens ?? 0) +
+    (usage.cacheReadTokens ?? 0) +
+    (usage.cacheWriteTokens ?? 0);
+  usage.estimatedTokens =
+    totalRealTokens > 0
+      ? totalRealTokens
+      : Math.round(cleanedContent.length / 4);
+
+  usage.totalCost = (usage.llmCost ?? 0) + (usage.toolCost ?? 0);
+
+  return {
+    content: cleanedContent,
+    aborted,
+    toolCalls: [...toolCalls.values()],
+    totalDurationMs: executionTimer.getDurationMs(),
+    error,
+    runId,
+    usage,
+  };
+}
+
+/**
+ * Filter out <thinking>...</thinking> tags from text.
+ */
+export function filterThinkingTags(text: string): string {
+  return text.replace(/<thinking>[\s\S]*?<\/thinking>\s*/g, "");
+}

package/src/lib/index.ts

@@ -0,0 +1,16 @@
+export { executeSubagent, filterThinkingTags } from "./executor";
+export { resolveModel } from "./model-resolver";
+export {
+  createExecutionTimer,
+  markExecutionEnd,
+  markExecutionStart,
+  type TimedExecution,
+} from "./timing";
+export type {
+  OnTextUpdate,
+  OnToolUpdate,
+  SubagentConfig,
+  SubagentResult,
+  SubagentToolCall,
+  SubagentUsage,
+} from "./types";

package/src/lib/model-resolver.ts

@@ -0,0 +1,47 @@
+/**
+ * Model resolution helper for subagents.
+ *
+ * Resolves a model by provider + ID from the model registry.
+ */
+
+import type { Model } from "@mariozechner/pi-ai";
+import type { ExtensionContext } from "@mariozechner/pi-coding-agent";
+
+/**
+ * Find a model by provider and ID.
+ *
+ * @param provider - Provider name (e.g., "openrouter", "anthropic", "openai-codex")
+ * @param modelId - Model ID (e.g., "anthropic/claude-haiku-4.5")
+ * @param ctx - Extension context with modelRegistry
+ * @returns The resolved model
+ * @throws Error if model not found or API key not configured
+ */
+export function resolveModel(
+  provider: string,
+  modelId: string,
+  ctx: ExtensionContext,
+  // biome-ignore lint/suspicious/noExplicitAny: Model type requires any for generic API
+): Model<any> {
+  const available = ctx.modelRegistry.getAvailable();
+  const model = available.find(
+    (m) => m.id === modelId && m.provider === provider,
+  );
+
+  if (model) {
+    return model;
+  }
+
+  // Check if the model exists but the API key is missing
+  const all = ctx.modelRegistry.getAll();
+  const existsWithoutKey = all.some(
+    (m) => m.id === modelId && m.provider === provider,
+  );
+
+  if (existsWithoutKey) {
+    throw new Error(
+      `Model "${modelId}" exists on ${provider} but no valid API key is configured.`,
+    );
+  }
+
+  throw new Error(`Model "${modelId}" not found on provider "${provider}".`);
+}