1// SPDX-FileCopyrightText: Amolith <amolith@secluded.site>
2// SPDX-FileCopyrightText: Petr Baudis <pasky@ucw.cz>
3//
4// SPDX-License-Identifier: MIT
5
6import { complete, type Message, type Tool } from "@mariozechner/pi-ai";
7import type { ExtensionContext } from "@mariozechner/pi-coding-agent";
8import { Type } from "@sinclair/typebox";
9
10const SYSTEM_PROMPT = `You're helping transfer context between coding sessions. The next session starts fresh with no memory of this conversation, so extract what matters.
11
12Consider these questions:
13- What was just done or implemented?
14- What decisions were made and why?
15- What technical details were discovered (APIs, methods, patterns)?
16- What constraints, limitations, or caveats were found?
17- What patterns or approaches are being followed?
18- What is still unfinished or unresolved?
19- What open questions or risks are worth flagging?
20
21Rules:
22- Be concrete: prefer file paths, specific decisions, and actual commands over vague summaries.
23- Only include files from the provided candidate list.
24- Only include skills from the provided loaded skills list.
25- If something wasn't explicitly discussed, don't include it.`;
26
27/**
28 * Tool definition passed to the extraction LLM call. Models produce
29 * structured tool-call arguments far more reliably than free-form JSON
30 * in a text response, so we use a single tool instead of responseFormat.
31 */
32const HANDOFF_EXTRACTION_TOOL: Tool = {
33 name: "extract_handoff_context",
34 description: "Extract handoff context from the conversation for a new session.",
35 parameters: Type.Object({
36 relevantFiles: Type.Array(Type.String(), {
37 description: "File paths relevant to continuing this work, chosen from the provided candidate list",
38 }),
39 skillsInUse: Type.Array(Type.String(), {
40 description: "Skills relevant to the goal, chosen from the provided list of loaded skills",
41 }),
42 context: Type.String({
43 description:
44 "Key context: what was done, decisions made, technical details discovered, constraints, and patterns being followed",
45 }),
46 openItems: Type.Array(Type.String(), {
47 description: "Unfinished work, open questions, known risks, and things to watch out for",
48 }),
49 }),
50};
51
52export type HandoffExtraction = {
53 files: string[];
54 skills: string[];
55 context: string;
56 openItems: string[];
57};
58
59/**
60 * Run the extraction LLM call and return structured context, or null if
61 * aborted or the model didn't produce a valid tool call.
62 */
63export async function extractHandoffContext(
64 model: NonNullable<ExtensionContext["model"]>,
65 apiKey: string | undefined,
66 headers: Record<string, string> | undefined,
67 conversationText: string,
68 goal: string,
69 candidateFiles: string[],
70 loadedSkills: string[],
71 signal?: AbortSignal,
72): Promise<HandoffExtraction | null> {
73 const filesContext =
74 candidateFiles.length > 0
75 ? `\n\n## Candidate Files\n\nThese files were touched or mentioned during the session. Return only the ones relevant to the goal.\n\n${candidateFiles.map((f) => `- ${f}`).join("\n")}`
76 : "";
77
78 const skillsContext =
79 loadedSkills.length > 0
80 ? `\n\n## Skills Loaded During This Session\n\n${loadedSkills.map((s) => `- ${s}`).join("\n")}\n\nReturn only the skills from this list that are relevant to the goal.`
81 : "";
82
83 const userMessage: Message = {
84 role: "user",
85 content: [
86 {
87 type: "text",
88 text: `## Conversation\n\n${conversationText}\n\n## Goal for Next Session\n\n${goal}${filesContext}${skillsContext}`,
89 },
90 ],
91 timestamp: Date.now(),
92 };
93
94 const response = await complete(
95 model,
96 {
97 systemPrompt: SYSTEM_PROMPT,
98 messages: [userMessage],
99 tools: [HANDOFF_EXTRACTION_TOOL],
100 },
101 { apiKey, headers, signal },
102 );
103
104 if (response.stopReason === "aborted") return null;
105
106 const toolCall = response.content.find((c) => c.type === "toolCall" && c.name === "extract_handoff_context");
107
108 if (!toolCall || toolCall.type !== "toolCall") {
109 console.error("Model did not call extract_handoff_context:", response.content);
110 return null;
111 }
112
113 const args = toolCall.arguments as Record<string, unknown>;
114
115 if (
116 !Array.isArray(args.relevantFiles) ||
117 !Array.isArray(args.skillsInUse) ||
118 typeof args.context !== "string" ||
119 !Array.isArray(args.openItems)
120 ) {
121 console.error("Unexpected tool call arguments shape:", args);
122 return null;
123 }
124
125 const candidateSet = new Set(candidateFiles);
126 const loadedSkillSet = new Set(loadedSkills);
127
128 return {
129 files: (args.relevantFiles as string[]).filter((f) => typeof f === "string" && candidateSet.has(f)),
130 skills: (args.skillsInUse as string[]).filter((s) => typeof s === "string" && loadedSkillSet.has(s)),
131 context: args.context as string,
132 openItems: (args.openItems as string[]).filter((item) => typeof item === "string"),
133 };
134}
135
136/**
137 * Assemble the handoff draft. The user's goal goes last so it has the
138 * most weight in the new session (recency bias).
139 */
140export function assembleHandoffDraft(
141 result: HandoffExtraction,
142 goal: string,
143 parentSessionFile: string | undefined,
144): string {
145 const parentBlock = parentSessionFile
146 ? `**Parent session:** \`${parentSessionFile}\`\n\nUse the \`session_query\` tool with this path if you need details from the prior thread.\n\n`
147 : "";
148
149 const filesSection = result.files.length > 0 ? `## Files\n\n${result.files.map((f) => `- ${f}`).join("\n")}\n\n` : "";
150
151 const skillsSection =
152 result.skills.length > 0 ? `## Skills in Use\n\n${result.skills.map((s) => `- ${s}`).join("\n")}\n\n` : "";
153
154 const contextSection = `## Context\n\n${result.context}\n\n`;
155
156 const openItemsSection =
157 result.openItems.length > 0 ? `## Open Items\n\n${result.openItems.map((item) => `- ${item}`).join("\n")}\n\n` : "";
158
159 return `${parentBlock}${filesSection}${skillsSection}${contextSection}${openItemsSection}${goal}`.trim();
160}