oh-my-pi/packages/coding-agent/src/utils/shell-snapshot.ts

/**
 * Shell environment snapshot for preserving user aliases, functions, and options.
 *
 * Creates a snapshot file that captures the user's shell environment from their
 * .bashrc/.zshrc, which can be sourced before each command to provide a familiar
 * shell experience.
 */
import * as fs from "node:fs";
import * as os from "node:os";
import * as path from "node:path";
import { postmortem } from "@oh-my-pi/pi-utils";

const cachedSnapshotPaths = new Map<string, string>();
const SNAPSHOT_TIMEOUT_MS = 2_000;

function sanitizeSnapshotEnv(env: Record<string, string | undefined>): Record<string, string | undefined> {
	const sanitized = { ...env };
	delete sanitized.BASH_ENV;
	delete sanitized.ENV;
	return sanitized;
}

/**
 * Get the user's shell config file path.
 */
function getShellConfigFile(shell: string): string {
	const home = os.homedir();
	if (shell.includes("zsh")) return path.join(home, ".zshrc");
	if (shell.includes("bash")) return path.join(home, ".bashrc");
	return path.join(home, ".profile");
}

/**
 * Generate the snapshot creation script.
 * This script sources the user's rc file and extracts functions, aliases, and options.
 * Matches Claude Code's snapshot generation logic.
 */
function generateSnapshotScript(shell: string, snapshotPath: string, rcFile: string): string {
	const hasRcFile = fs.existsSync(rcFile);
	const isZsh = shell.includes("zsh");
	const commonToolsRegex =
		"^(ls|dir|vdir|cat|head|tail|less|more|grep|egrep|fgrep|rg|find|fd|locate|sed|awk|perl|cp|mv|rm|mkdir|rmdir|touch|chmod|chown|ln|pwd|readlink|stat|cut|sort|uniq|xargs|tee|tr|basename|dirname)$";

	// Escape the snapshot path for shell
	const escapedPath = snapshotPath.replace(/'/g, "'\\''");

	// Function extraction differs between bash and zsh
	const functionScript = isZsh
		? `
echo "# Functions" >> "$SNAPSHOT_FILE"
# Force autoload all functions first
typeset -f > /dev/null 2>&1
# Get user function names - filter system/private ones
typeset +f 2>/dev/null | grep -vE '^(_|__)' | grep -vE '${commonToolsRegex}' | while read func; do
   typeset -f "$func" >> "$SNAPSHOT_FILE" 2>/dev/null
done
`
		: `
echo "# Functions" >> "$SNAPSHOT_FILE"
# Force autoload all functions first
declare -f > /dev/null 2>&1
# Get user function names - filter system/private ones, use base64 for special chars
declare -F 2>/dev/null | cut -d' ' -f3 | grep -vE '^(_|__)' | grep -vE '${commonToolsRegex}' | while read func; do
   encoded_func=$(declare -f "$func" | base64)
   echo "eval \\"\\$(echo '$encoded_func' | base64 -d)\\" > /dev/null 2>&1" >> "$SNAPSHOT_FILE"
done
`;

	// Shell options extraction
	const optionsScript = isZsh
		? `
echo "# Shell Options" >> "$SNAPSHOT_FILE"
setopt 2>/dev/null | sed 's/^/setopt /' | head -n 1000 >> "$SNAPSHOT_FILE"
`
		: `
echo "# Shell Options" >> "$SNAPSHOT_FILE"
shopt -p 2>/dev/null | head -n 1000 >> "$SNAPSHOT_FILE"
set -o 2>/dev/null | awk '$2 == "on" && $1 !~ /^(onecmd|monitor|restricted)$/ {print "set -o " $1}' | head -n 1000 >> "$SNAPSHOT_FILE"
echo "shopt -s expand_aliases" >> "$SNAPSHOT_FILE"
`;

	return `
SNAPSHOT_FILE='${escapedPath}'

# Source user's rc file if it exists
${hasRcFile ? `source "${rcFile}" < /dev/null 2>/dev/null` : "# No user config file to source"}

# Create/clear the snapshot file
echo "# Shell snapshot - generated by omp agent" >| "$SNAPSHOT_FILE"

# Unalias everything first to avoid conflicts when sourced
echo "unalias -a 2>/dev/null || true" >> "$SNAPSHOT_FILE"

${functionScript}

${optionsScript}

# Export aliases (limit to 1000)
echo "# Aliases" >> "$SNAPSHOT_FILE"
# Filter out winpty aliases on Windows to avoid "stdin is not a tty" errors
if [[ "$OSTYPE" == "msys" ]] || [[ "$OSTYPE" == "cygwin" ]]; then
   alias 2>/dev/null | grep -v "='winpty " | grep -vE '^alias (${commonToolsRegex})=' | sed 's/^alias //g' | sed 's/^/alias -- /' | head -n 1000 >> "$SNAPSHOT_FILE"
else
   alias 2>/dev/null | grep -vE '^alias (${commonToolsRegex})=' | sed 's/^alias //g' | sed 's/^/alias -- /' | head -n 1000 >> "$SNAPSHOT_FILE"
fi

# Export PATH
echo "export PATH='$PATH'" >> "$SNAPSHOT_FILE"

# Verify snapshot was created
if [ ! -f "$SNAPSHOT_FILE" ]; then
   echo "Error: Snapshot file was not created" >&2
   exit 1
fi
`.trim();
}

/**
 * Create a shell snapshot, caching the result.
 * Returns the path to the snapshot file, or null if creation failed.
 */
export async function getOrCreateSnapshot(
	shell: string,
	env: Record<string, string | undefined>,
): Promise<string | null> {
	const cacheKey = shell;
	// Return cached snapshot if valid
	const cached = cachedSnapshotPaths.get(cacheKey);
	if (cached && fs.existsSync(cached)) {
		return cached;
	}
	if (cached) {
		cachedSnapshotPaths.delete(cacheKey);
	}

	// Skip on Windows (no .bashrc in standard location)
	if (process.platform === "win32") {
		return null;
	}

	const rcFile = getShellConfigFile(shell);

	// Create snapshot directory
	const snapshotDir = path.join(os.tmpdir(), "omp-shell-snapshots");
	fs.mkdirSync(snapshotDir, { recursive: true });

	// Generate unique snapshot path
	const shellName = shell.includes("zsh") ? "zsh" : shell.includes("bash") ? "bash" : "sh";
	const snapshotPath = path.join(snapshotDir, `snapshot-${shellName}-${crypto.randomUUID()}.sh`);

	// Generate and execute snapshot script
	const script = generateSnapshotScript(shell, snapshotPath, rcFile);

	try {
		const snapshotEnv = sanitizeSnapshotEnv(env);
		const spawnEnv: Record<string, string> = {};
		for (const [key, value] of Object.entries(snapshotEnv)) {
			if (value !== undefined) {
				spawnEnv[key] = value;
			}
		}
		const child = Bun.spawn([shell, "-c", script], {
			env: spawnEnv,
			stdin: "ignore",
			stdout: "ignore",
			stderr: "ignore",
			timeout: SNAPSHOT_TIMEOUT_MS,
			killSignal: "SIGKILL",
		});

		await child.exited;
		if (child.exitCode === 0 && fs.existsSync(snapshotPath)) {
			cachedSnapshotPaths.set(cacheKey, snapshotPath);
			return snapshotPath;
		}
	} catch {
		// Snapshot creation failed, proceed without it
	}

	return null;
}

/**
 * Get the command prefix to source the snapshot.
 * Returns empty string if no snapshot available.
 */
export function getSnapshotSourceCommand(snapshotPath: string | null): string {
	if (!snapshotPath) return "";
	// Escape for shell
	const escaped = snapshotPath.replace(/'/g, "'\\''");
	return `source '${escaped}' 2>/dev/null && `;
}

postmortem.register("shell-snapshot", () => {
	for (const snapshotPath of cachedSnapshotPaths.values()) {
		fs.unlinkSync(snapshotPath);
	}
	cachedSnapshotPaths.clear();
});