202cbe7366a557ac9a8d8a106fee3755db14e713

Message

Implement /allow-bash command to allow patterns in current session

Diff

This diff is truncated to protect this page.

  1diff --git a/README.md b/README.md
  2index 0d0ccd58a4625252fb2337e05e638be97674dfa3..ce3847c3502344f9ca8b12c60ce3d19467fdf406 100644
  3--- a/README.md
  4+++ b/README.md
  5@@ -62,6 +62,16 @@ The extension handles Pi's `tool_call` event and evaluates each tool call before
  6 
  7 `webfetch`, `websearch`, and MCP tools are deterministically allowed. MCP tools bypass policy evaluation. Shell syntax prevents deterministic allow matches, so compound Bash commands normally proceed to the LLM stage.
  8 
  9+### Temporary Bash override
 10+
 11+Use `/allow-bash <regular expression>` to allow matching Bash commands in the current session before static permissions, deterministic rules, cache lookup, and LLM review. For example:
 12+
 13+```text
 14+/allow-bash ^databricks\b
 15+```
 16+
 17diff --git a/src/index.ts b/src/index.ts
 18index ea587b919a641c955f8ab9bab4c18b477548a479..85ac14d96b0628596eb85a712c00009218efb076 100644
 19--- a/src/index.ts
 20+++ b/src/index.ts
 21@@ -6,6 +6,13 @@ import { evaluateWithLLM } from "./llm"
 22 import { checkStaticPermission } from "./static-permissions"
 23 import { logDecision } from "./logger"
 24 import { cacheKey, normalizeRequest } from "./normalizer"
 25+import {
 26+  createSessionBashAllowOverride,
 27+  matchesSessionBashAllowOverride,
 28+  restoreSessionBashAllowOverride,
 29+  SESSION_BASH_ALLOW_ENTRY,
 30+  type SessionBashAllowOverride,
 31+} from "./session-override"
 32 import type { Decision } from "./types"
 33 
 34 type ToolInput = Record<string, unknown>
 35@@ -46,11 +53,34 @@ async function runPipelineSingle(
 36   input: ToolInput,
 37   sessionId: string,
 38   cwd?: string,
 39+  sessionBashAllowOverride?: SessionBashAllowOverride,
 40 ): Promise<PipelineResult> {
 41   const start = performance.now()
 42+  const normalized = normalizeRequest(toolType, input)
 43+
 44+  if (
 45+    toolType === "bash"
 46+    && typeof input.command === "string"
 47+    && matchesSessionBashAllowOverride(sessionBashAllowOverride, input.command)
 48+  ) {
 49+    const decision: Decision = {
 50+      decision: "allow",
 51+      reason: "Matched session Bash allow override",
 52+      category: "session_allow",
 53+    }
 54+    logDecision({
 55+      toolType,
 56+      normalized,
 57+      decision,
 58+      source: "session",
 59+      timingMs: performance.now() - start,
 60+      sessionId,
 61+    })
 62+    return { decision, source: "session" }
 63+  }
 64+
 65   const staticPermission = cwd ? await checkStaticPermission(toolType, input, cwd) : undefined
 66   const deterministic = staticPermission ?? checkDeterministic(toolType, input)
 67-  const normalized = normalizeRequest(toolType, input)
 68 
 69   if (deterministic) {
 70     logDecision({
 71@@ -122,11 +152,67 @@ export async function runPipeline(
 72 }
 73 
 74 export default function (pi: ExtensionAPI) {
 75-  pi.on("tool_call", async (event, ctx) => {
 76-    await capturePiCredentials(ctx)
 77+  let sessionBashAllowOverride: SessionBashAllowOverride | undefined
 78+
 79+  pi.on("session_start", (_event, ctx) => {
 80+    sessionBashAllowOverride = restoreSessionBashAllowOverride(
 81+      ctx.sessionManager.getBranch(),
 82+      ctx.sessionManager.getSessionId(),
 83+    )
 84+  })
 85+
 86+  pi.registerCommand("allow-bash", {
 87+    description: "Temporarily allow Bash commands that match a regular expression",
 88+    handler: async (args, ctx) => {
 89+      const source = args.trim()
 90+      if (!source) {
 91+        const status = sessionBashAllowOverride
 92+          ? `Temporary Bash allow override: /${sessionBashAllowOverride.source}/`
 93+          : "No temporary Bash allow override. Use /allow-bash <regular expression>."
 94+        ctx.ui.notify(status, "info")
 95+        return
 96+      }
 97+
 98+      if (source === "clear") {
 99+        sessionBashAllowOverride = undefined
100+        pi.appendEntry(SESSION_BASH_ALLOW_ENTRY, {
101+          pattern: null,
102+          sessionId: ctx.sessionManager.getSessionId(),
103+        })
104+        ctx.ui.notify("Temporary Bash allow override cleared.", "info")
105+        return
106+      }
107 
108+      try {
109+        sessionBashAllowOverride = createSessionBashAllowOverride(source)
110+      } catch (error) {
111+        const message = error instanceof Error ? error.message : "Invalid regular expression"
112+        ctx.ui.notify(`Invalid regular expression: ${message}`, "error")
113+        return
114+      }
115+
116+      pi.appendEntry(SESSION_BASH_ALLOW_ENTRY, {
117+        pattern: source,
118+        sessionId: ctx.sessionManager.getSessionId(),
119+      })
120+      ctx.ui.notify(`Temporary Bash allow override enabled: /${source}/`, "warning")
121diff --git a/src/session-override.ts b/src/session-override.ts
122new file mode 100644
123index 0000000000000000000000000000000000000000..e009a7b2b629ba7117562594e6248bfbb0a4ed8a
124--- /dev/null
125+++ b/src/session-override.ts
126@@ -0,0 +1,50 @@
127+export const SESSION_BASH_ALLOW_ENTRY = "policy-engine-session-bash-allow"
128+
129+export type SessionBashAllowOverride = {
130+  source: string
131+  pattern: RegExp
132+}
133+
134+type SessionEntry = {
135+  type?: unknown
136+  customType?: unknown
137+  data?: unknown
138+}
139+
140+function sessionPattern(data: unknown, sessionId: string): { matches: boolean; source?: string } {
141+  if (!data || typeof data !== "object" || Array.isArray(data)) return { matches: false }
142+  const entry = data as { pattern?: unknown; sessionId?: unknown }
143+  if (entry.sessionId !== sessionId) return { matches: false }
144+  return { matches: true, source: typeof entry.pattern === "string" ? entry.pattern : undefined }
145+}
146+
147+export function createSessionBashAllowOverride(source: string): SessionBashAllowOverride {
148+  return { source, pattern: new RegExp(source) }
149+}
150+
151+export function matchesSessionBashAllowOverride(
152+  override: SessionBashAllowOverride | undefined,
153+  command: string,
154+): boolean {
155+  return override?.pattern.test(command) ?? false
156+}
157+
158+export function restoreSessionBashAllowOverride(
159+  entries: readonly unknown[],
160+  sessionId: string,
161+): SessionBashAllowOverride | undefined {
162+  let source: string | undefined
163+  for (const entry of entries) {
164+    const { type, customType, data } = entry as SessionEntry
165+    if (type !== "custom" || customType !== SESSION_BASH_ALLOW_ENTRY) continue
166+    const storedPattern = sessionPattern(data, sessionId)
167+    if (storedPattern.matches) source = storedPattern.source
168+  }
169+
170+  if (source === undefined) return undefined
171+  try {
172+    return createSessionBashAllowOverride(source)
173+  } catch {
174+    return undefined
175+  }
176+}
177diff --git a/test/extension.test.ts b/test/extension.test.ts
178index 1120cfda462d6a8323488545435faccf3fd1ed90..ccd1df19dd14897cc876bb6675297732e5d72f7b 100644
179--- a/test/extension.test.ts
180+++ b/test/extension.test.ts
181@@ -20,13 +20,33 @@ afterEach(() => {
182 })
183 
184 type ToolCallHandler = (event: { toolName: string; input: Record<string, unknown> }, ctx: any) => Promise<unknown>
185+type SessionStartHandler = (event: { reason: string }, ctx: any) => unknown
186+type CommandHandler = (args: string, ctx: any) => Promise<void>
187 type EmittedEvent = { name: string; data: unknown }
188+type SessionEntry = { customType: string; data: unknown }
189 
190-function loadToolCallHandler(emittedEvents: EmittedEvent[] = []): ToolCallHandler {
191-  let handler: ToolCallHandler | undefined
192+type ExtensionHandlers = {
193+  toolCall: ToolCallHandler
194+  sessionStart: SessionStartHandler
195+  allowBash: CommandHandler
196+  entries: SessionEntry[]
197+}
198+
199+function loadExtension(emittedEvents: EmittedEvent[] = []): ExtensionHandlers {
200+  let toolCall: ToolCallHandler | undefined
201+  let sessionStart: SessionStartHandler | undefined
202+  let allowBash: CommandHandler | undefined
203+  const entries: SessionEntry[] = []
204   PolicyEngine({
205-    on(event: string, callback: ToolCallHandler) {
206-      if (event === "tool_call") handler = callback
207+    on(event: string, callback: ToolCallHandler | SessionStartHandler) {
208+      if (event === "tool_call") toolCall = callback as ToolCallHandler
209+      if (event === "session_start") sessionStart = callback as SessionStartHandler
210+    },
211+    registerCommand(name: string, command: { handler: CommandHandler }) {
212+      if (name === "allow-bash") allowBash = command.handler
213+    },
214+    appendEntry(customType: string, data: unknown) {
215+      entries.push({ customType, data })
216     },
217     events: {
218       emit(name: string, data: unknown) {
219@@ -34,8 +54,8 @@ function loadToolCallHandler(emittedEvents: EmittedEvent[] = []): ToolCallHandle
220       },
221     },
222   } as any)
223-  if (!handler) throw new Error("tool_call handler was not registered")
224-  return handler
225+  if (!toolCall || !sessionStart || !allowBash) throw new Error("policy handlers were not registered")
226+  return { toolCall, sessionStart, allowBash, entries }
227 }
228 
229 function context(hasUI: boolean, confirm = async () => false) {
230@@ -43,15 +63,16 @@ function context(hasUI: boolean, confirm = async () => false) {
231     hasUI,
232     model: undefined,
233     modelRegistry: {},
234-    sessionManager: { getSessionId: () => "test-session" },
235-    ui: { confirm },
236+    cwd: tempDir,
237+    sessionManager: { getSessionId: () => "test-session", getBranch: () => [] },
238+    ui: { confirm, notify() {} },
239   }
240 }
241 
242 describe("Pi policy extension", () => {
243   test("allows deterministically safe bash commands", async () => {
244-    const handler = loadToolCallHandler()
245-    await expect(handler({ toolName: "bash", input: { command: "git status" } }, context(false)))
246+    const { toolCall } = loadExtension()
247+    await expect(toolCall({ toolName: "bash", input: { command: "git status" } }, context(false)))
248       .resolves.toBeUndefined()
249   })
250 
251@@ -67,26 +88,26 @@ describe("Pi policy extension", () => {
252     writeFileSync(process.env.PI_POLICY_ENGINE_CONFIG!, JSON.stringify({
253       permission: { bash: { "*": "check" } },
254     }))
255-    const handler = loadToolCallHandler()
256-    await expect(handler({ toolName: "bash", input: { command: "git status" } }, context(false)))
257+    const { toolCall } = loadExtension()
258+    await expect(toolCall({ toolName: "bash", input: { command: "git status" } }, context(false)))
259       .resolves.toBeUndefined()
260   })
261 
262   test("blocks approval-required commands without a UI", async () => {
263-    const handler = loadToolCallHandler()
264-    await expect(handler({ toolName: "bash", input: { command: "sudo apt install foo" } }, context(false)))
265+    const { toolCall } = loadExtension()
266+    await expect(toolCall({ toolName: "bash", input: { command: "sudo apt install foo" } }, context(false)))
267       .resolves.toMatchObject({ block: true })
268   })
269 
270   test("allows a user-approved command", async () => {
271     const emittedEvents: EmittedEvent[] = []
272-    const handler = loadToolCallHandler(emittedEvents)
273+    const { toolCall } = loadExtension(emittedEvents)
274     let prompted = false
275     const confirm = async () => {
276       prompted = true
277       return true
278     }
279-    await expect(handler({ toolName: "bash", input: { command: "sudo apt install foo" } }, context(true, confirm)))
280+    await expect(toolCall({ toolName: "bash", input: { command: "sudo apt install foo" } }, context(true, confirm)))