pi-code 0.2.1 → 0.2.2

package/extensions/context-imports.ts

@@ -5,16 +5,18 @@
  * Claude Code's `@path` imports inside them. This fills that one gap: on
  * before_agent_start it reads the already-loaded context files from
  * systemPromptOptions, resolves any `@path` imports (recursive, depth-capped,
- * cycle-safe; ~ expands to home, relative paths resolve against the importing
- * file), and appends ONLY the imported content. pi already injected the base
- * files, so nothing is duplicated.
+ * cycle-safe, budget-capped; ~ expands to home, relative paths resolve against
+ * the importing file), and appends ONLY the imported content. pi already
+ * injected the base files, so nothing is duplicated.
  *
  * Security: context files can come from an untrusted project, so imports are
  * confined (after resolving symlinks) to the working directory and the user's
  * own ~/.claude and ~/.pi config roots. An import that escapes those roots
  * (absolute paths, ~/.ssh, ../.. traversal, symlinks) is ignored, so a hostile
  * CLAUDE.md cannot read arbitrary files into the prompt. Imports inside fenced
- * code blocks are also skipped.
+ * code blocks are also skipped. One byte-and-file budget is shared by the whole
+ * run, so a context file cannot flood the prompt by importing breadth-first;
+ * what the budget refused is stated in the prompt rather than dropped silently.
  *
  * Docs: https://code.claude.com/docs/en/memory.md (imports)
  */
@@ -25,6 +27,8 @@ import * as path from 'node:path'
 import type { ExtensionAPI } from '@earendil-works/pi-coding-agent'
 
 const MAX_IMPORT_DEPTH = 5
+export const MAX_IMPORT_FILES = 50
+export const MAX_IMPORT_BYTES = 256 * 1024
 
 export function expandHome(target: string, home: string): string {
   if (target === '~') return home
@@ -54,6 +58,18 @@ export interface ImportedFile {
   body: string
 }
 
+/** Appended to the last body the byte budget could only partly pay for. */
+export const IMPORT_TRUNCATED_MARKER = '[truncated: import byte budget exhausted]'
+
+/** Remaining import allowance, shared across every context file of one run. */
+export interface ImportBudget {
+  files: number
+  bytes: number
+  dropped: number
+}
+
+export const createImportBudget = (): ImportBudget => ({ files: MAX_IMPORT_FILES, bytes: MAX_IMPORT_BYTES, dropped: 0 })
+
 /** The `@path` targets of a context file, in document order, skipping fenced code blocks. */
 function importTargets(content: string): string[] {
   const targets: string[] = []
@@ -94,13 +110,22 @@ function readImport(target: string, fromDir: string, home: string, allowedRoots:
  * discovery order. Imports are resolved through symlinks and kept within
  * `allowedRoots` (which must already be realpath'd).
  */
-export function collectImports(content: string, fromDir: string, home: string, allowedRoots: string[], seen: Set<string>, depth = 0): ImportedFile[] {
+export function collectImports(content: string, fromDir: string, home: string, allowedRoots: string[], seen: Set<string>, budget: ImportBudget = createImportBudget(), depth = 0): ImportedFile[] {
   if (depth >= MAX_IMPORT_DEPTH) return []
   const out: ImportedFile[] = []
   for (const target of importTargets(content)) {
+    // Checked before the read so an exhausted budget costs no I/O.
+    if (budget.files === 0 || budget.bytes === 0) {
+      budget.dropped += 1
+      continue
+    }
     const file = readImport(target, fromDir, home, allowedRoots, seen)
     if (!file) continue
-    out.push({ path: file.real, body: file.body.trim() }, ...collectImports(file.body, path.dirname(file.real), home, allowedRoots, seen, depth + 1))
+    budget.files -= 1
+    const kept = file.body.slice(0, budget.bytes)
+    budget.bytes -= kept.length
+    const body = kept.length < file.body.length ? `${kept.trim()}\n${IMPORT_TRUNCATED_MARKER}` : kept.trim()
+    out.push({ path: file.real, body }, ...collectImports(kept, path.dirname(file.real), home, allowedRoots, seen, budget, depth + 1))
   }
   return out
 }
@@ -132,14 +157,17 @@ export default function contextImportsExtension(pi: ExtensionAPI) {
     const seenSet = new Set(seen)
 
     const imported: ImportedFile[] = []
+    // One budget for the whole run, so N context files cannot each spend a full one.
+    const budget = createImportBudget()
     for (const file of contextFiles) {
       // Roots are scoped per importing file: a project file never reaches user config.
       const allowedRoots = rootsForImporter(file.path, home, cwd)
-      imported.push(...collectImports(file.content, path.dirname(file.path), home, allowedRoots, seenSet))
+      imported.push(...collectImports(file.content, path.dirname(file.path), home, allowedRoots, seenSet, budget))
     }
     if (imported.length === 0) return
 
     const section = imported.map((entry) => `### ${entry.path}\n\n${entry.body}`).join('\n\n')
-    return { systemPrompt: `${event.systemPrompt}\n\n## Imported context (@)\n\n${section}` }
+    const notice = budget.dropped === 0 ? '' : `\n\n${budget.dropped} further @imports were skipped: the import budget (${MAX_IMPORT_FILES} files, ${MAX_IMPORT_BYTES} bytes) is spent.`
+    return { systemPrompt: `${event.systemPrompt}\n\n## Imported context (@)\n\n${section}${notice}` }
   })
 }

package/extensions/hooks.ts

@@ -21,7 +21,7 @@
  * Docs: https://code.claude.com/docs/en/hooks.md
  */
 
-import { spawn } from 'node:child_process'
+import { type ChildProcess, spawn } from 'node:child_process'
 import * as fs from 'node:fs'
 import * as os from 'node:os'
 import * as path from 'node:path'
@@ -50,6 +50,8 @@ export interface HookRunResult {
   code: number
   stdout: string
   stderr: string
+  /** The hook was killed at its timeout, so its exit code carries no verdict. */
+  timedOut: boolean
 }
 export type HookRunner = (command: string, payload: unknown, timeoutMs: number) => Promise<HookRunResult>
 
@@ -112,27 +114,64 @@ export function interpretHookResult(code: number, stdout: string, stderr: string
   return { block: false }
 }
 
+/** Memory backstop for a runaway hook. A decision payload is orders of magnitude smaller. */
+const MAX_HOOK_OUTPUT = 1_000_000
+
+/** Conventional exit code for a killed-on-timeout command, as `timeout(1)` reports it. */
+const TIMEOUT_EXIT_CODE = 124
+
+/**
+ * Kill the shell and everything it spawned. `sh -c 'a; b'` forks, so signalling the
+ * direct child alone leaves a grandchild alive holding stdout/stderr.
+ */
+function killTree(child: ChildProcess): void {
+  try {
+    // Negative pid targets the whole process group, which `detached` gave the shell.
+    if (child.pid) {
+      process.kill(-child.pid, 'SIGKILL')
+      return
+    }
+  } catch {
+    // Group already reaped, or the platform refused it; fall through to the direct kill.
+  }
+  child.kill('SIGKILL')
+}
+
 export const runHookCommand: HookRunner = (command, payload, timeoutMs) =>
   new Promise((resolve) => {
     // Absolute path so the shell can't be resolved through an attacker-controlled PATH.
-    const child = spawn('/bin/sh', ['-c', command], { stdio: ['pipe', 'pipe', 'pipe'] })
+    // `detached` makes the shell its own process group leader so the timeout can kill
+    // the descendants too.
+    const child = spawn('/bin/sh', ['-c', command], { stdio: ['pipe', 'pipe', 'pipe'], detached: true })
     let stdout = ''
     let stderr = ''
-    const timer = setTimeout(() => child.kill('SIGKILL'), timeoutMs)
-    child.stdout?.on('data', (chunk) => {
-      stdout += chunk
-    })
-    child.stderr?.on('data', (chunk) => {
-      stderr += chunk
-    })
-    child.on('close', (code) => {
+    let settled = false
+    const finish = (result: HookRunResult): void => {
+      if (settled) return
+      settled = true
       clearTimeout(timer)
-      resolve({ code: code ?? 0, stdout, stderr })
+      resolve(result)
+    }
+    // Resolve from the timer itself rather than waiting for `close`: `close` fires only
+    // once every stdio pipe is closed, and a grandchild that inherited them can hold the
+    // promise pending long past the timeout, stalling the tool call that awaits it.
+    const timer = setTimeout(() => {
+      killTree(child)
+      finish({ code: TIMEOUT_EXIT_CODE, stdout, stderr, timedOut: true })
+    }, timeoutMs)
+    // Decode on the stream: concatenating Buffers as strings mangles a multi-byte
+    // character split across chunks, and a mangled byte in a hook's deny decision makes
+    // it unparseable, which reads as an allow.
+    child.stdout?.setEncoding('utf8')
+    child.stderr?.setEncoding('utf8')
+    child.stdout?.on('data', (chunk: string) => {
+      if (stdout.length < MAX_HOOK_OUTPUT) stdout += chunk
     })
-    child.on('error', () => {
-      clearTimeout(timer)
-      resolve({ code: 0, stdout, stderr })
+    child.stderr?.on('data', (chunk: string) => {
+      if (stderr.length < MAX_HOOK_OUTPUT) stderr += chunk
     })
+    child.on('close', (code) => finish({ code: code ?? 0, stdout, stderr, timedOut: false }))
+    child.on('error', () => finish({ code: 0, stdout, stderr, timedOut: false }))
     // A hook that exits without reading stdin (e.g. `exit 2`) closes the pipe first,
     // so ignore EPIPE on this write rather than crashing the host process.
     child.stdin?.on('error', () => {})
@@ -147,6 +186,9 @@ function timeoutMs(command: HookCommand): number {
 export async function runPreToolUse(config: HooksConfig, toolName: string, toolInput: unknown, runner: HookRunner): Promise<HookDecision> {
   for (const command of matchingCommands(config.PreToolUse, toolName)) {
     const result = await runner(command.command, { hook_event_name: 'PreToolUse', tool_name: toolName, tool_input: toolInput }, timeoutMs(command))
+    // A killed hook never reached its verdict, and SIGKILL leaves a null exit code that
+    // would otherwise read as a clean allow. Fail closed instead.
+    if (result.timedOut) return { block: true, reason: `Hook timed out after ${timeoutMs(command)}ms: ${command.command}` }
     const decision = interpretHookResult(result.code, result.stdout, result.stderr)
     if (decision.block) return decision
   }

package/extensions/plan-mode/utils.ts

@@ -102,6 +102,14 @@ const SUBSTITUTION = /\$\(|`|<\(|>\(/
  *
  * A shell AST would be exact; this is the honest approximation for a quoting-only concern.
  */
+/** Length of the separator at `i`, or 0 when there is none. */
+function separatorAt(command: string, i: number): number {
+  const pair = command.slice(i, i + 2)
+  if (pair === '&&' || pair === '||' || pair === '|&') return 2
+  const ch = command[i]
+  return ch === ';' || ch === '|' || ch === '&' || ch === '\n' ? 1 : 0
+}
+
 function splitSegments(command: string): string[] {
   const segments: string[] = []
   let current = ''
@@ -123,16 +131,11 @@ function splitSegments(command: string): string[] {
       current += ch + command[++i]
       continue
     }
-    const pair = command.slice(i, i + 2)
-    if (pair === '&&' || pair === '||' || pair === '|&') {
-      segments.push(current)
-      current = ''
-      i++
-      continue
-    }
-    if (ch === ';' || ch === '|' || ch === '&' || ch === '\n') {
+    const separator = separatorAt(command, i)
+    if (separator > 0) {
       segments.push(current)
       current = ''
+      i += separator - 1
       continue
     }
     current += ch

package/extensions/subagent/agents.ts

@@ -99,12 +99,35 @@ function isDirectory(p: string): boolean {
   }
 }
 
+/** Project root at or above `from`. `.git` is a file in worktrees and submodules. */
+const ROOT_MARKERS = ['.git', 'package.json']
+
+function repoRoot(from: string): string | undefined {
+  let currentDir = from
+  while (true) {
+    if (ROOT_MARKERS.some((marker) => fs.existsSync(path.join(currentDir, marker)))) return currentDir
+    const parentDir = path.dirname(currentDir)
+    if (parentDir === currentDir) return undefined
+    currentDir = parentDir
+  }
+}
+
+/**
+ * Nearest `relative` directory at or above `cwd`, stopping at the repository root.
+ *
+ * Without the boundary the search runs to the filesystem root, so an agent planted in a
+ * world-writable ancestor such as /tmp is offered as a project agent for every session
+ * beneath it. With no project marker (.git, package.json) the extent is unknown, so only
+ * `cwd` is considered.
+ */
 function findNearestDir(cwd: string, relative: string): string | null {
+  const boundary = repoRoot(cwd) ?? cwd
   let currentDir = cwd
   while (true) {
     const candidate = path.join(currentDir, relative)
     if (isDirectory(candidate)) return candidate
 
+    if (currentDir === boundary) return null
     const parentDir = path.dirname(currentDir)
     if (parentDir === currentDir) return null
     currentDir = parentDir

package/extensions/subagent/index.ts

@@ -488,12 +488,16 @@ async function checkProjectAgentGate(params: SubagentParamsStatic, agents: Agent
   // isProjectTrusted alone is true for a repo pi never asked about; see project-approval.
   const approved = await isProjectApproved(ctx)
   const gate = projectAgentGate(requestedProjectAgents.length, approved, ctx.hasUI, params.confirmProjectAgents ?? true)
-  const names = requestedProjectAgents.map((a) => a.name).join(', ')
+  // Agent names come from repo-controlled frontmatter; a newline in one would otherwise
+  // let it write its own "Source:" line into the prompt body.
+  const names = requestedProjectAgents.map((a) => a.name.replace(/\s+/g, ' ').trim()).join(', ')
   if (gate === 'refuse') {
     return { content: [{ type: 'text', text: `Project-local agents (${names}) require a trusted project; refusing in non-interactive mode.` }], details: makeDetails(gateMode)([]) }
   }
   if (gate === 'confirm') {
-    const dir = projectAgentsDir ?? '(unknown)'
+    // Each agent knows where it was loaded from; projectAgentsDir only ever held .pi/agents.
+    const dirs = [...new Set(requestedProjectAgents.map((a) => path.dirname(a.filePath)))]
+    const dir = dirs.join(', ') || projectAgentsDir || '(unknown)'
     const ok = await ctx.ui.confirm('Run project-local agents?', `Agents: ${names}\nSource: ${dir}\n\nProject agents are repo-controlled. Only continue for trusted repositories.`)
     if (!ok) return { content: [{ type: 'text', text: 'Canceled: project-local agents not approved.' }], details: makeDetails(gateMode)([]) }
   }

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pi-code",
-  "version": "0.2.1",
+  "version": "0.2.2",
   "description": "Claude Code experience for the pi coding agent: reads your .claude config (rules, commands, skills, hooks, output styles, MCP servers, agents) and adds todo, checkpoints, memory, web, and subagents",
   "keywords": [
     "pi-package"