@hheei/omp-optimizer 0.1.0

package/src/rtk.test.ts

@@ -0,0 +1,277 @@
+import { describe, expect, it, mock } from "bun:test";
+import type { ExtensionAPI } from "@oh-my-pi/pi-coding-agent";
+import {
+	applyRtkRewrite,
+	type BashCallEvent,
+	buildSudoBlockReason,
+	detectSudoSegments,
+	probeRtkAvailability,
+	rewriteChain,
+	splitChain,
+} from "./rtk.ts";
+
+/** Build a fresh bash tool_call event for hook tests. */
+function bashEvent(command: string): BashCallEvent {
+	return { toolName: "bash", input: { command } };
+}
+
+describe("probeRtkAvailability", () => {
+	it("probes rtk directly and accepts a successful version check", async () => {
+		const exec = mock(async () => ({
+			stdout: "rtk 0.1.0",
+			stderr: "",
+			code: 0,
+			killed: false,
+		}));
+
+		expect(await probeRtkAvailability({ exec } as Pick<ExtensionAPI, "exec">)).toBe(true);
+		expect(exec).toHaveBeenCalledWith("rtk", ["--version"], { timeout: 3000 });
+	});
+
+	it("rejects an unsuccessful version check", async () => {
+		const exec = mock(async () => ({ stdout: "", stderr: "missing", code: 1, killed: false }));
+
+		expect(await probeRtkAvailability({ exec } as Pick<ExtensionAPI, "exec">)).toBe(false);
+	});
+
+	it("treats a spawn failure as unavailable", async () => {
+		const exec = mock(async () => {
+			throw new Error("ENOENT");
+		});
+
+		expect(await probeRtkAvailability({ exec } as Pick<ExtensionAPI, "exec">)).toBe(false);
+	});
+});
+
+describe("splitChain", () => {
+	it("returns single segment for plain command", () => {
+		expect(splitChain("git status")).toEqual(["git status"]);
+	});
+
+	it("splits on && keeping operator", () => {
+		expect(splitChain("git add . && git push")).toEqual(["git add . ", "&&", " git push"]);
+	});
+
+	it("splits on ||, ;, |", () => {
+		expect(splitChain("a || b")).toEqual(["a ", "||", " b"]);
+		expect(splitChain("a ; b")).toEqual(["a ", ";", " b"]);
+		expect(splitChain("a | b")).toEqual(["a ", "|", " b"]);
+	});
+
+	it("ignores operators inside double quotes", () => {
+		expect(splitChain('git commit -m "a && b"')).toEqual(['git commit -m "a && b"']);
+	});
+
+	it("ignores operators inside single quotes", () => {
+		expect(splitChain("echo 'x | y'")).toEqual(["echo 'x | y'"]);
+	});
+
+	it("returns null on unbalanced quotes", () => {
+		expect(splitChain('git commit -m "oops')).toBeNull();
+	});
+});
+
+describe("rewriteChain", () => {
+	it("prefixes a single known command", () => {
+		expect(rewriteChain("git status")).toBe("rtk git status");
+	});
+
+	it("prefixes every segment in a chain", () => {
+		expect(rewriteChain("git add . && git commit -m x && git push")).toBe(
+			"rtk git add . && rtk git commit -m x && rtk git push",
+		);
+	});
+
+	it("prefixes mixed known commands", () => {
+		expect(rewriteChain("cargo build && npm test")).toBe("rtk cargo build && rtk npm test");
+	});
+
+	it("leaves unknown commands alone", () => {
+		expect(rewriteChain("echo hi && mkdir x")).toBe("echo hi && mkdir x");
+	});
+
+	it("does not prefix find (rtk find rejects -not/-exec)", () => {
+		const cmd = "find . -type f -not -path '*/node_modules/*'";
+		expect(rewriteChain(cmd)).toBe(cmd);
+	});
+
+	it("only prefixes known segments in a mixed chain", () => {
+		expect(rewriteChain("cd /tmp && git status")).toBe("cd /tmp && rtk git status");
+	});
+
+	it("does not double-prefix already-rtk commands", () => {
+		expect(rewriteChain("rtk git status")).toBe("rtk git status");
+		expect(rewriteChain("rtk git add . && git push")).toBe("rtk git add . && rtk git push");
+	});
+
+	it("does not touch operators inside quotes", () => {
+		expect(rewriteChain('git commit -m "a && b"')).toBe('rtk git commit -m "a && b"');
+	});
+
+	it("returns original on unbalanced quotes", () => {
+		const cmd = 'git commit -m "oops';
+		expect(rewriteChain(cmd)).toBe(cmd);
+	});
+
+	it("prefixes known commands across a pipe (ls, wc)", () => {
+		expect(rewriteChain("ls -la | wc -l")).toBe("rtk ls -la | rtk wc -l");
+	});
+
+	it("truly leaves a chain of only-unknown commands untouched", () => {
+		expect(rewriteChain("cd /tmp | sort | uniq")).toBe("cd /tmp | sort | uniq");
+	});
+
+	it("handles pipes between known commands", () => {
+		expect(rewriteChain("git log | grep fix")).toBe("rtk git log | rtk grep fix");
+	});
+});
+
+// ── detectSudoSegments ───────────────────────────────────────────────────────
+
+describe("detectSudoSegments", () => {
+	it("returns empty for command with no sudo", () => {
+		expect(detectSudoSegments(["git status"])).toEqual([]);
+	});
+
+	it("detects a plain sudo segment", () => {
+		expect(detectSudoSegments(["sudo apt-get install foo"])).toEqual(["sudo apt-get install foo"]);
+	});
+
+	it("detects sudo in a chain (operators excluded)", () => {
+		const parts = ["git status ", "&&", " sudo make install"];
+		expect(detectSudoSegments(parts)).toEqual(["sudo make install"]);
+	});
+
+	it("detects multiple sudo segments", () => {
+		const parts = ["sudo rm -rf /tmp ", ";", " sudo reboot"];
+		expect(detectSudoSegments(parts)).toEqual(["sudo rm -rf /tmp", "sudo reboot"]);
+	});
+
+	it("does not match 'sudoer' or 'pseudo'", () => {
+		expect(detectSudoSegments(["sudoers-check", "pseudo sudo"])).toEqual([]);
+	});
+
+	it("returns empty for operators-only parts", () => {
+		expect(detectSudoSegments(["&&", "||", ";"])).toEqual([]);
+	});
+});
+
+// ── buildSudoBlockReason ──────────────────────────────────────────────────────
+
+describe("buildSudoBlockReason", () => {
+	it("with sudo_run available: mentions sudo_run tool", () => {
+		const reason = buildSudoBlockReason(["sudo apt install curl"], true);
+		expect(reason).toContain("sudo_run");
+		expect(reason).toContain("sudo apt install curl");
+		expect(reason).not.toContain("not available");
+	});
+
+	it("with sudo_run available: instructs to strip sudo prefix", () => {
+		const reason = buildSudoBlockReason(["sudo make install"], true);
+		expect(reason).toContain("Strip the leading");
+	});
+
+	it("without sudo_run: explains restriction and asks user to run manually", () => {
+		const reason = buildSudoBlockReason(["sudo reboot"], false);
+		expect(reason).toContain("no sudo_run tool is available");
+		expect(reason).toContain("manually");
+		expect(reason).toContain("sudo reboot");
+	});
+
+	it("lists all blocked commands", () => {
+		const reason = buildSudoBlockReason(["sudo rm -rf /tmp", "sudo reboot"], true);
+		expect(reason).toContain("sudo rm -rf /tmp");
+		expect(reason).toContain("sudo reboot");
+	});
+});
+
+// Integration tests for the `tool_call` hook step. These guard the bug that
+// silently disabled rewriting: wrong event name + wrong field + wrong patch
+// mechanism. They assert on the IN-PLACE mutation contract the SDK requires.
+describe("applyRtkRewrite (tool_call hook step)", () => {
+	it("mutates event.input.command in place for a known bash command", () => {
+		const event = bashEvent("git status");
+		const changed = applyRtkRewrite(event, {
+			enabled: true,
+			rtkAvailable: true,
+		});
+		expect(changed).toBe(true);
+		expect(event.input.command).toBe("rtk git status");
+	});
+
+	it("rewrites every segment of a chain in place", () => {
+		const event = bashEvent("git add . && git push");
+		applyRtkRewrite(event, { enabled: true, rtkAvailable: true });
+		expect(event.input.command).toBe("rtk git add . && rtk git push");
+	});
+
+	it("does not mutate when disabled", () => {
+		const event = bashEvent("git status");
+		const changed = applyRtkRewrite(event, {
+			enabled: false,
+			rtkAvailable: true,
+		});
+		expect(changed).toBe(false);
+		expect(event.input.command).toBe("git status");
+	});
+
+	it("does not mutate when rtk binary is unavailable", () => {
+		const event = bashEvent("git status");
+		const changed = applyRtkRewrite(event, {
+			enabled: true,
+			rtkAvailable: false,
+		});
+		expect(changed).toBe(false);
+		expect(event.input.command).toBe("git status");
+	});
+
+	it("ignores non-bash tools", () => {
+		const event: BashCallEvent = {
+			toolName: "grep",
+			input: { command: "git status" },
+		};
+		const changed = applyRtkRewrite(event, {
+			enabled: true,
+			rtkAvailable: true,
+		});
+		expect(changed).toBe(false);
+		expect(event.input.command).toBe("git status");
+	});
+
+	it("leaves unknown commands untouched", () => {
+		const event = bashEvent("mkdir build && cd build");
+		const changed = applyRtkRewrite(event, {
+			enabled: true,
+			rtkAvailable: true,
+		});
+		expect(changed).toBe(false);
+		expect(event.input.command).toBe("mkdir build && cd build");
+	});
+
+	it("does not double-prefix an already-rtk command", () => {
+		const event = bashEvent("rtk git status");
+		const changed = applyRtkRewrite(event, {
+			enabled: true,
+			rtkAvailable: true,
+		});
+		expect(changed).toBe(false);
+		expect(event.input.command).toBe("rtk git status");
+	});
+
+	it("handles missing / non-string command safely", () => {
+		const event: BashCallEvent = { toolName: "bash", input: {} };
+		expect(applyRtkRewrite(event, { enabled: true, rtkAvailable: true })).toBe(false);
+		const event2: BashCallEvent = { toolName: "bash", input: { command: 123 } };
+		expect(applyRtkRewrite(event2, { enabled: true, rtkAvailable: true })).toBe(false);
+	});
+
+	it("leaves command unchanged on unbalanced quotes", () => {
+		const event = bashEvent('git commit -m "oops');
+		const changed = applyRtkRewrite(event, {
+			enabled: true,
+			rtkAvailable: true,
+		});
+		expect(changed).toBe(false);
+		expect(event.input.command).toBe('git commit -m "oops');
+	});
+});

package/src/rtk.ts

@@ -0,0 +1,355 @@
+/**
+ * Simple RTK Integration
+ *
+ * 1. Injects RTK system prompt (tells model to prefix commands with rtk)
+ * 2. Rewrites bash commands to add rtk prefix when model forgets
+ * 3. Falls back gracefully if rtk binary is missing
+ */
+
+import type {
+	ExtensionAPI,
+	ExtensionCommandContext,
+	ExtensionContext,
+} from "@oh-my-pi/pi-coding-agent";
+import { canExecute } from "./capability.ts";
+import { loadOptValue, saveOptValue } from "./persist.ts";
+import type { OptimizerHandle, OptimizerStatus } from "./status.ts";
+
+/**
+ * Minimal structural shape of the SDK `tool_call` event we care about.
+ * Mirrors `BashToolCallEvent` from the SDK without importing it, so the
+ * rewrite logic stays unit-testable with plain objects.
+ */
+export interface BashCallEvent {
+	toolName: string;
+	input: { command?: unknown };
+}
+
+/**
+ * Pure decision + mutation step for the `tool_call` hook.
+ *
+ * Given a tool-call event and whether RTK is available, mutate `event.input`
+ * in place (the SDK's only supported way to patch tool args) when the command
+ * is a rewritable bash command. Returns true if the command was rewritten.
+ *
+ * Extracted from the hook closure so the integration is directly testable
+ * without a live ExtensionAPI.
+ */
+/**
+ * Return the list of sudo sub-commands found in parsed chain segments.
+ * Each entry is the full segment body (trimmed) that starts with `sudo`.
+ * Operators are excluded. Returns empty array if none found.
+ */
+export function detectSudoSegments(parts: string[]): string[] {
+	return parts
+		.filter((p) => !CHAIN_OPERATORS.has(p.trim()))
+		.map((p) => p.trim())
+		.filter((p) => /^sudo\b/.test(p));
+}
+
+/**
+ * Build the block reason string shown to the model when a sudo command is
+ * intercepted. Directs the model to `sudo_run` when available, otherwise
+ * explains the restriction clearly.
+ */
+export function buildSudoBlockReason(sudoCmds: string[], hasSudoRunTool: boolean): string {
+	const list = sudoCmds.map((c) => `  - ${c}`).join("\n");
+	if (hasSudoRunTool) {
+		return (
+			`bash cannot run sudo commands directly. ` +
+			`Use the \`sudo_run\` tool instead — it shows the user a confirmation dialog ` +
+			`and handles authentication securely.\n` +
+			`Blocked command(s):\n${list}\n` +
+			`Strip the leading \`sudo\` from the command and pass the rest to \`sudo_run\` ` +
+			`with a clear \`reason\` parameter.`
+		);
+	}
+	return (
+		`bash cannot run sudo commands in this session — ` +
+		`no sudo_run tool is available and direct sudo is not permitted.\n` +
+		`Blocked command(s):\n${list}\n` +
+		`Ask the user to run the command manually with elevated privileges.`
+	);
+}
+
+export function applyRtkRewrite(
+	event: BashCallEvent,
+	opts: { enabled: boolean; rtkAvailable: boolean },
+): boolean {
+	if (!opts.enabled) return false;
+	if (!opts.rtkAvailable) return false;
+	if (event.toolName !== "bash") return false;
+
+	const command = event.input?.command;
+	if (typeof command !== "string" || !command) return false;
+
+	const rewritten = rewriteChain(command);
+	if (rewritten === command) return false;
+
+	event.input.command = rewritten;
+	return true;
+}
+
+const RTK_SYSTEM_PROMPT = `# RTK — token-optimized command wrapper
+
+Prefix shell commands with \`rtk\` (e.g. \`rtk git status\`). RTK compacts output for git, gh, cargo, npm/pnpm/yarn/bun, tsc, lint, vitest/jest/playwright, docker, kubectl, ls, grep, prisma — and passes anything else through unchanged, so it's always safe.
+
+Prefix EVERY segment in a chain, not just the first:
+\`rtk git add . && rtk git commit -m "msg" && rtk git push\`
+
+RTK also has filtering subcommands the auto-rewriter won't add — reach for these yourself when useful: \`rtk err <cmd>\` (errors only), \`rtk summary <cmd>\`, \`rtk log <file>\` (dedup), \`rtk json <file>\` (structure), \`rtk test <cmd>\` (failures only), \`rtk gain\` (savings stats).`;
+
+// Commands that should be prefixed with rtk
+const RTK_COMMANDS = new Set([
+	"git",
+	"gh",
+	"ls",
+	"tree",
+	"grep",
+	"cat",
+	"head",
+	"tail",
+	"tsc",
+	"lint",
+	"eslint",
+	"prettier",
+	"next",
+	"cargo",
+	"rustc",
+	"vitest",
+	"playwright",
+	"jest",
+	"test",
+	"pnpm",
+	"npm",
+	"npx",
+	"yarn",
+	"bun",
+	"docker",
+	"kubectl",
+	"aws",
+	"psql",
+	"wc",
+	"prisma",
+	"dotnet",
+]);
+
+interface RtkStatus {
+	available: boolean;
+	checkedAt: number;
+}
+
+/** Probe the command we actually use instead of relying on a platform-specific locator. */
+export function probeRtkAvailability(pi: Pick<ExtensionAPI, "exec">): Promise<boolean> {
+	return canExecute(pi, "rtk", ["--version"]);
+}
+
+/**
+ * Split a command line into segments at top-level shell operators
+ * (&&, ||, ;, |), keeping the operators as their own tokens. Operators
+ * inside single/double quotes are ignored.
+ *
+ * Returns null if the parser hits something it can't safely reason about
+ * (unbalanced quotes), so the caller can skip rewriting.
+ */
+export function splitChain(command: string): string[] | null {
+	const out: string[] = [];
+	let buf = "";
+	let quote: "'" | '"' | null = null;
+
+	for (let i = 0; i < command.length; i++) {
+		const c = command[i] ?? "";
+		const next = command[i + 1];
+
+		if (quote) {
+			buf += c;
+			if (c === quote) quote = null;
+			continue;
+		}
+
+		if (c === "'" || c === '"') {
+			quote = c;
+			buf += c;
+			continue;
+		}
+
+		// two-char operators
+		if ((c === "&" && next === "&") || (c === "|" && next === "|")) {
+			out.push(buf, c + c);
+			buf = "";
+			i++;
+			continue;
+		}
+
+		// single-char operators
+		if (c === ";" || c === "|") {
+			out.push(buf, c);
+			buf = "";
+			continue;
+		}
+
+		buf += c;
+	}
+
+	if (quote) return null; // unbalanced quote — bail out
+	out.push(buf);
+	return out;
+}
+
+const CHAIN_OPERATORS = new Set(["&&", "||", ";", "|"]);
+
+/**
+ * Prefix each command segment with `rtk` when its first word is a known
+ * RTK command and it is not already prefixed. Operators are preserved.
+ * Returns the rewritten command, or the original if nothing changed.
+ */
+export function rewriteChain(command: string): string {
+	const parts = splitChain(command);
+	if (!parts) return command; // unparseable — leave untouched
+
+	let changed = false;
+	const rewritten = parts.map((part) => {
+		if (CHAIN_OPERATORS.has(part.trim())) return part;
+
+		const leading = part.match(/^\s*/)?.[0] ?? "";
+		const body = part.slice(leading.length);
+		if (!body) return part;
+
+		const firstWord = body.split(/\s+/)[0] ?? "";
+		if (firstWord === "rtk") return part;
+		if (!RTK_COMMANDS.has(firstWord)) return part;
+
+		changed = true;
+		return `${leading}rtk ${body}`;
+	});
+
+	return changed ? rewritten.join("") : command;
+}
+
+export function rtk(pi: ExtensionAPI, status: OptimizerStatus): OptimizerHandle {
+	let rtkStatus: RtkStatus | null = null;
+	let warnedMissing = false;
+	let enabled = true;
+	// Tracks whether pix-sudo's sudo_run tool is active this session.
+	// Set from before_agent_start selectedTools; defaults to false until known.
+	let hasSudoRunTool = false;
+
+	// Report into the shared optimizer indicator. RTK counts as "on" only when
+	// enabled AND the binary is actually available.
+	function syncStatus(ctx: Pick<ExtensionContext, "ui">) {
+		status.set("rtk", enabled && rtkStatus?.available === true, ctx);
+	}
+
+	// Check if rtk binary is available
+	const checkRtkAvailability = async (): Promise<RtkStatus> => {
+		// Cache for 60 seconds
+		if (rtkStatus && Date.now() - rtkStatus.checkedAt < 60000) {
+			return rtkStatus;
+		}
+
+		const available = await probeRtkAvailability(pi);
+		rtkStatus = {
+			available,
+			checkedAt: Date.now(),
+		};
+		if (available) warnedMissing = false;
+		return rtkStatus;
+	};
+
+	// Inject RTK system prompt + detect sudo_run tool availability.
+	pi.on("before_agent_start", async (event) => {
+		// The active tool list is exposed by OMP's public ExtensionAPI.
+		hasSudoRunTool = pi.getActiveTools().includes("sudo_run");
+
+		if (!enabled) return undefined;
+		return { systemPrompt: [RTK_SYSTEM_PROMPT, ...event.systemPrompt] };
+	});
+
+
+	// Keep the status indicator in sync across the agent lifecycle. Probe
+	// availability on session start so the icon reflects reality immediately.
+	pi.on("session_start", async (_event, ctx) => {
+		// Restore the user's on/off choice from disk (survives quit/restart).
+		const saved = loadOptValue("rtk");
+		if (saved === "on" || saved === "off") enabled = saved === "on";
+		const probe = await checkRtkAvailability();
+		if (!probe.available && !warnedMissing) {
+			ctx.ui.notify(
+				"rtk not found — RTK rewriting disabled. Install: cargo install rtk-ai",
+				"warning",
+			);
+			warnedMissing = true;
+		}
+		syncStatus(ctx);
+	});
+	pi.on("agent_start", async (_event, ctx) => {
+		syncStatus(ctx);
+	});
+	pi.on("agent_end", async (_event, ctx) => {
+		syncStatus(ctx);
+	});
+
+	// -- Overlay value handler (called by the /optimizer overlay) --
+
+	async function run(value: string, ctx: ExtensionCommandContext): Promise<void> {
+		enabled = value === "on";
+		saveOptValue("rtk", enabled ? "on" : "off");
+
+		await checkRtkAvailability();
+		syncStatus(ctx);
+		ctx.ui.notify(`RTK rewriting ${enabled ? "on" : "off"}.`, "info");
+	}
+
+	// Rewrite bash commands to add rtk prefix.
+	//
+	// The SDK fires a single `tool_call` event for every tool. The bash variant
+	// carries `event.toolName === "bash"` and a mutable `event.input` of shape
+	// `{ command: string; timeout?: number }`. Arguments are patched by mutating
+	// `event.input` IN PLACE — returning `{ toolInput: ... }` does nothing.
+	pi.on("tool_call", async (event, ctx) => {
+		if (!enabled) {
+			return undefined;
+		}
+
+		if (event.toolName !== "bash") {
+			return undefined;
+		}
+
+		const probe = await checkRtkAvailability();
+
+		if (!probe.available) {
+			return undefined; // Don't rewrite if rtk not available
+		}
+
+		// First confirmed-available probe may have flipped state — refresh icon.
+		syncStatus(ctx);
+
+		// Block sudo segments before rtk rewriting.
+		// splitChain is safe to call here — same parser used by rewriteChain.
+		const command = event.input?.command;
+		if (typeof command === "string" && command) {
+			const parts = splitChain(command);
+			if (parts) {
+				const sudoCmds = detectSudoSegments(parts);
+				if (sudoCmds.length > 0) {
+					const reason = buildSudoBlockReason(sudoCmds, hasSudoRunTool);
+					return { block: true, reason };
+				}
+			}
+		}
+
+		// Rewrite every segment in the command chain that uses a known RTK
+		// command (e.g. `git add . && git push` -> `rtk git add . && rtk git push`).
+		// Mutates `event.input.command` in place — the SDK's supported patch path.
+		applyRtkRewrite(event, { enabled, rtkAvailable: probe.available });
+		return undefined;
+	});
+
+	return {
+		name: "rtk",
+		help: "rtk — prefix shell commands with rtk (token-optimized)",
+		values: ["off", "on"],
+		current: () => (enabled ? "on" : "off"),
+		run,
+	};
+}