index.ts
12066 bytes
1import type { UserMessage } from "@earendil-works/pi-ai";
2import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
3import { createJsonlDecisionAudit } from "../core/audit";
4import { createJsonlDecisionCache } from "../core/cache";
5import { parseBash, type BashParseResult } from "../core/bash";
6import { checkDeterministic, evaluateParsedBash } from "../core/deterministic";
7import { auditInputSummary, cacheKey, normalizeRequest } from "../core/normalize";
8import { createPolicyPipeline } from "../core/pipeline";
9import { createPolicyReviewer, type PiReviewerCaller } from "../core/review";
10import type { PiClassifierCaller } from "../core/classify";
11import type { PolicyContext, PolicyEvaluation, PolicyRequest } from "../core/types";
12import { loadPiPolicyConfig, piPolicyPaths } from "./config";
13import { PolicyApprovalDialog } from "./policy-approval";
14import {
15 createSessionBashAllowOverride,
16 matchesSessionBashAllowOverride,
17 restoreSessionBashAllowOverride,
18 SESSION_BASH_ALLOW_ENTRY,
19 type SessionBashAllowOverride,
20} from "./session-override";
21import { checkStaticPermission } from "./static-permissions";
22
23type ToolInput = Record<string, unknown>;
24
25function toolInput(input: unknown): ToolInput {
26 return input && typeof input === "object" && !Array.isArray(input) ? { ...(input as ToolInput) } : {};
27}
28
29function inputSummary(toolName: string, input: ToolInput): string {
30 if (toolName === "bash" && typeof input.command === "string") return input.command;
31 return JSON.stringify(input).slice(0, 500);
32}
33
34function approvalMessage(toolName: string, input: ToolInput, result: PolicyEvaluation): string {
35 const rawResponse =
36 result.error === "Failed to parse response JSON" && result.rawResponse
37 ? `\n\nLLM response:\n${result.rawResponse}`
38 : "";
39 if (toolName !== "bash" || typeof input.command !== "string") {
40 return `${toolName}: ${inputSummary(toolName, input)}\n\n${result.decision.reason}${rawResponse}`;
41 }
42
43 const commands = result.approvalRequests
44 ?.map((request) => request.input.command)
45 .filter((command): command is string => typeof command === "string");
46 const requiringReview = commands?.length
47 ? `\n\nCommands requiring review:\n${commands.map((command) => `- ${command}`).join("\n")}`
48 : "";
49 return `Entire Bash input:\n${input.command}${requiringReview}\n\n${result.decision.reason}${rawResponse}`;
50}
51
52function piReviewerCaller(ctx: ExtensionContext): PiReviewerCaller {
53 return async (
54 config: {
55 provider: string;
56 model: string;
57 reasoningEffort?: "none" | "minimal" | "low" | "medium" | "high" | "xhigh" | "max";
58 promptCacheKey?: string;
59 },
60 messages: readonly { role: "system" | "user"; content: string }[],
61 _maxTokens: number,
62 context: PolicyContext,
63 ): Promise<string> => {
64 const model = ctx.modelRegistry.find(config.provider, config.model);
65 if (!model) throw new Error(`Pi model is unavailable: ${config.provider}/${config.model}`);
66 const systemPrompt = messages.find((message) => message.role === "system")?.content;
67 const userMessages: UserMessage[] = messages
68 .filter((message) => message.role === "user")
69 .map((message) => ({
70 role: "user",
71 content: [{ type: "text", text: message.content }],
72 timestamp: Date.now(),
73 }));
74 const response = await ctx.modelRegistry.complete(
75 model,
76 { systemPrompt, messages: userMessages },
77 {
78 signal: context.signal,
79 reasoningEffort: config.reasoningEffort,
80 sessionId: config.promptCacheKey,
81 },
82 );
83 if (response.stopReason !== "stop")
84 throw new Error(response.errorMessage ?? `Pi reviewer stopped: ${response.stopReason}`);
85 const text = response.content
86 .filter((content): content is { type: "text"; text: string } => content.type === "text")
87 .map((content) => content.text)
88 .join("\n");
89 if (!text.trim()) throw new Error("Pi reviewer returned no text");
90 return text;
91 };
92}
93
94function piClassifierCaller(ctx: ExtensionContext): PiClassifierCaller {
95 return async (config, input, context) => {
96 const model = ctx.modelRegistry.getModelOfType("classifier", config.provider, config.model);
97 if (!model) throw new Error(`Pi classifier is unavailable: ${config.provider}/${config.model}`);
98 return ctx.modelRegistry.classify(model, input, { signal: context.signal });
99 };
100}
101
102function createPiPipeline(
103 sessionBashAllowOverride: SessionBashAllowOverride | undefined,
104 cwd: string,
105 ctx: ExtensionContext,
106 parsedBash: WeakMap<object, BashParseResult>,
107 bashReviewCommands: WeakMap<object, string[]>,
108) {
109 const config = loadPiPolicyConfig(cwd);
110 const paths = piPolicyPaths();
111 return createPolicyPipeline({
112 prechecks: [
113 {
114 source: "session",
115 decide: async (request) => {
116 if (
117 request.toolName !== "bash" ||
118 typeof request.input.command !== "string" ||
119 !matchesSessionBashAllowOverride(sessionBashAllowOverride, request.input.command)
120 )
121 return undefined;
122 return {
123 decision: "allow",
124 reason: "Matched session Bash allow override",
125 category: "session_allow",
126 };
127 },
128 },
129 {
130 source: "static",
131 decide: async (request) => checkStaticPermission(request.toolName, request.input, cwd, config),
132 },
133 ],
134 deterministic: (request, context) => {
135 const parsed = parsedBash.get(request.input);
136 if (!parsed) return checkDeterministic(request.toolName, request.input, context.cwd ?? cwd);
137 const result = evaluateParsedBash(parsed, context.cwd ?? cwd);
138 if (result.reviewCommands) bashReviewCommands.set(request.input, result.reviewCommands);
139 return result.decision;
140 },
141 cache: createJsonlDecisionCache(paths.cacheFile),
142 audit: createJsonlDecisionAudit(paths.auditFile),
143 reviewer: createPolicyReviewer(
144 config.reviewer,
145 piReviewerCaller(ctx),
146 config.externalDirectories,
147 piClassifierCaller(ctx),
148 ),
149 normalize: (request, context) => `${context.cwd ?? ""}\n${normalizeRequest(request.toolName, request.input)}`,
150 cacheKey,
151 inputSummary: auditInputSummary,
152 });
153}
154
155export async function evaluatePiToolCall(
156 toolName: string,
157 input: ToolInput,
158 context: PolicyContext,
159 sessionBashAllowOverride: SessionBashAllowOverride | undefined,
160 parse: (command: string) => Promise<BashParseResult> = parseBash,
161 skipPermissions = false,
162 ctx?: ExtensionContext,
163): Promise<PolicyEvaluation> {
164 if (!ctx) throw new Error("Pi extension context is required");
165 const parsedBash = new WeakMap<object, BashParseResult>();
166 const bashReviewCommands = new WeakMap<object, string[]>();
167 const pipeline = createPiPipeline(
168 sessionBashAllowOverride,
169 context.cwd ?? process.cwd(),
170 ctx,
171 parsedBash,
172 bashReviewCommands,
173 );
174 const evaluationOptions = {
175 skipPrechecks: true,
176 skipLLMReview: skipPermissions,
177 };
178 const request: PolicyRequest = { toolName, input };
179 const prechecked = await pipeline.precheck(request, context);
180 if (prechecked) return prechecked;
181
182 if (toolName !== "bash" || typeof input.command !== "string") {
183 return pipeline.evaluate(request, context, evaluationOptions);
184 }
185
186 const parsed = await parse(input.command);
187 if (parsed.parserUnavailable) {
188 return {
189 decision: {
190 decision: "deny",
191 reason: "Bash parser is unavailable",
192 category: "bash",
193 },
194 source: "deterministic",
195 };
196 }
197
198 parsedBash.set(input, parsed);
199 const result = await pipeline.evaluate(request, context, evaluationOptions);
200 const reviewCommands = bashReviewCommands.get(input);
201 return reviewCommands?.length
202 ? {
203 ...result,
204 approvalRequests: reviewCommands.map((command) => ({ toolName, input: { ...input, command } })),
205 }
206 : result;
207}
208
209export default function policyEngine(pi: ExtensionAPI) {
210 let sessionBashAllowOverride: SessionBashAllowOverride | undefined;
211 let skipPermissions = false;
212
213 pi.on("session_start", (_event, ctx) => {
214 sessionBashAllowOverride = restoreSessionBashAllowOverride(
215 ctx.sessionManager.getBranch(),
216 ctx.sessionManager.getSessionId(),
217 );
218 });
219
220 pi.registerCommand("pe-allow-bash", {
221 description: "Temporarily allow Bash commands that match a regular expression",
222 handler: async (args, ctx) => {
223 const source = args.trim();
224 if (!source) {
225 const status = sessionBashAllowOverride
226 ? `Temporary Bash allow override: /${sessionBashAllowOverride.source}/`
227 : "No temporary Bash allow override. Use /pe-allow-bash <regular expression>.";
228 ctx.ui.notify(status, "info");
229 return;
230 }
231
232 if (source === "clear") {
233 sessionBashAllowOverride = undefined;
234 pi.appendEntry(SESSION_BASH_ALLOW_ENTRY, {
235 pattern: null,
236 sessionId: ctx.sessionManager.getSessionId(),
237 });
238 ctx.ui.notify("Temporary Bash allow override cleared.", "info");
239 return;
240 }
241
242 try {
243 sessionBashAllowOverride = createSessionBashAllowOverride(source);
244 } catch (error) {
245 const message = error instanceof Error ? error.message : "Invalid regular expression";
246 ctx.ui.notify(`Invalid regular expression: ${message}`, "error");
247 return;
248 }
249
250 pi.appendEntry(SESSION_BASH_ALLOW_ENTRY, {
251 pattern: source,
252 sessionId: ctx.sessionManager.getSessionId(),
253 });
254 ctx.ui.notify(`Temporary Bash allow override enabled: /${source}/`, "warning");
255 },
256 });
257
258 pi.registerCommand("pe-skip-permissions", {
259 description: "Temporarily skip LLM permission checks",
260 handler: async (args, ctx) => {
261 if (args.trim() === "clear") {
262 skipPermissions = false;
263 ctx.ui.notify("LLM permission checks enabled.", "info");
264 return;
265 }
266 skipPermissions = true;
267 ctx.ui.notify(
268 "LLM permission checks disabled until this extension reloads. Use /pe-skip-permissions clear to re-enable them.",
269 "warning",
270 );
271 },
272 });
273
274 pi.on("tool_call", async (event, ctx) => {
275 if (event.toolName.startsWith("mcp")) return undefined;
276
277 const input = toolInput(event.input);
278 const result = await evaluatePiToolCall(
279 event.toolName,
280 input,
281 {
282 sessionId: ctx.sessionManager.getSessionId(),
283 cwd: ctx.cwd,
284 signal: ctx.signal,
285 },
286 sessionBashAllowOverride,
287 parseBash,
288 skipPermissions,
289 ctx,
290 );
291
292 if (result.decision.decision === "allow") return undefined;
293 if (result.decision.decision === "deny") {
294 return {
295 block: true,
296 reason: `Denied by policy: ${result.decision.reason}`,
297 };
298 }
299 if (!ctx.hasUI) {
300 return {
301 block: true,
302 reason: `Policy approval required: ${result.decision.reason}`,
303 };
304 }
305
306 pi.events.emit("herdr:blocked", {
307 active: true,
308 label: "Policy approval required",
309 });
310 let approved: boolean;
311 try {
312 const title = "Policy approval required";
313 const message = approvalMessage(event.toolName, input, result);
314 approved =
315 ctx.mode === "tui"
316 ? await ctx.ui.custom(
317 (tui, theme, keybindings, done) =>
318 new PolicyApprovalDialog(tui, theme, keybindings, title, message, done),
319 {
320 overlay: true,
321 overlayOptions: { width: "90%", margin: 1 },
322 },
323 )
324 : await ctx.ui.confirm(title, message);
325 } finally {
326 pi.events.emit("herdr:blocked", { active: false });
327 }
328 if (approved) return undefined;
329 const feedback = await ctx.ui.input(
330 "Why was this rejected?",
331 "Tell the agent why it was rejected and what to do differently",
332 );
333 return {
334 block: true,
335 reason: `Blocked by user: ${feedback?.trim() || result.decision.reason}`,
336 };
337 });
338}