pi-ext

// ultra settings: validates the `ultra` block from Pi's user and trusted project settings.

import type {
	ExtensionAPI,
	ExtensionContext,
} from "@earendil-works/pi-coding-agent";
import {
	resolveSetting,
	type SettingDeclaration,
} from "./src/pi-ext-settings.ts";

export const THINKING_LEVELS = [
	"off",
	"minimal",
	"low",
	"medium",
	"high",
	"xhigh",
	"max",
] as const satisfies readonly Parameters<ExtensionAPI["setThinkingLevel"]>[0][];

export const DEFAULT_FANOUT_TOOLS = ["read", "grep", "find", "ls", "bash"];

export interface ModelTier {
	/** null selects the session model. */
	model: string | null;
	thinkingLevel: (typeof THINKING_LEVELS)[number];
	/** Selection guidance for the orchestrator, not instructions for the child. */
	instructions: string;
}

export type ModelTiers = Record<string, ModelTier>;
/** Provider name → tier name → per-field overrides for when that provider is active. */
export type ProviderModelTiers = Record<
	string,
	Record<string, Partial<ModelTier>>
>;

export const DEFAULT_MODEL_TIERS: ModelTiers = {
	tiny: {
		model: null,
		thinkingLevel: "low",
		instructions: "No unresolved judgement.",
	},
	small: {
		model: null,
		thinkingLevel: "low",
		instructions: "Few local decisions; low consequence.",
	},
	medium: {
		model: null,
		thinkingLevel: "medium",
		instructions:
			"Several related decisions; moderate ambiguity or consequence.",
	},
	large: {
		model: null,
		thinkingLevel: "high",
		instructions:
			"Many dependent decisions; substantial ambiguity or consequence.",
	},
	huge: {
		model: null,
		thinkingLevel: "max",
		instructions:
			"Exceptional residual burden after decomposition, context, and verification.",
	},
};

export interface UltraSettings {
	/** When false, the tool + command refuse to run. */
	enabled: boolean;
	/** Activate `run_workflow` when the session starts. */
	autoEnable: boolean;
	/** Max concurrent sub-agents in a `fanout` phase. */
	concurrency: number;
	/** Result-submission reminders after an agent stops without output. */
	maxRetries: number;
	/** Extensions loaded into every sub-agent: all discovered extensions (`"all"`,
	 * default), none (`"none"`), or one static list of sibling pi-ext names.
	 * `step.tools` selects the initial tools; loaded extensions may change them. */
	subagentExtensions: "none" | "all" | string[];
	/** Additional initial tools permitted in fanout without writeIsolation, merged with defaults. */
	fanoutToolAllowlist: string[];
	/** Load host skills into sub-agents (`"all"`, default) or omit them (`"none"`). */
	subagentSkills: "none" | "all";
	/** Retain parent-linked sub-agent sessions, or run them entirely in memory. */
	subagentSessionRetention: "all" | "none";
	/** Settings-defined model, thinking, and selection policy used by step.model. */
	modelTiers: ModelTiers;
	/** Per-provider tier overrides folded over `modelTiers` when the session model's
	 * provider matches; providers without a block keep the base tiers. */
	providerModelTiers: ProviderModelTiers;
	/** Fast model used to turn raw tool args into board action phrases; null disables it. */
	actionSummarizerModel: string | null;
}

/** Documented defaults (design §Settings table). */
export const ULTRA_DEFAULTS: UltraSettings = {
	enabled: true,
	autoEnable: false,
	concurrency: 5,
	maxRetries: 2,
	subagentExtensions: "all",
	fanoutToolAllowlist: [...DEFAULT_FANOUT_TOOLS],
	subagentSkills: "all",
	subagentSessionRetention: "all",
	modelTiers: DEFAULT_MODEL_TIERS,
	providerModelTiers: {},
	actionSummarizerModel: null,
};

function isPlainObject(v: unknown): v is Record<string, unknown> {
	return v !== null && typeof v === "object" && !Array.isArray(v);
}

function pickModelTiers(
	raw: unknown,
	prefix = "ultra.modelTiers",
): Record<string, Partial<ModelTier>> | undefined {
	if (raw === undefined) return undefined;
	if (!isPlainObject(raw)) throw new Error(`${prefix} must be an object.`);
	const out: Record<string, Partial<ModelTier>> = {};
	for (const [name, value] of Object.entries(raw)) {
		const key = `${prefix}.${name}`;
		if (!/^[a-z][a-z0-9_-]*$/.test(name) || name in Object.prototype)
			throw new Error(`${key}: invalid tier name.`);
		if (!isPlainObject(value)) throw new Error(`${key} must be an object.`);
		const tier: Partial<ModelTier> = {};
		for (const field of Object.keys(value)) {
			if (!["model", "thinkingLevel", "instructions"].includes(field))
				throw new Error(`${key}: unknown field ${field}.`);
		}
		if ("model" in value) {
			if (
				value.model !== null &&
				(typeof value.model !== "string" ||
					!/^[^\s/:]+\/\S+$/.test(value.model) ||
					/:(off|minimal|low|medium|high|xhigh|max)$/.test(value.model))
			)
				throw new Error(
					`${key}.model must be provider/model or null; set thinkingLevel separately.`,
				);
			tier.model = value.model as string | null;
		}
		if ("thinkingLevel" in value) {
			if (
				!THINKING_LEVELS.includes(
					value.thinkingLevel as ModelTier["thinkingLevel"],
				)
			)
				throw new Error(
					`${key}.thinkingLevel must be ${THINKING_LEVELS.join(", ")}.`,
				);
			tier.thinkingLevel = value.thinkingLevel as ModelTier["thinkingLevel"];
		}
		if ("instructions" in value) {
			if (typeof value.instructions !== "string" || !value.instructions.trim())
				throw new Error(`${key}.instructions must be a non-empty string.`);
			tier.instructions = value.instructions;
		}
		out[name] = tier;
	}
	return out;
}

function pickProviderModelTiers(raw: unknown): ProviderModelTiers | undefined {
	if (raw === undefined) return undefined;
	if (!isPlainObject(raw))
		throw new Error("ultra.providerModelTiers must be an object.");
	const out: ProviderModelTiers = {};
	for (const [provider, tiers] of Object.entries(raw)) {
		if (!/^[^\s/:]+$/.test(provider) || provider in Object.prototype)
			throw new Error(
				`ultra.providerModelTiers.${provider}: invalid provider name.`,
			);
		out[provider] =
			pickModelTiers(tiers, `ultra.providerModelTiers.${provider}`) ?? {};
	}
	return out;
}

export function fanoutToolPolicyPrompt(allowed: readonly string[]): string {
	return `Fanout without writeIsolation: initial tools allowed: ${allowed.join(", ")}. Other initial tools require genuine disjoint ownership declared in writeIsolation. Without it, child prompts must forbid mutation. Tool availability is checked separately.`;
}

/** The sole tier rubric, rendered from the effective settings for each agent turn. */
export function modelTierPrompt(tiers: ModelTiers): string {
	return [
		"# ultra model tier policy",
		"Orchestrator only. Rows: tier: model; default thinkingLevel; selection guidance.",
		"Default small; tiny only for determined outcomes. Choose lowest adequate tier; lower wins ties.",
		"Judge residual decisions, dependencies, ambiguity, and consequence; credit upstream findings.",
		"Decompose only when reduced decisions, context, or error propagation outweigh handoffs.",
		"Use minimum sufficient context; verify proportionately to risk, adding a check when one child is insufficient.",
		"Ask user for missing decisions; report missing child evidence or authority as blockers, not reasons to escalate tiers.",
		...Object.entries(tiers).map(
			([name, tier]) =>
				`- ${name}: ${tier.model ?? "session model"}; ${tier.thinkingLevel}; ${tier.instructions}`,
		),
	].join("\n");
}

function oneOf<T extends string>(
	key: string,
	value: unknown,
	options: readonly T[],
): T {
	if (!options.includes(value as T))
		throw new Error(
			`${key} must be ${options.map((option) => `"${option}"`).join(" or ")}.`,
		);
	return value as T;
}

/** Validate the merged `ultra` block and fold it over the defaults.
 * Tiers merge per field over DEFAULT_MODEL_TIERS; fanoutToolAllowlist extends the default tools. */
export function parseUltraSettings(raw: unknown): UltraSettings {
	if (!isPlainObject(raw)) throw new Error("ultra must be an object.");
	const settings: UltraSettings = { ...ULTRA_DEFAULTS };
	for (const [field, value] of Object.entries(raw)) {
		const key = `ultra.${field}`;
		switch (field) {
			case "enabled":
			case "autoEnable":
				if (typeof value !== "boolean")
					throw new Error(`${key} must be a boolean.`);
				settings[field] = value;
				break;
			case "concurrency":
			case "maxRetries":
				if (typeof value !== "number")
					throw new Error(`${key} must be a number.`);
				settings[field] = value;
				break;
			case "subagentExtensions":
				if (value === "all" || value === "none")
					settings.subagentExtensions = value;
				else if (
					Array.isArray(value) &&
					value.every((name) => typeof name === "string" && name.trim())
				)
					settings.subagentExtensions = [
						...new Set(value.map((name: string) => name.trim())),
					];
				else
					throw new Error(
						`${key} must be "all", "none", or an array of extension names.`,
					);
				break;
			case "fanoutToolAllowlist":
				if (
					!Array.isArray(value) ||
					value.some(
						(tool) => typeof tool !== "string" || !tool || /\s/u.test(tool),
					)
				)
					throw new Error(
						`${key} must be an array of non-empty tool names without whitespace.`,
					);
				settings.fanoutToolAllowlist = [
					...new Set([...DEFAULT_FANOUT_TOOLS, ...value]),
				];
				break;
			case "subagentSkills":
			case "subagentSessionRetention":
				settings[field] = oneOf(key, value, ["all", "none"] as const);
				break;
			case "modelTiers":
			case "providerModelTiers":
				break;
			case "actionSummarizerModel":
				if (value !== null && (typeof value !== "string" || !value.trim()))
					throw new Error(`${key} must be a provider/model string or null.`);
				settings.actionSummarizerModel = value?.trim() ?? null;
				break;
			default:
				throw new Error(`unknown key ${key}.`);
		}
	}
	const modelTiers: Record<string, Partial<ModelTier>> = {};
	for (const tiers of [DEFAULT_MODEL_TIERS, pickModelTiers(raw.modelTiers)])
		for (const [name, tier] of Object.entries(tiers ?? {}))
			modelTiers[name] = { ...modelTiers[name], ...tier };
	for (const [name, tier] of Object.entries(modelTiers))
		if (
			tier.model === undefined ||
			tier.thinkingLevel === undefined ||
			tier.instructions === undefined
		)
			throw new Error(
				`ultra.modelTiers.${name} requires model, thinkingLevel, and instructions.`,
			);
	settings.modelTiers = modelTiers as ModelTiers;
	settings.providerModelTiers =
		pickProviderModelTiers(raw.providerModelTiers) ?? {};
	for (const [provider, tiers] of Object.entries(settings.providerModelTiers))
		for (const name of Object.keys(tiers))
			if (!Object.hasOwn(modelTiers, name))
				throw new Error(
					`ultra.providerModelTiers.${provider}.${name}: unknown tier.`,
				);
	return settings;
}

// Precedence: trusted project, user, defaults; project and user blocks deep-merge.
export const ULTRA_SETTING: SettingDeclaration<UltraSettings> = {
	key: "ultra",
	parse: parseUltraSettings,
	default: ULTRA_DEFAULTS,
};

/** Resolve effective settings; invalid configuration throws through Pi's extension error path. */
export function loadUltraSettings(
	pi: Pick<ExtensionAPI, "getFlag">,
	ctx: Pick<ExtensionContext, "cwd" | "isProjectTrusted">,
): UltraSettings {
	const result = resolveSetting(pi, ctx, ULTRA_SETTING);
	if (!result.ok) throw new Error(result.error);
	return result.value;
}

/** Fold provider-scoped overrides over the base tiers; unknown providers keep the base tiers. */
export function effectiveModelTiers(
	settings: UltraSettings,
	provider: string | undefined,
): ModelTiers {
	const overrides =
		provider === undefined ? undefined : settings.providerModelTiers[provider];
	if (!overrides) return settings.modelTiers;
	const out: ModelTiers = {};
	for (const [name, tier] of Object.entries(settings.modelTiers))
		out[name] = { ...tier, ...overrides[name] };
	return out;
}