From 47a8c3e4a9c19883ecb482673052dd18fa54e26a Mon Sep 17 00:00:00 2001 From: YeonGyu-Kim Date: Tue, 17 Feb 2026 14:43:27 +0900 Subject: [PATCH] fix: harden run completion checks and graceful timeout --- src/cli/run/completion.test.ts | 19 +++++++ src/cli/run/completion.ts | 9 +++- src/cli/run/poll-for-completion.test.ts | 67 ++++++++++++++++++++++++- src/cli/run/poll-for-completion.ts | 28 +++++++++++ src/cli/run/runner.ts | 14 ++---- 5 files changed, 124 insertions(+), 13 deletions(-) diff --git a/src/cli/run/completion.test.ts b/src/cli/run/completion.test.ts index a763f68b..d54ae9ea 100644 --- a/src/cli/run/completion.test.ts +++ b/src/cli/run/completion.test.ts @@ -143,6 +143,25 @@ describe("checkCompletionConditions", () => { expect(result).toBe(false) }) + it("returns false when child status is missing", async () => { + // given + spyOn(console, "log").mockImplementation(() => {}) + const ctx = createMockContext({ + childrenBySession: { + "test-session": [{ id: "child-1" }], + "child-1": [], + }, + statuses: {}, + }) + const { checkCompletionConditions } = await import("./completion") + + // when + const result = await checkCompletionConditions(ctx) + + // then + expect(result).toBe(false) + }) + it("returns true when all descendants idle (recursive)", async () => { // given spyOn(console, "log").mockImplementation(() => {}) diff --git a/src/cli/run/completion.ts b/src/cli/run/completion.ts index cb4759d5..c2807f4e 100644 --- a/src/cli/run/completion.ts +++ b/src/cli/run/completion.ts @@ -65,7 +65,14 @@ async function areAllDescendantsIdle( for (const child of children) { const status = allStatuses[child.id] - if (status && status.type !== "idle") { + if (!status) { + console.log( + pc.dim(` Waiting: session ${child.id.slice(0, 8)}... status unknown`) + ) + return false + } + + if (status.type !== "idle") { console.log( pc.dim(` Waiting: session ${child.id.slice(0, 8)}... is ${status.type}`) ) diff --git a/src/cli/run/poll-for-completion.test.ts b/src/cli/run/poll-for-completion.test.ts index 46d8a0f8..47db01e8 100644 --- a/src/cli/run/poll-for-completion.test.ts +++ b/src/cli/run/poll-for-completion.test.ts @@ -310,7 +310,7 @@ describe("pollForCompletion", () => { //#then - returns 1 (not 130/timeout), error message printed expect(result).toBe(1) const errorCalls = (console.error as ReturnType).mock.calls - expect(errorCalls.some((call) => call[0]?.includes("Session ended with error"))).toBe(true) + expect(errorCalls.some((call: unknown[]) => String(call[0] ?? "").includes("Session ended with error"))).toBe(true) }) it("returns 1 when session errors while tool is active (error not masked by tool gate)", async () => { @@ -335,4 +335,69 @@ describe("pollForCompletion", () => { //#then - returns 1 expect(result).toBe(1) }) + + it("returns 130 after graceful timeout window expires", async () => { + //#given + spyOn(console, "log").mockImplementation(() => {}) + spyOn(console, "error").mockImplementation(() => {}) + const ctx = createMockContext({ + statuses: { + "test-session": { type: "busy" }, + }, + }) + const eventState = createEventState() + eventState.mainSessionIdle = false + eventState.hasReceivedMeaningfulWork = true + const abortController = new AbortController() + + //#when + const result = await pollForCompletion(ctx, eventState, abortController, { + pollIntervalMs: 10, + requiredConsecutive: 1, + minStabilizationMs: 0, + timeoutMs: 30, + timeoutGraceMs: 40, + }) + + //#then + expect(result).toBe(130) + expect(abortController.signal.aborted).toBe(true) + }) + + it("allows completion during graceful timeout window", async () => { + //#given + spyOn(console, "log").mockImplementation(() => {}) + spyOn(console, "error").mockImplementation(() => {}) + const ctx = createMockContext() + const eventState = createEventState() + eventState.mainSessionIdle = true + eventState.hasReceivedMeaningfulWork = true + const abortController = new AbortController() + let todoCalls = 0 + + ;(ctx.client.session as unknown as { + todo: ReturnType + children: ReturnType + status: ReturnType + }).todo = mock(async () => { + todoCalls++ + if (todoCalls === 1) { + return { data: [{ id: "1", content: "wip", status: "in_progress", priority: "high" }] } + } + return { data: [] } + }) + + //#when + const result = await pollForCompletion(ctx, eventState, abortController, { + pollIntervalMs: 10, + requiredConsecutive: 1, + minStabilizationMs: 0, + timeoutMs: 20, + timeoutGraceMs: 80, + }) + + //#then + expect(result).toBe(0) + expect(abortController.signal.aborted).toBe(false) + }) }) diff --git a/src/cli/run/poll-for-completion.ts b/src/cli/run/poll-for-completion.ts index f1ba6cd3..a1ea7415 100644 --- a/src/cli/run/poll-for-completion.ts +++ b/src/cli/run/poll-for-completion.ts @@ -8,11 +8,14 @@ const DEFAULT_POLL_INTERVAL_MS = 500 const DEFAULT_REQUIRED_CONSECUTIVE = 3 const ERROR_GRACE_CYCLES = 3 const MIN_STABILIZATION_MS = 10_000 +const DEFAULT_TIMEOUT_GRACE_MS = 15_000 export interface PollOptions { pollIntervalMs?: number requiredConsecutive?: number minStabilizationMs?: number + timeoutMs?: number + timeoutGraceMs?: number } export async function pollForCompletion( @@ -26,14 +29,35 @@ export async function pollForCompletion( options.requiredConsecutive ?? DEFAULT_REQUIRED_CONSECUTIVE const minStabilizationMs = options.minStabilizationMs ?? MIN_STABILIZATION_MS + const timeoutMs = options.timeoutMs ?? 0 + const timeoutGraceMs = options.timeoutGraceMs ?? DEFAULT_TIMEOUT_GRACE_MS let consecutiveCompleteChecks = 0 let errorCycleCount = 0 let firstWorkTimestamp: number | null = null const pollStartTimestamp = Date.now() + let timeoutNoticePrinted = false while (!abortController.signal.aborted) { await new Promise((resolve) => setTimeout(resolve, pollIntervalMs)) + if (abortController.signal.aborted) { + return 130 + } + + if (timeoutMs > 0) { + const elapsedMs = Date.now() - pollStartTimestamp + if (elapsedMs >= timeoutMs && !timeoutNoticePrinted) { + console.log(pc.yellow("\nTimeout reached. Entering graceful shutdown window...")) + timeoutNoticePrinted = true + } + + if (elapsedMs >= timeoutMs + timeoutGraceMs) { + console.log(pc.yellow("Grace period expired. Aborting...")) + abortController.abort() + return 130 + } + } + // ERROR CHECK FIRST — errors must not be masked by other gates if (eventState.mainSessionError) { errorCycleCount++ @@ -91,6 +115,10 @@ export async function pollForCompletion( const shouldExit = await checkCompletionConditions(ctx) if (shouldExit) { + if (abortController.signal.aborted) { + return 130 + } + consecutiveCompleteChecks++ if (consecutiveCompleteChecks >= requiredConsecutive) { console.log(pc.green("\n\nAll tasks completed.")) diff --git a/src/cli/run/runner.ts b/src/cli/run/runner.ts index e56accc8..91e4cb6b 100644 --- a/src/cli/run/runner.ts +++ b/src/cli/run/runner.ts @@ -48,14 +48,6 @@ export async function run(options: RunOptions): Promise { const pluginConfig = loadPluginConfig(directory, { command: "run" }) const resolvedAgent = resolveRunAgent(options, pluginConfig) const abortController = new AbortController() - let timeoutId: ReturnType | null = null - - if (timeout > 0) { - timeoutId = setTimeout(() => { - console.log(pc.yellow("\nTimeout reached. Aborting...")) - abortController.abort() - }, timeout) - } try { const { client, cleanup: serverCleanup } = await createServerConnection({ @@ -65,7 +57,6 @@ export async function run(options: RunOptions): Promise { }) const cleanup = () => { - if (timeoutId) clearTimeout(timeoutId) serverCleanup() } @@ -108,7 +99,9 @@ export async function run(options: RunOptions): Promise { }) console.log(pc.dim("Waiting for completion...\n")) - const exitCode = await pollForCompletion(ctx, eventState, abortController) + const exitCode = await pollForCompletion(ctx, eventState, abortController, { + timeoutMs: timeout, + }) // Abort the event stream to stop the processor abortController.abort() @@ -144,7 +137,6 @@ export async function run(options: RunOptions): Promise { throw err } } catch (err) { - if (timeoutId) clearTimeout(timeoutId) if (jsonManager) jsonManager.restore() if (err instanceof Error && err.name === "AbortError") { return 130