can1357--oh-my-pi/packages/coding-agent/test/shell-snapshot.test.ts

import { describe, expect, it } from "bun:test";
import { existsSync } from "node:fs";
import * as fs from "node:fs/promises";
import * as os from "node:os";
import * as path from "node:path";
import { getOrCreateSnapshot, sanitizeSnapshotForBrush } from "@oh-my-pi/pi-coding-agent/utils/shell-snapshot";
import fnEnvHelper from "../src/utils/shell-snapshot-fn-env.sh" with { type: "text" };

// macOS ships bash at `/bin/bash`, not `/usr/bin/bash`; resolve a real bash the
// same way `bash-executor.test.ts` does so these e2e tests stay portable. The
// `getOrCreateSnapshot` tests below symlink this under a unique name and skip
// when no bash is present.
const REAL_BASH = Bun.env.SHELL?.includes("bash") ? Bun.env.SHELL : "/bin/bash";
// Likewise resolve `echo` (the stand-in for the mise binary the captured
// function invokes): macOS has no `/usr/bin/echo`, so hard-coding it makes the
// replay fail with `No such file or directory` even though the export landed.
const REAL_ECHO = Bun.which("echo") ?? "/bin/echo";

// `sanitizeSnapshotForBrush` is the snapshot-side mitigation for brush's
// whitespace-only alias expander (`crates/vendor/brush-core/src/interp.rs:1500`,
// brush issue reubeno/brush#57). Aliases whose body needs real shell parsing
// must be dropped or brush turns the first whitespace piece into the command
// name and the user sees `error: command not found: (alias;` (issue #3234).

describe("sanitizeSnapshotForBrush", () => {
	it("drops the Fedora `which` alias and reports the dropped name", () => {
		const snapshot = [
			"unalias -a 2>/dev/null || true",
			"alias -- which='(alias; declare -f) | /usr/bin/which --tty-only --read-alias --show-dot --show-tilde'",
			"alias -- ll='ls -l'",
			"",
		].join("\n");

		const result = sanitizeSnapshotForBrush(snapshot);

		expect(result.dropped).toEqual(["which"]);
		expect(result.content).not.toContain("alias -- which=");
		// Simple aliases must still pass through untouched.
		expect(result.content).toContain("alias -- ll='ls -l'");
	});

	it.each([
		["subshell", "alias -- a='(echo hi)'"],
		["pipe", "alias -- a='cat /etc/hostname | head -n1'"],
		["semicolon", "alias -- a='echo a; echo b'"],
		["redirect-out", "alias -- a='tee >foo'"],
		["redirect-in", "alias -- a='cat <foo'"],
		["background", "alias -- a='sleep 1 &'"],
		["command-substitution", "alias -- a='echo `date`'"],
		// Embedded single-quote `'\''` → `'` decodes to a body with a literal
		// single quote (and a pipe), so the filter still trips on the pipe.
		["decoded-pipe", "alias -- a='cat '\\''one'\\'' | head'"],
	])("drops aliases whose body needs shell parsing (%s)", (_label, line) => {
		const result = sanitizeSnapshotForBrush(`${line}\n`);
		expect(result.dropped).toEqual(["a"]);
		expect(result.content).not.toContain("alias -- a=");
	});

	it.each([
		["simple", "alias -- ll='ls -l'"],
		["flag-with-equals", "alias -- gc='git --color=auto commit'"],
		["multi-flag", "alias -- la='ls -lAh --group-directories-first'"],
		// A plain single quote escape that decodes to a metachar-free body
		// must survive — we only ban truly unparseable bodies.
		["embedded-quote", "alias -- say='echo '\\''hello'\\'''"],
	])("preserves aliases brush can handle by whitespace split (%s)", (_label, line) => {
		const result = sanitizeSnapshotForBrush(`${line}\n`);
		expect(result.dropped).toEqual([]);
		expect(result.content).toContain(line);
	});

	it("leaves non-alias lines (functions, exports, unalias) untouched", () => {
		const snapshot = [
			"# Shell snapshot - generated by omp agent",
			"unalias -a 2>/dev/null || true",
			"my_fn () { echo hi; }",
			"export PATH='/usr/bin:/bin'",
			"shopt -s expand_aliases",
			"alias -- which='(alias; declare -f) | /usr/bin/which'", // poisoned
		].join("\n");

		const result = sanitizeSnapshotForBrush(snapshot);

		expect(result.dropped).toEqual(["which"]);
		for (const keep of [
			"# Shell snapshot - generated by omp agent",
			"unalias -a 2>/dev/null || true",
			"my_fn () { echo hi; }", // function body contains `;` but is not an alias line
			"export PATH='/usr/bin:/bin'",
			"shopt -s expand_aliases",
		]) {
			expect(result.content).toContain(keep);
		}
	});
});

// `__omp_emit_referenced_exports` (shell-snapshot-fn-env.sh) re-exports env
// vars that snapshotted functions reference. mise activates a `mise()` shell
// function whose body expands `$__MISE_EXE`; the snapshot used to persist the
// function but discard the sidecar var, so the replay shell ran
// `command "" "$@"` and died with `command: command not found:` (issue #3470).

async function readStream(stream: ReadableStream<Uint8Array> | null): Promise<string> {
	if (!stream) return "";
	return await new Response(stream).text();
}

describe("shell-snapshot fn-env helper", () => {
	it("emits export lines for env vars referenced by captured functions, skips unset and shell-internal names", async () => {
		const funcs = [
			`mise () { command "$__MISE_EXE" "$@"; }`,
			// biome-ignore lint/suspicious/noTemplateCurlyInString: literal shell parameter expansion `${FOO_TEST_DIR}`
			'my_fn () { echo "$FOO_TEST_DIR ${FOO_TEST_DIR}"; }',
			`uses_path () { echo "$PATH"; }`,
			`uses_locale () { echo "$LC_ALL"; }`,
			`secret () { : "$NEVER_SET_TEST_VAR"; }`,
			``,
		].join("\n");

		const child = Bun.spawn(["bash", "-c", `${fnEnvHelper}\n__omp_emit_referenced_exports`], {
			env: {
				PATH: process.env.PATH ?? "/usr/bin:/bin",
				__MISE_EXE: "/opt/echo",
				FOO_TEST_DIR: "/opt/dir",
				LC_ALL: "C",
			},
			stdin: "pipe",
			stdout: "pipe",
			stderr: "pipe",
		});
		child.stdin.write(funcs);
		await child.stdin.end();
		const out = await readStream(child.stdout as ReadableStream<Uint8Array> | null);
		await child.exited;
		expect(child.exitCode).toBe(0);

		expect(out).toContain("export __MISE_EXE='/opt/echo'");
		expect(out).toContain("export FOO_TEST_DIR='/opt/dir'");
		// Shell-internal names must never be re-exported.
		expect(out).not.toMatch(/^export PATH=/m);
		expect(out).not.toMatch(/^export LC_ALL=/m);
		// Unset names produce no line.
		expect(out).not.toContain("NEVER_SET_TEST_VAR");
	});

	it("never emits export lines for likely-secret env var names", async () => {
		const funcs = [
			`deploy () { curl -H "Authorization: $GITHUB_TOKEN" .; }`,
			`call_openai () { curl -H "Authorization: Bearer $OPENAI_API_KEY" .; }`,
			`aws_sign () { echo "$AWS_SECRET_ACCESS_KEY"; }`,
			`db () { mysql --password="$DB_PASSWORD" -u root; }`,
			`legacy () { echo "$LDAP_PASSWD"; }`,
			`vault () { echo "$VAULT_PRIVATE_KEY"; }`,
			`session () { echo "$REDIS_SESSION_KEY"; }`,
			`creds () { echo "$AZURE_CREDENTIAL"; }`,
			// Control: non-secret-shaped var must still be emitted.
			`mise () { command "$__MISE_EXE" "$@"; }`,
			``,
		].join("\n");

		const child = Bun.spawn(["bash", "-c", `${fnEnvHelper}\n__omp_emit_referenced_exports`], {
			env: {
				PATH: process.env.PATH ?? "/usr/bin:/bin",
				GITHUB_TOKEN: "ghp_REDACTED",
				OPENAI_API_KEY: "sk-REDACTED",
				AWS_SECRET_ACCESS_KEY: "REDACTED",
				DB_PASSWORD: "hunter2",
				LDAP_PASSWD: "hunter2",
				VAULT_PRIVATE_KEY: "-----BEGIN-----",
				REDIS_SESSION_KEY: "abc",
				AZURE_CREDENTIAL: "xyz",
				__MISE_EXE: "/opt/echo",
			},
			stdin: "pipe",
			stdout: "pipe",
			stderr: "ignore",
		});
		child.stdin.write(funcs);
		await child.stdin.end();
		const out = await readStream(child.stdout as ReadableStream<Uint8Array> | null);
		await child.exited;
		expect(child.exitCode).toBe(0);

		for (const secret of [
			"GITHUB_TOKEN",
			"OPENAI_API_KEY",
			"AWS_SECRET_ACCESS_KEY",
			"DB_PASSWORD",
			"LDAP_PASSWD",
			"VAULT_PRIVATE_KEY",
			"REDIS_SESSION_KEY",
			"AZURE_CREDENTIAL",
		]) {
			expect(out).not.toContain(secret);
		}
		// And the secret VALUES — make sure nothing leaked through a different
		// quoting path.
		for (const value of ["ghp_REDACTED", "sk-REDACTED", "hunter2", "-----BEGIN-----"]) {
			expect(out).not.toContain(value);
		}
		// The non-secret helper var still goes through.
		expect(out).toContain("export __MISE_EXE='/opt/echo'");
	});

	it("single-quote-escapes values containing apostrophes and preserves newlines", async () => {
		const funcs = `shout () { echo "$TRICKY_VAL $NL_VAL"; }\n`;
		const child = Bun.spawn(["bash", "-c", `${fnEnvHelper}\n__omp_emit_referenced_exports`], {
			env: {
				PATH: process.env.PATH ?? "/usr/bin:/bin",
				TRICKY_VAL: "it's 'tricky'",
				NL_VAL: "line1\nline2",
			},
			stdin: "pipe",
			stdout: "pipe",
			stderr: "ignore",
		});
		child.stdin.write(funcs);
		await child.stdin.end();
		const out = await readStream(child.stdout as ReadableStream<Uint8Array> | null);
		await child.exited;

		expect(out).toContain(`export TRICKY_VAL='it'\\''s '\\''tricky'\\'''`);
		expect(out).toContain(`export NL_VAL='line1\nline2'`);

		// Eval the emitted lines and verify the round-trip values match.
		const round = Bun.spawn(
			["bash", "-c", `eval "$1"; printf '%s\\n' "$TRICKY_VAL"; printf '%s\\n' "$NL_VAL"`, "_", out],
			{ stdout: "pipe", stderr: "ignore" },
		);
		const echoed = await readStream(round.stdout as ReadableStream<Uint8Array> | null);
		await round.exited;
		expect(echoed).toBe("it's 'tricky'\nline1\nline2\n");
	});
});

describe("getOrCreateSnapshot", () => {
	it("re-exports env vars referenced by snapshotted functions (issue #3470)", async () => {
		const home = await fs.mkdtemp(path.join(os.tmpdir(), "omp-snap-3470-"));
		await fs.writeFile(
			path.join(home, ".bashrc"),
			[
				`export __MISE_EXE=${REAL_ECHO}`,
				`export FOO_TEST_DIR=/opt/foo`,
				`mise () { command "$__MISE_EXE" "$@"; }`,
				`shout () { echo "$FOO_TEST_DIR"; }`,
				``,
			].join("\n"),
		);

		// Symlink bash under a unique path so this call misses the process-wide
		// snapshot cache (`cachedSnapshotPaths` is keyed on the shell path, and
		// sibling tests in the same worker create a cached `/usr/bin/bash` entry
		// from a different HOME before this test runs).
		const realBash = REAL_BASH;
		if (!existsSync(realBash)) return;
		const shellLink = path.join(home, "bash-omp-3470");
		await fs.symlink(realBash, shellLink);

		// `getShellConfigFile` honours `env.HOME` so the spawned shell sources
		// our fixture bashrc rather than the test runner's home.
		const env = { ...process.env, HOME: home };
		const snapshotPath = await getOrCreateSnapshot(shellLink, env);
		expect(snapshotPath).not.toBeNull();
		const content = await fs.readFile(snapshotPath!, "utf8");

		expect(content).toContain(`export __MISE_EXE='${REAL_ECHO}'`);
		expect(content).toContain(`export FOO_TEST_DIR='/opt/foo'`);

		// Replay the snapshot in a fresh bash with `set -u` and confirm the
		// mise() function resolves cleanly instead of dying on the empty var.
		const replay = Bun.spawn(
			[realBash, "--noprofile", "--norc", "-c", `set -u; source "$1"; mise hello world; shout`, "_", snapshotPath!],
			{ stdout: "pipe", stderr: "pipe" },
		);
		const stdout = await readStream(replay.stdout as ReadableStream<Uint8Array> | null);
		const stderr = await readStream(replay.stderr as ReadableStream<Uint8Array> | null);
		await replay.exited;
		expect({ exitCode: replay.exitCode, stderr }).toEqual({ exitCode: 0, stderr: "" });
		expect(stdout).toBe("hello world\n/opt/foo\n");

		// PR-review hardening: snapshot file must be group/world-unreadable since
		// it now inlines env-var values. Directory must be 0700 for the same
		// reason — UUID filenames shouldn't leak via `ls /tmp/omp-shell-snapshots`.
		const fileStat = await fs.stat(snapshotPath!);
		expect(fileStat.mode & 0o077).toBe(0);
		const dirStat = await fs.stat(path.dirname(snapshotPath!));
		expect(dirStat.mode & 0o077).toBe(0);
	});

	it("keeps the snapshot file at 0600 even when the rc file resets umask to 022", async () => {
		// PR-review regression: previous revision ran `umask 077` BEFORE sourcing
		// the rc, so a typical `.bashrc` with `umask 022` reopened the world-read
		// window between the shell's first `>|` and the JS post-spawn chmod.
		// Fix: JS now pre-creates the file at 0600 (shell `>|`/`>>` preserve the
		// inode mode) AND the script re-applies `umask 077` after the source.
		const home = await fs.mkdtemp(path.join(os.tmpdir(), "omp-snap-umask-"));
		await fs.writeFile(
			path.join(home, ".bashrc"),
			[`umask 022`, `export __MISE_EXE=${REAL_ECHO}`, `mise () { command "$__MISE_EXE" "$@"; }`, ``].join("\n"),
		);

		const realBash = REAL_BASH;
		if (!existsSync(realBash)) return;
		const shellLink = path.join(home, "bash-omp-umask");
		await fs.symlink(realBash, shellLink);

		const env = { ...process.env, HOME: home };
		const snapshotPath = await getOrCreateSnapshot(shellLink, env);
		expect(snapshotPath).not.toBeNull();

		// Snapshot file must be 0600 — verifies the mode survives an rc-injected
		// `umask 022` AND that captured env was actually written (sanity: non-empty).
		const fileStat = await fs.stat(snapshotPath!);
		expect(fileStat.mode & 0o077).toBe(0);
		const content = await fs.readFile(snapshotPath!, "utf8");
		expect(content).toContain(`export __MISE_EXE='${REAL_ECHO}'`);
	});
	it("cleans up the empty snapshot file when the shell exits with a non-zero code", async () => {
		const testRoot = await fs.mkdtemp(path.join(os.tmpdir(), "omp-snap-fail-"));
		const originalTmpDir = process.env.TMPDIR;
		process.env.TMPDIR = testRoot;
		try {
			const fakeShell = path.join(testRoot, "fail-shell.sh");
			await fs.writeFile(fakeShell, `#!/bin/sh\nexit 1\n`);
			await fs.chmod(fakeShell, 0o755);

			const env = { ...process.env, HOME: testRoot };
			const snapshotDir = path.join(testRoot, "omp-shell-snapshots");

			const snapshotPath = await getOrCreateSnapshot(fakeShell, env);
			expect(snapshotPath).toBeNull();

			// Verify that no snapshot files were left behind in the isolated tmpdir
			if (existsSync(snapshotDir)) {
				const files = await fs.readdir(snapshotDir);
				expect(files).toEqual([]);
			}
		} finally {
			if (originalTmpDir === undefined) delete process.env.TMPDIR;
			else process.env.TMPDIR = originalTmpDir;
			await fs.rm(testRoot, { recursive: true, force: true });
		}
	});

	it("cleans up the empty snapshot file when the shell fails to spawn", async () => {
		const testRoot = await fs.mkdtemp(path.join(os.tmpdir(), "omp-snap-spawn-fail-"));
		const originalTmpDir = process.env.TMPDIR;
		process.env.TMPDIR = testRoot;
		try {
			const fakeShell = path.join(testRoot, "does-not-exist-shell");

			const env = { ...process.env, HOME: testRoot };
			const snapshotDir = path.join(testRoot, "omp-shell-snapshots");

			const snapshotPath = await getOrCreateSnapshot(fakeShell, env);
			expect(snapshotPath).toBeNull();

			if (existsSync(snapshotDir)) {
				const files = await fs.readdir(snapshotDir);
				expect(files).toEqual([]);
			}
		} finally {
			if (originalTmpDir === undefined) delete process.env.TMPDIR;
			else process.env.TMPDIR = originalTmpDir;
			await fs.rm(testRoot, { recursive: true, force: true });
		}
	});

	it("cleans up the empty snapshot file when the shell execution times out", async () => {
		const testRoot = await fs.mkdtemp(path.join(os.tmpdir(), "omp-snap-timeout-"));
		const originalTmpDir = process.env.TMPDIR;
		process.env.TMPDIR = testRoot;
		try {
			const fakeShell = path.join(testRoot, "timeout-shell.sh");
			// Sleep longer than SNAPSHOT_TIMEOUT_MS (2000)
			await fs.writeFile(fakeShell, `#!/bin/sh\nsleep 3\n`);
			await fs.chmod(fakeShell, 0o755);

			const env = { ...process.env, HOME: testRoot };
			const snapshotDir = path.join(testRoot, "omp-shell-snapshots");

			const snapshotPath = await getOrCreateSnapshot(fakeShell, env);
			expect(snapshotPath).toBeNull();

			if (existsSync(snapshotDir)) {
				const files = await fs.readdir(snapshotDir);
				expect(files).toEqual([]);
			}
		} finally {
			if (originalTmpDir === undefined) delete process.env.TMPDIR;
			else process.env.TMPDIR = originalTmpDir;
			await fs.rm(testRoot, { recursive: true, force: true });
		}
	}, 5000); // increase test timeout to 5s to accommodate the 2s snapshot timeout
});