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 } from "@mariozechner/pi-ai";
7import type { ExtensionAPI, SessionEntry } from "@mariozechner/pi-coding-agent";
8import { SessionManager, convertToLlm, serializeConversation } from "@mariozechner/pi-coding-agent";
9import { Type } from "@sinclair/typebox";
10import * as fs from "node:fs";
11import { getFallbackSessionsRoot, getSessionsRoot, normalizeSessionPath, sessionPathAllowed } from "./session-paths.js";
12
13const QUERY_SYSTEM_PROMPT = `You answer questions about a prior pi session.
14
15Rules:
16- Use only facts from the provided conversation.
17- Prefer concrete outputs: file paths, decisions, TODOs, errors.
18- If not present, say explicitly: "Not found in provided session.".
19- Keep answer concise.`;
20
21export function registerSessionQueryTool(pi: ExtensionAPI) {
22 pi.registerTool({
23 name: "session_query",
24 label: "Session Query",
25 description:
26 "Query a prior pi session file. Use when handoff prompt references a parent session and you need details.",
27 parameters: Type.Object({
28 sessionPath: Type.String({
29 description:
30 "Session .jsonl path. Absolute path, or relative to sessions root (e.g. 2026-02-16/foo/session.jsonl)",
31 }),
32 question: Type.String({ description: "Question about that session" }),
33 }),
34 async execute(_toolCallId, params, signal, onUpdate, ctx) {
35 const currentSessionFile = ctx.sessionManager.getSessionFile();
36 const sessionsRoot = getSessionsRoot(currentSessionFile) ?? getFallbackSessionsRoot();
37 const resolvedPath = normalizeSessionPath(params.sessionPath, sessionsRoot);
38
39 const error = (text: string) => ({
40 content: [{ type: "text" as const, text }],
41 details: { error: true } as const,
42 });
43
44 const cancelled = () => ({
45 content: [{ type: "text" as const, text: "Session query cancelled." }],
46 details: { cancelled: true } as const,
47 });
48
49 if (signal?.aborted) {
50 return cancelled();
51 }
52
53 if (!resolvedPath.endsWith(".jsonl")) {
54 return error(`Invalid session path (expected .jsonl): ${params.sessionPath}`);
55 }
56
57 if (!sessionPathAllowed(resolvedPath, sessionsRoot)) {
58 return error(`Session path outside allowed sessions directory: ${params.sessionPath}`);
59 }
60
61 if (!fs.existsSync(resolvedPath)) {
62 return error(`Session file not found: ${resolvedPath}`);
63 }
64
65 let fileStats: fs.Stats;
66 try {
67 fileStats = fs.statSync(resolvedPath);
68 } catch (err) {
69 return error(`Failed to stat session file: ${String(err)}`);
70 }
71
72 if (!fileStats.isFile()) {
73 return error(`Session path is not a file: ${resolvedPath}`);
74 }
75
76 onUpdate?.({
77 content: [{ type: "text", text: `Querying: ${resolvedPath}` }],
78 details: { status: "loading", sessionPath: resolvedPath },
79 });
80
81 let sessionManager: SessionManager;
82 try {
83 sessionManager = SessionManager.open(resolvedPath);
84 } catch (err) {
85 return error(`Failed to open session: ${String(err)}`);
86 }
87
88 const branch = sessionManager.getBranch();
89 const messages = branch
90 .filter((entry): entry is SessionEntry & { type: "message" } => entry.type === "message")
91 .map((entry) => entry.message);
92
93 if (messages.length === 0) {
94 return {
95 content: [{ type: "text" as const, text: "Session has no messages." }],
96 details: { empty: true, sessionPath: resolvedPath },
97 };
98 }
99
100 if (!ctx.model) {
101 return error("No model selected for session query.");
102 }
103
104 const conversationText = serializeConversation(convertToLlm(messages));
105 try {
106 const auth = await ctx.modelRegistry.getApiKeyAndHeaders(ctx.model);
107 if (!auth.ok) {
108 return error(`Failed to get API key: ${auth.error}`);
109 }
110 const userMessage: Message = {
111 role: "user",
112 content: [
113 {
114 type: "text",
115 text: `## Session\n\n${conversationText}\n\n## Question\n\n${params.question}`,
116 },
117 ],
118 timestamp: Date.now(),
119 };
120
121 const response = await complete(
122 ctx.model,
123 { systemPrompt: QUERY_SYSTEM_PROMPT, messages: [userMessage] },
124 { apiKey: auth.apiKey, headers: auth.headers, signal: signal as AbortSignal },
125 );
126
127 if (response.stopReason === "aborted") {
128 return cancelled();
129 }
130
131 const answer = response.content
132 .filter((c): c is { type: "text"; text: string } => c.type === "text")
133 .map((c) => c.text)
134 .join("\n")
135 .trim();
136
137 return {
138 content: [{ type: "text" as const, text: answer || "No answer generated." }],
139 details: {
140 sessionPath: resolvedPath,
141 question: params.question,
142 messageCount: messages.length,
143 },
144 };
145 } catch (err) {
146 if (signal?.aborted) {
147 return cancelled();
148 }
149 if (err instanceof Error && err.name === "AbortError") {
150 return cancelled();
151 }
152 return error(`Session query failed: ${String(err)}`);
153 }
154 },
155 });
156}