|
| 1 | +import { APIError } from "@cheatcode/observability"; |
| 2 | +import type { |
| 3 | + SandboxCompareAndSwapFileResult, |
| 4 | + SandboxWriteFileResult, |
| 5 | +} from "@cheatcode/sandbox-contracts"; |
| 6 | +import { decodeBase64, dirname, shellQuote } from "../sandbox-support"; |
| 7 | +import { assertMutableWorkspacePath, WORKSPACE_DIR } from "./project-sandbox-content-support"; |
| 8 | +import { timeoutSeconds } from "./project-sandbox-process-support"; |
| 9 | +import { |
| 10 | + type ProjectCompareAndSwapFileInput, |
| 11 | + ProjectCompareAndSwapFileInputSchema, |
| 12 | + type ProjectWriteFileInput, |
| 13 | + ProjectWriteFileInputSchema, |
| 14 | +} from "./project-sandbox-runtime"; |
| 15 | +import type { SandboxRuntime } from "./project-sandbox-runtime-handle"; |
| 16 | + |
| 17 | +const COMPARE_AND_SWAP_MISMATCH_EXIT = 73; |
| 18 | +const COMPARE_AND_SWAP_FILE_SCRIPT = ` |
| 19 | +import hashlib |
| 20 | +import os |
| 21 | +import stat |
| 22 | +import sys |
| 23 | +
|
| 24 | +target = sys.argv[1] |
| 25 | +candidate = sys.argv[2] |
| 26 | +expected = sys.argv[3] |
| 27 | +digest = hashlib.sha256() |
| 28 | +descriptor = os.open(target, os.O_RDONLY | os.O_NOFOLLOW) |
| 29 | +with os.fdopen(descriptor, "rb") as source: |
| 30 | + metadata = os.fstat(source.fileno()) |
| 31 | + if not stat.S_ISREG(metadata.st_mode): |
| 32 | + raise RuntimeError("Edit target is not a regular file") |
| 33 | + for chunk in iter(lambda: source.read(1024 * 1024), b""): |
| 34 | + digest.update(chunk) |
| 35 | +if digest.hexdigest() != expected: |
| 36 | + raise SystemExit(${COMPARE_AND_SWAP_MISMATCH_EXIT}) |
| 37 | +os.chmod(candidate, stat.S_IMODE(metadata.st_mode)) |
| 38 | +with open(candidate, "rb") as pending: |
| 39 | + os.fsync(pending.fileno()) |
| 40 | +os.replace(candidate, target) |
| 41 | +`; |
| 42 | + |
| 43 | +type FileApplyRuntime = Pick<SandboxRuntime, "client" | "ensureSandbox">; |
| 44 | + |
| 45 | +export async function writeFile( |
| 46 | + runtime: FileApplyRuntime, |
| 47 | + input: ProjectWriteFileInput, |
| 48 | +): Promise<SandboxWriteFileResult> { |
| 49 | + const parsed = ProjectWriteFileInputSchema.parse(input); |
| 50 | + assertMutableWorkspacePath(parsed.path); |
| 51 | + const id = await runtime.ensureSandbox(); |
| 52 | + await runtime.client().createFolder(id, dirname(parsed.path)); |
| 53 | + const bytes = |
| 54 | + parsed.encoding === "base64" |
| 55 | + ? decodeBase64(parsed.content) |
| 56 | + : new TextEncoder().encode(parsed.content); |
| 57 | + await runtime.client().uploadFile(id, parsed.path, bytes); |
| 58 | + return { path: parsed.path, success: true }; |
| 59 | +} |
| 60 | + |
| 61 | +export async function compareAndSwapFile( |
| 62 | + runtime: FileApplyRuntime, |
| 63 | + input: ProjectCompareAndSwapFileInput, |
| 64 | +): Promise<SandboxCompareAndSwapFileResult> { |
| 65 | + const parsed = ProjectCompareAndSwapFileInputSchema.parse(input); |
| 66 | + assertMutableWorkspacePath(parsed.path); |
| 67 | + const id = await runtime.ensureSandbox(); |
| 68 | + const stagingDir = `${WORKSPACE_DIR}/.cheatcode/runtime`; |
| 69 | + const candidatePath = `${stagingDir}/file-apply-${crypto.randomUUID()}`; |
| 70 | + await runtime.client().createFolder(id, stagingDir); |
| 71 | + await runtime.client().uploadFile(id, candidatePath, new TextEncoder().encode(parsed.content)); |
| 72 | + try { |
| 73 | + await executeCompareAndSwap(runtime, id, parsed, candidatePath); |
| 74 | + return { path: parsed.path, success: true }; |
| 75 | + } finally { |
| 76 | + await runtime |
| 77 | + .client() |
| 78 | + .deleteFilePath(id, candidatePath, false) |
| 79 | + .catch(() => undefined); |
| 80 | + } |
| 81 | +} |
| 82 | + |
| 83 | +async function executeCompareAndSwap( |
| 84 | + runtime: FileApplyRuntime, |
| 85 | + sandboxId: string, |
| 86 | + input: ProjectCompareAndSwapFileInput, |
| 87 | + candidatePath: string, |
| 88 | +): Promise<void> { |
| 89 | + const completed = await runtime.client().execute(sandboxId, { |
| 90 | + command: [ |
| 91 | + "python3", |
| 92 | + "-c", |
| 93 | + COMPARE_AND_SWAP_FILE_SCRIPT, |
| 94 | + input.path, |
| 95 | + candidatePath, |
| 96 | + input.expectedSha256, |
| 97 | + ] |
| 98 | + .map(shellQuote) |
| 99 | + .join(" "), |
| 100 | + cwd: WORKSPACE_DIR, |
| 101 | + timeout: timeoutSeconds(30_000), |
| 102 | + }); |
| 103 | + assertCompareAndSwapSucceeded(completed.exitCode, completed.result); |
| 104 | +} |
| 105 | + |
| 106 | +function assertCompareAndSwapSucceeded(exitCode: number, result: string | null | undefined): void { |
| 107 | + if (exitCode === 0) { |
| 108 | + return; |
| 109 | + } |
| 110 | + if (exitCode === COMPARE_AND_SWAP_MISMATCH_EXIT) { |
| 111 | + throw new APIError(409, "conflict_state_invalid", "File changed while the edit was prepared", { |
| 112 | + hint: "Read the latest file and retry the edit.", |
| 113 | + retriable: true, |
| 114 | + }); |
| 115 | + } |
| 116 | + throw new APIError(502, "sandbox_command_failed", "Sandbox file edit failed", { |
| 117 | + hint: result?.trim().slice(-300) || "Check that the target is a regular workspace file.", |
| 118 | + retriable: false, |
| 119 | + }); |
| 120 | +} |
0 commit comments