1import { constants } from "node:fs";

2import { access as fsAccess } from "node:fs/promises";

3import { constants as osConstants } from "node:os";

4import type { AgentTool } from "@earendil-works/pi-agent-core";

5import { spawn } from "child_process";

6import { type Static, Type } from "typebox";

7import { waitForChildProcess } from "../../utils/child-process.ts";

8import {

9 getShellConfig,

10 getShellEnv,

11 killProcessTree,

12 type ShellConfig,

13 trackDetachedChildPid,

14 untrackDetachedChildPid,

15} from "../../utils/shell.ts";

16import type { ExtensionContext, ToolDefinition } from "../extensions/types.ts";

17import { OutputAccumulator } from "./output-accumulator.ts";

18import { BASH_UPDATE_THROTTLE_MS, createShellRenderers } from "./renderers/bash.ts";

19import { wrapToolDefinition } from "./tool-definition-wrapper.ts";

20import { DEFAULT_MAX_BYTES, DEFAULT_MAX_LINES, formatSize, type TruncationResult } from "./truncate.ts";

21

22const MAX_TIMEOUT_MS = 2_147_483_647;

23/** Output limit of `structuredContent.output`, which programmatic callers such as codemode scripts receive. */

24const STRUCTURED_OUTPUT_MAX_BYTES = 1024 * 1024;

25const MAX_TIMEOUT_SECONDS = MAX_TIMEOUT_MS / 1000;

26

27function resolveTimeoutMs(timeout: number | undefined): number | undefined {

28 if (timeout === undefined) return undefined;

29 if (!Number.isFinite(timeout) || timeout <= 0) {

30 throw new Error("Invalid timeout: must be a finite number of seconds");

31 }

32

33 const timeoutMs = timeout * 1000;

34 if (timeoutMs > MAX_TIMEOUT_MS) {

35 throw new Error(`Invalid timeout: maximum is ${MAX_TIMEOUT_SECONDS} seconds`);

36 }

37 return timeoutMs;

38}

39

40const bashSchema = Type.Object({

41 command: Type.String({ description: "Shell command to execute" }),

42 timeout: Type.Optional(Type.Number({ description: "Timeout in seconds (optional, no default timeout)" })),

43});

44

45export const bashToolSystemPromptContribution = {

46 snippet: "Execute bash commands (ls, grep, find, etc.)",

47 guidelines: ["You can inspect PI_* environment variables for current model and session details."],

48} as const;

49

50export type BashToolInput = Static<typeof bashSchema>;

51

52/**

53 * Result for programmatic callers such as codemode scripts. A non-zero exit code is an error result for the model, but scripts still resolve to this value.

54 * `output` is not limited like the model-facing output: callers decide how much of it reaches the model.

55 */

56const bashOutputSchema = Type.Object({

57 output: Type.String({ description: "Combined stdout and stderr, possibly truncated" }),

58 truncated: Type.Boolean(),

59 full_output_path: Type.Optional(Type.String({ description: "Full output, when truncated" })),

60 exit_code: Type.Number(),

61 wall_time_seconds: Type.Number(),

62});

63

64export type BashToolOutput = Static<typeof bashOutputSchema>;

65

66export interface BashToolDetails {

67 truncation?: TruncationResult;

68 fullOutputPath?: string;

69}

70

71/**

72 * Pluggable operations for the bash tool.

73 * Override these to delegate command execution to remote systems (for example SSH).

74 */

75export interface BashOperations {

76 /**

77 * Execute a command and stream output.

78 * @param command The command to execute

79 * @param cwd Working directory

80 * @param options Execution options

81 * @returns Promise resolving to the exit code. Report signal terminations as 128 + signal number;

82 * a null exit code is treated as a failed command.

83 */

84 exec: (

85 command: string,

86 cwd: string,

87 options: {

88 onData: (data: Buffer) => void;

89 signal?: AbortSignal;

90 timeout?: number;

91 env?: NodeJS.ProcessEnv;

92 },

93 ) => Promise<{ exitCode: number | null }>;

94}

95

96/** Shared process execution used by the built-in shell tools. */

97export function createLocalShellOperations(shellName: string, resolveShellConfig: () => ShellConfig): BashOperations {

98 return {

99 exec: async (command, cwd, { onData, signal, timeout, env }) => {

100 const timeoutMs = resolveTimeoutMs(timeout);

101 if (signal?.aborted) {

102 throw new Error("aborted");

103 }

104 const shellConfig = resolveShellConfig();

105 try {

106 await fsAccess(cwd, constants.F_OK);

107 } catch {

108 throw new Error(`Working directory does not exist: ${cwd}\nCannot execute ${shellName} commands.`);

109 }

110

111 const commandFromStdin = shellConfig.commandTransport === "stdin";

112 const child = spawn(shellConfig.shell, commandFromStdin ? shellConfig.args : [...shellConfig.args, command], {

113 cwd,

114 detached: process.platform !== "win32",

115 env: env ?? getShellEnv(),

116 stdio: [commandFromStdin ? "pipe" : "ignore", "pipe", "pipe"],

117 windowsHide: true,

118 });

119 if (commandFromStdin) {

120 child.stdin?.on("error", () => {});

121 child.stdin?.end(command);

122 }

123 if (child.pid) trackDetachedChildPid(child.pid);

124 let timedOut = false;

125 let timeoutHandle: NodeJS.Timeout | undefined;

126 const onAbort = () => {

127 if (child.pid) killProcessTree(child.pid);

128 };

129

130 try {

131 // Set timeout if provided.

132 if (timeoutMs !== undefined) {

133 timeoutHandle = setTimeout(() => {

134 timedOut = true;

135 if (child.pid) killProcessTree(child.pid);

136 }, timeoutMs);

137 }

138 // Stream stdout and stderr.

139 child.stdout?.on("data", onData);

140 child.stderr?.on("data", onData);

141 // Handle abort signal by killing the entire process tree.

142 if (signal) {

143 if (signal.aborted) onAbort();

144 else signal.addEventListener("abort", onAbort, { once: true });

145 }

146 // Handle shell spawn errors and wait for the process to terminate without hanging

147 // on inherited stdio handles held by detached descendants.

148 const exitCode = await waitForChildProcess(child);

149 if (signal?.aborted) {

150 throw new Error("aborted");

151 }

152 if (timedOut) {

153 throw new Error(`timeout:${timeout}`);

154 }

155 // A signal-killed shell has no exit code. Use the standard shell convention so

156 // callers do not mistake the termination for a successful command.

157 const signalCode = child.signalCode;

158 return { exitCode: exitCode ?? (signalCode ? 128 + (osConstants.signals[signalCode] ?? 0) : 1) };

159 } finally {

160 if (child.pid) untrackDetachedChildPid(child.pid);

161 if (timeoutHandle) clearTimeout(timeoutHandle);

162 if (signal) signal.removeEventListener("abort", onAbort);

163 }

164 },

165 };

166}

167

168/**

169 * Create bash operations using pi's built-in local shell execution backend.

170 *

171 * This is useful for extensions that intercept user_bash and still want pi's

172 * standard local shell behavior while wrapping or rewriting commands.

173 */

174export function createLocalBashOperations(options?: { shellPath?: string }): BashOperations {

175 return createLocalShellOperations("bash", () => getShellConfig(options?.shellPath));

176}

177

178export interface BashSpawnContext {

179 command: string;

180 cwd: string;

181 env: NodeJS.ProcessEnv;

182}

183

184export type BashSpawnHook = (context: BashSpawnContext) => BashSpawnContext;

185

186function resolveSpawnContext(

187 command: string,

188 cwd: string,

189 spawnHook: BashSpawnHook | undefined,

190 exposeSessionEnvironment: boolean,

191 ctx: ExtensionContext | undefined,

192): BashSpawnContext {

193 const env = { ...getShellEnv() };

194 delete env.PI_SESSION_ID;

195 delete env.PI_SESSION_FILE;

196 delete env.PI_PROVIDER;

197 delete env.PI_MODEL;

198 delete env.PI_REASONING_LEVEL;

199 if (exposeSessionEnvironment && ctx) {

200 const model = ctx.model;

201 env.PI_SESSION_ID = ctx.sessionManager.getSessionId();

202 const sessionFile = ctx.sessionManager.getSessionFile();

203 if (sessionFile) env.PI_SESSION_FILE = sessionFile;

204 if (model) {

205 env.PI_PROVIDER = model.provider;

206 env.PI_MODEL = model.id;

207 }

208 if (ctx.thinkingLevel) env.PI_REASONING_LEVEL = ctx.thinkingLevel;

209 }

210 const baseContext: BashSpawnContext = { command, cwd, env };

211 return spawnHook ? spawnHook(baseContext) : baseContext;

212}

213

214export interface BashToolOptions {

215 /** Custom operations for command execution. Default: local shell */

216 operations?: BashOperations;

217 /** Command prefix prepended to every command (for example shell setup commands) */

218 commandPrefix?: string;

219 /** Optional explicit shell path from settings */

220 shellPath?: string;

221 /** Expose current Pi session metadata as PI_* environment variables. Default: true */

222 exposeSessionEnvironment?: boolean;

223 /** Hook to adjust command, cwd, or env before execution */

224 spawnHook?: BashSpawnHook;

225}

226

227export type BashRenderState = {

228 startedAt: number | undefined;

229 endedAt: number | undefined;

230 interval: NodeJS.Timeout | undefined;

231};

232

233export interface ShellToolConfig {

234 name: string;

235 label: string;

236 shellName: string;

237 prompt: string;

238 promptSnippet: string;

239 promptGuidelines?: readonly string[];

240 tempFilePrefix: string;

241}

242

243export function createShellToolDefinition(

244 cwd: string,

245 config: ShellToolConfig,

246 options?: BashToolOptions,

247): ToolDefinition<typeof bashSchema, BashToolDetails | undefined, BashRenderState> {

248 const ops = options?.operations ?? createLocalBashOperations({ shellPath: options?.shellPath });

249 const commandPrefix = options?.commandPrefix;

250 const exposeSessionEnvironment = options?.exposeSessionEnvironment ?? true;

251 const spawnHook = options?.spawnHook;

252 return {

253 name: config.name,

254 label: config.label,

255 description: `Execute a ${config.shellName} command in the current working directory. Returns stdout and stderr. Output is truncated to last ${DEFAULT_MAX_LINES} lines or ${DEFAULT_MAX_BYTES / 1024}KB (whichever is hit first). If truncated, full output is saved to a temp file. Optionally provide a timeout in seconds.`,

256 promptSnippet: config.promptSnippet,

257 promptGuidelines: exposeSessionEnvironment && config.promptGuidelines ? [...config.promptGuidelines] : undefined,

258 parameters: bashSchema,

259 outputSchema: bashOutputSchema,

260 constrainedSampling: { type: "json_schema", strict: "prefer" },

261 async execute(

262 _toolCallId,

263 { command, timeout }: { command: string; timeout?: number },

264 signal?: AbortSignal,

265 onUpdate?,

266 ctx?: ExtensionContext,

267 ) {

268 const resolvedCommand = commandPrefix ? `${commandPrefix}\n${command}` : command;

269 const spawnContext = resolveSpawnContext(

270 resolvedCommand,

271 ctx?.cwd || cwd,

272 spawnHook,

273 exposeSessionEnvironment,

274 ctx,

275 );

276 const output = new OutputAccumulator({ tempFilePrefix: config.tempFilePrefix });

277 let acceptingOutput = true;

278 let updateTimer: NodeJS.Timeout | undefined;

279 let updateDirty = false;

280 let lastUpdateAt = 0;

281

282 const emitOutputUpdate = () => {

283 if (!onUpdate || !updateDirty) return;

284 updateDirty = false;

285 lastUpdateAt = Date.now();

286 const snapshot = output.snapshot({ persistIfTruncated: true });

287 onUpdate({

288 content: [{ type: "text", text: snapshot.content || "" }],

289 details: {

290 truncation: snapshot.truncation.truncated ? snapshot.truncation : undefined,

291 fullOutputPath: snapshot.fullOutputPath,

292 },

293 });

294 };

295

296 const clearUpdateTimer = () => {

297 if (updateTimer) {

298 clearTimeout(updateTimer);

299 updateTimer = undefined;

300 }

301 };

302

303 const scheduleOutputUpdate = () => {

304 if (!onUpdate) return;

305 updateDirty = true;

306 const delay = BASH_UPDATE_THROTTLE_MS - (Date.now() - lastUpdateAt);

307 if (delay <= 0) {

308 clearUpdateTimer();

309 emitOutputUpdate();

310 return;

311 }

312 updateTimer ??= setTimeout(() => {

313 updateTimer = undefined;

314 emitOutputUpdate();

315 }, delay);

316 };

317

318 if (onUpdate) {

319 onUpdate({ content: [], details: undefined });

320 }

321

322 const handleData = (data: Buffer) => {

323 if (!acceptingOutput) return;

324 output.append(data);

325 scheduleOutputUpdate();

326 };

327

328 const finishOutput = async () => {

329 acceptingOutput = false;

330 output.finish();

331 clearUpdateTimer();

332 emitOutputUpdate();

333 const snapshot = output.snapshot({ persistIfTruncated: true });

334 await output.closeTempFile();

335 return snapshot;

336 };

337

338 const formatOutput = (snapshot: Awaited<ReturnType<typeof finishOutput>>, emptyText = "(no output)") => {

339 const truncation = snapshot.truncation;

340 let text = snapshot.content || emptyText;

341 let details: BashToolDetails | undefined;

342 if (truncation.truncated) {

343 details = { truncation, fullOutputPath: snapshot.fullOutputPath };

344 const startLine = truncation.totalLines - truncation.outputLines + 1;

345 const endLine = truncation.totalLines;

346 if (truncation.lastLinePartial) {

347 const lastLineSize = formatSize(output.getLastLineBytes());

348 text += `\n\n[Showing last ${formatSize(truncation.outputBytes)} of line ${endLine} (line is ${lastLineSize}). Full output: ${snapshot.fullOutputPath}]`;

349 } else if (truncation.truncatedBy === "lines") {

350 text += `\n\n[Showing lines ${startLine}-${endLine} of ${truncation.totalLines}. Full output: ${snapshot.fullOutputPath}]`;

351 } else {

352 text += `\n\n[Showing lines ${startLine}-${endLine} of ${truncation.totalLines} (${formatSize(DEFAULT_MAX_BYTES)} limit). Full output: ${snapshot.fullOutputPath}]`;

353 }

354 }

355 return { text, details };

356 };

357

358 const appendStatus = (text: string, status: string) => `${text ? `${text}\n\n` : ""}${status}`;

359 const startedAt = performance.now();

360

361 try {

362 let exitCode: number | null;

363 try {

364 const result = await ops.exec(spawnContext.command, spawnContext.cwd, {

365 onData: handleData,

366 signal,

367 timeout,

368 env: spawnContext.env,

369 });

370 exitCode = result.exitCode;

371 } catch (err) {

372 const snapshot = await finishOutput();

373 const { text } = formatOutput(snapshot, "");

374 if (err instanceof Error && err.message === "aborted") {

375 throw new Error(appendStatus(text, "Command aborted"));

376 }

377 if (err instanceof Error && err.message.startsWith("timeout:")) {

378 const timeoutSecs = err.message.split(":")[1];

379 throw new Error(appendStatus(text, `Command timed out after ${timeoutSecs} seconds`));

380 }

381 throw err;

382 }

383

384 const snapshot = await finishOutput();

385 const { text: outputText, details } = formatOutput(snapshot);

386 if (exitCode === null) {

387 throw new Error(appendStatus(outputText, "Command terminated without an exit code"));

388 }

389 const wallTimeSeconds = Math.round((performance.now() - startedAt) / 100) / 10;

390 const fullOutput = await output.readFullOutput(STRUCTURED_OUTPUT_MAX_BYTES);

391 const structuredContent: BashToolOutput = {

392 output: fullOutput.content,

393 truncated: fullOutput.truncated,

394 ...(fullOutput.truncated && snapshot.fullOutputPath

395 ? { full_output_path: snapshot.fullOutputPath }

396 : {}),

397 exit_code: exitCode,

398 wall_time_seconds: wallTimeSeconds,

399 };

400 if (exitCode !== 0) {

401 return {

402 content: [{ type: "text", text: appendStatus(outputText, `Command exited with code ${exitCode}`) }],

403 details,

404 structuredContent,

405 isError: true,

406 };

407 }

408 return { content: [{ type: "text", text: outputText }], details, structuredContent };

409 } finally {

410 clearUpdateTimer();

411 }

412 },

413 ...createShellRenderers(config.prompt),

414 };

415}

416

417const bashToolConfig: ShellToolConfig = {

418 name: "bash",

419 label: "bash",

420 shellName: "bash",

421 prompt: "$",

422 promptSnippet: bashToolSystemPromptContribution.snippet,

423 promptGuidelines: bashToolSystemPromptContribution.guidelines,

424 tempFilePrefix: "pi-bash",

425};

426

427export function createBashToolDefinition(

428 cwd: string,

429 options?: BashToolOptions,

430): ToolDefinition<typeof bashSchema, BashToolDetails | undefined, BashRenderState> {

431 return createShellToolDefinition(cwd, bashToolConfig, options);

432}

433

434export function createBashTool(cwd: string, options?: BashToolOptions): AgentTool<typeof bashSchema> {

435 const definition = createBashToolDefinition(cwd, options);

436 const tool = wrapToolDefinition(definition);

437 Object.assign(tool, {

438 promptSnippet: definition.promptSnippet,

439 promptGuidelines: definition.promptGuidelines,

440 });

441 return tool;

442}