Skip to content

Commit 802b34c

Browse files
saltbosaltbo
committed
fix(daemon): preserve worktree on any successful result, not just in_review
Previously finalize checked getTask status — if the task was rejected between review_requested and agent finish, getTask returned in_progress and finalize released the worktree, destroying the resume entry point. Now: any agent that produced a result (resultReceived=true, not crashed) gets its worktree preserved. The daemon loop's checkRejectedReviews handles cleanup for done/cancelled tasks. This eliminates the race between reject and finalize, and removes the getTask call from finalize.
1 parent e46b724 commit 802b34c

8 files changed

Lines changed: 58 additions & 56 deletions

‎packages/cli/src/daemon/loop.ts‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -280,9 +280,15 @@ export async function checkRejectedReviews(
280280
if (task.status === "in_progress") {
281281
const notes = (await client.getTaskNotes(s.taskId)) as Array<{ action?: string; detail?: string }>;
282282
const rejectLog = [...notes].reverse().find((l) => l.action === "rejected");
283-
const reason = rejectLog?.detail || "No reason provided";
284-
const message = `Task rejected. Reason: ${reason}\n\nPlease fix the issues and submit for review again.`;
285-
await resumeOne(s, message);
283+
if (rejectLog) {
284+
const reason = rejectLog.detail || "No reason provided";
285+
const message = `Task rejected. Reason: ${reason}\n\nPlease fix the issues and submit for review again.`;
286+
await resumeOne(s, message);
287+
} else {
288+
// Agent finished with result but never submitted for review — release.
289+
logger.info(`Task ${s.taskId} in_progress without reject, releasing orphan review session`);
290+
await completeTerminalFromReview(sessions, s, { type: "task_cancelled" });
291+
}
286292
}
287293
}
288294
}

‎packages/cli/src/daemon/runtimePool.ts‎

Lines changed: 9 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ import { createLogger } from "../logger.js";
1212
import type { AgentEvent, AgentHandle, AgentProvider } from "../providers/types.js";
1313
import { getSessionManager } from "../session/manager.js";
1414
import { classifyIteratorEnd, type SessionEvent } from "../session/stateMachine.js";
15-
import { apiCallOptional, apiFireAndForget, providerExecute } from "./boundaries.js";
15+
import { apiFireAndForget, providerExecute } from "./boundaries.js";
1616
import { classify } from "./errors.js";
1717

1818
const logger = createLogger("runtime-pool");
@@ -423,16 +423,14 @@ async function finalize(agent: AgentProcess, opts: { crashed: boolean; error?: u
423423
);
424424
}
425425

426-
// Check task status after the iterator ends — only now is the state final.
427-
let taskInReview = false;
428-
if (agent.resultReceived) {
429-
try {
430-
const task = (await apiCallOptional("getTask", () => agent.agentClient.getTask(taskId))) as { status?: string } | null;
431-
taskInReview = task?.status === "in_review";
432-
} catch (e) {
433-
logger.warn(`Failed to check task status for ${taskId}, assuming not in review: ${errMessage(e)}`);
434-
}
435-
}
426+
// Agent produced a result → preserve worktree. The work product lives in
427+
// the worktree and must survive for review, reject-resume, or completion
428+
// cleanup. The daemon loop handles all post-finalize lifecycle:
429+
// - in_review → wait for human review
430+
// - in_progress + rejected → resume agent
431+
// - in_progress + no reject → release (agent forgot to submit review)
432+
// - done/cancelled → cleanup
433+
const taskInReview = agent.resultReceived && !opts.crashed;
436434

437435
const event: SessionEvent = classifyIteratorEnd({
438436
resultReceived: agent.resultReceived,

‎packages/cli/tests/daemon-integration.test.ts‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -449,7 +449,7 @@ describe("review watcher — task rejected (in_progress)", () => {
449449
expect(message).toContain("wrong approach");
450450
});
451451

452-
it("uses fallback message when no rejection note exists", async () => {
452+
it("cleans up session when in_progress with no rejection note (agent never submitted review)", async () => {
453453
const sessionId = randomUUID();
454454
const taskId = `task-${randomUUID()}`;
455455
await sm.create(makeWorkerFile(sessionId, { taskId, status: "in_review" }));
@@ -463,9 +463,9 @@ describe("review watcher — task rejected (in_progress)", () => {
463463
const resumeOne = vi.fn().mockResolvedValue(undefined);
464464
await checkRejectedReviews(sm, pool, client as any, resumeOne, 5);
465465

466-
expect(resumeOne).toHaveBeenCalledOnce();
467-
const [, message] = resumeOne.mock.calls[0];
468-
expect(message).toContain("No reason provided");
466+
// No reject action → don't resume, clean up instead
467+
expect(resumeOne).not.toHaveBeenCalled();
468+
expect(sm.read(sessionId)).toBeNull();
469469
});
470470
});
471471

‎packages/cli/tests/processManager.test.ts‎

Lines changed: 12 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -900,12 +900,11 @@ describe("ProcessManager — onComplete skips cleanup when session is in_review"
900900
});
901901

902902
describe("ProcessManager — onComplete normal completion invokes cleanup", () => {
903-
it("invokes onCleanup when getTask returns done (normal completion)", async () => {
903+
it("preserves worktree when agent finishes with result (session stays in_review)", async () => {
904904
const tunnel = makeTunnel();
905905
const resultEvent: AgentEvent = { type: "turn.end", cost: 0.001 };
906906
const handle = makeHandle([resultEvent]);
907907
const onCleanup = vi.fn();
908-
// getTask returns "done" → taskInReview=false → state machine → completing → runCleanup
909908
const apiClient = makeApiClient({ getTask: vi.fn().mockResolvedValue({ status: "done" }) });
910909

911910
const pm = makePool(apiClient, makeCallbacks(), 0, tunnel);
@@ -925,33 +924,36 @@ describe("ProcessManager — onComplete normal completion invokes cleanup", () =
925924
await flushPromises();
926925
await flushPromises();
927926

928-
expect(onCleanup).toHaveBeenCalled();
927+
// resultReceived=true → worktree preserved, cleanup deferred to daemon loop
928+
expect(onCleanup).not.toHaveBeenCalled();
929929
});
930930

931-
it("removes session file when getTask returns done (normal completion)", async () => {
931+
it("session stays in in_review when agent finishes with result", async () => {
932932
const tunnel = makeTunnel();
933933
const resultEvent: AgentEvent = { type: "turn.end", cost: 0.001 };
934934
const handle = makeHandle([resultEvent]);
935935
const apiClient = makeApiClient({ getTask: vi.fn().mockResolvedValue({ status: "done" }) });
936936

937937
const pm = makePool(apiClient, makeCallbacks(), 0, tunnel);
938-
writeSession(makeWorkerSession("sess-normal-remove", { taskId: "task-normal-remove" }));
938+
writeSession(makeWorkerSession("sess-normal-keep", { taskId: "task-normal-keep" }));
939939

940940
await pm.spawnAgent({
941941
provider: makeProvider(handle),
942-
taskId: "task-normal-remove",
943-
sessionId: "sess-normal-remove",
942+
taskId: "task-normal-keep",
943+
sessionId: "sess-normal-keep",
944944
cwd: "/tmp",
945945
taskContext: "ctx",
946-
agentClient: makeAgentClient("a1", "sess-normal-remove"),
946+
agentClient: makeAgentClient("a1", "sess-normal-keep"),
947947
agentEnv: {},
948948
});
949949

950950
await flushPromises();
951951
await flushPromises();
952952

953-
// After cleanup_done the state machine removes the file
954-
expect(readSession("sess-normal-remove")).toBeNull();
953+
// Session preserved for daemon loop to clean up later
954+
const session = readSession("sess-normal-keep");
955+
expect(session).not.toBeNull();
956+
expect(session?.status).toBe("in_review");
955957
});
956958
});
957959

‎packages/cli/tests/rejectResumeFlow.integration.test.ts‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -799,10 +799,13 @@ describe("Scenario 7 (Fix 2): onComplete skips onCleanup when real session file
799799

800800
await flushPromises(8);
801801

802-
// Normal completion: onCleanup must be called
803-
expect(cleanupInvoked).toBe(true);
802+
// Agent produced result → worktree preserved for potential reject-resume.
803+
// Cleanup is deferred to daemon loop (checkRejectedReviews detects done → cleans up).
804+
expect(cleanupInvoked).toBe(false);
804805

805-
// Session file must be removed
806-
expect(readSession(sessionId)).toBeNull();
806+
// Session stays in in_review (daemon loop handles terminal cleanup)
807+
const session = readSession(sessionId);
808+
expect(session).not.toBeNull();
809+
expect(session?.status).toBe("in_review");
807810
});
808811
});

‎packages/cli/tests/scheduler.test.ts‎

Lines changed: 4 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -343,7 +343,7 @@ describe("DaemonLoop tick — in_review session resumption", () => {
343343
expect(sm.read("sess-notfound")).toBeNull();
344344
});
345345

346-
it("falls back to 'No reason provided' when no rejected note exists", async () => {
346+
it("cleans up session when in_progress with no rejected note", async () => {
347347
writeSession(makeWorkerSession("sess-noreason", "task-noreason", "in_review"));
348348

349349
const stubs = makeStubs();
@@ -356,12 +356,9 @@ describe("DaemonLoop tick — in_review session resumption", () => {
356356
await new Promise((r) => setTimeout(r, 100));
357357
loop.stop();
358358

359-
expect(mockResumeOneSession).toHaveBeenCalledWith(
360-
expect.objectContaining({ sessionId: "sess-noreason" }),
361-
expect.stringContaining("No reason provided"),
362-
expect.anything(),
363-
expect.anything(),
364-
);
359+
// No reject action → don't resume, clean up instead
360+
expect(mockResumeOneSession).not.toHaveBeenCalled();
361+
expect(sm.read("sess-noreason")).toBeNull();
365362
});
366363
});
367364

‎tests/runtimePool-coverage.test.ts‎

Lines changed: 8 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -604,7 +604,7 @@ describe("routeTurnEnd — per-segment token reporting", () => {
604604
expect(costCall[2].cost_micro_usd).toBe(Math.round(secondCost * 1_000_000));
605605
});
606606

607-
it("sets resultReceived=true so finalize calls getTask", async () => {
607+
it("resultReceived=true causes finalize to preserve worktree (no release)", async () => {
608608
const taskId = randomUUID();
609609
const sessionId = randomUUID();
610610
await seedActiveSession(sessions, sessionId, taskId);
@@ -614,21 +614,21 @@ describe("routeTurnEnd — per-segment token reporting", () => {
614614

615615
await spawnAndWait(apiClient, { events: [turnEnd], taskId, sessionId, agentClient });
616616

617-
// getTask is only called when resultReceived=true
618-
expect(agentClient.getTask as ReturnType<typeof vi.fn>).toHaveBeenCalledWith(taskId);
617+
// resultReceived=true → worktree preserved, no release
618+
expect(apiClient.releaseTask as ReturnType<typeof vi.fn>).not.toHaveBeenCalled();
619619
});
620620

621-
it("does NOT call getTask when no turn.end is received (resultReceived=false)", async () => {
621+
it("resultReceived=false causes finalize to release task", async () => {
622622
const taskId = randomUUID();
623623
const sessionId = randomUUID();
624624
await seedActiveSession(sessions, sessionId, taskId);
625625

626-
// No turn.end → resultReceived stays false → getTask not called
626+
// No turn.end → resultReceived stays false → release
627627
const agentClient = makeAgentClient(null);
628628

629629
await spawnAndWait(apiClient, { events: [], taskId, sessionId, agentClient });
630630

631-
expect(agentClient.getTask as ReturnType<typeof vi.fn>).not.toHaveBeenCalled();
631+
expect(apiClient.releaseTask as ReturnType<typeof vi.fn>).toHaveBeenCalledWith(taskId);
632632
});
633633
});
634634

@@ -959,8 +959,7 @@ describe("consumeEvents — syncs flags back to AgentProcess", () => {
959959

960960
await spawnAndWait(apiClient, { events: [makeTurnEndEvent(0.001)], taskId, sessionId, agentClient });
961961

962-
// getTask is called because resultReceived=true was synced back from flags
963-
expect(agentClient.getTask as ReturnType<typeof vi.fn>).toHaveBeenCalledWith(taskId);
962+
// resultReceived=true → worktree preserved, no release
964963
expect(apiClient.releaseTask as ReturnType<typeof vi.fn>).not.toHaveBeenCalled();
965964
});
966965

@@ -1217,10 +1216,8 @@ describe("multi-result: taskInReview reflects final state, not intermediate", ()
12171216
await spawnAndWait(apiClient, { handle, taskId, sessionId, agentClient });
12181217
await new Promise((r) => setTimeout(r, 10));
12191218

1220-
// in_review path: releaseTask is NOT called, worktree is preserved
1219+
// resultReceived=true → worktree preserved, no release
12211220
expect(apiClient.releaseTask as ReturnType<typeof vi.fn>).not.toHaveBeenCalled();
1222-
// getTask was called once after the iterator fully ended
1223-
expect(agentClient.getTask as ReturnType<typeof vi.fn>).toHaveBeenCalledWith(taskId);
12241221
});
12251222
});
12261223

‎tests/runtimePool-finalize.test.ts‎

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -299,17 +299,16 @@ describe("RuntimePool finalize() — releaseTask on completing", () => {
299299
});
300300

301301
// --------------------------------------------------------------------------
302-
// Case 2: agent finishes normally and task is NOT in_review (e.g. still in_progress)
303-
// → nextStatus === "completing", opts.crashed === false, taskInReview === false
304-
// → releaseTask MUST be called (the new behavior)
302+
// Case 2: agent finishes normally with result but task is in_progress
303+
// (e.g. rejected before finalize ran). Since agent produced a result,
304+
// worktree is preserved for resume — releaseTask must NOT be called.
305305
// --------------------------------------------------------------------------
306306

307-
it("calls releaseTask when agent finishes normally but task is NOT in_review", async () => {
307+
it("does NOT call releaseTask when agent produces result even if task is in_progress", async () => {
308308
const taskId = randomUUID();
309309
const sessionId = randomUUID();
310310
await seedActiveSession(sessions, sessionId, taskId);
311311

312-
// turn.end fires; getTask returns in_progress → agent.taskInReview = false
313312
const agentClient = makeAgentClient({ status: "in_progress" });
314313
const handle = makeHandle([makeTurnEndEvent()]);
315314
const provider = makeProvider(handle);
@@ -319,7 +318,7 @@ describe("RuntimePool finalize() — releaseTask on completing", () => {
319318
pool.spawnAgent({ provider, taskId, sessionId, cwd: "/tmp", taskContext: "test", agentClient, agentEnv: {} });
320319
});
321320

322-
expect(apiClient.releaseTask as ReturnType<typeof vi.fn>).toHaveBeenCalledWith(taskId);
321+
expect(apiClient.releaseTask as ReturnType<typeof vi.fn>).not.toHaveBeenCalled();
323322
});
324323

325324
// --------------------------------------------------------------------------

0 commit comments

Comments
 (0)