pi-code 1.0.6 → 1.0.8

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -46,6 +46,7 @@ import { Type } from 'typebox'
46
46
  import { matchesBashRules } from './internal/bash-rules.js'
47
47
  import { type CommandExec, type DiscoveredCommand, discoverCommandFiles, expandDynamicContent, type ParsedCommand, parseCommandFile, resolvePowershellBinary, spanExec, substituteArgsDetailed, substituteVars } from './internal/command-file.js'
48
48
  import { readManagedSettings } from './internal/managed-settings.js'
49
+ import { capForContext } from './internal/output-guard.js'
49
50
  import { matchesPathRules } from './internal/path-rules.js'
50
51
  import { type InstalledPlugin, installedPlugins } from './internal/plugins.js'
51
52
  import { isProjectApproved } from './internal/project-approval.js'
@@ -315,12 +316,17 @@ export default function commandsExtension(pi: ExtensionAPI) {
315
316
  pendingBashRules = undefined
316
317
  pendingPathRules = undefined
317
318
  if (pendingModelRestore) {
318
- void pi.setModel(pendingModelRestore as Parameters<typeof pi.setModel>[0])
319
+ const restore = pendingModelRestore as Parameters<typeof pi.setModel>[0]
319
320
  pendingModelRestore = undefined
321
+ // setModel can reject (e.g. auth resolution fails), and a floated rejection would
322
+ // escape as unhandled; surface it instead of leaving the session silently on the
323
+ // command's override model.
324
+ void pi.setModel(restore).catch(() => {})
325
+ }
326
+ if (pendingRestore) {
327
+ pi.setActiveTools(pendingRestore)
328
+ pendingRestore = undefined
320
329
  }
321
- if (!pendingRestore) return
322
- pi.setActiveTools(pendingRestore)
323
- pendingRestore = undefined
324
330
  })
325
331
 
326
332
  // The active-tool set has no argument dimension, so a scoped grant hands the turn
@@ -416,6 +422,16 @@ export default function commandsExtension(pi: ExtensionAPI) {
416
422
  // restored when that run ends. Restoring inline does not work: sendUserMessage is
417
423
  // fire-and-forget, so the restore would land before the agent ever read the tool
418
424
  // list, leaving the command running with everything enabled.
425
+ // A command invoked while the agent is streaming must not narrow the in-flight run's
426
+ // tools or switch its model (that would corrupt a run it does not own), and a bare
427
+ // sendUserMessage throws mid-stream and would be silently dropped. Queue it as a
428
+ // follow-up through pi's own queue, which is abort-aware and shown to the user; its
429
+ // frontmatter scoping is not applied in that case, since it cannot land on a run that
430
+ // has not started yet.
431
+ if (!ctx.isIdle()) {
432
+ pi.sendUserMessage(expanded, { deliverAs: 'followUp' })
433
+ return
434
+ }
419
435
  applyAllowedTools(parsed, vars)
420
436
  applyDisallowedTools(parsed)
421
437
  await applyModelOverride(parsed, varCtx)
@@ -500,7 +516,10 @@ export default function commandsExtension(pi: ExtensionAPI) {
500
516
  // the tool result is the channel, and frontmatter scoping stays user-path
501
517
  // territory (see the header).
502
518
  const expanded = await expandCommand(pi, current, args, execCtx, command.filePath, command.plugin, { allowShell: false })
503
- return { content: [{ type: 'text' as const, text: `Contents of /${name} (expanded):\n\n${expanded}` }], details: {} }
519
+ // Cap the tool result: a command body can inline an arbitrarily large @file, and
520
+ // an uncapped tool result overflows the model's context (every other pi-code tool
521
+ // routes its output through capForContext). The user-invoked path stays uncapped.
522
+ return { content: [{ type: 'text' as const, text: capForContext(`Contents of /${name} (expanded):\n\n${expanded}`) }], details: {} }
504
523
  },
505
524
  })
506
525
  })
@@ -13,12 +13,15 @@
13
13
  * the checkpoint (files created after the checkpoint are left in place).
14
14
  */
15
15
 
16
+ import { createHash } from 'node:crypto'
16
17
  import * as fs from 'node:fs'
17
18
  import * as os from 'node:os'
18
19
  import * as path from 'node:path'
19
20
  import type { ExtensionAPI, ExtensionCommandContext, ExtensionContext } from '@earendil-works/pi-coding-agent'
20
21
 
21
22
  const CUSTOM_TYPE = 'git-checkpoint'
23
+ /** Sidecar inside the bare shadow repo recording the work tree it snapshots. */
24
+ const WORK_TREE_FILE = 'pi-work-tree'
22
25
  const PROMPT_SNIPPET_LENGTH = 60
23
26
  const RESTORE_MODES = ['Code and conversation', 'Conversation only', 'Code only']
24
27
 
@@ -72,6 +75,31 @@ export function sessionSlug(sessionFile: string | undefined): string {
72
75
  return path.basename(sessionFile).replace(/[^\w.-]+/g, '_')
73
76
  }
74
77
 
78
+ /** A stable per-directory key, so a session resumed elsewhere gets its own shadow. */
79
+ function cwdSlug(cwd: string): string {
80
+ const resolved = path.resolve(cwd)
81
+ const hash = createHash('sha256').update(resolved).digest('hex').slice(0, 8)
82
+ return `${path.basename(resolved).replace(/[^\w.-]+/g, '_')}-${hash}`
83
+ }
84
+
85
+ /** The work tree a shadow repo was created against, or undefined for a repo that
86
+ * predates the sidecar or does not exist yet. */
87
+ function recordedWorkTree(shadowDir: string): string | undefined {
88
+ try {
89
+ return fs.readFileSync(path.join(shadowDir, WORK_TREE_FILE), 'utf8').trim() || undefined
90
+ } catch {
91
+ return undefined
92
+ }
93
+ }
94
+
95
+ function rememberWorkTree(shadowDir: string, cwd: string): void {
96
+ try {
97
+ fs.writeFileSync(path.join(shadowDir, WORK_TREE_FILE), `${cwd}\n`)
98
+ } catch {
99
+ // best effort: without the marker the next resume simply cannot detect a move
100
+ }
101
+ }
102
+
75
103
  function extractText(content: unknown): string {
76
104
  if (typeof content === 'string') return content
77
105
  if (!Array.isArray(content)) return ''
@@ -134,6 +162,16 @@ export default function gitCheckpointExtension(pi: ExtensionAPI) {
134
162
  const sessionFile = (ctx.sessionManager as { getSessionFile?: () => string | undefined }).getSessionFile?.()
135
163
  const checkpointsRoot = path.join(os.homedir(), '.pi', 'agent', 'checkpoints')
136
164
  shadowDir = path.join(checkpointsRoot, sessionSlug(sessionFile))
165
+ // A resumed session can arrive from a different directory than the one the shadow
166
+ // snapshotted; restoring those commits here would silently overwrite unrelated
167
+ // same-named files. Key a fresh shadow to this directory instead of ever checking
168
+ // one tree out into another. Resuming back in the recorded directory takes the
169
+ // original shadow again, so its checkpoints stay restorable there.
170
+ const recorded = recordedWorkTree(shadowDir)
171
+ if (recorded && path.resolve(recorded) !== path.resolve(ctx.cwd)) {
172
+ shadowDir = path.join(checkpointsRoot, `${sessionSlug(sessionFile)}-${cwdSlug(ctx.cwd)}`)
173
+ ctx.ui.notify(`Checkpoints for this session were recorded in ${recorded}; starting fresh checkpoints for ${ctx.cwd} (earlier ones are not restorable here)`, 'warning')
174
+ }
137
175
  pruneCheckpointRepos(checkpointsRoot, CHECKPOINT_RETENTION_DAYS, shadowDir)
138
176
  const check = await pi.exec('git', ['--git-dir', shadowDir, 'rev-parse', '--git-dir'], { cwd: ctx.cwd })
139
177
  if (check.code !== 0) {
@@ -147,6 +185,8 @@ export default function gitCheckpointExtension(pi: ExtensionAPI) {
147
185
  await pi.exec('git', ['--git-dir', shadowDir, 'config', 'user.email', 'checkpoint@pi-code'], { cwd: ctx.cwd })
148
186
  await pi.exec('git', ['--git-dir', shadowDir, 'config', 'user.name', 'pi-code-checkpoint'], { cwd: ctx.cwd })
149
187
  }
188
+ // Written on every start, so repos that predate the sidecar pick it up too.
189
+ rememberWorkTree(shadowDir, ctx.cwd)
150
190
  }
151
191
 
152
192
  /** `checkout -f <ref> -- .` errors when the ref's tree holds no files, so an empty
@@ -156,6 +156,42 @@ export function readDisableAllHooks(files: string[], managed: Record<string, unk
156
156
  return false
157
157
  }
158
158
 
159
+ /** Claude's `allowedHttpHookUrls` setting: URL patterns http hooks may target, with
160
+ * `*` as a wildcard. Per Claude's documentation: undefined (no source sets the key)
161
+ * means no restrictions, an empty array blocks every http hook, and arrays merge
162
+ * across settings sources. Merging is a union of managed settings plus every file in
163
+ * the chain; the chain already gates project files on trust (see hookFiles), and a
164
+ * trusted project can run arbitrary shell hooks anyway, so letting it extend the
165
+ * allowlist is no escalation. */
166
+ export function readAllowedHttpHookUrls(files: string[], managed: Record<string, unknown> = readManagedSettings()): string[] | undefined {
167
+ let found: string[] | undefined
168
+ const collect = (value: unknown): void => {
169
+ if (!Array.isArray(value)) return
170
+ found = [...(found ?? []), ...value.filter((entry): entry is string => typeof entry === 'string')]
171
+ }
172
+ collect(managed.allowedHttpHookUrls)
173
+ for (const file of files) {
174
+ try {
175
+ const parsed: unknown = JSON.parse(fs.readFileSync(file, 'utf-8'))
176
+ if (isRecord(parsed)) collect(parsed.allowedHttpHookUrls)
177
+ } catch {
178
+ // missing or invalid file: skip
179
+ }
180
+ }
181
+ return found
182
+ }
183
+
184
+ /** Whether an http hook may target `url`. `*` in an allowlist entry matches any run
185
+ * of characters; everything else is literal and the whole URL must match. An
186
+ * undefined allowlist means the setting is absent, so there are no restrictions. */
187
+ export function httpUrlAllowed(url: string, allowlist: string[] | undefined): boolean {
188
+ if (allowlist === undefined) return true
189
+ return allowlist.some((pattern) => {
190
+ const literal = pattern.split('*').map((part) => part.replace(/[.*+?^${}()|[\]\\]/g, String.raw`\$&`))
191
+ return new RegExp(`^${literal.join('.*')}$`).test(url)
192
+ })
193
+ }
194
+
159
195
  export function loadHooks(files: string[], sources?: Map<HookMatcher, string>): HooksConfig {
160
196
  const config: HooksConfig = {}
161
197
  for (const file of files) {
@@ -459,12 +495,15 @@ function interpolateHeaders(headers: Record<string, string> | undefined, allowed
459
495
  * with a valid JSON body renders a decision, read exactly like command stdout.
460
496
  * Everything else, including non-2xx statuses, connection failures and timeouts,
461
497
  * is a non-blocking error by contract, so none of these outcomes ever reports
462
- * `timedOut`, which PreToolUse fails closed on. The user wrote the URL into their
463
- * own settings, so it carries the same trust as a command hook's shell string and
464
- * gets no SSRF screening.
498
+ * `timedOut`, which PreToolUse fails closed on. Claude's `allowedHttpHookUrls`
499
+ * allowlist gates the fetch itself: a URL matching no entry is never contacted,
500
+ * so a settings file cannot point a hook at an arbitrary endpoint and exfiltrate
501
+ * the payload; when the setting is absent there are no restrictions, as Claude
502
+ * documents. A blocked hook renders no decision, like every other http failure.
465
503
  */
466
- export async function runHttpHook(hook: { type?: string; command: string; url?: string; headers?: Record<string, string>; allowedEnvVars?: string[] }, payload: unknown, timeoutMs: number): Promise<HookRunResult> {
504
+ export async function runHttpHook(hook: { type?: string; command: string; url?: string; headers?: Record<string, string>; allowedEnvVars?: string[] }, payload: unknown, timeoutMs: number, allowedUrls?: string[]): Promise<HookRunResult> {
467
505
  const url = hook.url ?? hook.command
506
+ if (!httpUrlAllowed(url, allowedUrls)) return { code: 1, stdout: '', stderr: `${url} does not match allowedHttpHookUrls; the hook was not called`, timedOut: false }
468
507
  try {
469
508
  const response = await fetch(url, {
470
509
  method: 'POST',
@@ -526,7 +565,7 @@ export async function runPromptHook(hook: HookCommand, payload: unknown, model:
526
565
  const prompt = substituteArguments(hook.prompt, payload)
527
566
  const signal = AbortSignal.timeout(timeoutMs)
528
567
  try {
529
- const answer = await completeText(model, prompt, { system: PROMPT_HOOK_SYSTEM, maxTokens: 512, signal })
568
+ const { text: answer } = await completeText(model, prompt, { system: PROMPT_HOOK_SYSTEM, maxTokens: 512, signal })
530
569
  return { code: 0, stdout: answer, stderr: '', timedOut: false }
531
570
  } catch (error) {
532
571
  return abortAwareFailure(signal, error)
@@ -615,8 +654,9 @@ function replaceRecord(target: Record<string, unknown>, next: Record<string, unk
615
654
  Object.assign(target, next)
616
655
  }
617
656
 
618
- /** Claude surfaces a hook error notice and the action proceeds; silence would read a
619
- * guard that never ran as a clean allow. */
657
+ /** Claude surfaces a hook error notice; on ungated events the action proceeds, while
658
+ * PreToolUse and UserPromptSubmit additionally fail closed on the same results (see
659
+ * their spawnFailed checks). Silence would hide that a guard never ran. */
620
660
  function surfaceHookFailures(commands: HookCommand[], results: HookRunResult[], notify?: SystemMessageSink): void {
621
661
  if (!notify) return
622
662
  for (const [i, result] of results.entries()) {
@@ -648,6 +688,10 @@ export async function runPreToolUse(config: HooksConfig, toolName: string, toolI
648
688
  // A killed hook never reached its verdict, and SIGKILL leaves a null exit code that
649
689
  // would otherwise read as a clean allow. Fail closed instead.
650
690
  if (result.timedOut) return { block: true, reason: `Hook timed out after ${timeoutMs(commands[i])}ms: ${commands[i].command}` }
691
+ // A hook that never spawned (EMFILE, missing /bin/sh) reached no verdict either;
692
+ // its code 0 must fail closed like a timeout, not read as an allow exactly when
693
+ // the machine is degraded.
694
+ if (result.spawnFailed) return { block: true, reason: `Hook failed to run: ${commands[i].command}: ${result.stderr.trim() || 'unknown error'}` }
651
695
  }
652
696
  if (onSystemMessage) surfaceSystemMessages(results, onSystemMessage)
653
697
  // A hard deny wins over an ask, matching Claude's deny > ask > allow precedence:
@@ -698,6 +742,8 @@ export async function runUserPromptSubmit(config: HooksConfig, prompt: string, r
698
742
  surfaceHookFailures(commands, results, onSystemMessage)
699
743
  for (const [i, result] of results.entries()) {
700
744
  if (result.timedOut) return { block: true, reason: `Hook timed out after ${timeoutMs(commands[i])}ms: ${commands[i].command}`, context: '' }
745
+ // No verdict was delivered, so fail closed like a timeout (see runPreToolUse).
746
+ if (result.spawnFailed) return { block: true, reason: `Hook failed to run: ${commands[i].command}: ${result.stderr.trim() || 'unknown error'}`, context: '' }
701
747
  }
702
748
  if (onSystemMessage) surfaceSystemMessages(results, onSystemMessage)
703
749
  const contexts: string[] = []
@@ -742,6 +788,8 @@ function postToolFeedback(result: HookRunResult, eventName: string, isError: boo
742
788
  export default function hooksExtension(pi: ExtensionAPI) {
743
789
  let config: HooksConfig = {}
744
790
  let projectDir = ''
791
+ /** Claude's allowedHttpHookUrls allowlist, resolved from the settings chain. */
792
+ let allowedHttpHookUrls: string[] | undefined
745
793
  let pendingSessionContext: string[] = []
746
794
  let stopHookActive = false
747
795
  let sessionCtx: ExtensionContext | undefined
@@ -763,7 +811,7 @@ export default function hooksExtension(pi: ExtensionAPI) {
763
811
  (ctx: ExtensionContext, extra?: Record<string, unknown>): HookRunner =>
764
812
  (hook, payload, ms) => {
765
813
  const merged = { ...commonPayload(ctx), ...extra, ...(payload as Record<string, unknown>) }
766
- if (hook.type === 'http') return runHttpHook(hook, merged, ms)
814
+ if (hook.type === 'http') return runHttpHook(hook, merged, ms, allowedHttpHookUrls)
767
815
  if (hook.type === 'prompt') return runPromptHook(hook, merged, ctx.model, ms)
768
816
  if (hook.type === 'agent') return runAgentHook(hook, merged, ms, (ctx.model as { id?: string } | undefined)?.id)
769
817
  if (hook.type === 'mcp_tool') return runMcpToolHook(hook, merged, ms)
@@ -814,8 +862,14 @@ export default function hooksExtension(pi: ExtensionAPI) {
814
862
  const ctx = sessionCtx
815
863
  const eventName = data.phase === 'start' ? 'SubagentStart' : 'SubagentStop'
816
864
  const payload = { hook_event_name: eventName, agent_type: data.agentType, agent_id: data.agentId }
817
- const results = await runNotifyHooks(matchingCommands(config[eventName], data.agentType), payload, boundRunner(ctx))
818
- surfaceSystemMessages(results, (message) => ctx.ui.notify(message, 'warning'))
865
+ try {
866
+ const results = await runNotifyHooks(matchingCommands(config[eventName], data.agentType), payload, boundRunner(ctx))
867
+ surfaceSystemMessages(results, (message) => ctx.ui.notify(message, 'warning'))
868
+ } catch {
869
+ // The bus outlives the session: an event landing between /new disposing this
870
+ // ctx and the next session_start hits disposed getters, and nothing awaits a
871
+ // bus listener, so a throw here would escape as an unhandled rejection.
872
+ }
819
873
  })
820
874
 
821
875
  pi.on('session_start', async (event, ctx) => {
@@ -827,6 +881,7 @@ export default function hooksExtension(pi: ExtensionAPI) {
827
881
  projectDir = repoRoot(ctx.cwd) ?? ctx.cwd
828
882
  const files = hookFiles(ctx.cwd, os.homedir(), trusted)
829
883
  hookSources.clear()
884
+ allowedHttpHookUrls = readAllowedHttpHookUrls(files)
830
885
  // The disableAllHooks escape hatch, checked before any config loads: with no
831
886
  // config resolved, no event, plugin hooks included, can fire a hook.
832
887
  hooksDisabled = readDisableAllHooks(files)
@@ -916,6 +971,12 @@ export default function hooksExtension(pi: ExtensionAPI) {
916
971
  // turn, and stop_hook_active in the payload tells the next firing it is already
917
972
  // continuing from a stop hook, which is the hook script's documented loop guard.
918
973
  // Only exit 2 and decision:"block" continue; continue:false means "stay stopped".
974
+ //
975
+ // On agent_end rather than agent_settled: agent_settled is only emitted after every
976
+ // agent_end handler returns, and a peer extension (plan mode) blocks its agent_end
977
+ // handler on a UI dialog, which would starve the Stop hook and idle notification
978
+ // until the user answers it. agent_end can fire slightly early before a rare
979
+ // automatic retry or compaction; that is the better tradeoff.
919
980
  pi.on('agent_end', async (event, ctx) => {
920
981
  // Claude's Notification event, for the one type pi can honestly source: the
921
982
  // agent finished and is waiting for input (idle_prompt). Observational only;
@@ -75,7 +75,9 @@ export default function initExtension(pi: ExtensionAPI) {
75
75
  const existing = findExistingContextFile(root)
76
76
  const cursorRules = statOf(path.join(root, '.cursor', 'rules'))?.isDirectory() === true || statOf(path.join(root, '.cursorrules'))?.isFile() === true
77
77
  const copilotRules = statOf(path.join(root, '.github', 'copilot-instructions.md'))?.isFile() === true
78
- pi.sendUserMessage(buildInitPrompt({ ...(existing !== undefined ? { existingContextFile: existing } : {}), cursorRules, copilotRules }))
78
+ // A bare send throws (and is silently swallowed) while the agent is
79
+ // streaming, so mid-stream invocations queue as a follow-up turn.
80
+ pi.sendUserMessage(buildInitPrompt({ ...(existing !== undefined ? { existingContextFile: existing } : {}), cursorRules, copilotRules }), ctx.isIdle() ? {} : { deliverAs: 'followUp' })
79
81
  },
80
82
  })
81
83
  }
@@ -49,6 +49,11 @@ export class FileOAuthProvider implements OAuthClientProvider {
49
49
  private readonly data: StoredAuth
50
50
  private port = 0
51
51
  private readonly onRedirect: (authorizationUrl: URL) => void
52
+ // A fresh random CSRF token per login attempt. The SDK puts it in the authorization
53
+ // URL's `state` param, the server echoes it back on the redirect, and waitForAuthCode
54
+ // rejects any callback that does not carry it, so another local process or an open web
55
+ // page cannot inject an authorization code into this login (RFC 8252 8.9).
56
+ private readonly loginState = crypto.randomBytes(16).toString('hex')
52
57
 
53
58
  constructor(serverName: string, onRedirect: (authorizationUrl: URL) => void) {
54
59
  this.storePath = storeFileFor(serverName)
@@ -117,6 +122,12 @@ export class FileOAuthProvider implements OAuthClientProvider {
117
122
  return this.data.tokens !== undefined
118
123
  }
119
124
 
125
+ /** The CSRF token the SDK adds to the authorization URL as `state`; waitForAuthCode
126
+ * verifies the redirect echoes exactly this value. */
127
+ state(): string {
128
+ return this.loginState
129
+ }
130
+
120
131
  redirectToAuthorization(authorizationUrl: URL): void {
121
132
  this.onRedirect(authorizationUrl)
122
133
  }
@@ -147,11 +158,30 @@ export async function startCallbackServer(preferredPort?: number): Promise<{ ser
147
158
  return { server, port: (server.address() as { port: number }).port }
148
159
  }
149
160
 
150
- export function waitForAuthCode(server: http.Server, timeoutMs: number): Promise<string> {
161
+ export function waitForAuthCode(server: http.Server, timeoutMs: number, expectedState?: string): Promise<string> {
151
162
  return new Promise((resolve, reject) => {
152
163
  const timer = setTimeout(() => reject(new Error(`authorization timed out after ${timeoutMs}ms`)), timeoutMs)
164
+ // Do not let the pending timer keep the process alive on its own: if the login is
165
+ // abandoned or resolved out of band, the event loop can still drain.
166
+ timer.unref?.()
153
167
  server.on('request', (request, response) => {
154
168
  const url = new URL(request.url ?? '/', 'http://127.0.0.1')
169
+ // Only the redirect path settles the login. A stray request (a favicon fetch, a
170
+ // local port scan, or a forged redirect from another process or an open web page)
171
+ // is answered but ignored, so it can neither inject a code nor abort the login by
172
+ // rejecting the promise (a repeatable DoS on a stable, guessable loopback port).
173
+ if (url.pathname !== '/callback') {
174
+ response.writeHead(404, { 'content-type': 'text/plain' })
175
+ response.end('not found')
176
+ return
177
+ }
178
+ // The CSRF check: a callback that does not echo this login's state is rejected
179
+ // without settling, so an attacker who cannot read the state cannot complete it.
180
+ if (expectedState !== undefined && url.searchParams.get('state') !== expectedState) {
181
+ response.writeHead(400, { 'content-type': 'text/plain' })
182
+ response.end('state mismatch')
183
+ return
184
+ }
155
185
  const code = url.searchParams.get('code')
156
186
  const error = url.searchParams.get('error')
157
187
  response.writeHead(200, { 'content-type': 'text/html' })
@@ -14,7 +14,7 @@
14
14
  * throw and the caller falls back to its non-model behavior.
15
15
  */
16
16
 
17
- import type { Api, AssistantMessage, Context, Model, ModelsSimpleStreamOptions } from '@earendil-works/pi-ai'
17
+ import type { Api, AssistantMessage, Context, Model, ModelsSimpleStreamOptions, Usage } from '@earendil-works/pi-ai'
18
18
  import { ModelRuntime } from '@earendil-works/pi-coding-agent'
19
19
 
20
20
  /** The completion backend: model + context -> assistant message. Overridable for tests. */
@@ -52,11 +52,13 @@ export interface CompleteOptions {
52
52
  }
53
53
 
54
54
  /**
55
- * Run `prompt` through `model` as a single user turn and return the reply text.
55
+ * Run `prompt` through `model` as a single user turn and return the reply text plus
56
+ * the call's usage. A tool that makes a nested LLM call must return that usage on
57
+ * its tool result, or the call's tokens and cost vanish from pi's session totals.
56
58
  * Throws on any failure so the caller can fall back; never returns a partial or a
57
59
  * tool call, only assistant text.
58
60
  */
59
- export async function completeText(model: Model<Api>, prompt: string, options: CompleteOptions = {}): Promise<string> {
61
+ export async function completeText(model: Model<Api>, prompt: string, options: CompleteOptions = {}): Promise<{ text: string; usage: Usage }> {
60
62
  backend ??= realBackend()
61
63
  const complete = await backend
62
64
  const context: Context = {
@@ -64,5 +66,5 @@ export async function completeText(model: Model<Api>, prompt: string, options: C
64
66
  messages: [{ role: 'user', content: prompt, timestamp: Date.now() }],
65
67
  }
66
68
  const message = await complete(model, context, { maxTokens: options.maxTokens ?? 1024, signal: options.signal })
67
- return assistantText(message)
69
+ return { text: assistantText(message), usage: message.usage }
68
70
  }
@@ -21,8 +21,66 @@ export interface PathAnchors {
21
21
 
22
22
  const escapeRegExp = (text: string): string => text.replace(/[.*+?^${}()|[\]\\]/g, String.raw`\$&`)
23
23
 
24
- /** One gitignore-style pattern as an anchored regular expression source. */
25
- export function globToRegExpSource(pattern: string): string {
24
+ /** Cap on brace-expanded alternatives per pattern, mirroring Claude's ~1000
25
+ * budget; an over-budget pattern is used unexpanded. */
26
+ const BRACE_EXPANSION_LIMIT = 1000
27
+
28
+ interface BraceGroup {
29
+ start: number
30
+ end: number
31
+ options: string[]
32
+ }
33
+
34
+ /** The `{...}` group opening at `open`, or null when it is unmatched or carries no
35
+ * top-level comma (literal braces). Options are split on commas at the group's own
36
+ * depth so a nested group stays inside one option. */
37
+ function parseBraceGroup(pattern: string, open: number): BraceGroup | null {
38
+ let depth = 1
39
+ let optionStart = open + 1
40
+ const options: string[] = []
41
+ for (let i = open + 1; i < pattern.length; i += 1) {
42
+ const ch = pattern[i]
43
+ if (ch === '{') depth += 1
44
+ else if (ch === ',' && depth === 1) {
45
+ options.push(pattern.slice(optionStart, i))
46
+ optionStart = i + 1
47
+ } else if (ch === '}' && --depth === 0) {
48
+ if (options.length === 0) return null // no top-level comma: literal braces
49
+ options.push(pattern.slice(optionStart, i))
50
+ return { start: open, end: i, options }
51
+ }
52
+ }
53
+ return null
54
+ }
55
+
56
+ /** The first expandable `{...}` group. A comma-less or unmatched `{` is skipped as
57
+ * literal, so the scan can still find an expandable group nested inside it. */
58
+ function findBraceGroup(pattern: string): BraceGroup | null {
59
+ for (let open = pattern.indexOf('{'); open !== -1; open = pattern.indexOf('{', open + 1)) {
60
+ const group = parseBraceGroup(pattern, open)
61
+ if (group) return group
62
+ }
63
+ return null
64
+ }
65
+
66
+ /** Bash-style brace expansion of one pattern into its alternatives: each group
67
+ * multiplies out (Cartesian across groups, nested groups recurse). Returns null
68
+ * when the expansion would exceed the budget. */
69
+ function expandBraces(pattern: string): string[] | null {
70
+ const group = findBraceGroup(pattern)
71
+ if (group === null) return [pattern]
72
+ const expanded: string[] = []
73
+ for (const option of group.options) {
74
+ const branch = expandBraces(pattern.slice(0, group.start) + option + pattern.slice(group.end + 1))
75
+ if (branch === null) return null
76
+ expanded.push(...branch)
77
+ if (expanded.length > BRACE_EXPANSION_LIMIT) return null
78
+ }
79
+ return expanded
80
+ }
81
+
82
+ /** One glob pattern, braces already expanded, as a regular expression source. */
83
+ function translateGlob(pattern: string): string {
26
84
  let out = ''
27
85
  let i = 0
28
86
  while (i < pattern.length) {
@@ -54,6 +112,15 @@ export function globToRegExpSource(pattern: string): string {
54
112
  return out
55
113
  }
56
114
 
115
+ /** One gitignore-style pattern as an anchored regular expression source. Brace
116
+ * groups (`{ts,tsx}`, nested, Cartesian across groups) expand into ORed
117
+ * alternatives; an over-budget expansion falls back to the literal pattern. */
118
+ export function globToRegExpSource(pattern: string): string {
119
+ const alternatives = expandBraces(pattern) ?? [pattern]
120
+ if (alternatives.length === 1) return translateGlob(alternatives[0])
121
+ return `(?:${alternatives.map(translateGlob).join('|')})`
122
+ }
123
+
57
124
  /** A rule resolved to an absolute glob per its anchor form. */
58
125
  function resolveRule(rule: string, anchors: PathAnchors): string {
59
126
  if (rule.startsWith('//')) return rule.slice(1)
package/extensions/mcp.ts CHANGED
@@ -11,8 +11,10 @@
11
11
  * per-project `projects[cwd].mcpServers` local scope, and ~/.pi/agent/mcp.json) is the
12
12
  * user's own and loads on the first session. Project config (.mcp.json, .pi/mcp.json)
13
13
  * can run arbitrary commands on connect, so it loads only once the project is approved
14
- * (see project-approval). The two scopes are loaded separately, not merged; user config
15
- * connects first, so a project server cannot take the name of a user server that connected.
14
+ * (see project-approval). The two scopes are loaded separately, not merged. Claude's
15
+ * precedence is project over user for a duplicate name, so a project server the user has
16
+ * consented to (or an approved project's) wins; a merely-present untrusted project entry
17
+ * cannot shadow a user server, and a gated project server does not preempt it.
16
18
  * Values support ${VAR} / ${VAR:-default} interpolation, connect and per-call timeouts
17
19
  * honor MCP_TIMEOUT / MCP_TOOL_TIMEOUT, and a stdio server receives only the SDK's default
18
20
  * environment plus its own `env` block, not the whole process environment.
@@ -108,11 +110,18 @@ export type ServerConfig = StdioServerConfig | HttpServerConfig
108
110
 
109
111
  /** Claude's .mcp.json expansion: ${VAR}, and ${VAR:-default}. The syntax borrows
110
112
  * shell's `:-`, which substitutes when the variable is unset OR empty. */
111
- export function interpolateEnv(value: string, env: NodeJS.ProcessEnv = process.env): string {
112
- return value.replace(/\$\{(\w+)(:-([^}]*))?\}/g, (_, name, hasDefault, fallback) => {
113
+ export function interpolateEnv(value: string, env: NodeJS.ProcessEnv = process.env, onMissing?: (name: string) => void): string {
114
+ return value.replace(/\$\{(\w+)(:-([^}]*))?\}/g, (fullMatch, name, hasDefault, fallback) => {
113
115
  const current = env[name]
114
116
  if (hasDefault !== undefined) return current || fallback
115
- return current ?? ''
117
+ if (current === undefined) {
118
+ // A referenced variable with no value and no default: keep the literal ${VAR} and
119
+ // report it, matching Claude, rather than silently substituting an empty string that
120
+ // turns `Bearer ${TOKEN}` into a confusing `Bearer ` and a mystery 401.
121
+ onMissing?.(name)
122
+ return fullMatch
123
+ }
124
+ return current
116
125
  })
117
126
  }
118
127
 
@@ -387,11 +396,40 @@ export function promptMessageContent(messages: ReadonlyArray<{ content: unknown
387
396
  return mapContent(messages.map((message) => message.content as McpContentBlock)).filter((block) => block.type !== 'text' || block.text.trim() !== '')
388
397
  }
389
398
 
399
+ /** Merge the `properties` (and, for allOf, the `required`) of a root-level combinator's
400
+ * branches into one flat object schema. Without this a tool whose input schema is a bare
401
+ * anyOf/oneOf/allOf (no top-level `type`) would present no properties at all, so the model
402
+ * would be forced to call it with no arguments. */
403
+ function mergeCombinatorBranches(branches: unknown[]): { properties: Record<string, unknown>; required: string[] } {
404
+ const properties: Record<string, unknown> = {}
405
+ const required = new Set<string>()
406
+ for (const branch of branches) {
407
+ if (!branch || typeof branch !== 'object') continue
408
+ const b = branch as Record<string, unknown>
409
+ if (b.properties && typeof b.properties === 'object') Object.assign(properties, b.properties as Record<string, unknown>)
410
+ if (Array.isArray(b.required)) for (const name of b.required) if (typeof name === 'string') required.add(name)
411
+ }
412
+ return { properties, required: [...required] }
413
+ }
414
+
390
415
  export function normalizeSchema(schema: unknown): object {
391
416
  const base = (schema as Record<string, unknown>) ?? {}
392
417
  const { $schema: _dropSchema, additionalProperties: _dropAdditional, ...rest } = base
393
- if (!rest.type) return { type: 'object', properties: {} }
394
- return rest
418
+ if (rest.type) return rest
419
+ // A root-level combinator carries the real parameters in its branches; flatten them
420
+ // into one object schema rather than emptying it. allOf means every branch applies, so
421
+ // its required union is kept; anyOf/oneOf branches are alternatives, so required is left
422
+ // open (the server still enforces its own).
423
+ const allOf = Array.isArray(rest.allOf) ? rest.allOf : undefined
424
+ let branches = allOf
425
+ if (!branches && Array.isArray(rest.anyOf)) branches = rest.anyOf
426
+ if (!branches && Array.isArray(rest.oneOf)) branches = rest.oneOf
427
+ if (!branches) return { type: 'object', properties: {} }
428
+ const { properties, required } = mergeCombinatorBranches(branches)
429
+ const merged: Record<string, unknown> = { type: 'object', properties }
430
+ if (typeof rest.description === 'string') merged.description = rest.description
431
+ if (allOf && required.length > 0) merged.required = required
432
+ return merged
395
433
  }
396
434
 
397
435
  interface McpContentBlock {
@@ -494,23 +532,32 @@ async function withTimeout<T>(promise: Promise<T>, ms: number, label: string): P
494
532
 
495
533
  async function connect(name: string, config: ServerConfig, authUi?: AuthUi): Promise<Client> {
496
534
  const client = new Client({ name: 'pi-code-mcp', version: '0.1.0' })
535
+ // Names referenced by ${VAR} with no value and no default, gathered across this
536
+ // server's interpolated fields so the connect can warn once rather than fail with a
537
+ // mystery 401 or a command that lost an argument.
538
+ const missing = new Set<string>()
539
+ const fill = (value: string): string => interpolateEnv(value, process.env, (varName) => missing.add(varName))
540
+ const warnMissing = (): void => {
541
+ if (missing.size > 0) console.warn(`pi-code-mcp: server ${name} references undefined variable(s) ${[...missing].join(', ')}; leaving them unexpanded`)
542
+ }
497
543
  if (isStdio(config)) {
498
544
  // Start from the SDK's allowlist (PATH, HOME, SHELL, ...) rather than the whole
499
545
  // process env: a server should not receive ANTHROPIC_API_KEY or GITHUB_TOKEN just
500
546
  // for being launched. A server that needs a variable names it in its own env block.
501
547
  const env: Record<string, string> = { ...getDefaultEnvironment() }
502
- for (const [key, value] of Object.entries(config.env ?? {})) env[key] = interpolateEnv(value)
548
+ for (const [key, value] of Object.entries(config.env ?? {})) env[key] = fill(value)
503
549
  const transport = new StdioClientTransport({
504
- command: interpolateEnv(config.command),
505
- args: (config.args ?? []).map((arg) => interpolateEnv(arg)),
550
+ command: fill(config.command),
551
+ args: (config.args ?? []).map((arg) => fill(arg)),
506
552
  env,
507
553
  cwd: expandCwd(config.cwd),
508
554
  stderr: 'ignore',
509
555
  })
556
+ warnMissing()
510
557
  await connectWithTimeout(client, transport, `connect ${name}`)
511
558
  return client
512
559
  }
513
- const url = new URL(interpolateEnv(config.url))
560
+ const url = new URL(fill(config.url))
514
561
  if (config.type === 'ws' || config.type === 'websocket') {
515
562
  // The SDK's WebSocket transport takes only a url: it carries no headers, bearer
516
563
  // token, or headersHelper output. Warn rather than silently dropping configured
@@ -520,16 +567,18 @@ async function connect(name: string, config: ServerConfig, authUi?: AuthUi): Pro
520
567
  console.warn(`pi-code-mcp: server ${name} is a WebSocket server; the SDK ws transport is url-only, so its headers/bearerToken/headersHelper are ignored`)
521
568
  }
522
569
  const transport = new WebSocketClientTransport(url)
570
+ warnMissing()
523
571
  await connectWithTimeout(client, transport, `connect ${name} (ws)`)
524
572
  return client
525
573
  }
526
574
  const headers: Record<string, string> = {}
527
- for (const [key, value] of Object.entries(config.headers ?? {})) headers[key] = interpolateEnv(value)
575
+ for (const [key, value] of Object.entries(config.headers ?? {})) headers[key] = fill(value)
528
576
  const token = resolveBearerToken(config)
529
577
  if (token) headers.Authorization = `Bearer ${token}`
530
578
  // A headersHelper generates connect-time headers for non-OAuth auth schemes; its
531
579
  // JSON stdout merges over the static headers.
532
- if (config.headersHelper) Object.assign(headers, await runHeadersHelper(interpolateEnv(config.headersHelper)))
580
+ if (config.headersHelper) Object.assign(headers, await runHeadersHelper(fill(config.headersHelper)))
581
+ warnMissing()
533
582
  const sseTransport = (authProvider?: OAuthClientProvider) => new SSEClientTransport(url, { requestInit: { headers }, authProvider }) // NOSONAR: explicitly declared or deliberate legacy transport
534
583
  if (config.type === 'sse') {
535
584
  return await connectHttpFamily(name, config, sseTransport, `connect ${name} (sse)`, token, authUi)
@@ -645,7 +694,9 @@ async function runInteractiveOAuth(name: string, config: { url: string }, makeTr
645
694
  provider.bindRedirectPort(port)
646
695
  try {
647
696
  const transport = makeTransport(provider)
648
- const pendingCode = waitForAuthCode(server, OAUTH_FLOW_TIMEOUT_MS)
697
+ // Verify the redirect echoes this login's state, so a stray or forged callback to the
698
+ // loopback port cannot inject a code or abort the login (see waitForAuthCode).
699
+ const pendingCode = waitForAuthCode(server, OAUTH_FLOW_TIMEOUT_MS, provider.state())
649
700
  pendingCode.catch(() => {}) // consumed below; an abandoned login must not surface as unhandled
650
701
  const client = newClient()
651
702
  try {
@@ -790,7 +841,7 @@ export default async function mcpExtension(pi: ExtensionAPI) {
790
841
  const aliases: McpToolAlias[] = []
791
842
 
792
843
  /** Register every not-yet-registered tool of a server; returns how many were added. */
793
- function registerTools(name: string, config: ServerConfig, client: Client, tools: McpToolInfo[]): number {
844
+ function registerTools(name: string, config: ServerConfig, tools: McpToolInfo[]): number {
794
845
  let count = 0
795
846
  for (const tool of tools) {
796
847
  const toolName = formatToolName(name, tool.name)
@@ -809,12 +860,18 @@ export default async function mcpExtension(pi: ExtensionAPI) {
809
860
  description: tool.description ?? `MCP tool ${tool.name} from ${name}`,
810
861
  parameters: Type.Unsafe(normalizeSchema(tool.inputSchema)),
811
862
  async execute(_id, params) {
863
+ // Resolve the live client by name at call time rather than capturing the one
864
+ // present at registration: pi has no tool unregister, so after a server drops
865
+ // and a later session_start reconnects it, registerTools skips re-registration
866
+ // and this closure would otherwise keep calling the old, closed client.
867
+ const current = clients.get(name)
868
+ if (!current) throw new Error(`MCP server "${name}" is not connected`)
812
869
  // Pass the timeout to the SDK too: its own default request timeout is 60s and
813
870
  // would otherwise reject first, so the outer race at CALL_TIMEOUT_MS was dead.
814
871
  // Claude's per-server timeout wins over MCP_TOOL_TIMEOUT, with a 1s floor.
815
872
  const declared = typeof config.timeout === 'number' && config.timeout >= 1000 ? config.timeout : undefined
816
873
  const budget = declared ?? callTimeoutMs()
817
- const result = await withTimeout(client.callTool({ name: tool.name, arguments: params as Record<string, unknown> }, undefined, { timeout: budget }), budget, toolName)
874
+ const result = await withTimeout(current.callTool({ name: tool.name, arguments: params as Record<string, unknown> }, undefined, { timeout: budget }), budget, toolName)
818
875
  const content = mapContent(result.content as McpContentBlock[], result.structuredContent)
819
876
  const details: { error?: string } = {}
820
877
  if (result.isError) {
@@ -839,7 +896,7 @@ export default async function mcpExtension(pi: ExtensionAPI) {
839
896
  * no command unregister, so, like tools, a withdrawn prompt keeps its registration
840
897
  * and surfaces the server's own error when invoked; an edit to a prompt's declared
841
898
  * arguments only lands on new names, since an existing command keeps its binding. */
842
- function registerPrompts(name: string, client: Client, prompts: McpPromptInfo[]): void {
899
+ function registerPrompts(name: string, prompts: McpPromptInfo[]): void {
843
900
  for (const prompt of prompts) {
844
901
  const commandName = formatPromptCommandName(name, prompt.name)
845
902
  const owner = registeredPrompts.get(commandName)
@@ -855,11 +912,19 @@ export default async function mcpExtension(pi: ExtensionAPI) {
855
912
  description: hint ? `${base} ${hint}` : base,
856
913
  handler: async (args, ctx) => {
857
914
  try {
915
+ // Resolve the live client at call time, not the one captured at registration:
916
+ // pi has no command unregister, so after a reconnect this closure must not keep
917
+ // calling the old, closed client (see registerTools for the same reason).
918
+ const current = clients.get(name)
919
+ if (!current) {
920
+ ctx.ui.notify(`${commandName}: MCP server "${name}" is not connected`, 'error')
921
+ return
922
+ }
858
923
  const promptArgs = mapPromptArguments(prompt.arguments, args)
859
924
  const params: { name: string; arguments?: Record<string, string> } = { name: prompt.name }
860
925
  if (Object.keys(promptArgs).length > 0) params.arguments = promptArgs
861
926
  const budget = callTimeoutMs()
862
- const result = await withTimeout(client.getPrompt(params, { timeout: budget }), budget, commandName)
927
+ const result = await withTimeout(current.getPrompt(params, { timeout: budget }), budget, commandName)
863
928
  // The prompt drives a turn exactly the way a custom slash command does
864
929
  // (see commands.ts), carrying its image blocks through. A prompt that
865
930
  // yields no content is reported rather than sent as an empty turn.
@@ -868,7 +933,9 @@ export default async function mcpExtension(pi: ExtensionAPI) {
868
933
  ctx.ui.notify(`${commandName}: prompt returned no content`, 'info')
869
934
  return
870
935
  }
871
- pi.sendUserMessage(content)
936
+ // A bare send throws (and is silently swallowed) while the agent is
937
+ // streaming, so mid-stream invocations queue as a follow-up turn.
938
+ pi.sendUserMessage(content, ctx.isIdle() ? {} : { deliverAs: 'followUp' })
872
939
  } catch (error) {
873
940
  ctx.ui.notify(`${commandName}: ${error instanceof Error ? error.message : String(error)}`, 'error')
874
941
  }
@@ -882,7 +949,7 @@ export default async function mcpExtension(pi: ExtensionAPI) {
882
949
  async function connectPrompts(name: string, client: Client): Promise<void> {
883
950
  if (!client.getServerCapabilities()?.prompts) return
884
951
  try {
885
- registerPrompts(name, client, await withTimeout(listAllPrompts(client), connectTimeoutMs(), `list prompts ${name}`))
952
+ registerPrompts(name, await withTimeout(listAllPrompts(client), connectTimeoutMs(), `list prompts ${name}`))
886
953
  } catch (error) {
887
954
  console.warn(`pi-code-mcp: prompt listing failed for ${name}: ${error instanceof Error ? error.message : String(error)}`)
888
955
  }
@@ -894,7 +961,7 @@ export default async function mcpExtension(pi: ExtensionAPI) {
894
961
  try {
895
962
  client.setNotificationHandler(PromptListChangedNotificationSchema, async () => {
896
963
  try {
897
- registerPrompts(name, client, await withTimeout(listAllPrompts(client), connectTimeoutMs(), `list prompts ${name}`))
964
+ registerPrompts(name, await withTimeout(listAllPrompts(client), connectTimeoutMs(), `list prompts ${name}`))
898
965
  } catch (error) {
899
966
  console.warn(`pi-code-mcp: prompt refresh failed for ${name}: ${error instanceof Error ? error.message : String(error)}`)
900
967
  }
@@ -978,7 +1045,7 @@ export default async function mcpExtension(pi: ExtensionAPI) {
978
1045
  client.setNotificationHandler(ToolListChangedNotificationSchema, async () => {
979
1046
  try {
980
1047
  const refreshed = await withTimeout(listAllTools(client), connectTimeoutMs(), `list tools ${name}`)
981
- const added = registerTools(name, config, client, refreshed)
1048
+ const added = registerTools(name, config, refreshed)
982
1049
  if (added === 0) return
983
1050
  const current = status.get(name)
984
1051
  status.set(name, { state: current?.state ?? 'connected', tools: (current?.tools ?? 0) + added })
@@ -1014,7 +1081,7 @@ export default async function mcpExtension(pi: ExtensionAPI) {
1014
1081
  const client = await connect(name, config, authUi)
1015
1082
  clients.set(name, client)
1016
1083
  const tools = await withTimeout(listAllTools(client), connectTimeoutMs(), `list tools ${name}`)
1017
- const count = registerTools(name, config, client, tools)
1084
+ const count = registerTools(name, config, tools)
1018
1085
  subscribeToToolChanges(name, config, client)
1019
1086
  // Prompts and resources are additive surfaces: their failures warn (inside
1020
1087
  // connectPrompts) rather than flipping a tool-serving server to failed.
@@ -1075,7 +1142,15 @@ export default async function mcpExtension(pi: ExtensionAPI) {
1075
1142
  const pluginServers = loadPluginServers(installedPlugins(os.homedir()))
1076
1143
  const { allowed, denied } = mcpAllowDeny()
1077
1144
  const scoped = applyServerPolicy({ ...pluginServers, ...loadUserScope(os.homedir(), ctx.cwd) }, allowed, denied)
1078
- const userServers = Object.fromEntries(Object.entries(scoped).filter(([name]) => !clients.has(name)))
1145
+ // Claude's precedence is project over user for a duplicate name. A project .mcp.json
1146
+ // server only outranks the user's own when it will actually connect (the user already
1147
+ // consented to it, or an approved project's), so a merely-present untrusted project
1148
+ // entry cannot shadow a trusted user server by reusing its name. A gated project
1149
+ // server still awaiting the approval prompt does not preempt the user server: that is
1150
+ // a deliberate narrowing of Claude's rule to keep the safe default.
1151
+ const projectPolicy = projectServerPolicy(ctx.cwd, os.homedir(), isProjectApprovedSilently(ctx))
1152
+ const projectWinners = new Set(Object.keys(splitByPolicy(applyServerPolicy(loadConfigFrom(projectConfigPaths(ctx.cwd)), allowed, denied), projectPolicy).consented))
1153
+ const userServers = Object.fromEntries(Object.entries(scoped).filter(([name]) => !clients.has(name) && !projectWinners.has(name)))
1079
1154
  if (Object.keys(userServers).length > 0) await connectServers(userServers, authUiFor(ctx))
1080
1155
  // A project .mcp.json can run arbitrary commands on connect, so only honor it once
1081
1156
  // the project is trusted. Per-server settings refine that: disabled servers never
@@ -12,7 +12,7 @@ import * as fs from 'node:fs'
12
12
  import * as os from 'node:os'
13
13
  import * as path from 'node:path'
14
14
  import { StringEnum } from '@earendil-works/pi-ai'
15
- import type { ExtensionAPI } from '@earendil-works/pi-coding-agent'
15
+ import { type ExtensionAPI, withFileMutationQueue } from '@earendil-works/pi-coding-agent'
16
16
  import { Type } from 'typebox'
17
17
  import { capForContext } from './internal/output-guard.js'
18
18
  import { isProjectApprovedSilently } from './internal/project-approval.js'
@@ -138,29 +138,39 @@ export function indexWouldOverflow(index: string, name: string, description: str
138
138
  return next.split('\n').length > INDEX_MAX_LINES || Buffer.byteLength(next, 'utf-8') > INDEX_MAX_BYTES
139
139
  }
140
140
 
141
- /** Write a memory and its index line, or say why it cannot be written. */
142
- export function saveMemory(dir: string, indexPath: string, name: string | undefined, description: string | undefined, content: string | undefined, now: string = new Date().toISOString()): { content: Array<{ type: 'text'; text: string }>; details: Record<string, never> } {
141
+ type MemoryToolResult = { content: Array<{ type: 'text'; text: string }>; details: Record<string, never> }
142
+
143
+ /** Write a memory and its index line, or say why it cannot be written. The whole
144
+ * read-modify-write holds the index's mutation queue: tool calls run in parallel, so
145
+ * two unqueued saves both read the same index and the second silently drops the first's
146
+ * line. The queue keys ONLY on the index, the shared file every save touches, and never
147
+ * also on the memory file: a second nested queue self-deadlocks when a memory name
148
+ * canonicalizes to the same key as the index (e.g. `memory.md` and `MEMORY.md` under a
149
+ * case-insensitive filesystem, since the queue keys on realpath). */
150
+ export async function saveMemory(dir: string, indexPath: string, name: string | undefined, description: string | undefined, content: string | undefined, now: string = new Date().toISOString()): Promise<MemoryToolResult> {
143
151
  if (!name || !description || !content) {
144
152
  return { content: [{ type: 'text', text: 'save requires name, description, and content.' }], details: {} }
145
153
  }
146
- const index = readIndex(dir)
147
- // Claude reports an explicit error rather than writing a memory the next session
148
- // would never load, and says what to do about it.
149
- if (indexWouldOverflow(index, name, description)) {
150
- return {
151
- content: [{ type: 'text', text: `Memory index is full (${INDEX_MAX_LINES} entries or ${INDEX_MAX_BYTES} bytes). Delete or consolidate memories before saving ${name}.` }],
152
- details: {},
154
+ return withFileMutationQueue(indexPath, async (): Promise<MemoryToolResult> => {
155
+ const index = readIndex(dir)
156
+ // Claude reports an explicit error rather than writing a memory the next session
157
+ // would never load, and says what to do about it.
158
+ if (indexWouldOverflow(index, name, description)) {
159
+ return {
160
+ content: [{ type: 'text', text: `Memory index is full (${INDEX_MAX_LINES} entries or ${INDEX_MAX_BYTES} bytes). Delete or consolidate memories before saving ${name}.` }],
161
+ details: {},
162
+ }
153
163
  }
154
- }
155
- fs.mkdirSync(dir, { recursive: true })
156
- // A memory with frontmatter records its write time; one without is left as-is.
157
- fs.writeFileSync(path.join(dir, `${name}.md`), stampModified(content, now))
158
- writeIndex(indexPath, upsertIndexLine(index, name, description))
159
- return { content: [{ type: 'text', text: `Saved memory ${name}.` }], details: {} }
164
+ fs.mkdirSync(dir, { recursive: true })
165
+ // A memory with frontmatter records its write time; one without is left as-is.
166
+ fs.writeFileSync(path.join(dir, `${name}.md`), stampModified(content, now))
167
+ writeIndex(indexPath, upsertIndexLine(index, name, description))
168
+ return { content: [{ type: 'text', text: `Saved memory ${name}.` }], details: {} }
169
+ })
160
170
  }
161
171
 
162
172
  /** The read action: a memory's body, capped for context, or a not-found message. */
163
- function readMemory(dir: string, name: string): { content: Array<{ type: 'text'; text: string }>; details: Record<string, never> } {
173
+ function readMemory(dir: string, name: string): MemoryToolResult {
164
174
  try {
165
175
  const body = fs.readFileSync(path.join(dir, `${name}.md`), 'utf-8')
166
176
  return { content: [{ type: 'text', text: capForContext(body) }], details: {} }
@@ -169,21 +179,23 @@ function readMemory(dir: string, name: string): { content: Array<{ type: 'text';
169
179
  }
170
180
  }
171
181
 
172
- /** The delete action: remove a memory file and its index line. The index is read
173
- * before anything is removed: refusing on a failed read must leave both the memory
174
- * file and the index as they were. */
175
- function deleteMemory(dir: string, indexPath: string, name: string): { content: Array<{ type: 'text'; text: string }>; details: Record<string, never> } {
176
- let index: string
182
+ /** The delete action: remove a memory file and its index line, queued on the index
183
+ * like save (single key, no deadlock). The index is read before anything is removed,
184
+ * and any failure (a bad index read, or an unreadable store the queue key cannot
185
+ * realpath) leaves both the memory file and the index as they were. */
186
+ async function deleteMemory(dir: string, indexPath: string, name: string): Promise<MemoryToolResult> {
177
187
  try {
178
- index = readIndex(dir)
188
+ return await withFileMutationQueue(indexPath, async (): Promise<MemoryToolResult> => {
189
+ const index = readIndex(dir)
190
+ fs.rmSync(path.join(dir, `${name}.md`), { force: true })
191
+ const remaining = removeIndexLine(index, name)
192
+ if (remaining) writeIndex(indexPath, remaining)
193
+ else fs.rmSync(indexPath, { force: true })
194
+ return { content: [{ type: 'text', text: `Deleted memory ${name}.` }], details: {} }
195
+ })
179
196
  } catch (error) {
180
197
  return { content: [{ type: 'text', text: `Memory delete failed: ${error instanceof Error ? error.message : String(error)}. Nothing was deleted.` }], details: {} }
181
198
  }
182
- fs.rmSync(path.join(dir, `${name}.md`), { force: true })
183
- const remaining = removeIndexLine(index, name)
184
- if (remaining) writeIndex(indexPath, remaining)
185
- else fs.rmSync(indexPath, { force: true })
186
- return { content: [{ type: 'text', text: `Deleted memory ${name}.` }], details: {} }
187
199
  }
188
200
 
189
201
  /** The index as injected into the prompt, bounded like Claude's startup load. */
@@ -346,7 +358,8 @@ export default function memoryExtension(pi: ExtensionAPI) {
346
358
 
347
359
  if (params.action === 'save') {
348
360
  try {
349
- return saveMemory(dir, indexPath, name, params.description, params.content)
361
+ // Awaited here, not returned: the catch must see a queued write's rejection.
362
+ return await saveMemory(dir, indexPath, name, params.description, params.content)
350
363
  } catch (error) {
351
364
  return { content: [{ type: 'text' as const, text: `Memory save failed: ${error instanceof Error ? error.message : String(error)}. The index was left untouched.` }], details: {} }
352
365
  }
@@ -113,6 +113,11 @@ export default function notifyExtension(pi: ExtensionAPI) {
113
113
  lastInputAt = Date.now()
114
114
  })
115
115
 
116
+ // Fires on agent_end rather than agent_settled deliberately: agent_settled is only
117
+ // emitted after every agent_end handler returns, and a peer extension (plan mode)
118
+ // blocks its agent_end handler on a UI dialog, which would starve this notification
119
+ // exactly when the user has stepped away. agent_end can fire slightly early before a
120
+ // rare automatic retry or compaction, which is a better failure than never notifying.
116
121
  pi.on('agent_end', async () => {
117
122
  if (channel === 'off') return
118
123
  // Piped or headless stdout (pi -p, CI) must not receive raw escape bytes.
@@ -209,7 +209,9 @@ export default function planModeExtension(pi: ExtensionAPI): void {
209
209
  planFromTool = false
210
210
  const refinement = await ctx.ui.editor('Refine the plan:', '')
211
211
  if (refinement?.trim()) {
212
- pi.sendUserMessage(refinement.trim())
212
+ // A bare send throws (and is silently swallowed) while the agent is
213
+ // streaming, so mid-stream invocations queue as a follow-up turn.
214
+ pi.sendUserMessage(refinement.trim(), ctx.isIdle() ? {} : { deliverAs: 'followUp' })
213
215
  }
214
216
  }
215
217
  }
@@ -238,7 +238,39 @@ async function askOne(params: QuestionSpec, ctx: ExtensionContext): Promise<{ co
238
238
  // The free-text option does not compose with checkbox selection, so it is single-select only.
239
239
  const allOptions: DisplayOption[] = multiSelect ? [...params.options] : [...params.options, { label: 'Type something.', isOther: true }]
240
240
 
241
- const result = await ctx.ui.custom<{ answer: string; wasCustom: boolean; index?: number } | null>((tui: Parameters<Parameters<ExtensionContext['ui']['custom']>[0]>[0], theme: Theme, _kb: unknown, done: (value: { answer: string; wasCustom: boolean; index?: number } | null) => void) => {
241
+ // ui.custom() is terminal-only: with a UI but no terminal (RPC mode) it resolves
242
+ // undefined immediately, which would read as a cancel without ever asking. Ask
243
+ // through the dialog primitives there instead.
244
+ const result = ctx.mode === 'tui' ? await askViaOverlay(params, ctx, allOptions, multiSelect) : await askViaDialogs(params, ctx, allOptions, multiSelect)
245
+
246
+ // Build simple options list for details; header/multiSelect appear only when set,
247
+ // so single-select details are unchanged.
248
+ const simpleOptions = params.options.map((o) => o.label)
249
+ const base = { question: params.question, options: simpleOptions, ...(params.header ? { header: shortHeader(params.header) } : {}), ...(multiSelect ? { multiSelect: true } : {}) }
250
+
251
+ if (!result) {
252
+ return {
253
+ content: [{ type: 'text', text: 'User cancelled the selection' }],
254
+ details: { ...base, answer: null } as QuestionDetails,
255
+ }
256
+ }
257
+
258
+ if (result.wasCustom) {
259
+ return {
260
+ content: [{ type: 'text', text: `User wrote: ${result.answer}` }],
261
+ details: { ...base, answer: result.answer, wasCustom: true } as QuestionDetails,
262
+ }
263
+ }
264
+ const selectionText = multiSelect ? `User selected: ${result.answer || '(none)'}` : `User selected: ${result.index}. ${result.answer}`
265
+ return {
266
+ content: [{ type: 'text', text: selectionText }],
267
+ details: { ...base, answer: result.answer, wasCustom: false } as QuestionDetails,
268
+ }
269
+ }
270
+
271
+ /** Terminal path: the full custom overlay (options list, checkboxes, inline editor). */
272
+ function askViaOverlay(params: QuestionSpec, ctx: ExtensionContext, allOptions: DisplayOption[], multiSelect: boolean): Promise<{ answer: string; wasCustom: boolean; index?: number } | null> {
273
+ return ctx.ui.custom<{ answer: string; wasCustom: boolean; index?: number } | null>((tui: Parameters<Parameters<ExtensionContext['ui']['custom']>[0]>[0], theme: Theme, _kb: unknown, done: (value: { answer: string; wasCustom: boolean; index?: number } | null) => void) => {
242
274
  let optionIndex = 0
243
275
  let editMode = false
244
276
  const checked: boolean[] = allOptions.map(() => false)
@@ -339,28 +371,29 @@ async function askOne(params: QuestionSpec, ctx: ExtensionContext): Promise<{ co
339
371
  handleInput,
340
372
  }
341
373
  })
374
+ }
342
375
 
343
- // Build simple options list for details; header/multiSelect appear only when set,
344
- // so single-select details are unchanged.
345
- const simpleOptions = params.options.map((o) => o.label)
346
- const base = { question: params.question, options: simpleOptions, ...(params.header ? { header: shortHeader(params.header) } : {}), ...(multiSelect ? { multiSelect: true } : {}) }
347
-
348
- if (!result) {
349
- return {
350
- content: [{ type: 'text', text: 'User cancelled the selection' }],
351
- details: { ...base, answer: null } as QuestionDetails,
352
- }
353
- }
354
-
355
- if (result.wasCustom) {
356
- return {
357
- content: [{ type: 'text', text: `User wrote: ${result.answer}` }],
358
- details: { ...base, answer: result.answer, wasCustom: true } as QuestionDetails,
359
- }
360
- }
361
- const selectionText = multiSelect ? `User selected: ${result.answer || '(none)'}` : `User selected: ${result.index}. ${result.answer}`
362
- return {
363
- content: [{ type: 'text', text: selectionText }],
364
- details: { ...base, answer: result.answer, wasCustom: false } as QuestionDetails,
376
+ /** Dialog-primitive fallback for UI without a terminal (RPC mode supports
377
+ * select/input/notify but not custom components). Mirrors the overlay's result
378
+ * shape; a dismissed dialog reads as a cancel, same as Escape in the overlay. */
379
+ async function askViaDialogs(params: QuestionSpec, ctx: ExtensionContext, allOptions: DisplayOption[], multiSelect: boolean): Promise<{ answer: string; wasCustom: boolean; index?: number } | null> {
380
+ const header = shortHeader(params.header)
381
+ const title = header ? `[${header}] ${params.question}` : params.question
382
+ // Number the labels: ctx.ui.select returns the chosen label string, so duplicate
383
+ // labels (or a model-supplied option named like the free-text entry) would be
384
+ // ambiguous by text alone; the number is the unambiguous way back to the option.
385
+ const labels = allOptions.map((option, i) => `${i + 1}. ${option.label}`)
386
+ const choice = await ctx.ui.select(title, labels)
387
+ if (choice === undefined) return null
388
+ const index = labels.indexOf(choice)
389
+ const chosen = allOptions[index]
390
+ if (!multiSelect && chosen?.isOther === true) {
391
+ const typed = await ctx.ui.input(params.question, 'Your answer')
392
+ // A dismissed dialog cancels; a submitted empty answer is an (empty) answer, not a
393
+ // cancel, so one accidental blank Enter does not abort the rest of a question batch.
394
+ if (typed === undefined) return null
395
+ return { answer: typed.trim(), wasCustom: true }
365
396
  }
397
+ const answer = chosen?.label ?? choice
398
+ return multiSelect ? { answer, wasCustom: false } : { answer, wasCustom: false, index: index + 1 }
366
399
  }
@@ -26,6 +26,10 @@ export interface BackgroundRun {
26
26
  /** True until the child process actually closes: a cancelled child that ignores
27
27
  * SIGTERM is still alive and must keep holding its concurrency slot. */
28
28
  live?: boolean
29
+ /** Monotonic finish order, stamped when the run completes. Eviction drops the
30
+ * earliest-finished runs by this, not Map insertion (start) order: a long run
31
+ * started first but finished last must not vanish the instant it completes. */
32
+ finishedAt?: number
29
33
  /** pi session the child ran under, so a follow-up can continue its context. */
30
34
  sessionId: string
31
35
  /** How the child was spawned, so a follow-up can repeat it with a new task. */
@@ -63,8 +67,13 @@ export function activeBackgroundRuns(): number {
63
67
  return [...runs.values()].filter((run) => run.live || run.state === 'running').length
64
68
  }
65
69
 
70
+ /** Stamps BackgroundRun.finishedAt; a counter rather than a clock so two runs
71
+ * completing in the same millisecond still evict in their true finish order. */
72
+ let finishSequence = 0
73
+
66
74
  function evictFinishedRuns(): void {
67
75
  const finished = [...runs.values()].filter((run) => !run.live && run.state !== 'running')
76
+ finished.sort((a, b) => (a.finishedAt ?? 0) - (b.finishedAt ?? 0))
68
77
  for (const stale of finished.slice(0, Math.max(0, finished.length - MAX_FINISHED_RUNS))) runs.delete(stale.id)
69
78
  }
70
79
 
@@ -160,6 +169,7 @@ export function resumeBackgroundRun(id: string, task: string, onComplete: (run:
160
169
  run.output = undefined
161
170
  run.exitCode = undefined
162
171
  run.stderr = undefined
172
+ run.finishedAt = undefined
163
173
  driveRun(run, { ...run.spawn, args }, onComplete)
164
174
  return 'resumed'
165
175
  }
@@ -256,6 +266,7 @@ function driveRun(run: BackgroundRun, invocation: BackgroundSpawn, onComplete: (
256
266
  const complete = (): void => {
257
267
  if (completed) return
258
268
  completed = true
269
+ run.finishedAt = ++finishSequence
259
270
  evictFinishedRuns()
260
271
  // A run outlives the session that started it, and pi's loader wires assertActive()
261
272
  // into every runtime call, so notifying a disposed session throws. This fires from
@@ -84,6 +84,9 @@ const listMark = (status: TodoStatus): string => {
84
84
  return '[ ]'
85
85
  }
86
86
 
87
+ /** Plain-text list, shared by the list action and the non-terminal /todos path. */
88
+ const plainTodoList = (todos: Todo[]): string => (todos.length ? todos.map((t) => `${listMark(t.status)} #${t.id}: ${t.text}`).join('\n') : 'No todos')
89
+
87
90
  const overlayLabel = (todo: Todo, theme: Theme): string => {
88
91
  if (todo.status === 'completed') return theme.fg('dim', todo.text)
89
92
  if (todo.status === 'in_progress') return theme.fg('text', todo.activeForm ?? todo.text)
@@ -433,7 +436,7 @@ export default function todoExtension(pi: ExtensionAPI) {
433
436
  return ok('clear', `Cleared ${count} todos`)
434
437
  }
435
438
 
436
- const handleList = () => ok('list', todos.length ? todos.map((t) => `${listMark(t.status)} #${t.id}: ${t.text}`).join('\n') : 'No todos')
439
+ const handleList = () => ok('list', plainTodoList(todos))
437
440
 
438
441
  // Register the todo tool for the LLM
439
442
  pi.registerTool({
@@ -516,6 +519,13 @@ export default function todoExtension(pi: ExtensionAPI) {
516
519
  return
517
520
  }
518
521
 
522
+ // ui.custom() is terminal-only: with a UI but no terminal (RPC mode) it
523
+ // resolves undefined without showing anything, so notify the plain list.
524
+ if (ctx.mode !== 'tui') {
525
+ ctx.ui.notify(plainTodoList(todos), 'info')
526
+ return
527
+ }
528
+
519
529
  await ctx.ui.custom<void>((_tui, theme, _kb, done) => {
520
530
  return new TodoListComponent(todos, theme, () => done())
521
531
  })
package/extensions/web.ts CHANGED
@@ -9,6 +9,7 @@
9
9
  import type { LookupAddress } from 'node:dns'
10
10
  import { lookup } from 'node:dns/promises'
11
11
  import type { LookupFunction } from 'node:net'
12
+ import type { Usage } from '@earendil-works/pi-ai'
12
13
  import type { ExtensionAPI } from '@earendil-works/pi-coding-agent'
13
14
  import { Type } from 'typebox'
14
15
 
@@ -252,16 +253,16 @@ function rememberFetch(cache: FetchCache, url: string, body: string, now: number
252
253
  }
253
254
 
254
255
  /** Claude's WebFetch runs the prompt over the page with a fast model and returns
255
- * that answer, not the raw page. Best-effort: any failure (no model, provider error)
256
+ * that answer, not the raw page, along with the nested call's usage so the tool
257
+ * result can account for it. Best-effort: any failure (no model, provider error)
256
258
  * yields null so the caller falls back to the markdown. */
257
- async function answerFromPage(model: Parameters<typeof completeText>[0], prompt: string, url: string, body: string, signal?: AbortSignal): Promise<string | null> {
259
+ async function answerFromPage(model: Parameters<typeof completeText>[0], prompt: string, url: string, body: string, signal?: AbortSignal): Promise<{ text: string; usage: Usage } | null> {
258
260
  try {
259
- const answer = await completeText(model, `${prompt}\n\nAnswer using only the page content below, fetched from ${url}:\n\n${body}`, {
261
+ return await completeText(model, `${prompt}\n\nAnswer using only the page content below, fetched from ${url}:\n\n${body}`, {
260
262
  system: 'You extract and answer questions from a web page. Answer only from the provided content, concisely. If the content does not contain the answer, say so.',
261
263
  maxTokens: 1024,
262
264
  signal,
263
265
  })
264
- return answer || null
265
266
  } catch {
266
267
  return null
267
268
  }
@@ -318,14 +319,18 @@ export default function webExtension(pi: ExtensionAPI) {
318
319
  }
319
320
 
320
321
  // Best-effort prompt-over-page: a failure returns null, so web_fetch always
321
- // falls back to the raw markdown and returns something.
322
+ // falls back to the raw markdown and returns something. The nested call's
323
+ // usage rides on the result either way, so pi counts it in session totals.
324
+ let usage: Usage | undefined
322
325
  if (params.prompt && ctx?.model) {
323
326
  const answer = await answerFromPage(ctx.model, params.prompt, params.url, body, signal)
324
- if (answer) return { content: [{ type: 'text' as const, text: answer }], details: {} }
327
+ if (answer?.text) return { content: [{ type: 'text' as const, text: answer.text }], details: {}, usage: answer.usage }
328
+ // An empty answer still cost the completion; the fallback carries its usage.
329
+ usage = answer?.usage
325
330
  }
326
331
  // The char cap alone admits thousands of short lines; pi's tool-output budget
327
332
  // bounds lines too, which the shared guard enforces.
328
- return { content: [{ type: 'text' as const, text: capForContext(body) || '(empty response)' }], details: {} }
333
+ return { content: [{ type: 'text' as const, text: capForContext(body) || '(empty response)' }], details: {}, usage }
329
334
  },
330
335
  })
331
336
  }
package/package.json CHANGED
@@ -1,6 +1,6 @@
1
1
  {
2
2
  "name": "pi-code",
3
- "version": "1.0.6",
3
+ "version": "1.0.8",
4
4
  "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",
5
5
  "keywords": [
6
6
  "pi",