feat(engine): add opt-in RTK bash rewriting

This commit is contained in:
Phil Larson
2026-05-27 11:23:15 -07:00
parent 2f80c6ea7b
commit 199204989e
3 changed files with 273 additions and 3 deletions

View File

@@ -0,0 +1,5 @@
---
"@fusion/engine": minor
---
Add opt-in RTK command rewriting for Pi bash tools via `FUSION_RTK_REWRITE`.

View File

@@ -20,6 +20,10 @@ const settingsManagerInMemoryMock = vi.fn(() => ({ kind: "settings-manager" }));
const setFallbackResolverMock = vi.fn(); const setFallbackResolverMock = vi.fn();
const reloadMock = vi.fn(async () => {}); const reloadMock = vi.fn(async () => {});
const execSyncMock = vi.fn((_cmd?: any, _opts?: any) => ""); const execSyncMock = vi.fn((_cmd?: any, _opts?: any) => "");
const execFileMock = vi.fn((_file?: any, _args?: any, _opts?: any, cb?: any) => {
const callback = typeof _opts === "function" ? _opts : cb;
if (typeof callback === "function") callback(null, "", "");
});
const existsSyncMock = vi.fn((_path: PathLike) => false); const existsSyncMock = vi.fn((_path: PathLike) => false);
const readFileSyncMock = vi.fn((_path?: any) => "{}"); const readFileSyncMock = vi.fn((_path?: any) => "{}");
const readCustomProvidersMock = vi.fn(() => []); const readCustomProvidersMock = vi.fn(() => []);
@@ -59,7 +63,7 @@ vi.mock("node:child_process", () => {
} }
}); });
}); });
return { execSync: execSyncFn, exec: execFn }; return { execSync: execSyncFn, exec: execFn, execFile: execFileMock };
}); });
vi.mock("node:fs", async () => { vi.mock("node:fs", async () => {
@@ -133,6 +137,149 @@ vi.mock("@mariozechner/pi-coding-agent", () => ({
}, },
})); }));
describe("RTK bash rewrite wrapper", () => {
beforeEach(() => {
execFileMock.mockReset();
execFileMock.mockImplementation((_file?: any, _args?: any, _opts?: any, cb?: any) => {
const callback = typeof _opts === "function" ? _opts : cb;
if (typeof callback === "function") callback(null, "", "");
});
});
it("rewrites bash commands when rtk returns an accepted rewrite", async () => {
execFileMock.mockImplementation((_file: string, _args: string[], _opts: any, cb: any) => {
cb(null, "rtk git status\n", "");
});
const bashTool = {
name: "bash",
execute: vi.fn().mockResolvedValue({ ok: true }),
};
const { wrapToolsWithRtkRewrite } = await import("../pi.js");
const wrapped = wrapToolsWithRtkRewrite([bashTool as any], { mode: "rewrite", timeoutMs: 100 });
await (wrapped[0] as any).execute("call-1", { command: "git status", cwd: "/project" });
expect(execFileMock).toHaveBeenCalledWith("rtk", ["rewrite", "git status"], expect.objectContaining({ timeout: 100 }), expect.any(Function));
expect(bashTool.execute).toHaveBeenCalledWith("call-1", { command: "rtk git status", cwd: "/project" });
});
it("accepts rtk rewrite exit code 3", async () => {
execFileMock.mockImplementation((_file: string, _args: string[], _opts: any, cb: any) => {
const err = new Error("ask") as any;
err.code = 3;
cb(err, "rtk ls\n", "");
});
const bashTool = {
name: "bash",
execute: vi.fn().mockResolvedValue({ ok: true }),
};
const { wrapToolsWithRtkRewrite } = await import("../pi.js");
const wrapped = wrapToolsWithRtkRewrite([bashTool as any], { mode: "rewrite", timeoutMs: 100 });
await (wrapped[0] as any).execute("call-1", { command: "ls" });
expect(bashTool.execute).toHaveBeenCalledWith("call-1", { command: "rtk ls" });
});
it("fails open when rtk is unavailable or declines a rewrite", async () => {
execFileMock.mockImplementation((_file: string, _args: string[], _opts: any, cb: any) => {
const err = new Error("no equivalent") as any;
err.code = 1;
cb(err, "", "");
});
const bashTool = {
name: "bash",
execute: vi.fn().mockResolvedValue({ ok: true }),
};
const { wrapToolsWithRtkRewrite } = await import("../pi.js");
const wrapped = wrapToolsWithRtkRewrite([bashTool as any], { mode: "rewrite", timeoutMs: 100 });
await (wrapped[0] as any).execute("call-1", { command: "git status" });
expect(bashTool.execute).toHaveBeenCalledWith("call-1", { command: "git status" });
});
it("does not rewrite non-bash tools or when mode is off", async () => {
const readTool = {
name: "read",
execute: vi.fn().mockResolvedValue({ ok: true }),
};
const bashTool = {
name: "bash",
execute: vi.fn().mockResolvedValue({ ok: true }),
};
const { wrapToolsWithRtkRewrite } = await import("../pi.js");
const wrapped = wrapToolsWithRtkRewrite([readTool as any, bashTool as any], { mode: "off", timeoutMs: 100 });
await (wrapped[0] as any).execute("call-read", { command: "cat package.json" });
await (wrapped[1] as any).execute("call-bash", { command: "git status" });
expect(execFileMock).not.toHaveBeenCalled();
expect(readTool.execute).toHaveBeenCalledWith("call-read", { command: "cat package.json" });
expect(bashTool.execute).toHaveBeenCalledWith("call-bash", { command: "git status" });
});
it("passes the tool abort signal to the rtk subprocess", async () => {
execFileMock.mockImplementation((_file: string, _args: string[], _opts: any, cb: any) => {
cb(null, "rtk git status\n", "");
});
const signal = new AbortController().signal;
const bashTool = {
name: "bash",
execute: vi.fn().mockResolvedValue({ ok: true }),
};
const { wrapToolsWithRtkRewrite } = await import("../pi.js");
const wrapped = wrapToolsWithRtkRewrite([bashTool as any], { mode: "rewrite", timeoutMs: 100 });
await (wrapped[0] as any).execute("call-1", { command: "git status" }, signal);
expect(execFileMock).toHaveBeenCalledWith("rtk", ["rewrite", "git status"], expect.objectContaining({ signal }), expect.any(Function));
});
it("keeps action gating outside RTK rewriting so git policies see the original command", async () => {
execFileMock.mockImplementation((_file: string, _args: string[], _opts: any, cb: any) => {
cb(null, "rtk git push\n", "");
});
const bashTool = {
name: "bash",
execute: vi.fn().mockResolvedValue({ ok: true }),
};
const { wrapToolsWithActionGate, wrapToolsWithRtkRewrite } = await import("../pi.js");
const rtkWrapped = wrapToolsWithRtkRewrite([bashTool as any], { mode: "rewrite", timeoutMs: 100 });
const gated = wrapToolsWithActionGate(rtkWrapped, {
agentId: "agent-1",
agentName: "Agent",
isEphemeral: false,
taskId: "FN-1",
permissionPolicy: {
presetId: "custom",
rules: {
git_write: "block",
file_write_delete: "allow",
command_execution: "allow",
network_api: "allow",
task_agent_mutation: "allow",
},
},
createApprovalRequest: vi.fn(),
findApprovalByDedupeKey: vi.fn(),
});
const result = await (gated[0] as any).execute("call-1", { command: "git push" });
expect((result as any).isError).toBe(true);
expect((result as any).decision.category).toBe("git_write");
expect(execFileMock).not.toHaveBeenCalled();
expect(bashTool.execute).not.toHaveBeenCalled();
});
});
describe("worktree path boundary helpers", () => { describe("worktree path boundary helpers", () => {
// Test helper functions directly by importing them // Test helper functions directly by importing them
// Note: These tests verify the boundary logic without needing a full agent session // Note: These tests verify the boundary logic without needing a full agent session

View File

@@ -7,7 +7,7 @@
/* eslint-disable @typescript-eslint/no-explicit-any */ /* eslint-disable @typescript-eslint/no-explicit-any */
import { existsSync, readFileSync } from "node:fs"; import { existsSync, readFileSync } from "node:fs";
import { exec } from "node:child_process"; import { exec, execFile } from "node:child_process";
import { promisify } from "node:util"; import { promisify } from "node:util";
import { createRequire } from "node:module"; import { createRequire } from "node:module";
import { basename, dirname, join, relative, isAbsolute, resolve } from "node:path"; import { basename, dirname, join, relative, isAbsolute, resolve } from "node:path";
@@ -59,6 +59,84 @@ import { resolvePermanentAgentToolDecision } from "./permanent-agent-gating.js";
import type { SystemPromptLayers } from "./prompt-layers.js"; import type { SystemPromptLayers } from "./prompt-layers.js";
import { READONLY_ALLOWLIST, filterCustomToolsForReadonly, isReadonlyAllowed } from "./workflow-step-tool-policy.js"; import { READONLY_ALLOWLIST, filterCustomToolsForReadonly, isReadonlyAllowed } from "./workflow-step-tool-policy.js";
const RTK_ACCEPTED_REWRITE_EXIT_CODES = new Set([0, 3]);
const RTK_EXPECTED_PASSTHROUGH_EXIT_CODES = new Set([1, 2]);
const RTK_EXPECTED_FAIL_OPEN_ERROR_CODES = new Set(["ABORT_ERR", "ENOENT", "ETIMEDOUT"]);
const RTK_REWRITE_MAX_BUFFER_BYTES = 64 * 1024;
export type RtkRewriteMode = "off" | "rewrite";
export interface RtkRewriteOptions {
mode?: RtkRewriteMode;
timeoutMs?: number;
}
function normalizeRtkRewriteOptions(options?: RtkRewriteOptions): Required<RtkRewriteOptions> {
const modeEnv = process.env.FUSION_RTK_REWRITE?.toLowerCase();
const envMode: RtkRewriteMode = modeEnv === "1" || modeEnv === "true" || modeEnv === "rewrite" ? "rewrite" : "off";
const envTimeoutMs = Number.parseInt(process.env.FUSION_RTK_REWRITE_TIMEOUT_MS ?? "2000", 10);
const requestedTimeoutMs = options?.timeoutMs ?? envTimeoutMs;
return {
mode: options?.mode ?? envMode,
timeoutMs: Number.isFinite(requestedTimeoutMs) && requestedTimeoutMs > 0 ? requestedTimeoutMs : 2000,
};
}
function resolveRtkRewriteOptions(): Required<RtkRewriteOptions> {
return normalizeRtkRewriteOptions();
}
function getRtkErrorCode(error: Error | null): number | string | null {
if (!error) return 0;
const rawCode = (error as unknown as { code?: unknown }).code;
if (typeof rawCode === "number" || typeof rawCode === "string") return rawCode;
return null;
}
function shouldWarnForRtkFailure(code: number | string | null): boolean {
if (typeof code === "number") return !RTK_EXPECTED_PASSTHROUGH_EXIT_CODES.has(code);
if (typeof code === "string") return !RTK_EXPECTED_FAIL_OPEN_ERROR_CODES.has(code);
return true;
}
function rewriteCommandWithRtk(
command: string,
options: Required<RtkRewriteOptions>,
signal?: AbortSignal,
): Promise<string | null> {
if (signal?.aborted) {
return Promise.resolve(null);
}
return new Promise((resolve) => {
execFile(
"rtk",
["rewrite", command],
{
timeout: options.timeoutMs,
maxBuffer: RTK_REWRITE_MAX_BUFFER_BYTES,
signal,
windowsHide: true,
},
(error, stdout, stderr) => {
const code = getRtkErrorCode(error);
if (typeof code !== "number" || !RTK_ACCEPTED_REWRITE_EXIT_CODES.has(code)) {
if (shouldWarnForRtkFailure(code)) {
const reason = stderr?.toString().trim() || (error instanceof Error ? error.message : `exit ${String(code)}`);
piLog.warn(`[pi] rtk rewrite failed open: ${reason}`);
}
resolve(null);
return;
}
const rewritten = stdout.toString().trim();
resolve(rewritten && rewritten !== command ? rewritten : null);
},
);
});
}
export interface AgentResult { export interface AgentResult {
session: AgentSession; session: AgentSession;
/** Path to the persisted session file (undefined for in-memory sessions). */ /** Path to the persisted session file (undefined for in-memory sessions). */
@@ -1493,6 +1571,45 @@ export function wrapToolsWithBoundary(
}); });
} }
export function wrapToolsWithRtkRewrite(
tools: ToolDefinition[],
options: RtkRewriteOptions = resolveRtkRewriteOptions(),
): ToolDefinition[] {
const resolvedOptions = normalizeRtkRewriteOptions(options);
if (resolvedOptions.mode !== "rewrite") {
return tools;
}
return tools.map((tool) => {
if (tool.name !== "bash") {
return tool;
}
const originalExecute = tool.execute as any;
return {
...tool,
execute: async (...args: any[]) => {
const params = args[1] as Record<string, unknown> | undefined;
const command = params?.command;
if (typeof command !== "string" || !command.trim()) {
return originalExecute(...args);
}
const signal = args[2] as AbortSignal | undefined;
const rewrittenCommand = await rewriteCommandWithRtk(command, resolvedOptions, signal);
if (!rewrittenCommand) {
return originalExecute(...args);
}
const rewrittenArgs = [...args];
rewrittenArgs[1] = { ...(params ?? {}), command: rewrittenCommand };
return originalExecute(...rewrittenArgs);
},
};
});
}
export function wrapToolsWithPermanentAgentGating( export function wrapToolsWithPermanentAgentGating(
tools: ToolDefinition[], tools: ToolDefinition[],
gating: PermanentAgentGatingContext | undefined, gating: PermanentAgentGatingContext | undefined,
@@ -1873,8 +1990,9 @@ export async function createFnAgent(options: AgentOptions): Promise<AgentResult>
...(tools as ToolDefinition[]), ...(tools as ToolDefinition[]),
...readonlyFilteredCustomTools.allowed, ...readonlyFilteredCustomTools.allowed,
]; ];
const toolsWithRtkRewrite = wrapToolsWithRtkRewrite(toolChainStart);
const toolsWithPermanentGating = wrapToolsWithPermanentAgentGating( const toolsWithPermanentGating = wrapToolsWithPermanentAgentGating(
toolChainStart, toolsWithRtkRewrite,
options.permanentAgentGating, options.permanentAgentGating,
); );
const toolsWithActionGate = wrapToolsWithActionGate( const toolsWithActionGate = wrapToolsWithActionGate(