diff --git a/src/agents/bash-process-registry.ts b/src/agents/bash-process-registry.ts index 278b11199..b32c265db 100644 --- a/src/agents/bash-process-registry.ts +++ b/src/agents/bash-process-registry.ts @@ -1,4 +1,5 @@ import type { ChildProcessWithoutNullStreams } from "node:child_process"; +import type { ExecEventSessionState } from "../infra/exec-events.js"; import { createSessionSlug as createSessionSlugId } from "./session-slug.js"; const DEFAULT_JOB_TTL_MS = 30 * 60 * 1000; // 30 minutes @@ -47,6 +48,7 @@ export interface ProcessSession { exited: boolean; truncated: boolean; backgrounded: boolean; + execEvents?: ExecEventSessionState; } export interface FinishedSession { diff --git a/src/agents/bash-tools.exec.ts b/src/agents/bash-tools.exec.ts index b9de81872..a99acb8d2 100644 --- a/src/agents/bash-tools.exec.ts +++ b/src/agents/bash-tools.exec.ts @@ -4,6 +4,7 @@ import path from "node:path"; import type { AgentTool, AgentToolResult } from "@mariozechner/pi-agent-core"; import { Type } from "@sinclair/typebox"; +import { loadConfig } from "../config/config.js"; import { type ExecAsk, type ExecHost, @@ -19,6 +20,16 @@ import { resolveExecApprovals, resolveExecApprovalsFromFile, } from "../infra/exec-approvals.js"; +import { getExecEventContext } from "../infra/exec-events-context.js"; +import { + createThrottledExecOutputBuffer, + emitExecEvent, + matchExecCommandAgainstWhitelist, + resolveExecEventsConfig, + toExecEventRuntimeConfig, + type ExecEventContext, + type ExecEventSessionState, +} from "../infra/exec-events.js"; import { requestHeartbeatNow } from "../infra/heartbeat-wake.js"; import { buildNodeShellCommand } from "../infra/node-shell.js"; import { @@ -251,6 +262,21 @@ function normalizeNotifyOutput(value: string) { return value.replace(/\s+/g, " ").trim(); } +function normalizeExecEventContext( + context: ExecEventContext | undefined, +): ExecEventContext | undefined { + if (!context) return undefined; + const runId = context.runId?.trim(); + const toolCallId = context.toolCallId?.trim(); + const sessionKey = context.sessionKey?.trim(); + if (!runId && !toolCallId && !sessionKey) return undefined; + return { + ...(runId ? { runId } : {}), + ...(toolCallId ? { toolCallId } : {}), + ...(sessionKey ? { sessionKey } : {}), + } satisfies ExecEventContext; +} + function normalizePathPrepend(entries?: string[]) { if (!Array.isArray(entries)) return []; const seen = new Set(); @@ -361,6 +387,16 @@ async function runExecProcess(opts: { let child: ChildProcessWithoutNullStreams | null = null; let pty: PtyHandle | null = null; let stdin: SessionStdin | undefined; + const execEventsConfig = resolveExecEventsConfig(loadConfig()); + const execEventsMatch = + execEventsConfig.emitEvents === true + ? matchExecCommandAgainstWhitelist(opts.command, execEventsConfig) + : { matched: false }; + const execEventsEnabled = execEventsConfig.emitEvents === true && execEventsMatch.matched; + const execEventContext = normalizeExecEventContext(getExecEventContext()); + const execEventsRuntimeConfig = execEventsEnabled + ? toExecEventRuntimeConfig(execEventsConfig) + : undefined; if (opts.sandbox) { const { child: spawned } = await spawnWithFallback({ @@ -517,9 +553,88 @@ async function runExecProcess(opts: { exitSignal: undefined as NodeJS.Signals | number | null | undefined, truncated: false, backgrounded: false, + execEvents: { + enabled: execEventsEnabled, + commandName: execEventsMatch.commandName, + context: execEventContext, + config: execEventsRuntimeConfig, + startedEmitted: false, + completedEmitted: false, + } satisfies ExecEventSessionState, } satisfies ProcessSession; addSession(session); + const execEventsState = session.execEvents; + if (execEventsState?.enabled && execEventsState.config) { + execEventsState.outputBuffer = createThrottledExecOutputBuffer({ + throttleMs: execEventsState.config.outputThrottleMs, + maxChunkBytes: execEventsState.config.outputMaxChunkBytes, + maxBufferBytes: execEventsState.config.outputBufferMaxBytes, + onFlush: (chunks) => { + if (!execEventsState.enabled) return; + for (const chunk of chunks) { + emitExecEvent("exec.output", { + sessionId, + pid: session.pid, + command: session.command, + commandName: execEventsState.commandName, + context: execEventsState.context, + stream: chunk.stream, + output: chunk.output, + ...(chunk.truncated ? { truncated: true } : {}), + }); + } + }, + }); + } + + const emitExecStarted = () => { + if (!execEventsState?.enabled || execEventsState.startedEmitted) return; + execEventsState.startedEmitted = true; + emitExecEvent("exec.started", { + sessionId, + pid: session.pid, + command: session.command, + commandName: execEventsState.commandName, + context: execEventsState.context, + startedAt, + cwd: session.cwd, + }); + }; + + const finalizeExecEvents = (params: { + status: "completed" | "failed"; + exitCode: number | null; + exitSignal: NodeJS.Signals | number | null; + timedOut?: boolean; + reason?: string; + }) => { + if (!execEventsState?.enabled || execEventsState.completedEmitted) return; + execEventsState.completedEmitted = true; + execEventsState.outputBuffer?.flushAll(); + execEventsState.outputBuffer?.dispose(); + const completedAt = Date.now(); + emitExecEvent("exec.completed", { + sessionId, + pid: session.pid, + command: session.command, + commandName: execEventsState.commandName, + context: execEventsState.context, + startedAt, + completedAt, + durationMs: completedAt - startedAt, + status: params.status, + exitCode: params.exitCode, + exitSignal: params.exitSignal, + ...(params.timedOut ? { timedOut: true } : {}), + ...(params.reason ? { reason: params.reason } : {}), + }); + }; + + emitExecStarted(); + + const execOutputBuffer = execEventsState?.outputBuffer; + let settled = false; let timeoutTimer: NodeJS.Timeout | null = null; let timeoutFinalizeTimer: NodeJS.Timeout | null = null; @@ -539,6 +654,13 @@ async function runExecProcess(opts: { maybeNotifyOnExit(session, "failed"); const aggregated = session.aggregated.trim(); const reason = `Command timed out after ${opts.timeoutSec} seconds`; + finalizeExecEvents({ + status: "failed", + exitCode: null, + exitSignal: "SIGKILL", + timedOut: true, + reason, + }); settle({ status: "failed", exitCode: null, @@ -587,6 +709,9 @@ async function runExecProcess(opts: { const str = sanitizeBinaryOutput(data.toString()); for (const chunk of chunkString(str)) { appendOutput(session, "stdout", chunk); + if (execOutputBuffer) { + execOutputBuffer.append("stdout", chunk); + } emitUpdate(); } }; @@ -595,6 +720,9 @@ async function runExecProcess(opts: { const str = sanitizeBinaryOutput(data.toString()); for (const chunk of chunkString(str)) { appendOutput(session, "stderr", chunk); + if (execOutputBuffer) { + execOutputBuffer.append("stderr", chunk); + } emitUpdate(); } }; @@ -625,23 +753,32 @@ async function runExecProcess(opts: { const wasSignal = exitSignal != null; const isSuccess = code === 0 && !wasSignal && !timedOut; const status: "completed" | "failed" = isSuccess ? "completed" : "failed"; - markExited(session, code, exitSignal, status); - maybeNotifyOnExit(session, status); - if (!session.child && session.stdin) { - session.stdin.destroyed = true; - } - - if (settled) return; - const aggregated = session.aggregated.trim(); - if (!isSuccess) { - const reason = timedOut + const failureReason = !isSuccess + ? timedOut ? `Command timed out after ${opts.timeoutSec} seconds` : wasSignal && exitSignal ? `Command aborted by signal ${exitSignal}` : code === null ? "Command aborted before exit code was captured" - : `Command exited with code ${code}`; - const message = aggregated ? `${aggregated}\n\n${reason}` : reason; + : `Command exited with code ${code}` + : undefined; + markExited(session, code, exitSignal, status); + maybeNotifyOnExit(session, status); + if (!session.child && session.stdin) { + session.stdin.destroyed = true; + } + finalizeExecEvents({ + status, + exitCode: code ?? null, + exitSignal: exitSignal ?? null, + ...(timedOut ? { timedOut: true } : {}), + ...(failureReason ? { reason: failureReason } : {}), + }); + + if (settled) return; + const aggregated = session.aggregated.trim(); + if (!isSuccess) { + const message = aggregated ? `${aggregated}\n\n${failureReason}` : failureReason; settle({ status: "failed", exitCode: code ?? null, @@ -649,7 +786,7 @@ async function runExecProcess(opts: { durationMs, aggregated, timedOut, - reason: message, + reason: message ?? undefined, }); return; } @@ -679,8 +816,16 @@ async function runExecProcess(opts: { if (timeoutFinalizeTimer) clearTimeout(timeoutFinalizeTimer); markExited(session, null, null, "failed"); maybeNotifyOnExit(session, "failed"); + const reason = String(err); + finalizeExecEvents({ + status: "failed", + exitCode: null, + exitSignal: null, + ...(timedOut ? { timedOut: true } : {}), + reason, + }); const aggregated = session.aggregated.trim(); - const message = aggregated ? `${aggregated}\n\n${String(err)}` : String(err); + const message = aggregated ? `${aggregated}\n\n${reason}` : reason; settle({ status: "failed", exitCode: null, diff --git a/src/infra/exec-events.ts b/src/infra/exec-events.ts index 95a329912..2fc9e9e45 100644 --- a/src/infra/exec-events.ts +++ b/src/infra/exec-events.ts @@ -116,6 +116,13 @@ const WRAPPER_COMMANDS = new Set([ const WRAPPER_SUBCOMMANDS = new Set(["exec", "run", "run-script", "dlx", "x"]); +type ExecEventsConfigInput = { + emitEvents?: boolean; + commandWhitelist?: string[]; + outputThrottleMs?: number; + outputMaxChunkBytes?: number; +}; + type WhitelistMatchResult = { matched: boolean; commandName?: string; @@ -131,6 +138,21 @@ export type ExecEventsConfigResolved = { outputBufferMaxBytes: number; }; +export type ExecEventRuntimeConfig = Pick< + ExecEventsConfigResolved, + "outputThrottleMs" | "outputMaxChunkBytes" | "outputBufferMaxBytes" +>; + +export type ExecEventSessionState = { + enabled: boolean; + commandName?: string; + context?: ExecEventContext; + config?: ExecEventRuntimeConfig; + startedEmitted: boolean; + completedEmitted: boolean; + outputBuffer?: ExecOutputBuffer; +}; + function normalizeCommandName(value: string): string { return path.basename(value.trim()).toLowerCase(); } @@ -158,7 +180,8 @@ function resolvePositiveInt(value: unknown, fallback: number, min?: number): num } export function resolveExecEventsConfig(cfg?: MoltbotConfig): ExecEventsConfigResolved { - const execCfg = cfg?.hooks?.exec; + const hooks = cfg?.hooks as { exec?: ExecEventsConfigInput } | undefined; + const execCfg = hooks?.exec; const commandWhitelist = normalizeCommandWhitelist(execCfg?.commandWhitelist); const commandWhitelistSet = new Set(commandWhitelist); const outputThrottleMs = resolvePositiveInt(execCfg?.outputThrottleMs, DEFAULT_OUTPUT_THROTTLE_MS); @@ -183,6 +206,14 @@ export function resolveExecEventsConfig(cfg?: MoltbotConfig): ExecEventsConfigRe } satisfies ExecEventsConfigResolved; } +export function toExecEventRuntimeConfig(cfg: ExecEventsConfigResolved): ExecEventRuntimeConfig { + return { + outputThrottleMs: cfg.outputThrottleMs, + outputMaxChunkBytes: cfg.outputMaxChunkBytes, + outputBufferMaxBytes: cfg.outputBufferMaxBytes, + }; +} + function includesWhitelistToken(command: string, tokens: string[]): boolean { if (!command || tokens.length === 0) return false; const lower = command.toLowerCase(); @@ -381,19 +412,34 @@ function sliceByMaxBytes(text: string, maxBytes: number): string { return text.slice(0, low); } +function trimSuffixByMaxBytes(text: string, maxBytes: number): string { + if (!text) return ""; + if (maxBytes <= 0) return text.slice(-1); + if (byteLength(text) <= maxBytes) return text; + + let low = 0; + let high = text.length; + while (low < high) { + const mid = Math.floor((low + high) / 2); + const suffix = text.slice(mid); + const bytes = byteLength(suffix); + if (bytes > maxBytes) { + low = mid + 1; + } else { + high = mid; + } + } + + const trimmed = text.slice(low); + return trimmed || text.slice(-1); +} + function capBuffer(state: OutputStreamState, maxBufferBytes: number): void { const totalBytes = byteLength(state.buffer); if (totalBytes <= maxBufferBytes) return; state.truncated = true; - // Drop oldest data until we're within the cap. - let remaining = state.buffer; - while (remaining && byteLength(remaining) > maxBufferBytes) { - const overflowBytes = byteLength(remaining) - maxBufferBytes; - const dropChars = Math.min(remaining.length, Math.max(1, Math.floor(overflowBytes / 2))); - remaining = remaining.slice(dropChars); - } - state.buffer = remaining; + state.buffer = trimSuffixByMaxBytes(state.buffer, maxBufferBytes); } export function createThrottledExecOutputBuffer(options: OutputBufferOptions): ExecOutputBuffer { @@ -405,6 +451,7 @@ export function createThrottledExecOutputBuffer(options: OutputBufferOptions): E const stderr: OutputStreamState = { buffer: "", truncated: false }; let timer: NodeJS.Timeout | null = null; let disposed = false; + let flushingAll = false; const getState = (stream: OutputStream) => (stream === "stdout" ? stdout : stderr); @@ -440,25 +487,24 @@ export function createThrottledExecOutputBuffer(options: OutputBufferOptions): E if (chunks.length > 0) { options.onFlush(chunks); } - if (!disposed && throttleMs > 0 && (stdout.buffer || stderr.buffer)) { + if (!disposed && !flushingAll && throttleMs > 0 && (stdout.buffer || stderr.buffer)) { schedule(); } }; const flushAll = () => { if (disposed) return; - if (timer) { - clearTimeout(timer); - timer = null; - } - while (stdout.buffer || stderr.buffer) { - flush(); - if (throttleMs > 0 && (stdout.buffer || stderr.buffer)) { - // Prevent tight loops in extremely large output bursts by yielding back to the event loop. - // The next scheduled flush will continue draining. - schedule(); - return; + flushingAll = true; + try { + if (timer) { + clearTimeout(timer); + timer = null; } + while (!disposed && (stdout.buffer || stderr.buffer)) { + flush(); + } + } finally { + flushingAll = false; } };