486 lines
16 KiB
TypeScript
486 lines
16 KiB
TypeScript
/**
|
|
* Bash command execution with streaming support and cancellation.
|
|
*
|
|
* Uses brush-core via native bindings for shell execution.
|
|
*/
|
|
import { ExponentialYield } from "@oh-my-pi/pi-agent-core/utils/yield";
|
|
import { executeShell, type MinimizerOptions, Shell, type ShellRunResult } from "@oh-my-pi/pi-natives";
|
|
import { isExecutable, type ShellConfig } from "@oh-my-pi/pi-utils/procmgr";
|
|
import { Settings, type ShellMinimizerSettings } from "../config/settings";
|
|
import { OutputSink } from "../session/streaming-output";
|
|
import { resolveOutputMaxColumns, resolveOutputSinkHeadBytes } from "../tools/output-meta";
|
|
import { getOrCreateSnapshot } from "../utils/shell-snapshot";
|
|
import { buildNonInteractiveEnv } from "./non-interactive-env";
|
|
|
|
export interface BashExecutorOptions {
|
|
cwd?: string;
|
|
/** Milliseconds before aborting the command; 0 disables the executor deadline. */
|
|
timeout?: number;
|
|
onChunk?: (chunk: string) => void;
|
|
chunkThrottleMs?: number;
|
|
signal?: AbortSignal;
|
|
/** Session key suffix to isolate shell sessions per agent */
|
|
sessionKey?: string;
|
|
/** Additional environment variables to inject */
|
|
env?: Record<string, string>;
|
|
/** Run through the configured user shell instead of brush parsing directly. */
|
|
useUserShell?: boolean;
|
|
/** Artifact path/id for full output storage */
|
|
artifactPath?: string;
|
|
artifactId?: string;
|
|
/**
|
|
* Invoked when the native minimizer rewrote the command's output, giving
|
|
* the caller a chance to persist the lossless original capture (typically
|
|
* via the session's `ArtifactManager`). The returned id is spliced into
|
|
* the sink output as `artifact://<id>` so the agent can retrieve the raw
|
|
* bytes. Return `undefined` to skip the footer.
|
|
*/
|
|
onMinimizedSave?: (
|
|
originalText: string,
|
|
info: { filter: string; inputBytes: number; outputBytes: number },
|
|
) => Promise<string | undefined>;
|
|
}
|
|
|
|
export interface BashResult {
|
|
output: string;
|
|
exitCode: number | undefined;
|
|
cancelled: boolean;
|
|
truncated: boolean;
|
|
totalLines: number;
|
|
totalBytes: number;
|
|
outputLines: number;
|
|
outputBytes: number;
|
|
artifactId?: string;
|
|
workingDir?: string;
|
|
}
|
|
|
|
const shellSessions = new Map<string, Shell>();
|
|
const brokenShellSessions = new Set<string>();
|
|
const shellSessionQuarantines = new Map<string, Promise<unknown>>();
|
|
/** Session keys with a command currently in flight on the persistent Shell. */
|
|
const shellSessionsInUse = new Set<string>();
|
|
|
|
/**
|
|
* Shells retained past their turn because a background (`nohup`/`&`) job is
|
|
* still running. A per-call `:async:` Shell is normally dropped at teardown,
|
|
* which SIGKILLs its children via kill-on-drop. Keeping the reference alive lets
|
|
* the process survive across turns; the Shell is dropped once its last
|
|
* background job exits (reaped by the poll loop below). Children stay
|
|
* kill-on-drop, so they still die when the harness tears the Shell down on exit.
|
|
*/
|
|
const retainedShells = new Set<Shell>();
|
|
const RETAIN_REAP_INTERVAL_MS = 5_000;
|
|
|
|
async function retainShellWithLiveBackgroundJobs(shell: Shell): Promise<void> {
|
|
let live: number;
|
|
try {
|
|
live = await shell.liveBackgroundJobCount();
|
|
} catch {
|
|
return;
|
|
}
|
|
if (live <= 0) return;
|
|
retainedShells.add(shell);
|
|
const interval = setInterval(() => {
|
|
void shell
|
|
.liveBackgroundJobCount()
|
|
.then(remaining => {
|
|
if (remaining > 0) return;
|
|
clearInterval(interval);
|
|
retainedShells.delete(shell);
|
|
})
|
|
.catch(() => {
|
|
clearInterval(interval);
|
|
retainedShells.delete(shell);
|
|
});
|
|
}, RETAIN_REAP_INTERVAL_MS);
|
|
interval.unref?.();
|
|
}
|
|
|
|
function quarantineShellSession(
|
|
sessionKey: string,
|
|
runPromise: Promise<ShellRunResult>,
|
|
abortCleanupPromise: Promise<void> | undefined,
|
|
): void {
|
|
brokenShellSessions.add(sessionKey);
|
|
const cleanup = abortCleanupPromise
|
|
? Promise.allSettled([runPromise, abortCleanupPromise])
|
|
: Promise.allSettled([runPromise]);
|
|
shellSessionQuarantines.set(sessionKey, cleanup);
|
|
void cleanup
|
|
.finally(() => {
|
|
if (shellSessionQuarantines.get(sessionKey) === cleanup) {
|
|
shellSessionQuarantines.delete(sessionKey);
|
|
brokenShellSessions.delete(sessionKey);
|
|
}
|
|
})
|
|
.catch(() => undefined);
|
|
}
|
|
|
|
function resolveShellCwd(cwd: string | undefined): string | undefined {
|
|
// Preserve the caller's logical cwd string. Brush uses this value to update `PWD` and its
|
|
// internal working directory, so realpathing here collapses symlinks before the shell sees them.
|
|
return cwd;
|
|
}
|
|
|
|
/** Translate `ShellMinimizerSettings` into native `MinimizerOptions`, or `undefined` when disabled. */
|
|
export function buildMinimizerOptions(group: ShellMinimizerSettings): MinimizerOptions | undefined {
|
|
if (!group.enabled) return undefined;
|
|
return {
|
|
enabled: true,
|
|
settingsPath: group.settingsPath || undefined,
|
|
only: group.only.length > 0 ? group.only : undefined,
|
|
except: group.except.length > 0 ? group.except : undefined,
|
|
maxCaptureBytes: group.maxCaptureBytes,
|
|
sourceOutlineLevel: group.sourceOutlineLevel === "default" ? undefined : group.sourceOutlineLevel,
|
|
legacyFilters: group.legacyFilters,
|
|
};
|
|
}
|
|
|
|
function shellBasename(shell: string): string {
|
|
return shell.replace(/\\/g, "/").split("/").pop()?.toLowerCase() ?? "";
|
|
}
|
|
|
|
function isBashShell(shell: string): boolean {
|
|
const basename = shellBasename(shell);
|
|
return basename.includes("bash");
|
|
}
|
|
|
|
function needsInteractiveShellArg(shell: string): boolean {
|
|
const basename = shellBasename(shell);
|
|
return basename.includes("zsh");
|
|
}
|
|
|
|
function supportsAutoUserShell(shell: string): boolean {
|
|
const basename = shellBasename(shell);
|
|
return basename.includes("bash") || basename.includes("zsh") || basename.includes("fish");
|
|
}
|
|
|
|
function hasInteractiveShellArg(args: string[]): boolean {
|
|
return args.some(arg => arg === "--interactive" || /^-[^-]*i/.test(arg));
|
|
}
|
|
|
|
function ensureInteractiveShellArgs(shell: string, args: string[]): string[] {
|
|
if (!needsInteractiveShellArg(shell) || hasInteractiveShellArg(args)) return args;
|
|
|
|
const commandIndex = args.findIndex(arg => arg === "-c" || arg === "--command");
|
|
if (commandIndex !== -1) {
|
|
return [...args.slice(0, commandIndex), "-i", ...args.slice(commandIndex)];
|
|
}
|
|
|
|
const compactCommandIndex = args.findIndex(arg => /^-[^-]*c[^-]*$/.test(arg));
|
|
if (compactCommandIndex !== -1) {
|
|
return args.map((arg, index) => (index === compactCommandIndex ? arg.replace("c", "ic") : arg));
|
|
}
|
|
|
|
return [...args, "-i"];
|
|
}
|
|
|
|
function quoteShellArg(value: string): string {
|
|
return `'${value.replace(/'/g, "'\\''")}'`;
|
|
}
|
|
|
|
function buildUserShellCommand(shell: string, args: string[], command: string): string {
|
|
return [shell, ...ensureInteractiveShellArgs(shell, args), command].map(quoteShellArg).join(" ");
|
|
}
|
|
|
|
function resolveUserShellConfig(settings: Settings, baseConfig: ShellConfig): ShellConfig {
|
|
const customShellPath = settings.get("shellPath");
|
|
const envShell = Bun.env.SHELL;
|
|
if (customShellPath || process.platform === "win32" || !envShell || envShell === baseConfig.shell) {
|
|
return baseConfig;
|
|
}
|
|
if (!supportsAutoUserShell(envShell) || !isExecutable(envShell)) {
|
|
return baseConfig;
|
|
}
|
|
|
|
return {
|
|
...baseConfig,
|
|
shell: envShell,
|
|
env: {
|
|
...baseConfig.env,
|
|
SHELL: envShell,
|
|
},
|
|
};
|
|
}
|
|
|
|
export async function executeBash(command: string, options?: BashExecutorOptions): Promise<BashResult> {
|
|
const settings = await Settings.init();
|
|
const baseShellConfig = settings.getShellConfig();
|
|
const shellConfig =
|
|
options?.useUserShell === true ? resolveUserShellConfig(settings, baseShellConfig) : baseShellConfig;
|
|
const { shell, args, env: shellEnv, prefix } = shellConfig;
|
|
const bashShell = isBashShell(shell);
|
|
const snapshotPath = bashShell ? await getOrCreateSnapshot(shell, shellEnv) : null;
|
|
|
|
const minimizer = buildMinimizerOptions(settings.getGroup("shellMinimizer"));
|
|
|
|
const commandCwd = resolveShellCwd(options?.cwd);
|
|
const commandEnv = buildNonInteractiveEnv(options?.env);
|
|
|
|
// Apply command prefix if configured
|
|
const prefixedCommand = prefix ? `${prefix} ${command}` : command;
|
|
const finalCommand =
|
|
options?.useUserShell === true && !bashShell
|
|
? buildUserShellCommand(shell, args, prefixedCommand)
|
|
: prefixedCommand;
|
|
|
|
// Create output sink for truncation and artifact handling
|
|
const sink = new OutputSink({
|
|
onChunk: options?.onChunk,
|
|
artifactPath: options?.artifactPath,
|
|
artifactId: options?.artifactId,
|
|
headBytes: resolveOutputSinkHeadBytes(settings),
|
|
maxColumns: resolveOutputMaxColumns(settings),
|
|
chunkThrottleMs: options?.onChunk ? (options.chunkThrottleMs ?? 50) : 0,
|
|
});
|
|
|
|
// sink.push() is synchronous — buffer management, counters, and onChunk
|
|
// all run inline. File writes (artifact path) are handled asynchronously
|
|
// inside the sink. No promise chain needed.
|
|
let acceptingChunks = true;
|
|
const enqueueChunk = (chunk: string) => {
|
|
if (acceptingChunks) sink.push(chunk);
|
|
};
|
|
|
|
if (options?.signal?.aborted) {
|
|
return {
|
|
exitCode: undefined,
|
|
cancelled: true,
|
|
...(await sink.dump("Command cancelled")),
|
|
};
|
|
}
|
|
|
|
const sessionKey = buildSessionKey(shell, prefix, snapshotPath, shellEnv, options?.sessionKey, minimizer);
|
|
const persistentSessionBroken = brokenShellSessions.has(sessionKey);
|
|
if (persistentSessionBroken) {
|
|
shellSessions.delete(sessionKey);
|
|
}
|
|
|
|
// A persistent Shell runs one command at a time (the native session is a
|
|
// mutex-guarded queue and `abort()` kills every in-flight run on it). When
|
|
// parallel bash calls overlap on the same key, the first one owns the
|
|
// persistent session; the rest degrade to isolated one-shot shells — the
|
|
// same path quarantined sessions take.
|
|
const sessionBusy = shellSessionsInUse.has(sessionKey);
|
|
let shellSession = persistentSessionBroken || sessionBusy ? undefined : shellSessions.get(sessionKey);
|
|
if (!shellSession && !persistentSessionBroken && !sessionBusy) {
|
|
shellSession = new Shell({
|
|
sessionEnv: shellEnv,
|
|
snapshotPath: snapshotPath ?? undefined,
|
|
minimizer,
|
|
});
|
|
shellSessions.set(sessionKey, shellSession);
|
|
}
|
|
const ownsPersistentSession = shellSession !== undefined;
|
|
if (ownsPersistentSession) {
|
|
shellSessionsInUse.add(sessionKey);
|
|
}
|
|
const userSignal = options?.signal;
|
|
const runAbortController = new AbortController();
|
|
let abortCleanupPromise: Promise<void> | undefined;
|
|
const abortCurrentExecution = () => {
|
|
if (!runAbortController.signal.aborted) {
|
|
runAbortController.abort();
|
|
}
|
|
if (shellSession && !abortCleanupPromise) {
|
|
abortCleanupPromise = shellSession.abort().catch(() => undefined);
|
|
}
|
|
};
|
|
const abortDeferred = Promise.withResolvers<"abort">();
|
|
const abortHandler = () => {
|
|
abortCurrentExecution();
|
|
abortDeferred.resolve("abort");
|
|
};
|
|
if (userSignal) {
|
|
userSignal.addEventListener("abort", abortHandler, { once: true });
|
|
}
|
|
|
|
let timeoutTimer: NodeJS.Timeout | undefined;
|
|
const timeoutDeferred = Promise.withResolvers<"timeout">();
|
|
const requestedTimeoutMs = options?.timeout;
|
|
const deadlineTimeoutMs = requestedTimeoutMs === 0 ? undefined : Math.max(1_000, requestedTimeoutMs ?? 300_000);
|
|
const nativeTimeoutMs = requestedTimeoutMs !== undefined && requestedTimeoutMs > 0 ? requestedTimeoutMs : undefined;
|
|
const nativeOwnsTimeout = nativeTimeoutMs !== undefined;
|
|
if (deadlineTimeoutMs !== undefined) {
|
|
timeoutTimer = setTimeout(() => {
|
|
// Explicit timeouts are already enforced inside pi-natives via
|
|
// `timeoutMs`. Do not also abort the JS AbortSignal here: on Windows,
|
|
// aborting that signal while a piped command is still forwarding output
|
|
// can terminate the Bun host before the native timeout result resolves.
|
|
if (!nativeOwnsTimeout) {
|
|
abortCurrentExecution();
|
|
}
|
|
timeoutDeferred.resolve("timeout");
|
|
}, deadlineTimeoutMs);
|
|
}
|
|
|
|
let resetSession = false;
|
|
|
|
try {
|
|
const runPromise = shellSession
|
|
? shellSession.run(
|
|
{
|
|
command: finalCommand,
|
|
cwd: commandCwd,
|
|
env: commandEnv,
|
|
timeoutMs: nativeTimeoutMs,
|
|
signal: runAbortController.signal,
|
|
},
|
|
(err, chunk) => {
|
|
if (!err) {
|
|
enqueueChunk(chunk);
|
|
}
|
|
},
|
|
)
|
|
: executeShell(
|
|
{
|
|
command: finalCommand,
|
|
cwd: commandCwd,
|
|
env: commandEnv,
|
|
sessionEnv: shellEnv,
|
|
snapshotPath: snapshotPath ?? undefined,
|
|
minimizer,
|
|
timeoutMs: nativeTimeoutMs,
|
|
signal: runAbortController.signal,
|
|
},
|
|
(err, chunk) => {
|
|
if (!err) {
|
|
enqueueChunk(chunk);
|
|
}
|
|
},
|
|
);
|
|
|
|
const ey = new ExponentialYield();
|
|
const winner = await ey.race<
|
|
{ kind: "result"; result: ShellRunResult } | { kind: "timeout" } | { kind: "abort" }
|
|
>([
|
|
runPromise.then(result => ({ kind: "result" as const, result })),
|
|
timeoutDeferred.promise.then(kind => ({ kind })),
|
|
abortDeferred.promise.then(kind => ({ kind })),
|
|
]);
|
|
|
|
if (winner.kind === "timeout" || winner.kind === "abort") {
|
|
acceptingChunks = false;
|
|
if (shellSession) {
|
|
resetSession = true;
|
|
quarantineShellSession(sessionKey, runPromise, abortCleanupPromise);
|
|
} else {
|
|
void runPromise.catch(() => undefined);
|
|
}
|
|
return {
|
|
exitCode: undefined,
|
|
cancelled: true,
|
|
...(await sink.dump(
|
|
winner.kind === "timeout" && deadlineTimeoutMs !== undefined
|
|
? `Command timed out after ${Math.round(deadlineTimeoutMs / 1000)} seconds`
|
|
: "Command cancelled",
|
|
)),
|
|
};
|
|
}
|
|
if (timeoutTimer) {
|
|
clearTimeout(timeoutTimer);
|
|
timeoutTimer = undefined;
|
|
}
|
|
|
|
// Handle timeout
|
|
if (winner.result.timedOut) {
|
|
const annotation = options?.timeout
|
|
? `Command timed out after ${Math.round(options.timeout / 1000)} seconds`
|
|
: "Command timed out";
|
|
resetSession = true;
|
|
if (shellSession) {
|
|
quarantineShellSession(sessionKey, runPromise, abortCleanupPromise);
|
|
}
|
|
return {
|
|
exitCode: undefined,
|
|
cancelled: true,
|
|
...(await sink.dump(annotation)),
|
|
};
|
|
}
|
|
|
|
// Handle cancellation
|
|
if (winner.result.cancelled) {
|
|
resetSession = true;
|
|
if (shellSession) {
|
|
quarantineShellSession(sessionKey, runPromise, abortCleanupPromise);
|
|
}
|
|
return {
|
|
exitCode: undefined,
|
|
cancelled: true,
|
|
...(await sink.dump("Command cancelled")),
|
|
};
|
|
}
|
|
|
|
// When the native minimizer rewrote the output, swap the sink's accumulated
|
|
// raw stream for the minimized text, persist the original as a session
|
|
// artifact, and splice an `artifact://<id>` footer into the visible text so
|
|
// the agent can retrieve the raw bytes losslessly.
|
|
const minimized = winner.result.minimized;
|
|
if (minimized && minimized.text !== minimized.originalText) {
|
|
sink.replace(minimized.text);
|
|
if (options?.onMinimizedSave) {
|
|
const artifactId = await options.onMinimizedSave(minimized.originalText, {
|
|
filter: minimized.filter,
|
|
inputBytes: minimized.inputBytes,
|
|
outputBytes: minimized.outputBytes,
|
|
});
|
|
if (artifactId) {
|
|
const sep = minimized.text.endsWith("\n") ? "" : "\n";
|
|
sink.push(`${sep}[raw output: artifact://${artifactId}]\n`);
|
|
}
|
|
}
|
|
}
|
|
|
|
// Normal completion
|
|
return {
|
|
exitCode: winner.result.exitCode,
|
|
cancelled: false,
|
|
workingDir: winner.result.workingDir,
|
|
...(await sink.dump()),
|
|
};
|
|
} catch (err) {
|
|
resetSession = true;
|
|
throw err;
|
|
} finally {
|
|
if (timeoutTimer) {
|
|
clearTimeout(timeoutTimer);
|
|
}
|
|
if (userSignal) {
|
|
userSignal.removeEventListener("abort", abortHandler);
|
|
}
|
|
if (ownsPersistentSession) {
|
|
shellSessionsInUse.delete(sessionKey);
|
|
if (resetSession || options?.sessionKey?.includes(":async:")) {
|
|
// `:async:` keys are per-job (jobId is unique), so the Shell would
|
|
// otherwise stay in the process-global map forever after completion.
|
|
shellSessions.delete(sessionKey);
|
|
// Dropping the only reference to a per-call `:async:` Shell SIGKILLs
|
|
// any `nohup`/`&` children (kill-on-drop). If the command left a live
|
|
// background job, retain the Shell so the process survives across
|
|
// turns; it is reaped once its last job exits and still dies with the
|
|
// harness. Skip on resetSession (cancel/error) — those tear down.
|
|
if (!resetSession && shellSession) {
|
|
await retainShellWithLiveBackgroundJobs(shellSession);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
function buildSessionKey(
|
|
shell: string,
|
|
prefix: string | undefined,
|
|
snapshotPath: string | null,
|
|
env: Record<string, string>,
|
|
agentSessionKey?: string,
|
|
minimizer?: MinimizerOptions,
|
|
): string {
|
|
const entries = Object.entries(env);
|
|
entries.sort(([a], [b]) => a.localeCompare(b));
|
|
const envSerialized = entries.map(([key, value]) => `${key}=${value}`).join("\n");
|
|
const minimizerSerialized = minimizer ? JSON.stringify(minimizer) : "";
|
|
return [agentSessionKey ?? "", shell, prefix ?? "", snapshotPath ?? "", envSerialized, minimizerSerialized].join(
|
|
"\n",
|
|
);
|
|
}
|