singularity-forge/src/resources/extensions/sf/bootstrap/register-hooks.ts

import { join, resolve, relative } from "node:path";

import type {
	ExtensionAPI,
	ExtensionContext,
} from "@singularity-forge/pi-coding-agent";
import { isToolCallEventType } from "@singularity-forge/pi-coding-agent";
import { resetAskUserQuestionsCache } from "../../ask-user-questions.js";
import { formatTokenCount } from "../../shared/format-utils.js";
import { saveActivityLog } from "../activity-log.js";
import {
	getAutoDashboardData,
	isAutoActive,
	isAutoPaused,
	markToolEnd,
	markToolStart,
	recordToolInvocationError,
} from "../auto.js";
import {
	applyCompletionNudgeTemperature,
	maybeInjectCompletionNudgeMessage,
	recordCompletionNudgeToolCall,
} from "../auto-completion-nudge.js";
import { recordToolCallName } from "../auto-tool-tracking.js";
import { loadToolApiKeys } from "../commands-config.js";
import { getEcosystemReadyPromise } from "../ecosystem/loader.js";
import type { SFEcosystemBeforeAgentStartHandler } from "../ecosystem/sf-extension-api.js";
import { updateSnapshot } from "../ecosystem/sf-extension-api.js";
import { formatContinue, loadFile, saveFile } from "../files.js";
import { getDiscussionMilestoneId } from "../guided-flow.js";
import { initHealthWidget } from "../health-widget.js";
import {
	initializeLearningRuntime,
	resetLearningRuntime,
	selectLearnedModel,
} from "../learning/runtime.js";
import { initNotificationStore } from "../notification-store.js";
import { initNotificationWidget } from "../notification-widget.js";
import {
	isParallelActive,
	shutdownParallel,
} from "../parallel-orchestrator.js";
import {
	buildMilestoneFileName,
	resolveMilestonePath,
	resolveSliceFile,
	resolveSlicePath,
} from "../paths.js";
import { cleanupQuickBranch } from "../quick.js";
import { classifyCommand } from "../safety/destructive-guard.js";
import {
	recordToolCall as safetyRecordToolCall,
	recordToolResult as safetyRecordToolResult,
} from "../safety/evidence-collector.js";
import { deriveState } from "../state.js";
import { countGoogleGeminiCliTokens } from "../token-counter.js";
import { logWarning as safetyLogWarning } from "../workflow-logger.js";
import {
	BLOCKED_WRITE_ERROR,
	isBashWriteToStateFile,
	isBlockedStateFile,
} from "../write-intercept.js";
import { handleAgentEnd } from "./agent-end-recovery.js";
import { installNotifyInterceptor } from "./notify-interceptor.js";
import { buildBeforeAgentStartResult } from "./system-context.js";
import {
	checkToolCallLoop,
	resetToolCallLoopGuard,
} from "./tool-call-loop-guard.js";
import {
	clearDiscussionFlowState,
	clearPendingGate,
	extractDepthVerificationMilestoneId,
	getPendingGate,
	getSelectedGateAnswer,
	isDepthConfirmationAnswer,
	isGateQuestionId,
	isQueuePhaseActive,
	markDepthVerified,
	resetWriteGateState,
	setPendingGate,
	shouldBlockContextWrite,
	shouldBlockPendingGate,
	shouldBlockPendingGateBash,
	shouldBlockQueueExecution,
} from "./write-gate.js";

// Skip the welcome screen on the very first session_start — cli.ts already
// printed it before the TUI launched. Only re-print on /clear (subsequent sessions).
let isFirstSession = true;
let lastGeminiPreflightWarning: string | undefined;

async function syncServiceTierStatus(ctx: ExtensionContext): Promise<void> {
	const {
		getEffectiveServiceTier,
		formatServiceTierFooterStatus,
		isServiceTierDisabled,
	} = await import("../service-tier.js");
	// Skip the footer event entirely when the feature is explicitly disabled —
	// no setStatus call, no RPC traffic, no leak into headless stderr even if
	// the TUI_FOOTER_STATUS_KEYS filter is bypassed.
	if (isServiceTierDisabled()) return;
	ctx.ui.setStatus(
		"sf-fast",
		formatServiceTierFooterStatus(getEffectiveServiceTier(), ctx.model?.id),
	);
}

export function registerHooks(
	pi: ExtensionAPI,
	ecosystemHandlers: SFEcosystemBeforeAgentStartHandler[] = [],
): void {
	pi.on("session_start", async (_event, ctx) => {
		lastGeminiPreflightWarning = undefined;
		resetLearningRuntime();
		try {
			const sid = ctx.sessionManager?.getSessionId?.() ?? "";
			const sfile = ctx.sessionManager?.getSessionFile?.() ?? "";
			if (sid) {
				process.stderr.write(`[forge] session ${sid.slice(0, 8)} · ${sfile}\n`);
			}
		} catch {
			/* non-fatal */
		}
		initNotificationStore(process.cwd());
		installNotifyInterceptor(ctx);
		initNotificationWidget(ctx);
		initHealthWidget(ctx);
		resetWriteGateState();
		resetToolCallLoopGuard();
		resetAskUserQuestionsCache();
		await syncServiceTierStatus(ctx);
		const { prepareWorkflowMcpForProject } = await import(
			"../workflow-mcp-auto-prep.js"
		);
		prepareWorkflowMcpForProject(ctx, process.cwd());
		await initializeLearningRuntime();

		// Apply show_token_cost preference (#1515)
		try {
			const { loadEffectiveSFPreferences } = await import("../preferences.js");
			const prefs = loadEffectiveSFPreferences();
			process.env.SF_SHOW_TOKEN_COST = prefs?.preferences.show_token_cost
				? "1"
				: "";
		} catch {
			/* non-fatal */
		}
		if (isFirstSession) {
			isFirstSession = false;
		} else {
			try {
				const sfBinPath = process.env.SF_BIN_PATH;
				if (sfBinPath) {
					const { dirname } = await import("node:path");
					const { printWelcomeScreen } = (await import(
						join(dirname(sfBinPath), "welcome-screen.js")
					)) as {
						printWelcomeScreen: (opts: {
							version: string;
							modelName?: string;
							provider?: string;
							remoteChannel?: string;
						}) => void;
					};

					let remoteChannel: string | undefined;
					try {
						const { resolveRemoteConfig } = await import(
							"../../remote-questions/config.js"
						);
						const rc = resolveRemoteConfig();
						if (rc) remoteChannel = rc.channel;
					} catch {
						/* non-fatal */
					}

					printWelcomeScreen({
						version: process.env.SF_VERSION || "0.0.0",
						remoteChannel,
					});
				}
			} catch {
				/* non-fatal */
			}
		}
		loadToolApiKeys();
		// Drain self-feedback backlog: auto-resolve entries whose blocking
		// sf-version constraint has been satisfied by the current sf bump,
		// and surface entries that remain blocked to the operator. Done after
		// other init so notifications appear in the same session-start sweep.
		try {
			const { triageBlockedEntries, markResolved } = await import(
				"../self-feedback.js"
			);
			const triage = triageBlockedEntries(process.cwd());
			const currentSfVersion = process.env.SF_VERSION || "unknown";
			for (const e of triage.retry) {
				markResolved(
					e.id,
					{
						reason: `sf bumped past ${e.sfVersion} (was blocking on this version)`,
						evidence: {
							kind: "auto-version-bump",
							fromVersion: e.sfVersion,
							toVersion: currentSfVersion,
						},
					},
					process.cwd(),
				);
				const occ = e.occurredIn;
				const unit = occ
					? [occ.milestone, occ.slice, occ.task].filter(Boolean).join("/") ||
						occ.unitType ||
						"(unknown unit)"
					: "(unknown unit)";
				ctx.ui?.notify?.(
					`Self-feedback ${e.id} (${e.kind}) auto-resolved — sf bumped past ${e.sfVersion}. Originating unit ${unit} should be re-run.`,
					"info",
				);
			}
			if (triage.stillBlocked.length > 0) {
				ctx.ui?.notify?.(
					`${triage.stillBlocked.length} self-feedback entr${triage.stillBlocked.length === 1 ? "y" : "ies"} still blocked on prior sf versions. See .sf/BACKLOG.md or ~/.sf/agent/upstream-feedback.jsonl.`,
					"warning",
				);
			}
		} catch {
			/* non-fatal — self-feedback drain must never block session start */
		}
	});

	pi.on("session_switch", async (_event, ctx) => {
		lastGeminiPreflightWarning = undefined;
		resetLearningRuntime();
		initNotificationStore(process.cwd());
		installNotifyInterceptor(ctx);
		resetWriteGateState();
		resetToolCallLoopGuard();
		resetAskUserQuestionsCache();
		clearDiscussionFlowState();
		await syncServiceTierStatus(ctx);
		const { prepareWorkflowMcpForProject } = await import(
			"../workflow-mcp-auto-prep.js"
		);
		prepareWorkflowMcpForProject(ctx, process.cwd());
		await initializeLearningRuntime();
		loadToolApiKeys();
	});

	pi.on("before_agent_start", async (event, ctx: ExtensionContext) => {
		// Refresh the ecosystem snapshot BEFORE running ecosystem handlers so they
		// see current phase/unit state (#3338).
		try {
			const { ensureDbOpen } = await import("./dynamic-tools.js");
			await ensureDbOpen();
			const basePath = process.cwd();
			const state = await deriveState(basePath);
			updateSnapshot(state);
		} catch {
			updateSnapshot(null);
		}

		// Await ecosystem loading, then dispatch any registered handlers.
		await getEcosystemReadyPromise();
		for (const handler of ecosystemHandlers) {
			try {
				await handler(event, ctx as any);
			} catch {
				// Non-fatal: don't break the SF turn if a third-party handler throws.
			}
		}

		return buildBeforeAgentStartResult(event, ctx);
	});

	pi.on("agent_end", async (event, ctx: ExtensionContext) => {
		resetToolCallLoopGuard();
		resetAskUserQuestionsCache();
		await handleAgentEnd(pi, event, ctx);
	});

	// Squash-merge quick-task branch back to the original branch after the
	// agent turn completes (#2668). cleanupQuickBranch is a no-op when no
	// quick-return state is pending, so this is safe to call on every turn.
	pi.on("turn_end", async () => {
		try {
			cleanupQuickBranch();
		} catch {
			// Best-effort: don't break the turn lifecycle if cleanup fails.
		}
	});

	pi.on("session_before_compact", async () => {
		// Only cancel compaction while auto-mode is actively running.
		// Paused auto-mode should allow compaction — the user may be doing
		// interactive work (#3165).
		if (isAutoActive()) {
			return { cancel: true };
		}
		const basePath = process.cwd();
		const { ensureDbOpen } = await import("./dynamic-tools.js");
		await ensureDbOpen();
		const state = await deriveState(basePath);
		if (!state.activeMilestone || !state.activeSlice || !state.activeTask)
			return;
		if (state.phase !== "executing") return;

		const sliceDir = resolveSlicePath(
			basePath,
			state.activeMilestone.id,
			state.activeSlice.id,
		);
		if (!sliceDir) return;

		const existingFile = resolveSliceFile(
			basePath,
			state.activeMilestone.id,
			state.activeSlice.id,
			"CONTINUE",
		);
		if (existingFile && (await loadFile(existingFile))) return;
		const legacyContinue = join(sliceDir, "continue.md");
		if (await loadFile(legacyContinue)) return;

		const continuePath = join(sliceDir, `${state.activeSlice.id}-CONTINUE.md`);
		await saveFile(
			continuePath,
			formatContinue({
				frontmatter: {
					milestone: state.activeMilestone.id,
					slice: state.activeSlice.id,
					task: state.activeTask.id,
					step: 0,
					totalSteps: 0,
					status: "compacted" as const,
					savedAt: new Date().toISOString(),
				},
				completedWork: `Task ${state.activeTask.id} (${state.activeTask.title}) was in progress when compaction occurred.`,
				remainingWork: "Check the task plan for remaining steps.",
				decisions: "Check task summary files for prior decisions.",
				context: "Session was auto-compacted by Pi. Resume with /sf.",
				nextAction: `Resume task ${state.activeTask.id}: ${state.activeTask.title}.`,
			}),
		);
	});

	pi.on("session_shutdown", async (_event, ctx: ExtensionContext) => {
		resetLearningRuntime();
		if (isParallelActive()) {
			try {
				await shutdownParallel(process.cwd());
			} catch {
				// best-effort
			}
		}
		if (!isAutoActive() && !isAutoPaused()) return;
		const dash = getAutoDashboardData();
		if (dash.currentUnit) {
			saveActivityLog(
				ctx,
				dash.basePath,
				dash.currentUnit.type,
				dash.currentUnit.id,
			);
		}
	});

	pi.on("tool_call", async (event) => {
		const discussionBasePath = process.cwd();
		// ── Loop guard: block repeated identical tool calls ──
		const loopCheck = checkToolCallLoop(
			event.toolName,
			event.input as Record<string, unknown>,
		);
		if (loopCheck.block) {
			return { block: true, reason: loopCheck.reason };
		}

		// ── Discussion gate enforcement: track pending gate questions ─────────
		// Only gate-shaped ask_user_questions calls should block execution.
		// The gate stays pending until the user selects the approval option.
		if (event.toolName === "ask_user_questions") {
			const questions: any[] = (event.input as any)?.questions ?? [];
			const questionId = questions.find(
				(question) =>
					typeof question?.id === "string" && isGateQuestionId(question.id),
			)?.id;
			if (typeof questionId === "string") {
				setPendingGate(questionId);
			}
		}

		// ── Discussion gate enforcement: block tool calls while gate is pending ──
		// If ask_user_questions was called with a gate ID but hasn't been confirmed,
		// block all non-read-only tool calls to prevent the model from skipping gates.
		if (getPendingGate()) {
			const milestoneId = getDiscussionMilestoneId(discussionBasePath);
			if (isToolCallEventType("bash", event)) {
				const bashGuard = shouldBlockPendingGateBash(
					event.input.command,
					milestoneId,
					isQueuePhaseActive(),
				);
				if (bashGuard.block) return bashGuard;
			} else {
				const gateGuard = shouldBlockPendingGate(
					event.toolName,
					milestoneId,
					isQueuePhaseActive(),
				);
				if (gateGuard.block) return gateGuard;
			}
		}

		// ── Queue-mode execution guard (#2545): block source-code mutations ──
		// When /sf queue is active, the agent should only create milestones,
		// not execute work. Block write/edit to non-.sf/ paths and bash commands
		// that would modify files.
		if (isQueuePhaseActive()) {
			let queueInput = "";
			if (isToolCallEventType("write", event)) {
				queueInput = event.input.path;
			} else if (isToolCallEventType("edit", event)) {
				queueInput = event.input.path;
			} else if (isToolCallEventType("bash", event)) {
				queueInput = event.input.command;
			}
			const queueGuard = shouldBlockQueueExecution(
				event.toolName,
				queueInput,
				true,
			);
			if (queueGuard.block) return queueGuard;
		}

		// ── Single-writer engine: block direct writes to STATE.md ──────────
		// Covers write, edit, and bash tools to prevent bypass vectors.
		if (isToolCallEventType("write", event)) {
			if (isBlockedStateFile(event.input.path)) {
				return { block: true, reason: BLOCKED_WRITE_ERROR };
			}
		}

		if (isToolCallEventType("edit", event)) {
			if (isBlockedStateFile(event.input.path)) {
				return { block: true, reason: BLOCKED_WRITE_ERROR };
			}
		}

		if (isToolCallEventType("bash", event)) {
			if (isBashWriteToStateFile(event.input.command)) {
				return { block: true, reason: BLOCKED_WRITE_ERROR };
			}
		}

		if (!isToolCallEventType("write", event)) return;

		// ── Worktree isolation: block writes outside the worktree and main .sf/ ──
		// Only enforced in auto-mode — interactive sessions skip this check.
		// When SF_WORKTREE is set, process.cwd() is the worktree directory.
		// The agent should only write inside the worktree OR inside the main repo's .sf/.
		if (isAutoActive() && process.env.SF_WORKTREE) {
			const worktreeRoot = process.cwd();
			const mainRepoRoot =
				process.env.SF_PROJECT_ROOT ??
				(resolve(worktreeRoot, ".."));
			const targetPath = resolve(event.input.path);
			const worktreeRel = relative(worktreeRoot, targetPath);
			const mainSfRel = relative(join(mainRepoRoot, ".sf"), targetPath);
			const worktreeOk =
				!worktreeRel.startsWith("..") && !worktreeRel.startsWith("/");
			const mainSfOk =
				!mainSfRel.startsWith("..") && !mainSfRel.startsWith("/");
			if (!worktreeOk && !mainSfOk) {
				return {
					block: true,
					reason:
						`HARD BLOCK: Worktree isolation is active. Cannot write to "${event.input.path}" — ` +
						`path is outside the worktree (${worktreeRoot}) and outside the main repo's .sf/ directory. ` +
						`Write only inside the worktree or inside ${join(mainRepoRoot, ".sf")}/milestones/ for planning artifacts.`,
				};
			}
		}

		const result = shouldBlockContextWrite(
			event.toolName,
			event.input.path,
			getDiscussionMilestoneId(discussionBasePath),
			isQueuePhaseActive(),
		);
		if (result.block) return result;
	});

	// ── Safety harness: evidence collection + destructive command warnings ──
	pi.on("tool_call", async (event, ctx) => {
		if (!isAutoActive()) return;
		safetyRecordToolCall(
			event.toolName,
			event.input as Record<string, unknown>,
		);

		// Destructive command classification (warn only, never block)
		if (isToolCallEventType("bash", event)) {
			const classification = classifyCommand(event.input.command);
			if (classification.destructive) {
				safetyLogWarning(
					"safety",
					`destructive command: ${classification.labels.join(", ")}`,
					{
						command: String(event.input.command).slice(0, 200),
					},
				);
				ctx.ui.notify(
					`Destructive command detected: ${classification.labels.join(", ")}`,
					"warning",
				);
			}
		}
	});

	pi.on("tool_result", async (event) => {
		if (event.toolName !== "ask_user_questions") return;
		const milestoneId = getDiscussionMilestoneId(process.cwd());
		const queueActive = isQueuePhaseActive();

		const details = event.details as any;

		// ── Discussion gate enforcement: handle gate question responses ──
		// Single consolidated loop: finds depth_verification questions, verifies the answer,
		// marks the milestone as depth-verified, and clears the pending gate.
		// Also handles the legacy pending-gate path (set by tool_call) for robustness.
		const questions: any[] = (event.input as any)?.questions ?? [];
		const currentPendingGate = getPendingGate();

		if (details?.cancelled || !details?.response) return;

		for (const question of questions) {
			if (typeof question.id !== "string") continue;

			// Check if this is a depth_verification question (either directly or via pending gate)
			const isDepthQ = question.id.includes("depth_verification");
			const isPendingQ = question.id === currentPendingGate;
			if (!isDepthQ && !isPendingQ) continue;

			const answer = details.response?.answers?.[question.id];
			if (
				isDepthConfirmationAnswer(getSelectedGateAnswer(answer), question.options)
			) {
				// Always mark depth-verified AND clear the gate
				if (isDepthQ) {
					const inferredMilestoneId =
						extractDepthVerificationMilestoneId(question.id) ?? milestoneId;
					markDepthVerified(inferredMilestoneId);
				}
				clearPendingGate();
				break;
			}
		}

		if (!milestoneId && !queueActive) return;
		if (!milestoneId) return;

		const basePath = process.cwd();
		const milestoneDir = resolveMilestonePath(basePath, milestoneId);
		if (!milestoneDir) return;

		const discussionPath = join(
			milestoneDir,
			buildMilestoneFileName(milestoneId, "DISCUSSION"),
		);
		const timestamp = new Date().toISOString();
		const lines: string[] = [`## Exchange — ${timestamp}`, ""];
		for (const question of questions) {
			lines.push(
				`### ${question.header ?? "Question"}`,
				"",
				question.question ?? "",
			);
			if (Array.isArray(question.options)) {
				lines.push("");
				for (const opt of question.options) {
					lines.push(`- **${opt.label}** — ${opt.description ?? ""}`);
				}
			}
			const answer = details.response?.answers?.[question.id];
			if (answer) {
				lines.push("");
				const selectedValue = getSelectedGateAnswer(answer);
				const selected = Array.isArray(selectedValue)
					? selectedValue.join(", ")
					: selectedValue;
				lines.push(`**Selected:** ${selected}`);
				if (answer.notes) {
					lines.push(`**Notes:** ${answer.notes}`);
				}
			}
			lines.push("");
		}
		lines.push("---", "");
		const existing =
			(await loadFile(discussionPath)) ?? `# ${milestoneId} Discussion Log\n\n`;
		await saveFile(discussionPath, existing + lines.join("\n"));
	});

	pi.on("tool_execution_start", async (event) => {
		if (!isAutoActive()) return;
		markToolStart(event.toolCallId, event.toolName);
		recordToolCallName(event.toolName);
		recordCompletionNudgeToolCall(event.toolName);
	});

	pi.on("tool_execution_end", async (event) => {
		markToolEnd(event.toolCallId);
		// #2883: Capture tool invocation errors (malformed/truncated JSON arguments)
		// so postUnitPreVerification can break the retry loop instead of re-dispatching.
		if (event.isError && event.toolName.startsWith("sf_")) {
			const errorText =
				typeof event.result === "string"
					? event.result
					: typeof event.result?.content?.[0]?.text === "string"
						? event.result.content[0].text
						: String(event.result);
			recordToolInvocationError(event.toolName, errorText);
		}
		// Safety harness: record tool execution results for evidence cross-referencing
		if (isAutoActive()) {
			safetyRecordToolResult(
				event.toolCallId,
				event.toolName,
				event.result,
				event.isError,
			);
		}
	});

	pi.on("model_select", async (_event, ctx) => {
		await syncServiceTierStatus(ctx);
	});

	pi.on("context", async (event) => {
		if (!isAutoActive()) return;
		const messages = maybeInjectCompletionNudgeMessage(event.messages);
		if (messages === event.messages) return;
		return { messages };
	});

	pi.on("before_provider_request", async (event, ctx) => {
		const payload = event.payload as Record<string, unknown> | null;
		if (!payload || typeof payload !== "object") return;
		applyCompletionNudgeTemperature(payload);

		// ── Observation Masking ─────────────────────────────────────────────
		// Replace old tool results with placeholders to reduce context bloat.
		// Only active during auto-mode when context_management.observation_masking is enabled.
		if (isAutoActive()) {
			try {
				const { loadEffectiveSFPreferences } = await import(
					"../preferences.js"
				);
				const prefs = loadEffectiveSFPreferences();
				const cmConfig = prefs?.preferences.context_management;

				// Observation masking: replace old tool results with placeholders
				if (cmConfig?.observation_masking !== false) {
					const keepTurns = cmConfig?.observation_mask_turns ?? 8;
					const { createObservationMask } = await import(
						"../context-masker.js"
					);
					const mask = createObservationMask(keepTurns);
					const messages = payload.messages;
					if (Array.isArray(messages)) {
						payload.messages = mask(messages);
					}
				}

				// Tool result truncation: cap individual tool result content length.
				// In pi-ai format, toolResult messages have role: "toolResult" and content: TextContent[].
				// Creates new objects to avoid mutating shared conversation state.
				const maxChars = cmConfig?.tool_result_max_chars ?? 800;
				const msgs = payload.messages;
				if (Array.isArray(msgs)) {
					payload.messages = msgs.map((msg: Record<string, unknown>) => {
						// Match toolResult messages (role: "toolResult", content is array of content blocks)
						if (msg?.role === "toolResult" && Array.isArray(msg.content)) {
							const blocks = msg.content as Array<Record<string, unknown>>;
							const totalLen = blocks.reduce(
								(sum: number, b) =>
									sum + (typeof b.text === "string" ? b.text.length : 0),
								0,
							);
							if (totalLen > maxChars) {
								const truncated = blocks.map((b) => {
									if (typeof b.text === "string" && b.text.length > maxChars) {
										return {
											...b,
											text: b.text.slice(0, maxChars) + "\n…[truncated]",
										};
									}
									return b;
								});
								return { ...msg, content: truncated };
							}
						}
						return msg;
					});
				}
			} catch {
				/* non-fatal */
			}
		}

		// ── Service Tier ────────────────────────────────────────────────────
		const modelId = event.model?.id;
		if (!modelId) {
			ctx.ui.setStatus("sf-gemini-tokens", undefined);
			return payload;
		}
		const {
			getEffectiveServiceTier,
			supportsServiceTier,
			isServiceTierDisabled,
		} = await import("../service-tier.js");
		// Short-circuit on explicit disable — never inject service_tier on any
		// setup that has opted out, regardless of model.
		if (!isServiceTierDisabled()) {
			const tier = getEffectiveServiceTier();
			if (tier && supportsServiceTier(modelId)) {
				payload.service_tier = tier;
			}
		}

		if (event.model?.provider !== "google-gemini-cli") {
			ctx.ui.setStatus("sf-gemini-tokens", undefined);
			return payload;
		}

		try {
			const resolvedModel =
				ctx.model &&
				ctx.model.provider === event.model.provider &&
				ctx.model.id === event.model.id
					? ctx.model
					: ctx.modelRegistry
							.getAvailable()
							.find(
								(m) =>
									m.provider === event.model?.provider &&
									m.id === event.model?.id,
							);
			if (!resolvedModel) {
				ctx.ui.setStatus("sf-gemini-tokens", undefined);
				return payload;
			}

			const apiKey = await ctx.modelRegistry.getApiKey(resolvedModel);
			const totalTokens = await countGoogleGeminiCliTokens(payload, apiKey);
			if (typeof totalTokens !== "number") {
				ctx.ui.setStatus("sf-gemini-tokens", undefined);
				return payload;
			}

			const contextWindow = resolvedModel.contextWindow ?? 0;
			const pct =
				contextWindow > 0
					? Math.round((totalTokens / contextWindow) * 100)
					: undefined;
			ctx.ui.setStatus(
				"sf-gemini-tokens",
				pct !== undefined
					? `gemini ${formatTokenCount(totalTokens)} (${pct}%)`
					: `gemini ${formatTokenCount(totalTokens)}`,
			);

			if (contextWindow > 0 && totalTokens >= Math.floor(contextWindow * 0.8)) {
				const warningKey = `${resolvedModel.id}:${totalTokens}:${contextWindow}`;
				if (lastGeminiPreflightWarning !== warningKey) {
					lastGeminiPreflightWarning = warningKey;
					ctx.ui.notify(
						`Gemini preflight: ${formatTokenCount(totalTokens)} tokens (${pct}% of ${formatTokenCount(contextWindow)} context).`,
						"warning",
					);
				}
			}
		} catch {
			ctx.ui.setStatus("sf-gemini-tokens", undefined);
		}

		return payload;
	});

	// Capability-aware model routing hook (ADR-004)
	// Extensions can override model selection by returning { modelId: "..." }
	// Return undefined to let the built-in capability scoring proceed.
	pi.on("before_model_select", async (event) => {
		return selectLearnedModel({
			unitType: event.unitType,
			eligibleModels: event.eligibleModels,
			phaseConfig: event.phaseConfig,
		});
	});

	// Tool set adaptation hook (ADR-005 Phase 4)
	// Extensions can override tool set after model selection by returning { toolNames: [...] }
	// Return undefined to let the built-in provider compatibility filtering proceed.
	pi.on("adjust_tool_set", async (_event) => {
		// Default: no override — let provider capability filtering handle tool set
		return undefined;
	});
}