1// SPDX-FileCopyrightText: Amolith <amolith@secluded.site>
2// SPDX-FileCopyrightText: Petr Baudis <pasky@ucw.cz>
3//
4// SPDX-License-Identifier: MIT
5
6import { existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs";
7import * as path from "node:path";
8import {
9 type ExtensionAPI,
10 type ExtensionCommandContext,
11 type ExtensionContext,
12 getAgentDir,
13 type SessionEntry,
14 SettingsManager,
15} from "@earendil-works/pi-coding-agent";
16import { pickModel } from "./model-utils.js";
17
18/**
19 * Model override sources for handoff extraction calls. Use provider/modelId
20 * values (e.g. "anthropic/claude-haiku-4-5"). Precedence is:
21 * --handoff-extraction-model, PI_HANDOFF_EXTRACTION_MODEL, user settings
22 * handoffExtractionModel, session model.
23 */
24const HANDOFF_EXTRACTION_MODEL_FLAG = "handoff-extraction-model";
25const HANDOFF_EXTRACTION_MODEL_ENV = "PI_HANDOFF_EXTRACTION_MODEL";
26const HANDOFF_EXTRACTION_MODEL_SETTING = "handoffExtractionModel";
27
28type JsonValue = string | number | boolean | null | JsonObject | JsonValue[];
29
30interface JsonObject {
31 [key: string]: JsonValue;
32}
33
34export function registerHandoffExtractionModelFlag(pi: ExtensionAPI): void {
35 pi.registerFlag(HANDOFF_EXTRACTION_MODEL_FLAG, {
36 description: "Model for handoff extraction as provider/modelId",
37 type: "string",
38 });
39}
40
41export function registerHandoffExtractionModelCommand(pi: ExtensionAPI): void {
42 pi.registerCommand("handoff:set-extraction-model", {
43 description: "Choose the model used to extract handoff context",
44 handler: async (args: string, ctx: ExtensionCommandContext) => {
45 if (!ctx.hasUI || ctx.mode !== "tui") {
46 ctx.ui.notify("/handoff:set-extraction-model requires interactive mode", "error");
47 return;
48 }
49
50 const currentModel = configuredExtractionModel(pi, ctx);
51 const selectedModel = await pickModel(ctx, currentModel, nonEmptyString(args));
52 if (!selectedModel) {
53 ctx.ui.notify("Handoff extraction model unchanged", "info");
54 return;
55 }
56
57 const modelSpec = `${selectedModel.provider}/${selectedModel.id}`;
58 try {
59 setUserSettingsHandoffExtractionModel(modelSpec);
60 } catch (err) {
61 const message = err instanceof Error ? err.message : String(err);
62 ctx.ui.notify(`Failed to save handoff extraction model: ${message}`, "error");
63 return;
64 }
65
66 process.env[HANDOFF_EXTRACTION_MODEL_ENV] = modelSpec;
67
68 if (configuredFlagHandoffExtractionModel(pi)) {
69 ctx.ui.notify(
70 `Saved handoff extraction model: ${modelSpec}. The --${HANDOFF_EXTRACTION_MODEL_FLAG} flag still takes precedence for this run.`,
71 "warning",
72 );
73 return;
74 }
75
76 ctx.ui.notify(`Handoff extraction model: ${modelSpec}`, "info");
77 },
78 });
79}
80
81function nonEmptyString(value: string | undefined): string | undefined {
82 const trimmed = value?.trim();
83 return trimmed ? trimmed : undefined;
84}
85
86function stringProperty(source: object, key: string): string | undefined {
87 const descriptor = Object.getOwnPropertyDescriptor(source, key);
88 if (!descriptor || !("value" in descriptor)) return undefined;
89 return typeof descriptor.value === "string" ? nonEmptyString(descriptor.value) : undefined;
90}
91
92function userSettingsHandoffExtractionModel(cwd: string): string | undefined {
93 const settingsManager = SettingsManager.create(cwd, getAgentDir());
94 return stringProperty(settingsManager.getGlobalSettings(), HANDOFF_EXTRACTION_MODEL_SETTING);
95}
96
97function configuredFlagHandoffExtractionModel(pi: ExtensionAPI): string | undefined {
98 const flagValue = pi.getFlag(HANDOFF_EXTRACTION_MODEL_FLAG);
99 const flagModel = typeof flagValue === "string" ? nonEmptyString(flagValue) : undefined;
100 return flagModel;
101}
102
103function configuredHandoffExtractionModel(pi: ExtensionAPI, ctx: { cwd: string }): string | undefined {
104 const flagModel = configuredFlagHandoffExtractionModel(pi);
105 return (
106 flagModel ??
107 nonEmptyString(process.env[HANDOFF_EXTRACTION_MODEL_ENV]) ??
108 userSettingsHandoffExtractionModel(ctx.cwd)
109 );
110}
111
112function parseJsonObject(content: string, source: string): JsonObject {
113 const parsed: JsonValue = JSON.parse(content);
114 if (parsed === null || typeof parsed !== "object" || Array.isArray(parsed)) {
115 throw new Error(`${source} must contain a JSON object`);
116 }
117 return parsed;
118}
119
120function readSettingsObject(settingsPath: string): JsonObject {
121 if (!existsSync(settingsPath)) return {};
122 return parseJsonObject(readFileSync(settingsPath, "utf-8"), settingsPath);
123}
124
125function setUserSettingsHandoffExtractionModel(modelSpec: string): void {
126 const settingsPath = path.join(getAgentDir(), "settings.json");
127 const settings = readSettingsObject(settingsPath);
128 settings[HANDOFF_EXTRACTION_MODEL_SETTING] = modelSpec;
129
130 mkdirSync(path.dirname(settingsPath), { recursive: true });
131 writeFileSync(settingsPath, `${JSON.stringify(settings, null, 2)}\n`, "utf-8");
132}
133
134function parseModelSpec(spec: string): { provider: string; modelId: string } | undefined {
135 const slashIdx = spec.indexOf("/");
136 if (slashIdx <= 0 || slashIdx === spec.length - 1) return undefined;
137 return {
138 provider: spec.slice(0, slashIdx),
139 modelId: spec.slice(slashIdx + 1),
140 };
141}
142
143export function configuredExtractionModel(
144 pi: ExtensionAPI,
145 ctx: {
146 cwd: string;
147 modelRegistry: ExtensionContext["modelRegistry"];
148 },
149): ExtensionContext["model"] {
150 const modelSpec = configuredHandoffExtractionModel(pi, ctx);
151 if (!modelSpec) return undefined;
152 const parsed = parseModelSpec(modelSpec);
153 if (!parsed) return undefined;
154 return ctx.modelRegistry.find(parsed.provider, parsed.modelId);
155}
156
157/**
158 * Build a candidate file set from two sources:
159 * 1. Primary: actual tool calls (read, write, edit, create) in the session
160 * 2. Secondary: file-like patterns in the conversation text (catches files
161 * that were discussed but never opened)
162 */
163export function extractCandidateFiles(entries: SessionEntry[], conversationText: string): Set<string> {
164 const files = new Set<string>();
165 const fileToolNames = new Set(["read", "write", "edit", "create"]);
166
167 // Primary: files from actual tool calls
168 for (const entry of entries) {
169 if (entry.type !== "message") continue;
170 const msg = entry.message;
171 if (msg.role !== "assistant") continue;
172
173 for (const block of msg.content) {
174 if (typeof block !== "object" || block === null || block.type !== "toolCall") continue;
175 if (!fileToolNames.has(block.name)) continue;
176
177 const args = block.arguments as Record<string, unknown>;
178 const filePath =
179 typeof args.path === "string" ? args.path : typeof args.file === "string" ? args.file : undefined;
180 if (!filePath) continue;
181 if (filePath.endsWith("/SKILL.md")) continue;
182
183 files.add(filePath);
184 }
185 }
186
187 // Secondary: file-like patterns from conversation text.
188 // Trailing lookahead so the boundary isn't consumed — otherwise adjacent
189 // files separated by a single space (e.g. "file1.txt file2.txt") get skipped.
190 const filePattern = /(?:^|\s)([a-zA-Z0-9._\-/]+\.[a-zA-Z0-9]+)(?=\s|$|[,;:)])/gm;
191 for (const match of conversationText.matchAll(filePattern)) {
192 const candidate = match[1];
193 if (candidate && !candidate.startsWith(".") && candidate.length > 2) {
194 files.add(candidate);
195 }
196 }
197
198 return files;
199}
200
201/**
202 * Extract skill names that were actually loaded during the conversation.
203 * Looks for read() tool calls targeting SKILL.md files and derives the
204 * skill name from the parent directory (the convention for pi skills).
205 */
206export function extractLoadedSkills(entries: SessionEntry[]): string[] {
207 const skills = new Set<string>();
208 for (const entry of entries) {
209 if (entry.type !== "message") continue;
210 const msg = entry.message;
211 if (msg.role !== "assistant") continue;
212
213 for (const block of msg.content) {
214 if (typeof block !== "object" || block === null || block.type !== "toolCall") continue;
215
216 // read() calls where the path ends in SKILL.md
217 if (block.name !== "read") continue;
218 const args = block.arguments as Record<string, unknown>;
219 const filePath = typeof args.path === "string" ? args.path : undefined;
220 if (!filePath?.endsWith("/SKILL.md")) continue;
221
222 // Skill name is the parent directory name:
223 // .../skills/backing-up-with-keld/SKILL.md → backing-up-with-keld
224 const parent = path.basename(path.dirname(filePath));
225 if (parent && parent !== "skills") {
226 skills.add(parent);
227 }
228 }
229 }
230 return [...skills].sort();
231}