pi-agent-config/shared/extensions/pi-guard-base.ts

/**
 * pi-guard-base — reusable building blocks for a dedicated Pi agent extension.
 *
 * Extracted from `pi-workspaces/memo-inbox/.pi/extensions/memo-guard.ts`, which
 * has enforced these patterns in production since 2026-07.
 *
 * Provides:
 *   - path containment that is safe against symlink escape
 *   - a path-restricted `read` that shadows the built-in tool
 *   - the two-layer capability guard (setActiveTools + tool_call block)
 *   - output truncation helpers
 *   - a loopback HTTP bridge for exposing host capabilities as tools
 *
 * Import from a scenario extension:
 *
 *   import { installGuard, registerRestrictedRead } from "../../../shared/extensions/pi-guard-base.ts";
 *
 * Rationale and mechanism references: `docs/isolation-baseline.md`,
 * `docs/pi-runtime-notes.md` sections 1, 4, 6, 7, 13, 14.
 */

import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import {
  existsSync,
  lstatSync,
  readFileSync,
  realpathSync,
} from "node:fs";
import { basename, dirname, isAbsolute, join, relative, resolve, sep } from "node:path";

// Pi truncates tool output at these limits itself, but a tool that returns more
// wastes a full round trip. Truncate before returning. (§13)
export const DEFAULT_MAX_BYTES = 50_000;
export const DEFAULT_MAX_LINES = 2_000;

// ---------------------------------------------------------------------------
// Results
// ---------------------------------------------------------------------------

export type ToolText = { type: "text"; text: string };

/** Build a text result. `details` is persisted in the session but is NOT sent
 *  to the model, so put structured payloads there and prose in `text`. (§13) */
export function textResult(text: string, details: Record<string, unknown> = {}) {
  return { content: [{ type: "text" as const, text }], details };
}

/**
 * Truncate from the tail, reporting what was dropped.
 *
 * Byte-aware rather than character-aware because the caps pi applies are byte
 * caps and CJK text is 3 bytes per character in UTF-8.
 */
export function truncate(
  text: string,
  maxBytes: number = DEFAULT_MAX_BYTES,
  maxLines: number = DEFAULT_MAX_LINES,
): { text: string; truncated: boolean } {
  let out = text;
  let truncated = false;

  const lines = out.split("\n");
  if (lines.length > maxLines) {
    out = lines.slice(0, maxLines).join("\n");
    truncated = true;
  }

  let bytes = Buffer.byteLength(out, "utf8");
  if (bytes > maxBytes) {
    // Cut on a character boundary by shrinking until it fits.
    let end = Math.floor((out.length * maxBytes) / bytes);
    while (end > 0 && Buffer.byteLength(out.slice(0, end), "utf8") > maxBytes) {
      end -= Math.max(1, Math.floor(end / 64));
    }
    out = out.slice(0, Math.max(0, end));
    truncated = true;
    bytes = Buffer.byteLength(out, "utf8");
  }

  if (truncated) {
    out += `\n\n[truncated: ${bytes} of ${Buffer.byteLength(text, "utf8")} bytes shown]`;
  }
  return { text: out, truncated };
}

// ---------------------------------------------------------------------------
// Path containment
// ---------------------------------------------------------------------------

/** Whether `candidate` is `root` or lies beneath it. Both must be absolute. */
export function inside(root: string, candidate: string): boolean {
  const rel = relative(root, candidate);
  return rel === "" || (!rel.startsWith(`..${sep}`) && rel !== ".." && !isAbsolute(rel));
}

/**
 * Canonicalise a path, resolving symlinks.
 *
 * Symlinks are resolved *before* any containment check so that a link inside an
 * allowed root cannot point outside it.
 */
export function safeRealPath(candidate: string, allowMissing = false): string {
  const absolute = resolve(candidate);
  if (existsSync(absolute)) return realpathSync(absolute);
  if (!allowMissing) throw new Error(`Path does not exist: ${candidate}`);
  return join(realpathSync(dirname(absolute)), basename(absolute));
}

export interface PathPolicy {
  /** Absolute roots the agent may reach. Resolved and canonicalised on build. */
  roots: string[];
  /** Base for relative inputs. Defaults to the first root. */
  base?: string;
  /** If set, only these lowercase extensions are permitted, e.g. [".md"]. */
  extensions?: string[];
  /** Message shown when a path is outside every root. */
  denyMessage?: string;
}

export interface PathResolver {
  (input: string): string;
  roots: string[];
}

/** Build a resolver that accepts only files inside the policy's roots. */
export function makePathResolver(policy: PathPolicy): PathResolver {
  const roots = policy.roots.map((r) => realpathSync(resolve(r)));
  const base = realpathSync(resolve(policy.base ?? roots[0]));
  const allowed = policy.extensions?.map((e) => e.toLowerCase());
  const deny = policy.denyMessage ?? "Access denied: path is outside the permitted roots.";

  const resolver = ((input: string): string => {
    if (typeof input !== "string" || input.trim() === "") {
      throw new Error("Access denied: empty path.");
    }
    const candidate = safeRealPath(isAbsolute(input) ? input : resolve(base, input));
    if (!roots.some((root) => inside(root, candidate))) throw new Error(deny);
    if (!lstatSync(candidate).isFile()) throw new Error("Access denied: path is not a file.");
    if (allowed) {
      const dot = candidate.lastIndexOf(".");
      const ext = dot === -1 ? "" : candidate.slice(dot).toLowerCase();
      if (!allowed.includes(ext)) {
        throw new Error(`Access denied: only ${allowed.join(", ")} files are readable.`);
      }
    }
    return candidate;
  }) as PathResolver;

  resolver.roots = roots;
  return resolver;
}

// ---------------------------------------------------------------------------
// Restricted `read`
// ---------------------------------------------------------------------------

export interface RestrictedReadOptions extends PathPolicy {
  /** Maximum characters returned in one call. */
  maxChars?: number;
  label?: string;
  description?: string;
}

/**
 * Register a path-restricted `read` that shadows the built-in tool.
 *
 * Two reasons a dedicated agent needs this rather than simply dropping `read`:
 *
 *   1. pi emits the `<available_skills>` block only when a tool named `read` is
 *      active, and skill bodies are loaded through it (progressive disclosure).
 *      Without `read`, every SKILL.md is unreachable. (§1, §2)
 *   2. The built-in `read` accepts absolute paths, so it can reach the service's
 *      own credential files.
 *
 * The policy roots MUST therefore include the skill directory.
 */
export function registerRestrictedRead(pi: ExtensionAPI, options: RestrictedReadOptions): void {
  const resolvePath = makePathResolver(options);
  const maxChars = options.maxChars ?? 80_000;
  const rootList = resolvePath.roots.join(", ");

  pi.registerTool({
    name: "read",
    label: options.label ?? "Read Permitted File",
    description:
      options.description ??
      `Read a UTF-8 text file from the permitted roots (${rootList}). ` +
        `Any other host path is denied. Use this to load a skill file when a task matches its description.`,
    promptSnippet: "read: read a text file from the agent's permitted directories",
    parameters: {
      type: "object",
      properties: {
        path: { type: "string", description: "Absolute path, or relative to the workspace." },
        offset: { type: "number", minimum: 0, description: "Character offset to start at." },
        limit: { type: "number", minimum: 1, maximum: maxChars, description: "Characters to return." },
      },
      required: ["path"],
      additionalProperties: false,
    } as any,
    async execute(_id: string, params: any) {
      // Denials are returned rather than thrown so the model can correct course
      // instead of the turn being marked failed.
      let path: string;
      try {
        path = resolvePath(String(params?.path ?? ""));
      } catch (error) {
        return textResult(`ERROR: ${(error as Error).message}`, { error: true });
      }
      const content = readFileSync(path, "utf8");
      const offset = Math.max(0, Math.floor(Number(params?.offset) || 0));
      const limit = Math.min(maxChars, Math.max(1, Math.floor(Number(params?.limit) || maxChars)));
      const slice = content.slice(offset, offset + limit);
      const { text } = truncate(slice);
      return textResult(text, {
        path,
        offset,
        returned: slice.length,
        truncated: offset + limit < content.length,
      });
    },
  });
}

// ---------------------------------------------------------------------------
// Capability guard
// ---------------------------------------------------------------------------

export interface GuardOptions {
  /** Scenario name, used in block messages. */
  scenario: string;
  /** The complete set of tool names this agent may use. */
  allowedTools: string[];
  /** Called when a tool call is blocked. */
  onBlocked?: (toolName: string) => void;
}

/**
 * Install the two independent capability layers.
 *
 * `setActiveTools` governs which tools the model is offered; the `tool_call`
 * hook blocks execution regardless. Both are needed: the first is the contract,
 * the second survives a mistake in the CLI flags or a later resource discovery
 * that re-activates something. (§4, §14)
 *
 * Re-asserted on `resources_discover` as well as `session_start` because
 * discovery can change the registry after startup.
 */
export function installGuard(pi: ExtensionAPI, options: GuardOptions): void {
  const allowed = new Set(options.allowedTools);
  const restrict = () => pi.setActiveTools([...allowed]);

  pi.on("session_start", async () => restrict());
  pi.on("resources_discover", async () => restrict());

  pi.on("tool_call", async (event: any) => {
    if (!allowed.has(event.toolName)) {
      options.onBlocked?.(event.toolName);
      return {
        block: true,
        reason: `${options.scenario} blocks tool: ${event.toolName}`,
      };
    }
    return undefined;
  });
}

// ---------------------------------------------------------------------------
// Loopback bridge
// ---------------------------------------------------------------------------

export interface BridgeOptions {
  /** Base URL, e.g. http://127.0.0.1:8767. Must be loopback. */
  baseUrl: string;
  /** Shared secret sent as `x-pi-bridge-token`. */
  token?: string;
  timeoutMs?: number;
}

export interface BridgeToolSpec {
  name: string;
  label?: string;
  description: string;
  /** JSON Schema. registerTool accepts a plain schema object — verified. (§5) */
  parameters: Record<string, unknown>;
  promptSnippet?: string;
  promptGuidelines?: string[];
  /** Path appended to baseUrl. Defaults to `/tools/<name>`. */
  path?: string;
  /** Provider-side schema enforcement for structured-output tools. (§12) */
  constrainedSampling?: { type: "json_schema"; strict: "prefer" | "require" };
  /** End the turn on success — for terminating structured-output tools. */
  terminate?: boolean;
  maxBytes?: number;
}

function assertLoopback(baseUrl: string): void {
  const url = new URL(baseUrl);
  const host = url.hostname;
  const ok =
    host === "127.0.0.1" ||
    host === "::1" ||
    host === "localhost" ||
    host.startsWith("127.");
  if (!ok) {
    throw new Error(
      `pi-guard-base: bridge baseUrl must be loopback, got ${host}. ` +
        `A dedicated agent must not be given a route to the wider network.`,
    );
  }
}

/**
 * Register tools that proxy to a loopback HTTP backend.
 *
 * There is no RPC command for the host to inject a tool result, so host
 * capabilities must be exposed this way. (§15, `docs/gateway-patterns.md` #9)
 *
 * Because `registerTool` accepts a plain JSON Schema object, the backend can own
 * the schema and serve it, giving one source of truth instead of a TypeScript
 * copy that drifts.
 */
export function registerBridgeTools(
  pi: ExtensionAPI,
  bridge: BridgeOptions,
  specs: BridgeToolSpec[],
): void {
  assertLoopback(bridge.baseUrl);
  const base = bridge.baseUrl.replace(/\/$/, "");
  const timeoutMs = bridge.timeoutMs ?? 30_000;

  for (const spec of specs) {
    const path = spec.path ?? `/tools/${spec.name}`;
    pi.registerTool({
      name: spec.name,
      label: spec.label ?? spec.name,
      description: spec.description,
      // Without promptSnippet a tool is absent from the prose tool list, even
      // though it stays callable through the provider API. (§7)
      promptSnippet: spec.promptSnippet ?? `${spec.name}: ${spec.description.slice(0, 80)}`,
      promptGuidelines: spec.promptGuidelines,
      parameters: spec.parameters as any,
      constrainedSampling: spec.constrainedSampling,
      async execute(_id: string, params: any, signal?: AbortSignal) {
        const controller = new AbortController();
        const timer = setTimeout(() => controller.abort(), timeoutMs);
        const onAbort = () => controller.abort();
        signal?.addEventListener("abort", onAbort, { once: true });
        try {
          const headers: Record<string, string> = { "content-type": "application/json" };
          if (bridge.token) headers["x-pi-bridge-token"] = bridge.token;

          const response = await fetch(`${base}${path}`, {
            method: "POST",
            headers,
            body: JSON.stringify(params ?? {}),
            signal: controller.signal,
          });

          const body = await response.text();
          if (!response.ok) {
            // Throwing marks the call as an error, which is correct for a
            // backend failure: it is not a valid answer. (§13)
            throw new Error(`${spec.name} failed: HTTP ${response.status} ${body.slice(0, 400)}`);
          }

          let parsed: any;
          try {
            parsed = JSON.parse(body);
          } catch {
            const { text } = truncate(body, spec.maxBytes);
            return { content: [{ type: "text" as const, text }], details: { raw: true } };
          }

          const rendered =
            typeof parsed?.text === "string"
              ? parsed.text
              : JSON.stringify(parsed, null, 2);
          const { text } = truncate(rendered, spec.maxBytes);
          const result: any = {
            content: [{ type: "text" as const, text }],
            details: parsed,
          };
          if (spec.terminate) result.terminate = true;
          return result;
        } finally {
          clearTimeout(timer);
          signal?.removeEventListener("abort", onAbort);
        }
      },
    });
  }
}

/** Fetch tool specifications from the backend so the schema has one owner. */
export async function fetchBridgeSpecs(bridge: BridgeOptions): Promise<BridgeToolSpec[]> {
  assertLoopback(bridge.baseUrl);
  const headers: Record<string, string> = {};
  if (bridge.token) headers["x-pi-bridge-token"] = bridge.token;
  const response = await fetch(`${bridge.baseUrl.replace(/\/$/, "")}/tools`, { headers });
  if (!response.ok) {
    throw new Error(`pi-guard-base: cannot load tool specs: HTTP ${response.status}`);
  }
  const payload = await response.json();
  const specs = Array.isArray(payload) ? payload : payload?.tools;
  if (!Array.isArray(specs)) {
    throw new Error("pi-guard-base: /tools must return an array or { tools: [...] }");
  }
  return specs as BridgeToolSpec[];
}