pi-auto-reinject/src/state.ts

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

/** How a skill was first observed in the session (SPEC §6.2). */
export type SkillSource = "slash" | "skill-block" | "read";

/** Compaction that triggered or skipped re-inject (SPEC §8). */
export type CompactionSource = "auto" | "manual";

/** pi-auto-compact delivery mode (SPEC §6.5, §16). */
export type AutoCompactIntegration = "auto" | "defer" | "immediate" | "off";

export interface TrackedSkill {
	name: string;
	filePath: string;
	baseDir: string;
	firstSeenAt: number;
	lastSeenAt: number;
	sources: SkillSource[];
}

export interface ExtensionState {
	version: 1;
	sessionOverride: boolean | null;
	/** Session override for autoCompactIntegration (SPEC §7.1, §16.4). */
	sessionIntegrationOverride: AutoCompactIntegration | null;
	skills: TrackedSkill[];
	lastCompactionSource: CompactionSource | null;
	/** Skill names awaiting re-inject on the next before_agent_start (SPEC §6.5). */
	pendingReinject: string[];
}

/** Runtime-only; not persisted via appendEntry (SPEC §6.1). */
export interface RuntimeFlags {
	autoCompactDetected: boolean;
	autoCompactIntegration: AutoCompactIntegration;
	/** One-time hint when Pi default compaction coexists with pi-auto-compact (SPEC §16.7). */
	compactionCoexistenceHintShown: boolean;
}

export const STATE_ENTRY_TYPE = "skill-reinject:state";

export function saveState(pi: ExtensionAPI, state: ExtensionState): void {
	pi.appendEntry<ExtensionState>(STATE_ENTRY_TYPE, state);
}

function isExtensionState(data: unknown): data is ExtensionState {
	if (!data || typeof data !== "object") {
		return false;
	}
	const candidate = data as ExtensionState;
	return (
		candidate.version === 1 &&
		(candidate.sessionOverride === null || typeof candidate.sessionOverride === "boolean") &&
		(candidate.sessionIntegrationOverride === null ||
			candidate.sessionIntegrationOverride === "auto" ||
			candidate.sessionIntegrationOverride === "defer" ||
			candidate.sessionIntegrationOverride === "immediate" ||
			candidate.sessionIntegrationOverride === "off" ||
			candidate.sessionIntegrationOverride === undefined) &&
		Array.isArray(candidate.skills) &&
		(candidate.lastCompactionSource === null ||
			candidate.lastCompactionSource === "auto" ||
			candidate.lastCompactionSource === "manual") &&
		Array.isArray(candidate.pendingReinject)
	);
}

/** Latest persisted state on the branch, or null if none (SPEC §6.3). */
export function loadStateFromBranch(branch: SessionEntry[]): ExtensionState | null {
	for (let i = branch.length - 1; i >= 0; i--) {
		const entry = branch[i];
		if (entry.type !== "custom" || entry.customType !== STATE_ENTRY_TYPE || entry.data === undefined) {
			continue;
		}
		if (!isExtensionState(entry.data)) {
			continue;
		}
		return structuredClone(entry.data);
	}
	return null;
}

export function createInitialState(): ExtensionState {
	return {
		version: 1,
		sessionOverride: null,
		sessionIntegrationOverride: null,
		skills: [],
		lastCompactionSource: null,
		pendingReinject: [],
	};
}

/** Copy persisted fields into live session state (SPEC §6.3). */
export function applyExtensionState(target: ExtensionState, loaded: ExtensionState): void {
	target.sessionOverride = loaded.sessionOverride;
	target.sessionIntegrationOverride = loaded.sessionIntegrationOverride ?? null;
	target.skills = loaded.skills;
	target.lastCompactionSource = loaded.lastCompactionSource;
	target.pendingReinject = loaded.pendingReinject;
}

export function createRuntimeFlags(): RuntimeFlags {
	return {
		autoCompactDetected: false,
		autoCompactIntegration: "auto",
		compactionCoexistenceHintShown: false,
	};
}

export interface TrackSkillInput {
	name: string;
	filePath: string;
	baseDir: string;
	source: SkillSource;
	seenAt?: number;
}

/** Upsert by name, merge sources, preserve insertion order (SPEC §6.1). */
export function trackSkill(state: ExtensionState, input: TrackSkillInput): void {
	const seenAt = input.seenAt ?? Date.now();
	const existing = state.skills.find((skill) => skill.name === input.name);
	if (existing) {
		existing.filePath = input.filePath;
		existing.baseDir = input.baseDir;
		existing.lastSeenAt = seenAt;
		if (!existing.sources.includes(input.source)) {
			existing.sources.push(input.source);
		}
		return;
	}
	state.skills.push({
		name: input.name,
		filePath: input.filePath,
		baseDir: input.baseDir,
		firstSeenAt: seenAt,
		lastSeenAt: seenAt,
		sources: [input.source],
	});
}