pi-code 0.2.1 → 0.2.3

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/extensions/web-transport.ts

@@ -0,0 +1,51 @@
+/**
+ * Web transport
+ *
+ * A single HTTP(S) request pinned to a caller-supplied DNS resolution. Global `fetch`
+ * resolves the hostname itself, independently of any prior guard, so a validate-then-fetch
+ * SSRF check has a time-of-check/time-of-use gap: a zero-TTL record can answer public to
+ * the guard and private to fetch's own lookup. `node:http`/`node:https` accept a `lookup`
+ * option, which is the seam that closes the gap: the socket connects to exactly the address
+ * the guard validated, while `servername` (SNI, certificate validation) and the `Host`
+ * header stay the real hostname, so virtual hosts and TLS still work.
+ */
+
+import { request as httpRequest } from 'node:http'
+import { request as httpsRequest } from 'node:https'
+import type { LookupFunction } from 'node:net'
+import { Readable } from 'node:stream'
+
+export interface TransportOptions {
+  signal: AbortSignal
+  lookup: LookupFunction
+  userAgent: string
+}
+
+/** One request, no redirect following (the caller re-validates and re-pins per hop). */
+export function httpFetch(url: URL, opts: TransportOptions): Promise<Response> {
+  const request = url.protocol === 'https:' ? httpsRequest : httpRequest
+  return new Promise((resolve, reject) => {
+    const req = request(
+      url,
+      {
+        method: 'GET',
+        headers: { 'User-Agent': opts.userAgent },
+        signal: opts.signal,
+        lookup: opts.lookup,
+        // servername is left to default to url.hostname, so SNI and certificate
+        // validation use the real host even though the socket connects to the pinned IP.
+      },
+      (res) => {
+        const headers = new Headers()
+        for (const [key, value] of Object.entries(res.headers)) {
+          if (typeof value === 'string') headers.set(key, value)
+          else if (Array.isArray(value)) headers.set(key, value.join(', '))
+        }
+        const body = Readable.toWeb(res) as ReadableStream<Uint8Array>
+        resolve(new Response(body, { status: res.statusCode ?? 0, headers }))
+      },
+    )
+    req.on('error', reject)
+    req.end()
+  })
+}

package/extensions/web.ts

@@ -6,10 +6,14 @@
  * Honors the local-only setup: no cloud accounts, plain HTTPS to public web.
  */
 
+import type { LookupAddress } from 'node:dns'
 import { lookup } from 'node:dns/promises'
+import type { LookupFunction } from 'node:net'
 import type { ExtensionAPI } from '@earendil-works/pi-coding-agent'
 import { Type } from 'typebox'
 
+import { httpFetch } from './web-transport.js'
+
 const SEARCH_ENDPOINT = 'https://html.duckduckgo.com/html/?q='
 const USER_AGENT = 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) pi-code-web/0.1'
 const MAX_FETCH_CHARS = 30_000
@@ -128,15 +132,32 @@ export function isPrivateAddress(ip: string): boolean {
   return addr.includes(':') ? isPrivateIpv6(addr) : isPrivateIpv4(addr)
 }
 
-async function assertPublicHost(url: URL): Promise<void> {
+/** A lookup that always yields `addresses`, so the socket cannot resolve the host again. */
+export function pinnedLookup(addresses: LookupAddress[]): LookupFunction {
+  return (_hostname, options, callback) => {
+    const cb = (typeof options === 'function' ? options : callback) as (err: Error | null, address: unknown, family?: number) => void
+    if (typeof options !== 'function' && options.all) return cb(null, addresses)
+    const [first] = addresses
+    cb(null, first.address, first.family)
+  }
+}
+
+/**
+ * Resolve a host once, reject any private address, and return a lookup pinned to exactly
+ * those addresses. Passing that lookup to the transport is what closes the SSRF
+ * time-of-check/time-of-use gap: the connection reuses the validated resolution rather
+ * than issuing a second, unchecked DNS query that a rebinding record could answer privately.
+ */
+async function resolveAndPin(url: URL): Promise<LookupFunction> {
   const host = url.hostname.replace(/^\[|\]$/g, '')
   const addresses = await lookup(host, { all: true, verbatim: true })
-  // An empty list would leave nothing for the loop to reject, so the guard would pass
-  // vacuously. Schemes without a host (data:, file:) reach here the same way.
+  // An empty list would leave nothing to reject, so the guard would pass vacuously.
+  // Schemes without a host (data:, file:) reach here the same way.
   if (addresses.length === 0) throw new Error(`${url.hostname || url.protocol} did not resolve to any address`)
   for (const { address } of addresses) {
     if (isPrivateAddress(address)) throw new Error(`refusing to fetch private/internal address for ${url.hostname} (${address})`)
   }
+  return pinnedLookup(addresses)
 }
 
 const MAX_REDIRECTS = 5
@@ -156,14 +177,15 @@ async function readCapped(response: Response): Promise<string> {
   return text.slice(0, MAX_RAW_CHARS)
 }
 
-async function fetchText(rawUrl: string): Promise<{ text: string; contentType: string }> {
+async function fetchText(rawUrl: string, transport = httpFetch): Promise<{ text: string; contentType: string }> {
   let url = new URL(rawUrl)
   for (let hop = 0; hop <= MAX_REDIRECTS; hop++) {
-    await assertPublicHost(url)
-    const response = await fetch(url, {
-      headers: { 'User-Agent': USER_AGENT },
+    // Resolve, validate and pin per hop: a redirect target gets the same guarantee.
+    const lookup = await resolveAndPin(url)
+    const response = await transport(url, {
       signal: AbortSignal.timeout(FETCH_TIMEOUT_MS),
-      redirect: 'manual',
+      lookup,
+      userAgent: USER_AGENT,
     })
     if (response.status >= 300 && response.status < 400) {
       const location = response.headers.get('location')

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pi-code",
-  "version": "0.2.1",
+  "version": "0.2.3",
   "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"