pi-auto-reinject/src/reinject.ts

import type {
	BeforeAgentStartEventResult,
	ExtensionAPI,
	ExtensionContext,
	SessionCompactEvent,
	Skill,
} from "@earendil-works/pi-coding-agent";
import { existsSync } from "node:fs";
import { expandSkill } from "./expand.js";
import {
	filterSkillsNeedingReinject,
	getKeptEntries,
	skillsPresentInKeptWindow,
} from "./kept.js";
import type { SkillReinjectSettings } from "./settings.js";
import type { CompactionRuntime } from "./compaction.js";
import type { ExtensionState } from "./state.js";

export const DEFERRED_REINJECT_CUSTOM_TYPE = "skill-reinject:inject";

function notifyWarning(ctx: ExtensionContext | undefined, message: string): void {
	if (!ctx?.hasUI) {
		return;
	}
	ctx.ui.notify(message, "warning");
}

function notifySkippedSkill(ctx: ExtensionContext | undefined, skillName: string, reason: string): void {
	notifyWarning(ctx, `skill-reinject: skipped "${skillName}" — ${reason}`);
}

/** First registered skill per name; warn on resourceLoader collisions (SPEC §11). */
export function registeredSkillsByName<T extends Pick<Skill, "name">>(
	registeredSkills: readonly T[],
	ctx?: ExtensionContext,
): Map<string, T> {
	const byName = new Map<string, T>();
	const warned = new Set<string>();
	for (const skill of registeredSkills) {
		if (byName.has(skill.name)) {
			if (!warned.has(skill.name)) {
				warned.add(skill.name);
				notifyWarning(
					ctx,
					`skill-reinject: duplicate skill name "${skill.name}" — using first from resourceLoader`,
				);
			}
			continue;
		}
		byName.set(skill.name, skill);
	}
	return byName;
}

/** Names still registered in resourceLoader (SPEC §5.2). */
export function registeredSkillNames(skills: readonly Pick<Skill, "name">[]): ReadonlySet<string> {
	return new Set(skills.map((skill) => skill.name));
}

/**
 * Skill names to re-inject after compaction: tracked, absent from kept window, still registered (SPEC §5.2).
 * `registeredSkills` comes from resourceLoader — ExtensionContext has no getSkills(); wired in index.ts.
 */
export function planReinject(
	state: ExtensionState,
	_settings: SkillReinjectSettings,
	ctx: ExtensionContext,
	compactionEvent: SessionCompactEvent,
	registeredSkills: readonly Pick<Skill, "name">[],
): string[] {
	const branch = ctx.sessionManager.getBranch();
	const keptEntries = getKeptEntries(branch, compactionEvent.compactionEntry.firstKeptEntryId);
	const trackedNames = state.skills.map((skill) => skill.name);
	const keptPresent = skillsPresentInKeptWindow(keptEntries, trackedNames);
	return filterSkillsNeedingReinject(
		state.skills,
		keptPresent,
		registeredSkillNames(registeredSkills),
	);
}

/** Defer path on session_compact: queue planned skills without sendUserMessage (SPEC §6.5.1, §16.2). */
export function enqueueDeferredReinjectFromCompact(
	state: ExtensionState,
	settings: SkillReinjectSettings,
	ctx: ExtensionContext,
	compactionEvent: SessionCompactEvent,
	registeredSkills: readonly Pick<Skill, "name">[],
): void {
	state.pendingReinject = planReinject(
		state,
		settings,
		ctx,
		compactionEvent,
		registeredSkills,
	);
}

/**
 * Update pending queue after session_compact; each compact recalculates or clears (SPEC §16.6).
 * Manual compaction with default settings keeps stale pending until the next user prompt.
 */
export function applyPendingReinjectAfterCompact(
	state: ExtensionState,
	compactionRuntime: CompactionRuntime,
	shouldReinject: boolean,
	planned: readonly string[],
): void {
	if (shouldReinject) {
		state.pendingReinject = [...planned];
		return;
	}
	if (!compactionRuntime.clearPendingReinjectOnNextUserInput) {
		state.pendingReinject = [];
	}
}

/** Warn when re-injecting many skills; unlimited by default with soft warn above 3 (SPEC §15). */
export function maybeWarnManySkills(
	skillCount: number,
	settings: SkillReinjectSettings,
	ctx?: ExtensionContext,
): void {
	if (skillCount <= 0) {
		return;
	}
	const threshold = settings.maxSkills ?? 3;
	if (skillCount <= threshold) {
		return;
	}
	if (settings.maxSkills !== undefined) {
		notifyWarning(
			ctx,
			`skill-reinject: re-injecting ${skillCount} skills (maxSkills warn threshold: ${settings.maxSkills})`,
		);
		return;
	}
	notifyWarning(ctx, `skill-reinject: re-injecting ${skillCount} tracked skills (soft warn above 3)`);
}

/** Expanded skill-block messages in queue order (SPEC §5.3). */
export function buildReinjectBlocks(
	skillNames: readonly string[],
	state: ExtensionState,
	settings: SkillReinjectSettings,
	registeredSkills: readonly Pick<Skill, "name" | "filePath" | "baseDir">[],
	ctx?: ExtensionContext,
): string[] {
	maybeWarnManySkills(skillNames.length, settings, ctx);
	const registeredByName = registeredSkillsByName(registeredSkills, ctx);
	const blocks: string[] = [];
	for (const name of skillNames) {
		const tracked = state.skills.find((skill) => skill.name === name);
		if (!tracked) {
			continue;
		}
		const registered = registeredByName.get(name);
		if (!registered) {
			notifySkippedSkill(ctx, name, "no longer registered");
			continue;
		}
		if (!existsSync(registered.filePath)) {
			notifySkippedSkill(ctx, name, "SKILL.md not found on disk");
			continue;
		}
		try {
			blocks.push(
				expandSkill(
					{
						name: tracked.name,
						filePath: registered.filePath,
						baseDir: registered.baseDir,
					},
					settings.suffix,
				),
			);
		} catch {
			notifySkippedSkill(ctx, name, "SKILL.md not readable");
		}
	}
	return blocks;
}

/** Combined skill-block user text for pending names in queue order (SPEC §5.3). */
export function buildDeferredReinjectContent(
	pendingNames: readonly string[],
	state: ExtensionState,
	settings: SkillReinjectSettings,
	registeredSkills: readonly Pick<Skill, "name" | "filePath" | "baseDir">[],
	ctx?: ExtensionContext,
): string {
	return buildReinjectBlocks(pendingNames, state, settings, registeredSkills, ctx).join("\n\n");
}

/** Immediate path when agent is idle: first block triggers turn, rest queue as followUp (SPEC §6.5.2). */
export function sendImmediateReinjectIdle(
	pi: ExtensionAPI,
	skillNames: readonly string[],
	state: ExtensionState,
	settings: SkillReinjectSettings,
	registeredSkills: readonly Pick<Skill, "name" | "filePath" | "baseDir">[],
	ctx?: ExtensionContext,
): void {
	const blocks = buildReinjectBlocks(skillNames, state, settings, registeredSkills, ctx);
	blocks.forEach((block, index) => {
		pi.sendUserMessage(block, index === 0 ? undefined : { deliverAs: "followUp" });
	});
}

/** Immediate path when streaming or overflow willRetry: all blocks as followUp (SPEC §5.2, §6.5.2). */
export function sendImmediateReinjectAllFollowUp(
	pi: ExtensionAPI,
	skillNames: readonly string[],
	state: ExtensionState,
	settings: SkillReinjectSettings,
	registeredSkills: readonly Pick<Skill, "name" | "filePath" | "baseDir">[],
	ctx?: ExtensionContext,
): void {
	for (const block of buildReinjectBlocks(skillNames, state, settings, registeredSkills, ctx)) {
		pi.sendUserMessage(block, { deliverAs: "followUp" });
	}
}

/** Force re-inject all tracked registered skills for /skill-reinject now (SPEC §7.1). */
export function reinjectNow(
	pi: ExtensionAPI,
	state: ExtensionState,
	settings: SkillReinjectSettings,
	ctx: ExtensionContext,
	registeredSkills: readonly Pick<Skill, "name" | "filePath" | "baseDir">[],
): void {
	const registered = registeredSkillNames(registeredSkills);
	const skillNames = state.skills.map((skill) => skill.name).filter((name) => registered.has(name));
	if (skillNames.length === 0) {
		if (ctx.hasUI) {
			ctx.ui.notify("skill-reinject: no tracked skills to re-inject", "info");
		}
		return;
	}
	if (ctx.isIdle()) {
		sendImmediateReinjectIdle(pi, skillNames, state, settings, registeredSkills, ctx);
		return;
	}
	sendImmediateReinjectAllFollowUp(pi, skillNames, state, settings, registeredSkills, ctx);
}

/**
 * Manual `/compact` with default settings: drop stale pending on the next user prompt (SPEC §16.5, §12.3).
 * Returns true when pending was cleared.
 */
export function clearPendingReinjectOnUserPrompt(
	state: ExtensionState,
	compactionRuntime: CompactionRuntime,
): boolean {
	if (!compactionRuntime.clearPendingReinjectOnNextUserInput) {
		return false;
	}
	compactionRuntime.clearPendingReinjectOnNextUserInput = false;
	const hadPending = state.pendingReinject.length > 0;
	state.pendingReinject = [];
	return hadPending;
}

/**
 * Defer path on before_agent_start: inject one combined message, then clear queue (SPEC §6.5.1).
 * Returns undefined when pendingReinject is empty or manual compaction scheduled a clear (SPEC §16.5).
 */
export function tryConsumeDeferredReinject(
	state: ExtensionState,
	settings: SkillReinjectSettings,
	registeredSkills: readonly Pick<Skill, "name" | "filePath" | "baseDir">[],
	ctx?: ExtensionContext,
	compactionRuntime?: CompactionRuntime,
): BeforeAgentStartEventResult | undefined {
	if (compactionRuntime?.clearPendingReinjectOnNextUserInput) {
		return undefined;
	}
	if (state.pendingReinject.length === 0) {
		return undefined;
	}
	const pendingNames = [...state.pendingReinject];
	const content = buildDeferredReinjectContent(pendingNames, state, settings, registeredSkills, ctx);
	if (!content) {
		state.pendingReinject = [];
		return undefined;
	}
	state.pendingReinject = [];
	return {
		message: {
			customType: DEFERRED_REINJECT_CUSTOM_TYPE,
			content,
			display: true,
		},
	};
}