ThothII/backend/src/pi/managed-config.ts

import {
  chmodSync, closeSync, constants, fstatSync, lstatSync, mkdtempSync, openSync,
  readFileSync, readdirSync, rmSync, symlinkSync, writeFileSync, type Dirent,
} from "node:fs";
import { homedir, tmpdir } from "node:os";
import { join, resolve } from "node:path";

const MAX_AGENT_CONFIG_BYTES = 1024 * 1024;

export const PI_MANAGED_CONFIG_ERROR_CODE = "PI_MANAGED_CONFIG_INVALID";
export const PI_MANAGED_CONFIG_ERROR_MESSAGE = "Pi provider/model configuration is invalid";

export class PiManagedConfigError extends Error {
  readonly code = PI_MANAGED_CONFIG_ERROR_CODE;

  constructor() {
    super(PI_MANAGED_CONFIG_ERROR_MESSAGE);
  }
}

export function isPiManagedConfigError(error: unknown): boolean {
  return Boolean(
    error && typeof error === "object"
    && (error as { code?: unknown }).code === PI_MANAGED_CONFIG_ERROR_CODE,
  );
}

export function parsePiConfigJson(raw: string): unknown {
  try {
    return JSON.parse(raw);
  } catch {
    throw new PiManagedConfigError();
  }
}

/** Reject every Pi shell-backed configuration value, including unknown future nested fields. */
export function assertDeclarativePiConfig(value: unknown): void {
  const pending: unknown[] = [value];
  while (pending.length > 0) {
    const current = pending.pop();
    if (typeof current === "string") {
      if (current.startsWith("!")) throw new PiManagedConfigError();
      continue;
    }
    if (Array.isArray(current)) {
      for (const nested of current) pending.push(nested);
      continue;
    }
    if (current && typeof current === "object") {
      for (const nested of Object.values(current as Record<string, unknown>)) pending.push(nested);
    }
  }
}

export function validateDeclarativePiConfig(raw: string): void {
  assertDeclarativePiConfig(parsePiConfigJson(raw));
}

/** Return the selected provider's declarative apiKey value without knowing provider IDs in code. */
export function configuredPiProviderApiKey(
  raw: string | undefined,
  provider: string | undefined,
): string | undefined {
  if (raw === undefined || provider === undefined) return undefined;
  const parsed = parsePiConfigJson(raw);
  if (!parsed || typeof parsed !== "object" || Array.isArray(parsed)) {
    throw new PiManagedConfigError();
  }
  const providers = (parsed as { providers?: unknown }).providers;
  if (!providers || typeof providers !== "object" || Array.isArray(providers)) {
    throw new PiManagedConfigError();
  }
  const entry = Object.entries(providers as Record<string, unknown>)
    .find(([id]) => id.trim().toLowerCase() === provider.trim().toLowerCase());
  if (!entry) return undefined;
  const config = entry[1];
  if (!config || typeof config !== "object" || Array.isArray(config)) {
    throw new PiManagedConfigError();
  }
  assertDeclarativePiConfig(config);
  const apiKey = (config as { apiKey?: unknown }).apiKey;
  if (apiKey === undefined) return undefined;
  if (typeof apiKey !== "string" || apiKey.length === 0) throw new PiManagedConfigError();
  return apiKey;
}

function configuredPiAgentDir(): string {
  return resolve(process.env.PI_CODING_AGENT_DIR ?? join(homedir(), ".pi", "agent"));
}

function readPiAgentFile(
  configuredAgentDir: string,
  name: "auth.json" | "models.json",
  optional: boolean,
): string | undefined {
  const path = join(configuredAgentDir, name);
  let fd: number | undefined;
  try {
    const before = lstatSync(path);
    if (!before.isFile() || before.isSymbolicLink() || before.size > MAX_AGENT_CONFIG_BYTES) {
      throw new PiManagedConfigError();
    }
    fd = openSync(path, constants.O_RDONLY | constants.O_NOFOLLOW);
    const opened = fstatSync(fd);
    if (!opened.isFile() || opened.size > MAX_AGENT_CONFIG_BYTES
      || before.dev !== opened.dev || before.ino !== opened.ino) {
      throw new PiManagedConfigError();
    }
    return readFileSync(fd, "utf8");
  } catch (error) {
    if (optional && (error as NodeJS.ErrnoException)?.code === "ENOENT") return undefined;
    throw new PiManagedConfigError();
  } finally {
    if (fd !== undefined) {
      try { closeSync(fd); } catch { /* preserve the stable validation outcome */ }
    }
  }
}

export function readConfiguredPiAgentFile(name: "auth.json"): string;
export function readConfiguredPiAgentFile(name: "auth.json", optional: true): string | undefined;
export function readConfiguredPiAgentFile(name: "models.json", optional: true): string | undefined;
export function readConfiguredPiAgentFile(
  name: "auth.json" | "models.json",
  optional = false,
): string | undefined {
  return readPiAgentFile(configuredPiAgentDir(), name, optional);
}

export interface PiRuntimeAgentSnapshot {
  agentDir: string;
  sessionDir: string;
  models?: string;
  cleanup: () => void;
}

/**
 * Bind a session Pi process to the exact managed auth/model bytes validated at spawn time.
 * Other agent resources remain live through symlinks, while session storage stays persistent.
 */
export function createPiRuntimeAgentSnapshot(): PiRuntimeAgentSnapshot {
  const sourceAgentDir = configuredPiAgentDir();
  const auth = readPiAgentFile(sourceAgentDir, "auth.json", true);
  const models = readPiAgentFile(sourceAgentDir, "models.json", true);
  if (auth !== undefined) validateDeclarativePiConfig(auth);
  if (models !== undefined) validateDeclarativePiConfig(models);

  let snapshotDir: string | undefined;
  try {
    snapshotDir = mkdtempSync(join(tmpdir(), "thoth-pi-runtime-agent-"));
    chmodSync(snapshotDir, 0o700);
    let entries: Dirent[];
    try {
      entries = readdirSync(sourceAgentDir, { withFileTypes: true });
    } catch (error) {
      if ((error as NodeJS.ErrnoException)?.code !== "ENOENT") throw error;
      entries = [];
    }
    for (const entry of entries) {
      if (entry.name === "auth.json" || entry.name === "models.json") continue;
      symlinkSync(
        join(sourceAgentDir, entry.name),
        join(snapshotDir, entry.name),
        entry.isDirectory() ? (process.platform === "win32" ? "junction" : "dir") : "file",
      );
    }
    if (auth !== undefined) {
      writeFileSync(join(snapshotDir, "auth.json"), auth, { flag: "wx", mode: 0o600 });
    }
    if (models !== undefined) {
      writeFileSync(join(snapshotDir, "models.json"), models, { flag: "wx", mode: 0o600 });
    }
  } catch {
    if (snapshotDir !== undefined) {
      try { rmSync(snapshotDir, { recursive: true, force: true }); } catch { /* sanitized */ }
    }
    throw new PiManagedConfigError();
  }

  let cleaned = false;
  return {
    agentDir: snapshotDir,
    sessionDir: process.env.PI_CODING_AGENT_SESSION_DIR || join(sourceAgentDir, "sessions"),
    models,
    cleanup: () => {
      if (cleaned) return;
      cleaned = true;
      try { rmSync(snapshotDir, { recursive: true, force: true }); } catch { /* sanitized */ }
    },
  };
}