mirror of
https://github.com/badlogic/pi-mono.git
synced 2026-07-09 17:28:45 +00:00
456 lines
16 KiB
TypeScript
456 lines
16 KiB
TypeScript
import { constants } from "node:fs";
|
|
import { access as fsAccess } from "node:fs/promises";
|
|
import type { AgentTool } from "@earendil-works/pi-agent-core";
|
|
import { Container, Text, truncateToWidth } from "@earendil-works/pi-tui";
|
|
import { spawn } from "child_process";
|
|
import { type Static, Type } from "typebox";
|
|
import { keyHint } from "../../modes/interactive/components/keybinding-hints.ts";
|
|
import { truncateToVisualLines } from "../../modes/interactive/components/visual-truncate.ts";
|
|
import { theme } from "../../modes/interactive/theme/theme.ts";
|
|
import { waitForChildProcess } from "../../utils/child-process.ts";
|
|
import {
|
|
getShellConfig,
|
|
getShellEnv,
|
|
killProcessTree,
|
|
trackDetachedChildPid,
|
|
untrackDetachedChildPid,
|
|
} from "../../utils/shell.ts";
|
|
import type { ToolDefinition, ToolRenderResultOptions } from "../extensions/types.ts";
|
|
import { OutputAccumulator } from "./output-accumulator.ts";
|
|
import { getTextOutput, invalidArgText, str } from "./render-utils.ts";
|
|
import { wrapToolDefinition } from "./tool-definition-wrapper.ts";
|
|
import { DEFAULT_MAX_BYTES, DEFAULT_MAX_LINES, formatSize, type TruncationResult } from "./truncate.ts";
|
|
|
|
const bashSchema = Type.Object({
|
|
command: Type.String({ description: "Bash command to execute" }),
|
|
timeout: Type.Optional(Type.Number({ description: "Timeout in seconds (optional, no default timeout)" })),
|
|
});
|
|
|
|
export type BashToolInput = Static<typeof bashSchema>;
|
|
|
|
export interface BashToolDetails {
|
|
truncation?: TruncationResult;
|
|
fullOutputPath?: string;
|
|
}
|
|
|
|
/**
|
|
* Pluggable operations for the bash tool.
|
|
* Override these to delegate command execution to remote systems (for example SSH).
|
|
*/
|
|
export interface BashOperations {
|
|
/**
|
|
* Execute a command and stream output.
|
|
* @param command The command to execute
|
|
* @param cwd Working directory
|
|
* @param options Execution options
|
|
* @returns Promise resolving to exit code (null if killed)
|
|
*/
|
|
exec: (
|
|
command: string,
|
|
cwd: string,
|
|
options: {
|
|
onData: (data: Buffer) => void;
|
|
signal?: AbortSignal;
|
|
timeout?: number;
|
|
env?: NodeJS.ProcessEnv;
|
|
},
|
|
) => Promise<{ exitCode: number | null }>;
|
|
}
|
|
|
|
/**
|
|
* Create bash operations using pi's built-in local shell execution backend.
|
|
*
|
|
* This is useful for extensions that intercept user_bash and still want pi's
|
|
* standard local shell behavior while wrapping or rewriting commands.
|
|
*/
|
|
export function createLocalBashOperations(options?: { shellPath?: string }): BashOperations {
|
|
return {
|
|
exec: (command, cwd, { onData, signal, timeout, env }) => {
|
|
return new Promise((resolve, reject) => {
|
|
void (async () => {
|
|
const { shell, args } = getShellConfig(options?.shellPath);
|
|
try {
|
|
await fsAccess(cwd, constants.F_OK);
|
|
} catch {
|
|
reject(new Error(`Working directory does not exist: ${cwd}\nCannot execute bash commands.`));
|
|
return;
|
|
}
|
|
if (signal?.aborted) {
|
|
reject(new Error("aborted"));
|
|
return;
|
|
}
|
|
const child = spawn(shell, [...args, command], {
|
|
cwd,
|
|
detached: process.platform !== "win32",
|
|
env: env ?? getShellEnv(),
|
|
stdio: ["ignore", "pipe", "pipe"],
|
|
windowsHide: true,
|
|
});
|
|
if (child.pid) trackDetachedChildPid(child.pid);
|
|
let timedOut = false;
|
|
let timeoutHandle: NodeJS.Timeout | undefined;
|
|
// Set timeout if provided.
|
|
if (timeout !== undefined && timeout > 0) {
|
|
timeoutHandle = setTimeout(() => {
|
|
timedOut = true;
|
|
if (child.pid) killProcessTree(child.pid);
|
|
}, timeout * 1000);
|
|
}
|
|
// Stream stdout and stderr.
|
|
child.stdout?.on("data", onData);
|
|
child.stderr?.on("data", onData);
|
|
// Handle abort signal by killing the entire process tree.
|
|
const onAbort = () => {
|
|
if (child.pid) killProcessTree(child.pid);
|
|
};
|
|
if (signal) {
|
|
if (signal.aborted) onAbort();
|
|
else signal.addEventListener("abort", onAbort, { once: true });
|
|
}
|
|
// Handle shell spawn errors and wait for the process to terminate without hanging
|
|
// on inherited stdio handles held by detached descendants.
|
|
waitForChildProcess(child)
|
|
.then((code) => {
|
|
if (child.pid) untrackDetachedChildPid(child.pid);
|
|
if (timeoutHandle) clearTimeout(timeoutHandle);
|
|
if (signal) signal.removeEventListener("abort", onAbort);
|
|
if (signal?.aborted) {
|
|
reject(new Error("aborted"));
|
|
return;
|
|
}
|
|
if (timedOut) {
|
|
reject(new Error(`timeout:${timeout}`));
|
|
return;
|
|
}
|
|
resolve({ exitCode: code });
|
|
})
|
|
.catch((err) => {
|
|
if (child.pid) untrackDetachedChildPid(child.pid);
|
|
if (timeoutHandle) clearTimeout(timeoutHandle);
|
|
if (signal) signal.removeEventListener("abort", onAbort);
|
|
reject(err);
|
|
});
|
|
})().catch((err: unknown) => reject(err instanceof Error ? err : new Error(String(err))));
|
|
});
|
|
},
|
|
};
|
|
}
|
|
|
|
export interface BashSpawnContext {
|
|
command: string;
|
|
cwd: string;
|
|
env: NodeJS.ProcessEnv;
|
|
}
|
|
|
|
export type BashSpawnHook = (context: BashSpawnContext) => BashSpawnContext;
|
|
|
|
function resolveSpawnContext(command: string, cwd: string, spawnHook?: BashSpawnHook): BashSpawnContext {
|
|
const baseContext: BashSpawnContext = { command, cwd, env: { ...getShellEnv() } };
|
|
return spawnHook ? spawnHook(baseContext) : baseContext;
|
|
}
|
|
|
|
export interface BashToolOptions {
|
|
/** Custom operations for command execution. Default: local shell */
|
|
operations?: BashOperations;
|
|
/** Command prefix prepended to every command (for example shell setup commands) */
|
|
commandPrefix?: string;
|
|
/** Optional explicit shell path from settings */
|
|
shellPath?: string;
|
|
/** Hook to adjust command, cwd, or env before execution */
|
|
spawnHook?: BashSpawnHook;
|
|
}
|
|
|
|
const BASH_PREVIEW_LINES = 5;
|
|
const BASH_UPDATE_THROTTLE_MS = 100;
|
|
|
|
type BashRenderState = {
|
|
startedAt: number | undefined;
|
|
endedAt: number | undefined;
|
|
interval: NodeJS.Timeout | undefined;
|
|
};
|
|
|
|
type BashResultRenderState = {
|
|
cachedWidth: number | undefined;
|
|
cachedLines: string[] | undefined;
|
|
cachedSkipped: number | undefined;
|
|
};
|
|
|
|
class BashResultRenderComponent extends Container {
|
|
state: BashResultRenderState = {
|
|
cachedWidth: undefined,
|
|
cachedLines: undefined,
|
|
cachedSkipped: undefined,
|
|
};
|
|
}
|
|
|
|
function formatDuration(ms: number): string {
|
|
return `${(ms / 1000).toFixed(1)}s`;
|
|
}
|
|
|
|
function formatBashCall(args: { command?: string; timeout?: number } | undefined): string {
|
|
const command = str(args?.command);
|
|
const timeout = args?.timeout as number | undefined;
|
|
const timeoutSuffix = timeout ? theme.fg("muted", ` (timeout ${timeout}s)`) : "";
|
|
const commandDisplay = command === null ? invalidArgText(theme) : command ? command : theme.fg("toolOutput", "...");
|
|
return theme.fg("toolTitle", theme.bold(`$ ${commandDisplay}`)) + timeoutSuffix;
|
|
}
|
|
|
|
function rebuildBashResultRenderComponent(
|
|
component: BashResultRenderComponent,
|
|
result: {
|
|
content: Array<{ type: string; text?: string; data?: string; mimeType?: string }>;
|
|
details?: BashToolDetails;
|
|
},
|
|
options: ToolRenderResultOptions,
|
|
showImages: boolean,
|
|
startedAt: number | undefined,
|
|
endedAt: number | undefined,
|
|
): void {
|
|
const state = component.state;
|
|
component.clear();
|
|
|
|
let output = getTextOutput(result as any, showImages).trim();
|
|
const truncation = result.details?.truncation;
|
|
const fullOutputPath = result.details?.fullOutputPath;
|
|
if (!options.isPartial && truncation?.truncated && fullOutputPath && output.endsWith("]")) {
|
|
const footerStart = output.lastIndexOf("\n\n[");
|
|
if (footerStart !== -1 && output.slice(footerStart).includes(fullOutputPath)) {
|
|
output = output.slice(0, footerStart).trimEnd();
|
|
}
|
|
}
|
|
|
|
if (output) {
|
|
const styledOutput = output
|
|
.split("\n")
|
|
.map((line) => theme.fg("toolOutput", line))
|
|
.join("\n");
|
|
|
|
if (options.expanded) {
|
|
component.addChild(new Text(`\n${styledOutput}`, 0, 0));
|
|
} else {
|
|
component.addChild({
|
|
render: (width: number) => {
|
|
if (state.cachedLines === undefined || state.cachedWidth !== width) {
|
|
const preview = truncateToVisualLines(styledOutput, BASH_PREVIEW_LINES, width);
|
|
state.cachedLines = preview.visualLines;
|
|
state.cachedSkipped = preview.skippedCount;
|
|
state.cachedWidth = width;
|
|
}
|
|
if (state.cachedSkipped && state.cachedSkipped > 0) {
|
|
const hint =
|
|
theme.fg("muted", `... (${state.cachedSkipped} earlier lines,`) +
|
|
` ${keyHint("app.tools.expand", "to expand")})`;
|
|
return ["", truncateToWidth(hint, width, "..."), ...(state.cachedLines ?? [])];
|
|
}
|
|
return ["", ...(state.cachedLines ?? [])];
|
|
},
|
|
invalidate: () => {
|
|
state.cachedWidth = undefined;
|
|
state.cachedLines = undefined;
|
|
state.cachedSkipped = undefined;
|
|
},
|
|
});
|
|
}
|
|
}
|
|
|
|
if (truncation?.truncated || fullOutputPath) {
|
|
const warnings: string[] = [];
|
|
if (fullOutputPath) {
|
|
warnings.push(`Full output: ${fullOutputPath}`);
|
|
}
|
|
if (truncation?.truncated) {
|
|
if (truncation.truncatedBy === "lines") {
|
|
warnings.push(`Truncated: showing ${truncation.outputLines} of ${truncation.totalLines} lines`);
|
|
} else {
|
|
warnings.push(
|
|
`Truncated: ${truncation.outputLines} lines shown (${formatSize(truncation.maxBytes ?? DEFAULT_MAX_BYTES)} limit)`,
|
|
);
|
|
}
|
|
}
|
|
component.addChild(new Text(`\n${theme.fg("warning", `[${warnings.join(". ")}]`)}`, 0, 0));
|
|
}
|
|
|
|
if (startedAt !== undefined) {
|
|
const label = options.isPartial ? "Elapsed" : "Took";
|
|
const endTime = endedAt ?? Date.now();
|
|
component.addChild(new Text(`\n${theme.fg("muted", `${label} ${formatDuration(endTime - startedAt)}`)}`, 0, 0));
|
|
}
|
|
}
|
|
|
|
export function createBashToolDefinition(
|
|
cwd: string,
|
|
options?: BashToolOptions,
|
|
): ToolDefinition<typeof bashSchema, BashToolDetails | undefined, BashRenderState> {
|
|
const ops = options?.operations ?? createLocalBashOperations({ shellPath: options?.shellPath });
|
|
const commandPrefix = options?.commandPrefix;
|
|
const spawnHook = options?.spawnHook;
|
|
return {
|
|
name: "bash",
|
|
label: "bash",
|
|
description: `Execute a bash 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.`,
|
|
promptSnippet: "Execute bash commands (ls, grep, find, etc.)",
|
|
parameters: bashSchema,
|
|
async execute(
|
|
_toolCallId,
|
|
{ command, timeout }: { command: string; timeout?: number },
|
|
signal?: AbortSignal,
|
|
onUpdate?,
|
|
_ctx?,
|
|
) {
|
|
const resolvedCommand = commandPrefix ? `${commandPrefix}\n${command}` : command;
|
|
const spawnContext = resolveSpawnContext(resolvedCommand, cwd, spawnHook);
|
|
const output = new OutputAccumulator({ tempFilePrefix: "pi-bash" });
|
|
let updateTimer: NodeJS.Timeout | undefined;
|
|
let updateDirty = false;
|
|
let lastUpdateAt = 0;
|
|
|
|
const emitOutputUpdate = () => {
|
|
if (!onUpdate || !updateDirty) return;
|
|
updateDirty = false;
|
|
lastUpdateAt = Date.now();
|
|
const snapshot = output.snapshot({ persistIfTruncated: true });
|
|
onUpdate({
|
|
content: [{ type: "text", text: snapshot.content || "" }],
|
|
details: {
|
|
truncation: snapshot.truncation.truncated ? snapshot.truncation : undefined,
|
|
fullOutputPath: snapshot.fullOutputPath,
|
|
},
|
|
});
|
|
};
|
|
|
|
const clearUpdateTimer = () => {
|
|
if (updateTimer) {
|
|
clearTimeout(updateTimer);
|
|
updateTimer = undefined;
|
|
}
|
|
};
|
|
|
|
const scheduleOutputUpdate = () => {
|
|
if (!onUpdate) return;
|
|
updateDirty = true;
|
|
const delay = BASH_UPDATE_THROTTLE_MS - (Date.now() - lastUpdateAt);
|
|
if (delay <= 0) {
|
|
clearUpdateTimer();
|
|
emitOutputUpdate();
|
|
return;
|
|
}
|
|
updateTimer ??= setTimeout(() => {
|
|
updateTimer = undefined;
|
|
emitOutputUpdate();
|
|
}, delay);
|
|
};
|
|
|
|
if (onUpdate) {
|
|
onUpdate({ content: [], details: undefined });
|
|
}
|
|
|
|
const handleData = (data: Buffer) => {
|
|
output.append(data);
|
|
scheduleOutputUpdate();
|
|
};
|
|
|
|
const finishOutput = async () => {
|
|
output.finish();
|
|
clearUpdateTimer();
|
|
emitOutputUpdate();
|
|
const snapshot = output.snapshot({ persistIfTruncated: true });
|
|
await output.closeTempFile();
|
|
return snapshot;
|
|
};
|
|
|
|
const formatOutput = (snapshot: Awaited<ReturnType<typeof finishOutput>>, emptyText = "(no output)") => {
|
|
const truncation = snapshot.truncation;
|
|
let text = snapshot.content || emptyText;
|
|
let details: BashToolDetails | undefined;
|
|
if (truncation.truncated) {
|
|
details = { truncation, fullOutputPath: snapshot.fullOutputPath };
|
|
const startLine = truncation.totalLines - truncation.outputLines + 1;
|
|
const endLine = truncation.totalLines;
|
|
if (truncation.lastLinePartial) {
|
|
const lastLineSize = formatSize(output.getLastLineBytes());
|
|
text += `\n\n[Showing last ${formatSize(truncation.outputBytes)} of line ${endLine} (line is ${lastLineSize}). Full output: ${snapshot.fullOutputPath}]`;
|
|
} else if (truncation.truncatedBy === "lines") {
|
|
text += `\n\n[Showing lines ${startLine}-${endLine} of ${truncation.totalLines}. Full output: ${snapshot.fullOutputPath}]`;
|
|
} else {
|
|
text += `\n\n[Showing lines ${startLine}-${endLine} of ${truncation.totalLines} (${formatSize(DEFAULT_MAX_BYTES)} limit). Full output: ${snapshot.fullOutputPath}]`;
|
|
}
|
|
}
|
|
return { text, details };
|
|
};
|
|
|
|
const appendStatus = (text: string, status: string) => `${text ? `${text}\n\n` : ""}${status}`;
|
|
|
|
try {
|
|
let exitCode: number | null;
|
|
try {
|
|
const result = await ops.exec(spawnContext.command, spawnContext.cwd, {
|
|
onData: handleData,
|
|
signal,
|
|
timeout,
|
|
env: spawnContext.env,
|
|
});
|
|
exitCode = result.exitCode;
|
|
} catch (err) {
|
|
const snapshot = await finishOutput();
|
|
const { text } = formatOutput(snapshot, "");
|
|
if (err instanceof Error && err.message === "aborted") {
|
|
throw new Error(appendStatus(text, "Command aborted"));
|
|
}
|
|
if (err instanceof Error && err.message.startsWith("timeout:")) {
|
|
const timeoutSecs = err.message.split(":")[1];
|
|
throw new Error(appendStatus(text, `Command timed out after ${timeoutSecs} seconds`));
|
|
}
|
|
throw err;
|
|
}
|
|
|
|
const snapshot = await finishOutput();
|
|
const { text: outputText, details } = formatOutput(snapshot);
|
|
if (exitCode !== 0 && exitCode !== null) {
|
|
throw new Error(appendStatus(outputText, `Command exited with code ${exitCode}`));
|
|
}
|
|
return { content: [{ type: "text", text: outputText }], details };
|
|
} finally {
|
|
clearUpdateTimer();
|
|
}
|
|
},
|
|
renderCall(args, _theme, context) {
|
|
const state = context.state;
|
|
if (context.executionStarted && state.startedAt === undefined) {
|
|
state.startedAt = Date.now();
|
|
state.endedAt = undefined;
|
|
}
|
|
const text = (context.lastComponent as Text | undefined) ?? new Text("", 0, 0);
|
|
text.setText(formatBashCall(args));
|
|
return text;
|
|
},
|
|
renderResult(result, options, _theme, context) {
|
|
const state = context.state;
|
|
if (state.startedAt !== undefined && options.isPartial && !state.interval) {
|
|
state.interval = setInterval(() => context.invalidate(), 1000);
|
|
}
|
|
if (!options.isPartial || context.isError) {
|
|
state.endedAt ??= Date.now();
|
|
if (state.interval) {
|
|
clearInterval(state.interval);
|
|
state.interval = undefined;
|
|
}
|
|
}
|
|
const component =
|
|
(context.lastComponent as BashResultRenderComponent | undefined) ?? new BashResultRenderComponent();
|
|
rebuildBashResultRenderComponent(
|
|
component,
|
|
result as any,
|
|
options,
|
|
context.showImages,
|
|
state.startedAt,
|
|
state.endedAt,
|
|
);
|
|
component.invalidate();
|
|
return component;
|
|
},
|
|
};
|
|
}
|
|
|
|
export function createBashTool(cwd: string, options?: BashToolOptions): AgentTool<typeof bashSchema> {
|
|
return wrapToolDefinition(createBashToolDefinition(cwd, options));
|
|
}
|