1// SPDX-FileCopyrightText: Amolith <amolith@secluded.site>
2// SPDX-FileCopyrightText: Petr Baudis <pasky@ucw.cz>
3//
4// SPDX-License-Identifier: MIT
5
6import type { ExtensionContext, SessionEntry } from "@mariozechner/pi-coding-agent";
7import * as path from "node:path";
8
9/**
10 * Model used for handoff extraction calls. Set PI_HANDOFF_MODEL env var
11 * as "provider/modelId" (e.g. "anthropic/claude-haiku-4-5") to use a
12 * different model for extraction than the session's current model.
13 */
14const HANDOFF_MODEL_OVERRIDE = process.env.PI_HANDOFF_MODEL;
15
16/**
17 * Build a candidate file set from two sources:
18 * 1. Primary: actual tool calls (read, write, edit, create) in the session
19 * 2. Secondary: file-like patterns in the conversation text (catches files
20 * that were discussed but never opened)
21 */
22export function extractCandidateFiles(entries: SessionEntry[], conversationText: string): Set<string> {
23 const files = new Set<string>();
24 const fileToolNames = new Set(["read", "write", "edit", "create"]);
25
26 // Primary: files from actual tool calls
27 for (const entry of entries) {
28 if (entry.type !== "message") continue;
29 const msg = entry.message;
30 if (msg.role !== "assistant") continue;
31
32 for (const block of msg.content) {
33 if (typeof block !== "object" || block === null || block.type !== "toolCall") continue;
34 if (!fileToolNames.has(block.name)) continue;
35
36 const args = block.arguments as Record<string, unknown>;
37 const filePath =
38 typeof args.path === "string" ? args.path : typeof args.file === "string" ? args.file : undefined;
39 if (!filePath) continue;
40 if (filePath.endsWith("/SKILL.md")) continue;
41
42 files.add(filePath);
43 }
44 }
45
46 // Secondary: file-like patterns from conversation text.
47 // Trailing lookahead so the boundary isn't consumed — otherwise adjacent
48 // files separated by a single space (e.g. "file1.txt file2.txt") get skipped.
49 const filePattern = /(?:^|\s)([a-zA-Z0-9._\-/]+\.[a-zA-Z0-9]+)(?=\s|$|[,;:)])/gm;
50 for (const match of conversationText.matchAll(filePattern)) {
51 const candidate = match[1];
52 if (candidate && !candidate.startsWith(".") && candidate.length > 2) {
53 files.add(candidate);
54 }
55 }
56
57 return files;
58}
59
60/**
61 * Extract skill names that were actually loaded during the conversation.
62 * Looks for read() tool calls targeting SKILL.md files and derives the
63 * skill name from the parent directory (the convention for pi skills).
64 */
65export function extractLoadedSkills(entries: SessionEntry[]): string[] {
66 const skills = new Set<string>();
67 for (const entry of entries) {
68 if (entry.type !== "message") continue;
69 const msg = entry.message;
70 if (msg.role !== "assistant") continue;
71
72 for (const block of msg.content) {
73 if (typeof block !== "object" || block === null || block.type !== "toolCall") continue;
74
75 // read() calls where the path ends in SKILL.md
76 if (block.name !== "read") continue;
77 const args = block.arguments as Record<string, unknown>;
78 const filePath = typeof args.path === "string" ? args.path : undefined;
79 if (!filePath?.endsWith("/SKILL.md")) continue;
80
81 // Skill name is the parent directory name:
82 // .../skills/backing-up-with-keld/SKILL.md → backing-up-with-keld
83 const parent = path.basename(path.dirname(filePath));
84 if (parent && parent !== "skills") {
85 skills.add(parent);
86 }
87 }
88 }
89 return [...skills].sort();
90}
91
92/**
93 * Resolve the model to use for handoff extraction calls. Uses the
94 * PI_HANDOFF_MODEL env var if set, otherwise falls back to the session model.
95 */
96export function resolveExtractionModel(ctx: {
97 model: ExtensionContext["model"];
98 modelRegistry: ExtensionContext["modelRegistry"];
99}): ExtensionContext["model"] {
100 if (!HANDOFF_MODEL_OVERRIDE) return ctx.model;
101 const slashIdx = HANDOFF_MODEL_OVERRIDE.indexOf("/");
102 if (slashIdx <= 0) return ctx.model;
103 const provider = HANDOFF_MODEL_OVERRIDE.slice(0, slashIdx);
104 const modelId = HANDOFF_MODEL_OVERRIDE.slice(slashIdx + 1);
105 return ctx.modelRegistry.find(provider, modelId) ?? ctx.model;
106}