Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
5.8 kB · 186 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187import { access, chmod, mkdtemp, readFile, rm, stat, symlink, writeFile } from "node:fs/promises";import { tmpdir } from "node:os";import { join } from "node:path";import { afterEach, describe, expect, it } from "vitest";import { createEditTool } from "../src/core/tools/edit.js";import { withFileMutationQueue } from "../src/core/tools/file-mutation-queue.js";
function delay(ms: number): Promise<void> { return new Promise((resolve) => setTimeout(resolve, ms));}
const tempDirs: string[] = [];
async function createTempDir(): Promise<string> { const dir = await mkdtemp(join(tmpdir(), "pi-file-mutation-queue-")); tempDirs.push(dir); return dir;}
afterEach(async () => { await Promise.all(tempDirs.splice(0, tempDirs.length).map((dir) => rm(dir, { recursive: true, force: true })));});
describe("withFileMutationQueue", () => { it("serializes operations for the same file", async () => { const order: string[] = []; const path = "/tmp/file-mutation-queue-same";
const first = withFileMutationQueue(path, async () => { order.push("first:start"); await delay(30); order.push("first:end"); }); const second = withFileMutationQueue(path, async () => { order.push("second:start"); order.push("second:end"); });
await Promise.all([first, second]); expect(order).toEqual(["first:start", "first:end", "second:start", "second:end"]); });
it("allows different files to proceed in parallel", async () => { const order: string[] = [];
await Promise.all([ withFileMutationQueue("/tmp/file-mutation-queue-a", async () => { order.push("a:start"); await delay(30); order.push("a:end"); }), withFileMutationQueue("/tmp/file-mutation-queue-b", async () => { order.push("b:start"); await delay(30); order.push("b:end"); }), ]);
expect(order.indexOf("a:start")).toBeLessThan(order.indexOf("a:end")); expect(order.indexOf("b:start")).toBeLessThan(order.indexOf("b:end")); expect(order.indexOf("b:start")).toBeLessThan(order.indexOf("a:end")); });
it("uses the same queue for symlink aliases", async () => { const dir = await createTempDir(); const targetPath = join(dir, "target.txt"); const symlinkPath = join(dir, "alias.txt"); await writeFile(targetPath, "hello\n", "utf8"); await symlink(targetPath, symlinkPath);
const order: string[] = []; await Promise.all([ withFileMutationQueue(targetPath, async () => { order.push("target:start"); await delay(30); order.push("target:end"); }), withFileMutationQueue(symlinkPath, async () => { order.push("alias:start"); order.push("alias:end"); }), ]);
expect(order).toEqual(["target:start", "target:end", "alias:start", "alias:end"]); });});
describe("built-in edit tool", () => { it("refuses to overwrite a file changed after the initial read", async () => { const dir = await createTempDir(); const filePath = join(dir, "concurrent-edit.txt"); await writeFile(filePath, "alpha\nbeta\n", "utf8"); let reads = 0; let writes = 0;
const editTool = createEditTool(dir, { operations: { access, readFile: async (path) => { const content = await readFile(path); if (++reads === 1) await writeFile(path, "alpha\nchanged\n", "utf8"); return content; }, writeFile: async () => { writes++; }, }, });
await expect( editTool.execute("call-1", { path: filePath, edits: [{ oldText: "beta", newText: "BETA" }] }), ).rejects.toThrow("refusing to overwrite concurrent changes"); expect(writes).toBe(0); expect(await readFile(filePath, "utf8")).toBe("alpha\nchanged\n"); });
it("preserves the destination mode during atomic replacement", async () => { const dir = await createTempDir(); const filePath = join(dir, "mode-edit.txt"); await writeFile(filePath, "alpha\nbeta\n", "utf8"); await chmod(filePath, 0o751);
await createEditTool(dir).execute("call-1", { path: filePath, edits: [{ oldText: "beta", newText: "BETA" }], });
expect((await stat(filePath)).mode & 0o7777).toBe(0o751); expect(await readFile(filePath, "utf8")).toBe("alpha\nBETA\n"); });
it("reports success when cancellation races after the write commit starts", async () => { const dir = await createTempDir(); const filePath = join(dir, "committed-edit.txt"); await writeFile(filePath, "alpha\nbeta\n", "utf8"); const controller = new AbortController(); const editTool = createEditTool(dir, { operations: { access, readFile, writeFile: async (path, content) => { await writeFile(path, content, "utf8"); controller.abort(); }, }, });
await expect( editTool.execute( "call-1", { path: filePath, edits: [{ oldText: "beta", newText: "BETA" }] }, controller.signal, ), ).resolves.toMatchObject({ content: [{ text: expect.stringContaining("Successfully replaced") }] }); expect(await readFile(filePath, "utf8")).toBe("alpha\nBETA\n"); });
it("preserves both parallel edits on the same file", async () => { const dir = await createTempDir(); const filePath = join(dir, "parallel-edit.txt"); await writeFile(filePath, "alpha\nbeta\ngamma\n", "utf8");
const editTool = createEditTool(dir, { operations: { access, readFile: async (path) => { const buffer = await readFile(path); await delay(30); return buffer; }, writeFile: async (path, content) => { await delay(30); await writeFile(path, content, "utf8"); }, }, });
await Promise.all([ editTool.execute("call-1", { path: filePath, edits: [{ oldText: "alpha", newText: "ALPHA" }] }), editTool.execute("call-2", { path: filePath, edits: [{ oldText: "beta", newText: "BETA" }] }), ]);
const content = await readFile(filePath, "utf8"); expect(content).toBe("ALPHA\nBETA\ngamma\n"); });});