pi-auto-reinject/src/index.ts

import { dirname } from "node:path";
import {
	isToolCallEventType,
	type ExtensionAPI,
	type ExtensionContext,
	type SessionCompactEvent,
	type Skill,
} from "@earendil-works/pi-coding-agent";
import { detectAndCachePiAutoCompact, resolveDeliveryMode } from "./auto-compact.js";
import { registerSkillReinjectCommand, updateSkillReinjectStatusLine } from "./commands.js";
import {
	consumeCompactionOnSessionCompact,
	createCompactionRuntime,
	ensureCompactionSourceMarked,
	markAutoCompactionBeforeCompact,
	markManualCompactionFromInput,
} from "./compaction.js";
import { buildReinjectDiagSnapshot, notifyReinjectDiag } from "./diag.js";
import { detectSlashSkill, matchReadPathToSkillWhenEnabled, parseSkillBlocksFromText, userMessageText } from "./detect.js";
import {
	getKeptEntries,
	skillsPresentInKeptWindow,
} from "./kept.js";
import {
	applyPendingReinjectAfterCompact,
	clearPendingReinjectOnUserPrompt,
	planDeferredReinject,
	planReinject,
	sendImmediateReinjectAllFollowUp,
	sendImmediateReinjectIdle,
	tryConsumeDeferredReinject,
} from "./reinject.js";
import { readSettings } from "./settings.js";
import { rescanSkillsFromBranch } from "./rescan.js";
import { findRegisteredSkillByName, resolveRegisteredSkills } from "./skills-registry.js";
import {
	applyExtensionState,
	createInitialState,
	createRuntimeFlags,
	loadStateFromBranch,
	saveState,
	trackSkill,
	type TrackSkillInput,
} from "./state.js";

export default function skillReinject(pi: ExtensionAPI): void {
	const state = createInitialState();
	const runtime = createRuntimeFlags();
	const compactionRuntime = createCompactionRuntime();
	let registeredSkills: Skill[] = [];

	function persistState(): void {
		saveState(pi, state);
	}

	registerSkillReinjectCommand(pi, {
		pi,
		state,
		runtime,
		getRegisteredSkills: () => registeredSkills,
		persistState,
	});

	function trackSkillAndPersist(input: TrackSkillInput, ctx?: ExtensionContext): void {
		trackSkill(state, input);
		persistState();
		if (ctx) {
			updateSkillReinjectStatusLine(ctx, state, readSettings(ctx));
		}
	}

	function trackReadSkillPath(path: string, ctx: ExtensionContext): void {
		const settings = readSettings(ctx);
		const skills = resolveRegisteredSkills(ctx.cwd, registeredSkills);
		const matched = matchReadPathToSkillWhenEnabled(path, skills, settings.trackReadPaths);
		if (!matched) {
			return;
		}
		trackSkillAndPersist({
			name: matched.name,
			filePath: matched.filePath,
			baseDir: matched.baseDir,
			source: "read",
		}, ctx);
	}

	function restoreSessionState(ctx: ExtensionContext): void {
		detectAndCachePiAutoCompact(pi, runtime);
		const settings = readSettings(ctx);
		const branch = ctx.sessionManager.getBranch();
		const loaded = loadStateFromBranch(branch);
		applyExtensionState(state, loaded ?? createInitialState());
		if (!loaded) {
			rescanSkillsFromBranch(state, branch, ctx.cwd, registeredSkills, settings.trackReadPaths);
		}
		updateSkillReinjectStatusLine(ctx, state, settings);
	}

	function handleSessionCompact(event: SessionCompactEvent, ctx: ExtensionContext): void {
		const settings = readSettings(ctx);
		const skills = resolveRegisteredSkills(ctx.cwd, registeredSkills);
		const branch = ctx.sessionManager.getBranch();
		compactionRuntime.lastCompactionFirstKeptEntryId = event.compactionEntry.firstKeptEntryId;
		const trackedNames = state.skills.map((skill) => skill.name);
		const keptEntries = getKeptEntries(branch, event.compactionEntry.firstKeptEntryId);
		const keptPresent = skillsPresentInKeptWindow(keptEntries, trackedNames);
		ensureCompactionSourceMarked(compactionRuntime);
		const shouldReinject = consumeCompactionOnSessionCompact(
			compactionRuntime,
			state,
			state.sessionOverride,
			settings,
		);
		const deliveryMode = resolveDeliveryMode(settings, runtime, state.sessionIntegrationOverride);
		const planned =
			deliveryMode === "defer"
				? planDeferredReinject(state, ctx, event)
				: planReinject(state, settings, ctx, event, skills);

		applyPendingReinjectAfterCompact(state, compactionRuntime, shouldReinject, planned);
		notifyReinjectDiag(
			ctx,
			settings,
			"session_compact",
			buildReinjectDiagSnapshot(state, skills, keptPresent, planned),
		);

		if (deliveryMode === "defer") {
			persistState();
			return;
		}

		if (!shouldReinject) {
			persistState();
			return;
		}

		if (planned.length === 0) {
			persistState();
			return;
		}
		if (ctx.isIdle()) {
			sendImmediateReinjectIdle(pi, planned, state, settings, skills, ctx);
		} else {
			sendImmediateReinjectAllFollowUp(pi, planned, state, settings, skills, ctx);
		}
		persistState();
	}

	pi.on("session_start", async (_event, ctx) => {
		restoreSessionState(ctx);
	});

	pi.on("session_tree", async (_event, ctx) => {
		restoreSessionState(ctx);
	});

	pi.on("session_shutdown", async () => {
		persistState();
	});

	pi.on("session_before_compact", async () => {
		markAutoCompactionBeforeCompact(compactionRuntime);
	});

	pi.on("session_compact", async (event, ctx) => {
		handleSessionCompact(event, ctx);
	});

	pi.on("before_agent_start", async (event, ctx) => {
		registeredSkills = event.systemPromptOptions.skills ?? registeredSkills;

		const settings = readSettings(ctx);
		const skills = resolveRegisteredSkills(ctx.cwd, registeredSkills);
		const trackedNames = state.skills.map((skill) => skill.name);
		const keptPresent =
			compactionRuntime.lastCompactionFirstKeptEntryId === null
				? new Set<string>()
				: skillsPresentInKeptWindow(
						getKeptEntries(
							ctx.sessionManager.getBranch(),
							compactionRuntime.lastCompactionFirstKeptEntryId,
						),
						trackedNames,
					);
		notifyReinjectDiag(
			ctx,
			settings,
			"before_agent_start",
			buildReinjectDiagSnapshot(state, skills, keptPresent, []),
		);
		const pendingBefore = state.pendingReinject.length;
		const deferred = tryConsumeDeferredReinject(
			state,
			settings,
			registeredSkills,
			ctx,
			compactionRuntime,
		);
		if (deferred) {
			compactionRuntime.lastCompactionFirstKeptEntryId = null;
		}
		if (pendingBefore > 0) {
			persistState();
		}
		if (deferred) {
			return deferred;
		}
	});

	pi.on("input", async (event, ctx) => {
		if (event.source === "extension") {
			return { action: "continue" };
		}

		markManualCompactionFromInput(event.text, compactionRuntime);

		if (clearPendingReinjectOnUserPrompt(state, compactionRuntime)) {
			persistState();
		}

		const skillName = detectSlashSkill(event.text);
		if (skillName) {
			const skills = resolveRegisteredSkills(ctx.cwd, registeredSkills);
			const skill = findRegisteredSkillByName(skills, skillName);
			if (skill) {
				trackSkillAndPersist({
					name: skill.name,
					filePath: skill.filePath,
					baseDir: skill.baseDir,
					source: "slash",
				}, ctx);
			}
		}

		return { action: "continue" };
	});

	pi.on("message_end", async (event, ctx) => {
		if (event.message.role !== "user") {
			return;
		}

		const blocks = parseSkillBlocksFromText(userMessageText(event.message.content));
		if (blocks.length === 0) {
			return;
		}

		const skills = resolveRegisteredSkills(ctx.cwd, registeredSkills);
		for (const block of blocks) {
			const registered = findRegisteredSkillByName(skills, block.name);
			trackSkill(state, {
				name: block.name,
				filePath: registered?.filePath ?? block.location,
				baseDir: registered?.baseDir ?? dirname(block.location),
				source: "skill-block",
			});
		}
		persistState();
		updateSkillReinjectStatusLine(ctx, state, readSettings(ctx));
	});

	pi.on("tool_call", async (event, ctx) => {
		if (!isToolCallEventType("read", event)) {
			return;
		}
		trackReadSkillPath(event.input.path, ctx);
	});
}