import { createHash, randomUUID } from "node:crypto"; import type { ExtensionAPI, ExtensionCommandContext, ExtensionContext, SessionEntry, } from "@earendil-works/pi-coding-agent"; type LabelEntry = Extract; import { offerLimitIncrease, showConfig } from "./tree-labeler/config.ts"; import { applyModelAwareLimits } from "./tree-labeler/limits.ts"; import { generateProposals, LabelGenerationError, LabelingAbortedError } from "./tree-labeler/model.ts"; import { APPLY_ENTRY_TYPE, reconstructOwnership } from "./tree-labeler/ownership.ts"; import { reviewProposals } from "./tree-labeler/review.ts"; import { buildSemanticTree, serializeTree } from "./tree-labeler/semantic-tree.ts"; import { syncSessionName } from "./tree-labeler/session-name.ts"; import { loadSettings, settingSource } from "./tree-labeler/settings.ts"; import { latestPendingRun, loadRuns, saveRun } from "./tree-labeler/store.ts"; import { showStatusPanel } from "./tree-labeler/status-ui.ts"; import type { ApplyLedger, LabelProposal, LabelRun, ProposalConflict, SemanticNode, TreeLabelerSettings } from "./tree-labeler/types.ts"; type RunReason = "manual" | "turn-threshold" | "new-branch" | "tree-navigation" | "compaction"; function hash(value: string): string { return createHash("sha256").update(value).digest("hex"); } /** * Fingerprint only immutable semantic content that was actually exposed to the * labeler. Raw leaves and the complete tree are intentionally excluded because * session names, labels, custom entries, and unrelated branch growth are * append-only changes that do not invalidate proposals for existing entries. */ export function semanticSnapshot(nodes: SemanticNode[], exposedIds: Set): string { const byId = new Map(nodes.map((node) => [node.id, node])); const snapshot = [...exposedIds].sort().map((id) => { const node = byId.get(id); return node ? [node.id, node.parentId, node.role, node.excerpt] : [id, ""]; }); return hash(JSON.stringify(snapshot)); } function publicSettings(settings: TreeLabelerSettings): LabelRun["settings"] { const { prompt, ...rest } = settings; return { ...rest, promptHash: hash(prompt) }; } function notify(ctx: ExtensionContext, message: string, level: "info" | "warning" | "error" = "info") { if (ctx.hasUI) ctx.ui.notify(message, level); } function splitConflicts( proposals: LabelProposal[], entries: SessionEntry[], ): { safe: LabelProposal[]; conflicts: ProposalConflict[] } { const ownership = reconstructOwnership(entries); const ids = new Set(entries.map((entry) => entry.id)); const safe: LabelProposal[] = []; const conflicts: ProposalConflict[] = []; for (const proposal of proposals) { if (!ids.has(proposal.entryId)) { conflicts.push({ ...proposal, conflict: "Target entry no longer exists" }); continue; } const latest = ownership.latestLabelEntryByTarget.get(proposal.entryId); if (latest && !ownership.ownedLabelEntryIds.has(latest.id)) { conflicts.push({ ...proposal, conflict: latest.label === undefined ? "Target was manually cleared" : `Manual label is protected: ${latest.label}` }); continue; } if (latest?.label === proposal.label) continue; safe.push(proposal); } return { safe, conflicts }; } async function applyLabels( pi: ExtensionAPI, ctx: ExtensionContext, run: LabelRun, selected: LabelProposal[], ): Promise { const currentEntries = ctx.sessionManager.getEntries(); const { safe, conflicts } = splitConflicts(selected, currentEntries); run.conflicts.push(...conflicts); const changes: ApplyLedger["changes"] = []; for (const proposal of safe) { const before = ctx.sessionManager.getEntries(); const previousLabel = ctx.sessionManager.getLabel(proposal.entryId); pi.setLabel(proposal.entryId, proposal.label); const appended = ctx.sessionManager.getEntries().slice(before.length); const labelEntry = [...appended].reverse().find((entry): entry is LabelEntry => entry.type === "label" && entry.targetId === proposal.entryId && entry.label === proposal.label); if (!labelEntry) throw new Error(`Pi did not expose the label entry created for ${proposal.entryId}.`); const change = { targetId: proposal.entryId, previousLabel, label: proposal.label, labelEntryId: labelEntry.id }; changes.push(change); const ledger: ApplyLedger = { version: 1, runId: run.id, changes: [change] }; pi.appendEntry(APPLY_ENTRY_TYPE, ledger); } run.status = "applied"; run.applied = changes.map((change) => ({ entryId: change.targetId, label: change.label, labelEntryId: change.labelEntryId })); syncSessionName(pi, ctx, loadSettings(ctx.cwd, ctx.isProjectTrusted())); await saveRun(run); notify(ctx, `Applied ${changes.length} tree label${changes.length === 1 ? "" : "s"}${conflicts.length ? `; skipped ${conflicts.length} conflict${conflicts.length === 1 ? "" : "s"}` : ""}.`); } function formatHistory(runs: LabelRun[]): string { if (!runs.length) return "No tree-label runs for this session."; return runs.map((run) => `${run.id.slice(0, 8)} ${run.createdAt} ${run.status} ${run.proposals.length} proposal(s) ${run.model}`).join("\n"); } interface ErrorGuidance { code: string; rootCause: string; guidance: string[]; } function limitHitFromError(error: unknown): "maxInputChars" | "outputTokens" | "timeoutMs" | undefined { if (error instanceof LabelGenerationError && error.details.code === "generation-timeout") return "timeoutMs"; if (error instanceof LabelGenerationError && error.details.code === "output-truncated") return "outputTokens"; const message = error instanceof Error ? error.message : String(error); if (/maxInputChars|skeleton exceeds|context.{0,20}(?:length|window|overflow)|input.{0,20}too (?:large|long)/i.test(message)) return "maxInputChars"; return undefined; } function classifyRunError(error: unknown): ErrorGuidance { if (error instanceof LabelGenerationError) { return { code: error.details.code, rootCause: error.details.rootCause, guidance: error.details.guidance, }; } const message = error instanceof Error ? error.message : String(error); if (/api key|authenticate|authentication|credential|not configured/i.test(message)) { return { code: "model-authentication", rootCause: "The configured label model is not currently authenticated or available.", guidance: ["Open /tree-labels config and choose an authenticated model, or authenticate the configured provider, then retry."], }; } if (/maxInputChars|skeleton exceeds|above maxInput/i.test(message)) { return { code: "input-budget", rootCause: "The serialized tree exceeded the configured input budget.", guidance: ["In /tree-labels config, use balanced/branch exposure, lower Max nodes or Excerpt characters, or raise Max input characters."], }; } return { code: "unexpected-extension-error", rootCause: "The extension encountered an error it could not classify as a known configuration or model-output problem.", guidance: ["Retry once. If it repeats, report the failed run ID together with the extension version and error message."], }; } interface FailedRunContext { settings: TreeLabelerSettings; sessionId: string; leafId: string | null; treeDigest: string; exposedNodes: number; omittedNodes: number; } async function persistFailedRun(reason: RunReason, error: unknown, context: FailedRunContext): Promise { const generationError = error instanceof LabelGenerationError ? error : undefined; const classification = classifyRunError(error); const message = error instanceof Error ? error.message : String(error); const run: LabelRun = { version: 1, id: randomUUID(), sessionId: context.sessionId, createdAt: new Date().toISOString(), status: "failed", trigger: reason === "manual" ? "manual" : "automatic", model: generationError?.details.model ?? context.settings.model, settings: publicSettings(context.settings), treeDigest: context.treeDigest, leafId: context.leafId, exposedNodes: context.exposedNodes, omittedNodes: context.omittedNodes, proposals: [], conflicts: [], error: message, diagnostic: { stage: generationError?.details.stage ?? "run", errorType: error instanceof Error ? error.name : typeof error, code: classification.code, rootCause: classification.rootCause, guidance: classification.guidance, ...(generationError?.details.outputShape === undefined ? {} : { outputShape: generationError.details.outputShape }), }, ...(context.settings.persistPayloads && generationError ? { prompt: generationError.details.prompt, response: generationError.details.raw } : {}), }; await saveRun(run); return run; } export default function treeLabeler(pi: ExtensionAPI) { let epoch = 0; let inFlight: Promise | undefined; let activeRunController: AbortController | undefined; let lastUserTurns = 0; let lastBranchPoints = 0; let compactDirty = false; let pendingProposals = 0; let runningReason: RunReason | undefined; let lastOutcome = "not run in this session"; const topologyState = (ctx: ExtensionContext) => { const entries = ctx.sessionManager.getEntries(); const userTurns = entries.filter((entry) => entry.type === "message" && entry.message.role === "user").length; const settings = loadSettings(ctx.cwd, ctx.isProjectTrusted()); const ownership = reconstructOwnership(entries); const activeIds = new Set(ctx.sessionManager.getBranch().map((entry) => entry.id)); const branchPoints = buildSemanticTree(ctx.sessionManager.getTree(), activeIds, ownership, settings).filter((node) => node.branchPoint).length; return { userTurns, branchPoints, settings }; }; const updateStatus = (ctx: ExtensionContext) => { const state = topologyState(ctx); let text: string; if (runningReason) { text = `tree-labels: running (${runningReason})`; } else { const pending = pendingProposals ? `${pendingProposals} pending · ` : ""; if (!state.settings.automation.enabled) text = `tree-labels: ${pending}manual`; else { const progress = Math.min(state.settings.automation.afterTurns, Math.max(0, state.userTurns - lastUserTurns)); const extra = [ state.settings.automation.onBranchCreated ? "branch" : "", state.settings.automation.onTreeNavigation ? "tree" : "", state.settings.automation.onCompaction ? "compact" : "", ].filter(Boolean).join("/"); text = `tree-labels: ${pending}${progress}/${state.settings.automation.afterTurns} turns${extra ? ` +${extra}` : ""}`; } } const color = runningReason ? "accent" : pendingProposals ? "warning" : "muted"; ctx.ui.setStatus("tree-labeler", ctx.ui.theme.fg(color, text)); }; const refreshPending = async (ctx: ExtensionContext) => { const sessionId = ctx.sessionManager.getSessionId(); const runs = await loadRuns(sessionId, Number.MAX_SAFE_INTEGER); if (ctx.sessionManager.getSessionId() !== sessionId) return; pendingProposals = runs.filter((item) => item.status === "proposed").reduce((sum, item) => sum + item.proposals.length, 0); updateStatus(ctx); }; pi.on("session_start", (_event, ctx) => { epoch += 1; const state = topologyState(ctx); lastUserTurns = state.userTurns; lastBranchPoints = state.branchPoints; compactDirty = false; pendingProposals = 0; runningReason = undefined; lastOutcome = "not run in this session"; syncSessionName(pi, ctx, state.settings); updateStatus(ctx); void refreshPending(ctx); }); pi.on("session_shutdown", (_event, ctx) => { epoch += 1; activeRunController?.abort(); activeRunController = undefined; inFlight = undefined; compactDirty = false; runningReason = undefined; ctx.ui.setStatus("tree-labeler", undefined); }); const run = async (ctx: ExtensionContext, forceApply: boolean, reason: RunReason = "manual") => { if (inFlight) { notify(ctx, `A tree-label run is already in progress (${runningReason ?? "unknown trigger"}).`, "warning"); return inFlight; } const startedEpoch = epoch; const runController = new AbortController(); activeRunController = runController; runningReason = reason; updateStatus(ctx); let succeeded = false; let failedRunContext: FailedRunContext | undefined; let effectiveRunSettings: TreeLabelerSettings | undefined; const work = (async () => { if ("waitForIdle" in ctx) await (ctx as ExtensionCommandContext).waitForIdle(); const configuredSettings = loadSettings(ctx.cwd, ctx.isProjectTrusted()); const settings = applyModelAwareLimits(ctx, configuredSettings).settings; effectiveRunSettings = settings; if (!settings.enabled) throw new Error("treeLabeler.enabled is false."); const inputSessionId = ctx.sessionManager.getSessionId(); const inputLeafId = ctx.sessionManager.getLeafId(); failedRunContext = { settings, sessionId: inputSessionId, leafId: inputLeafId, treeDigest: "", exposedNodes: 0, omittedNodes: 0, }; const entries = ctx.sessionManager.getEntries(); const ownership = reconstructOwnership(entries); const activeIds = new Set(ctx.sessionManager.getBranch().map((entry) => entry.id)); const semantic = buildSemanticTree(ctx.sessionManager.getTree(), activeIds, ownership, settings); if (!semantic.length) throw new Error("This session has no labelable conversation entries."); const serialized = serializeTree(semantic, settings); const inputTreeDigest = hash(serialized.text); failedRunContext.treeDigest = inputTreeDigest; failedRunContext.exposedNodes = serialized.nodeCount; failedRunContext.omittedNodes = serialized.omittedCount; const inputSemanticSnapshot = semanticSnapshot(semantic, serialized.exposedIds); notify(ctx, `Curating ${serialized.nodeCount} tree entries…`); const generated = await generateProposals(ctx, serialized, settings, runController.signal); if (runController.signal.aborted) throw new LabelingAbortedError(); const currentOwnership = reconstructOwnership(ctx.sessionManager.getEntries()); const currentActiveIds = new Set(ctx.sessionManager.getBranch().map((entry) => entry.id)); const currentSemantic = buildSemanticTree(ctx.sessionManager.getTree(), currentActiveIds, currentOwnership, settings); const sameSession = epoch === startedEpoch && ctx.sessionManager.getSessionId() === inputSessionId; const sameExposedContent = semanticSnapshot(currentSemantic, serialized.exposedIds) === inputSemanticSnapshot; if (!sameSession || !sameExposedContent) { if (reason !== "manual") throw new LabelingAbortedError("Tree labeling was cancelled because its source entries changed."); throw new Error("The conversation entries exposed to the labeler changed during generation; discarded stale result."); } const current = splitConflicts(generated.proposals, ctx.sessionManager.getEntries()); const runRecord: LabelRun = { version: 1, id: randomUUID(), sessionId: inputSessionId, createdAt: new Date().toISOString(), status: "proposed", trigger: reason === "manual" ? "manual" : "automatic", model: generated.model, settings: publicSettings(settings), treeDigest: inputTreeDigest, leafId: inputLeafId, exposedNodes: serialized.nodeCount, omittedNodes: serialized.omittedCount, proposals: current.safe, conflicts: current.conflicts, ...(generated.warnings.length ? { warnings: generated.warnings } : {}), ...(settings.persistPayloads ? { prompt: generated.prompt, response: generated.raw } : {}), }; if (!current.safe.length) runRecord.status = "discarded"; await saveRun(runRecord); if (generated.limitsHit.length) { await offerLimitIncrease(ctx as ExtensionCommandContext, settings, generated.limitsHit); } if (generated.warnings.length) { notify(ctx, `Tree labeling kept ${current.safe.length} valid proposal${current.safe.length === 1 ? "" : "s"} but found malformed output. ${generated.warnings.join(" ")}`, "warning"); } lastUserTurns = topologyState(ctx).userTurns; succeeded = true; if (!current.safe.length) { lastOutcome = current.conflicts.length ? `no safe labels (${current.conflicts.length} conflicts)` : "no labels proposed"; notify(ctx, current.conflicts.length ? `No safe labels proposed; ${current.conflicts.length} conflicted with manual state.` : "The curator found no worthwhile new labels."); return; } if (forceApply || settings.applyMode === "auto") { if (!forceApply) { const source = settingSource(ctx.cwd, ctx.isProjectTrusted(), "applyMode"); notify(ctx, `Auto-applying because apply mode is “auto” in ${source} settings. Change ${source === "project" ? "/tree-labels config project" : "/tree-labels config global"} to require review.`, "warning"); } await applyLabels(pi, ctx, runRecord, current.safe); lastOutcome = `applied ${runRecord.applied?.length ?? 0} labels`; return; } if (settings.applyMode === "manual" || reason !== "manual") { lastOutcome = `saved ${current.safe.length} pending proposals`; notify(ctx, `Saved ${current.safe.length} proposals. Run /tree-labels review to inspect them.`); return; } const selected = await reviewProposals(current.safe, ctx as ExtensionCommandContext); if (selected === undefined) { lastOutcome = `saved ${current.safe.length} pending proposals`; notify(ctx, `Saved ${current.safe.length} pending proposals.`); return; } await applyLabels(pi, ctx, runRecord, selected); lastOutcome = `applied ${runRecord.applied?.length ?? 0} labels`; })().catch(async (error: unknown) => { if (error instanceof LabelingAbortedError) { lastOutcome = error.message; if (reason === "manual") notify(ctx, error.message, "info"); return; } const message = error instanceof Error ? error.message : String(error); const classification = classifyRunError(error); lastOutcome = `failed: ${classification.rootCause}`; let logSuffix = ""; if (failedRunContext) { try { const failedRun = await persistFailedRun(reason, error, failedRunContext); logSuffix = ` Failed run: ${failedRun.id.slice(0, 8)}.`; } catch (logError) { logSuffix = ` Logging also failed: ${logError instanceof Error ? logError.message : String(logError)}.`; } } const fix = classification.guidance[0] ? ` Fix: ${classification.guidance[0]}` : ""; notify(ctx, `Tree labeling failed: ${classification.rootCause} (${message})${fix}${logSuffix}`, "error"); const limitHit = limitHitFromError(error); if (limitHit && effectiveRunSettings) { await offerLimitIncrease(ctx as ExtensionCommandContext, effectiveRunSettings, [limitHit]); } }).finally(async () => { const isCurrentRun = activeRunController === runController; if (isCurrentRun) { activeRunController = undefined; runningReason = undefined; } if (inFlight === work) inFlight = undefined; if (!succeeded && reason !== "manual" && lastOutcome.startsWith("failed:")) lastOutcome = `automatic run ${lastOutcome}`; if (isCurrentRun && epoch === startedEpoch) await refreshPending(ctx); }); inFlight = work; return work; }; const queueAutomaticRun = async (ctx: ExtensionContext, reason: Exclude) => { // Never cancel a valid snapshot run because new chat/tree activity arrived. // Let it finish against the entries it was originally shown, then start at // most one follow-up run for newer content. const queuedEpoch = epoch; const queuedSessionId = ctx.sessionManager.getSessionId(); if (inFlight) await inFlight; if (inFlight || epoch !== queuedEpoch || ctx.sessionManager.getSessionId() !== queuedSessionId) return; await run(ctx, false, reason); }; pi.on("session_compact", (event, ctx) => { const settings = loadSettings(ctx.cwd, ctx.isProjectTrusted()); if (settings.automation.enabled && settings.automation.onCompaction) compactDirty = true; if (!event.willRetry && compactDirty) { compactDirty = false; void queueAutomaticRun(ctx, "compaction"); } else { updateStatus(ctx); } }); pi.on("session_tree", (_event, ctx) => { const state = topologyState(ctx); lastUserTurns = state.userTurns; lastBranchPoints = state.branchPoints; syncSessionName(pi, ctx, state.settings); updateStatus(ctx); if (state.settings.automation.enabled && state.settings.automation.onTreeNavigation) void queueAutomaticRun(ctx, "tree-navigation"); }); pi.on("agent_settled", (_event, ctx) => { const state = topologyState(ctx); if (!state.settings.automation.enabled) { lastUserTurns = state.userTurns; lastBranchPoints = state.branchPoints; compactDirty = false; updateStatus(ctx); return; } const enoughTurns = state.userTurns - lastUserTurns >= state.settings.automation.afterTurns; const newBranch = state.settings.automation.onBranchCreated && state.branchPoints > lastBranchPoints; const shouldRun = enoughTurns || newBranch || (compactDirty && state.settings.automation.onCompaction); if (shouldRun) { lastUserTurns = state.userTurns; lastBranchPoints = state.branchPoints; const reason: RunReason = newBranch ? "new-branch" : enoughTurns ? "turn-threshold" : "compaction"; compactDirty = false; void queueAutomaticRun(ctx, reason); } else { updateStatus(ctx); } }); pi.registerCommand("tree-labels", { description: "Curate, review, and apply semantic labels for the Pi session tree", getArgumentCompletions(prefix) { return ["run", "run --apply", "review", "apply-pending", "discard", "status", "history", "config", "config global", "config project"].filter((value) => value.startsWith(prefix)).map((value) => ({ value, label: value })); }, handler: async (args, ctx) => { const command = args.trim(); if (!command || command === "run") return run(ctx, false); if (command === "run --apply" || command === "apply") return run(ctx, true); if (command === "review") { await ctx.waitForIdle(); const pending = await latestPendingRun(ctx.sessionManager.getSessionId()); if (!pending) return notify(ctx, "No pending tree-label proposal.", "warning"); const { safe, conflicts } = splitConflicts(pending.proposals, ctx.sessionManager.getEntries()); pending.conflicts.push(...conflicts); if (!safe.length) { pending.status = "discarded"; await saveRun(pending); await refreshPending(ctx); return notify(ctx, "No pending proposals remain safe to apply.", "warning"); } const selected = await reviewProposals(safe, ctx); if (selected !== undefined) { await applyLabels(pi, ctx, pending, selected); lastOutcome = `applied ${pending.applied?.length ?? 0} pending labels`; await refreshPending(ctx); } else if (ctx.mode !== "tui") notify(ctx, "Review UI requires TUI mode. Use /tree-labels apply-pending to apply all safe pending labels.", "warning"); return; } if (command === "apply-pending") { await ctx.waitForIdle(); const pending = await latestPendingRun(ctx.sessionManager.getSessionId()); if (!pending) return notify(ctx, "No pending tree-label proposal.", "warning"); await applyLabels(pi, ctx, pending, pending.proposals); lastOutcome = `applied ${pending.applied?.length ?? 0} pending labels`; await refreshPending(ctx); return; } if (command === "discard") { const pending = await latestPendingRun(ctx.sessionManager.getSessionId()); if (!pending) return notify(ctx, "No pending tree-label proposal.", "warning"); pending.status = "discarded"; await saveRun(pending); lastOutcome = `discarded pending run ${pending.id.slice(0, 8)}`; await refreshPending(ctx); notify(ctx, `Discarded pending run ${pending.id.slice(0, 8)}.`); return; } if (command === "config" || command === "config global" || command === "config project") { await showConfig(ctx, command === "config global" ? "global" : command === "config project" ? "project" : undefined); updateStatus(ctx); return; } if (command === "history") { const history = formatHistory(await loadRuns(ctx.sessionManager.getSessionId())); if (ctx.mode === "tui") await ctx.ui.input("Tree-label history", history); else notify(ctx, history); return; } if (command === "status") { await refreshPending(ctx); const state = topologyState(ctx); const ownership = reconstructOwnership(ctx.sessionManager.getEntries()); const owned = [...ownership.latestLabelEntryByTarget.values()].filter((entry) => ownership.ownedLabelEntryIds.has(entry.id)).length; const progress = Math.min(state.settings.automation.afterTurns, Math.max(0, state.userTurns - lastUserTurns)); const remaining = Math.max(0, state.settings.automation.afterTurns - progress); const triggers = [ state.settings.automation.onBranchCreated ? "new branch" : "", state.settings.automation.onTreeNavigation ? "/tree checkout" : "", state.settings.automation.onCompaction ? "compaction" : "", ].filter(Boolean); await showStatusPanel(ctx, [ `Automation: ${state.settings.automation.enabled ? "enabled" : "disabled"}`, `Current activity: ${runningReason ? `running (${runningReason})` : "idle"}`, `Turn threshold: ${progress}/${state.settings.automation.afterTurns}${state.settings.automation.enabled ? ` — ${remaining} turn${remaining === 1 ? "" : "s"} remaining` : ""}`, `Event triggers: ${triggers.length ? triggers.join(", ") : "none"}`, `Pending proposals: ${pendingProposals}`, `Last outcome: ${lastOutcome}`, `Model: ${state.settings.model} · apply mode: ${state.settings.applyMode} (${settingSource(ctx.cwd, ctx.isProjectTrusted(), "applyMode")}) · exposure: ${state.settings.exposure} · limits: ${state.settings.autoLimits ? "auto" : "custom"}`, `Owned labels: ${owned} · protected manual clears: ${ownership.manuallyClearedTargets.size}`, ]); return; } notify(ctx, "Usage: /tree-labels [run [--apply] | review | apply-pending | discard | config [global|project] | status | history]", "warning"); }, }); }