Skip to content

Commit ac07599

Browse files
kays0xkays0x
committed
fix(agents): stop false-positive session-takeover on runner's own transcript appends
1 parent 94d8391 commit ac07599

3 files changed

Lines changed: 474 additions & 21 deletions

File tree

‎src/agents/pi-embedded-runner/run/attempt.session-lock.test.ts‎

Lines changed: 297 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -156,7 +156,7 @@ describe("embedded attempt session lock lifecycle", () => {
156156
expect(events).toEqual(["acquire-1", "release", "events-drained", "acquire-2", "release"]);
157157
});
158158

159-
it("rejects post-prompt writes when another owner advances the session file", async () => {
159+
it("rejects post-prompt writes when another owner queues a new user turn", async () => {
160160
const sessionFile = await createTempSessionFile();
161161
const release = vi.fn(async () => {});
162162
const acquireSessionWriteLock = vi.fn(async () => ({ release }));
@@ -166,7 +166,11 @@ describe("embedded attempt session lock lifecycle", () => {
166166
});
167167

168168
await controller.releaseForPrompt();
169-
await fs.appendFile(sessionFile, '{"type":"message","id":"takeover"}\n', "utf8");
169+
await fs.appendFile(
170+
sessionFile,
171+
'{"type":"message","id":"takeover","message":{"role":"user","content":"new turn"}}\n',
172+
"utf8",
173+
);
170174

171175
await expect(controller.withSessionWriteLock(() => "late-write")).rejects.toBeInstanceOf(
172176
EmbeddedAttemptSessionTakeoverError,
@@ -179,6 +183,297 @@ describe("embedded attempt session lock lifecycle", () => {
179183
expect(release).toHaveBeenCalledTimes(2);
180184
});
181185

186+
it("accepts post-prompt writes when only the runner's own transcript entries appear", async () => {
187+
const sessionFile = await createTempSessionFile();
188+
const release = vi.fn(async () => {});
189+
const acquireSessionWriteLock = vi.fn(async () => ({ release }));
190+
const controller = await createEmbeddedAttemptSessionLockController({
191+
acquireSessionWriteLock,
192+
lockOptions: { ...lockOptions, sessionFile },
193+
});
194+
195+
await controller.releaseForPrompt();
196+
// Runner-owned writers (e.g. appendSessionTranscriptMessageLocked via
197+
// allowReentrant:true) append assistant replies and persisted tool outputs
198+
// during the released-lock window. These must not trip the takeover fence.
199+
// Persisted role names follow isAgentMessage in transcript-file-state.ts.
200+
await fs.appendFile(
201+
sessionFile,
202+
'{"type":"message","id":"assistant-1","message":{"role":"assistant","content":"calling tool"}}\n',
203+
"utf8",
204+
);
205+
await fs.appendFile(
206+
sessionFile,
207+
'{"type":"message","id":"toolresult-1","message":{"role":"toolResult","toolCallId":"call-1","toolName":"exec","isError":false,"content":[{"type":"text","text":"ok"}]}}\n',
208+
"utf8",
209+
);
210+
await fs.appendFile(
211+
sessionFile,
212+
'{"type":"message","id":"bash-1","message":{"role":"bashExecution","command":"ls","output":"","exitCode":0,"cancelled":false,"truncated":false}}\n',
213+
"utf8",
214+
);
215+
216+
await expect(controller.withSessionWriteLock(() => "late-write")).resolves.toBe("late-write");
217+
expect(controller.hasSessionTakeover()).toBe(false);
218+
});
219+
220+
it("rejects post-prompt writes when an accepted runner-owned role is malformed", async () => {
221+
const sessionFile = await createTempSessionFile();
222+
const release = vi.fn(async () => {});
223+
const acquireSessionWriteLock = vi.fn(async () => ({ release }));
224+
const controller = await createEmbeddedAttemptSessionLockController({
225+
acquireSessionWriteLock,
226+
lockOptions: { ...lockOptions, sessionFile },
227+
});
228+
229+
await controller.releaseForPrompt();
230+
// toolResult requires toolCallId, toolName, isError, and array content
231+
// per the persisted message contract. An entry that carries the role
232+
// but omits required fields is not a real runner-owned write and must
233+
// trip the takeover fence.
234+
await fs.appendFile(
235+
sessionFile,
236+
'{"type":"message","id":"malformed-1","message":{"role":"toolResult","toolCallId":"call-1"}}\n',
237+
"utf8",
238+
);
239+
240+
await expect(controller.withSessionWriteLock(() => "late-write")).rejects.toBeInstanceOf(
241+
EmbeddedAttemptSessionTakeoverError,
242+
);
243+
expect(controller.hasSessionTakeover()).toBe(true);
244+
});
245+
246+
it("rejects post-prompt writes when a runner-owned message entry omits id", async () => {
247+
const sessionFile = await createTempSessionFile();
248+
const release = vi.fn(async () => {});
249+
const acquireSessionWriteLock = vi.fn(async () => ({ release }));
250+
const controller = await createEmbeddedAttemptSessionLockController({
251+
acquireSessionWriteLock,
252+
lockOptions: { ...lockOptions, sessionFile },
253+
});
254+
255+
await controller.releaseForPrompt();
256+
// Canonical persisted session entries require a non-empty string id. An
257+
// appended assistant message that omits the outer entry base is not a
258+
// valid runner-owned write and must trip the takeover fence, matching
259+
// the behavior current main exhibits via its stat-mismatch path.
260+
await fs.appendFile(
261+
sessionFile,
262+
'{"type":"message","message":{"role":"assistant","content":[{"type":"text","text":"hi"}]}}\n',
263+
"utf8",
264+
);
265+
266+
await expect(controller.withSessionWriteLock(() => "late-write")).rejects.toBeInstanceOf(
267+
EmbeddedAttemptSessionTakeoverError,
268+
);
269+
expect(controller.hasSessionTakeover()).toBe(true);
270+
});
271+
272+
it("rejects post-prompt writes when a non-message entry is appended", async () => {
273+
const sessionFile = await createTempSessionFile();
274+
const release = vi.fn(async () => {});
275+
const acquireSessionWriteLock = vi.fn(async () => ({ release }));
276+
const controller = await createEmbeddedAttemptSessionLockController({
277+
acquireSessionWriteLock,
278+
lockOptions: { ...lockOptions, sessionFile },
279+
});
280+
281+
await controller.releaseForPrompt();
282+
// Custom/compaction/branch_summary entries are session-state mutations
283+
// outside the runner-owned transcript-write path. Appending one during the
284+
// released-lock window must trip the takeover fence even though it grew
285+
// the file append-only and no user-role entry appeared.
286+
await fs.appendFile(
287+
sessionFile,
288+
'{"type":"custom","id":"custom-1","customType":"branch_marker","data":{}}\n',
289+
"utf8",
290+
);
291+
292+
await expect(controller.withSessionWriteLock(() => "late-write")).rejects.toBeInstanceOf(
293+
EmbeddedAttemptSessionTakeoverError,
294+
);
295+
expect(controller.hasSessionTakeover()).toBe(true);
296+
});
297+
298+
it("rejects post-prompt writes when a tail line parses to a non-object value", async () => {
299+
const sessionFile = await createTempSessionFile();
300+
const release = vi.fn(async () => {});
301+
const acquireSessionWriteLock = vi.fn(async () => ({ release }));
302+
const controller = await createEmbeddedAttemptSessionLockController({
303+
acquireSessionWriteLock,
304+
lockOptions: { ...lockOptions, sessionFile },
305+
});
306+
307+
await controller.releaseForPrompt();
308+
// A literal JSON null parses cleanly but is not a session entry. The
309+
// slow path must fail closed as takeover rather than letting the
310+
// canonical validator dereference a non-object.
311+
await fs.appendFile(sessionFile, "null\n", "utf8");
312+
313+
await expect(controller.withSessionWriteLock(() => "late-write")).rejects.toBeInstanceOf(
314+
EmbeddedAttemptSessionTakeoverError,
315+
);
316+
expect(controller.hasSessionTakeover()).toBe(true);
317+
});
318+
319+
it("keeps the fence consistent when a guarded operation throws after assertion", async () => {
320+
const sessionFile = await createTempSessionFile();
321+
const release = vi.fn(async () => {});
322+
const acquireSessionWriteLock = vi.fn(async () => ({ release }));
323+
const controller = await createEmbeddedAttemptSessionLockController({
324+
acquireSessionWriteLock,
325+
lockOptions: { ...lockOptions, sessionFile },
326+
});
327+
328+
const assistantEntry = (id: string) =>
329+
`{"type":"message","id":"${id}","parentId":null,"timestamp":"2026-05-19T00:00:00.000Z","message":{"role":"assistant","content":[{"type":"text","text":"x"}]}}\n`;
330+
331+
await controller.releaseForPrompt();
332+
// First append-then-assert advances the fence over a runner-owned tail.
333+
// The guarded operation then throws, which skips refreshSessionFileFence
334+
// and used to leave the snapshot describing the new size but holding
335+
// the old prefix hash. A second runner-owned append would then false-
336+
// positive on prefix-hash mismatch in the next slow path.
337+
await fs.appendFile(sessionFile, assistantEntry("a1"), "utf8");
338+
await expect(
339+
controller.withSessionWriteLock(() => {
340+
throw new Error("post-assert");
341+
}),
342+
).rejects.toThrow("post-assert");
343+
expect(controller.hasSessionTakeover()).toBe(false);
344+
345+
await fs.appendFile(sessionFile, assistantEntry("a2"), "utf8");
346+
await expect(controller.withSessionWriteLock(() => "ok")).resolves.toBe("ok");
347+
expect(controller.hasSessionTakeover()).toBe(false);
348+
});
349+
350+
it("refuses to bless a snapshot when the file is rewritten in place during snapshot creation", async () => {
351+
const sessionFile = await createTempSessionFile();
352+
const release = vi.fn(async () => {});
353+
const acquireSessionWriteLock = vi.fn(async () => ({ release }));
354+
const controller = await createEmbeddedAttemptSessionLockController({
355+
acquireSessionWriteLock,
356+
lockOptions: { ...lockOptions, sessionFile },
357+
});
358+
359+
// Simulate a same-size in-place rewrite landing between snapshot's stat
360+
// and its read. Spy on fs.readFile so the very first call (made during
361+
// releaseForPrompt's snapshot) atomically rewrites the file with new
362+
// content of identical byte length before delegating to the real read.
363+
// The post-read re-stat then sees a different mtime/ctime and the
364+
// snapshot returns prefixHashHex: null, which makes the next
365+
// withSessionWriteLock fail closed even on an otherwise valid append.
366+
const originalSize = (await fs.stat(sessionFile)).size;
367+
const rewritten = Buffer.alloc(originalSize, 0x78); // same length, different content
368+
const readFileSpy = vi.spyOn(fs, "readFile");
369+
readFileSpy.mockImplementationOnce(async (target) => {
370+
readFileSpy.mockRestore();
371+
await fs.writeFile(target as string, rewritten);
372+
return fs.readFile(target as string);
373+
});
374+
375+
try {
376+
await controller.releaseForPrompt();
377+
await fs.appendFile(
378+
sessionFile,
379+
'{"type":"message","id":"a1","parentId":null,"timestamp":"2026-05-19T00:00:00.000Z","message":{"role":"assistant","content":[{"type":"text","text":"x"}]}}\n',
380+
"utf8",
381+
);
382+
383+
await expect(controller.withSessionWriteLock(() => "late-write")).rejects.toBeInstanceOf(
384+
EmbeddedAttemptSessionTakeoverError,
385+
);
386+
expect(controller.hasSessionTakeover()).toBe(true);
387+
} finally {
388+
readFileSpy.mockRestore();
389+
}
390+
});
391+
392+
it("refuses to bless a slow-path verification when the file is rewritten in place during the slow-path read", async () => {
393+
const sessionFile = await createTempSessionFile();
394+
const release = vi.fn(async () => {});
395+
const acquireSessionWriteLock = vi.fn(async () => ({ release }));
396+
const controller = await createEmbeddedAttemptSessionLockController({
397+
acquireSessionWriteLock,
398+
lockOptions: { ...lockOptions, sessionFile },
399+
});
400+
401+
// Snapshot the fence cleanly first, then inject the race during the
402+
// slow path's read. The slow path stats, then reads, then re-stats; the
403+
// injected rewrite changes mtime between the two stats and the
404+
// verification refuses to bless the post-rewrite content as runner-owned.
405+
await controller.releaseForPrompt();
406+
await fs.appendFile(
407+
sessionFile,
408+
'{"type":"message","id":"a1","parentId":null,"timestamp":"2026-05-19T00:00:00.000Z","message":{"role":"assistant","content":[{"type":"text","text":"x"}]}}\n',
409+
"utf8",
410+
);
411+
412+
const sizeBeforeRace = (await fs.stat(sessionFile)).size;
413+
const rewritten = Buffer.alloc(sizeBeforeRace, 0x79);
414+
const readFileSpy = vi.spyOn(fs, "readFile");
415+
readFileSpy.mockImplementationOnce(async (target) => {
416+
readFileSpy.mockRestore();
417+
await fs.writeFile(target as string, rewritten);
418+
return fs.readFile(target as string);
419+
});
420+
421+
try {
422+
await expect(controller.withSessionWriteLock(() => "late-write")).rejects.toBeInstanceOf(
423+
EmbeddedAttemptSessionTakeoverError,
424+
);
425+
expect(controller.hasSessionTakeover()).toBe(true);
426+
} finally {
427+
readFileSpy.mockRestore();
428+
}
429+
});
430+
431+
it("rejects post-prompt writes when the session file is rewritten in place", async () => {
432+
const sessionFile = await createTempSessionFile();
433+
const release = vi.fn(async () => {});
434+
const acquireSessionWriteLock = vi.fn(async () => ({ release }));
435+
const controller = await createEmbeddedAttemptSessionLockController({
436+
acquireSessionWriteLock,
437+
lockOptions: { ...lockOptions, sessionFile },
438+
});
439+
440+
await controller.releaseForPrompt();
441+
// Concurrent compaction or another owner rewriting the transcript in place
442+
// would change the file's stat without an append-only growth. The fence
443+
// must still trip even though no user-role entry was added.
444+
await fs.writeFile(sessionFile, '{"type":"session","compacted":true}\n', "utf8");
445+
446+
await expect(controller.withSessionWriteLock(() => "late-write")).rejects.toBeInstanceOf(
447+
EmbeddedAttemptSessionTakeoverError,
448+
);
449+
expect(controller.hasSessionTakeover()).toBe(true);
450+
});
451+
452+
it("rejects post-prompt writes when the session file is replaced", async () => {
453+
const sessionFile = await createTempSessionFile();
454+
const release = vi.fn(async () => {});
455+
const acquireSessionWriteLock = vi.fn(async () => ({ release }));
456+
const controller = await createEmbeddedAttemptSessionLockController({
457+
acquireSessionWriteLock,
458+
lockOptions: { ...lockOptions, sessionFile },
459+
});
460+
461+
await controller.releaseForPrompt();
462+
// A different owner could atomically replace the session file (unlink +
463+
// recreate). dev/ino change must trip the fence regardless of content.
464+
await fs.rm(sessionFile);
465+
await fs.writeFile(
466+
sessionFile,
467+
'{"type":"session"}\n{"type":"message","id":"new-owner","message":{"role":"assistant","content":"hello"}}\n',
468+
"utf8",
469+
);
470+
471+
await expect(controller.withSessionWriteLock(() => "late-write")).rejects.toBeInstanceOf(
472+
EmbeddedAttemptSessionTakeoverError,
473+
);
474+
expect(controller.hasSessionTakeover()).toBe(true);
475+
});
476+
182477
it("returns a no-op cleanup lock after prompt lock reacquisition times out", async () => {
183478
const releases: string[] = [];
184479
const acquireSessionWriteLock = vi

0 commit comments

Comments
 (0)