fix(cli): fix session restore race that triggered premature exit

.changeset/cozy-shirts-try.md

@@ -0,0 +1,5 @@
+---
+"@kilocode/cli": patch
+---
+
+fix session restore race that triggered premature exit

cli/src/cli.ts

@@ -278,8 +278,15 @@ export class CLI {
 				logs.debug("SessionManager workspace directory set", "CLI", { workspace })
 
 				if (this.options.session) {
+					// Set flag BEFORE restoring session to prevent race condition
+					// The session restoration triggers async state updates that may contain
+					// historical completion_result messages. Without this flag set first,
+					// the CI exit logic may trigger before the prompt can execute.
+					this.store.set(taskResumedViaContinueOrSessionAtom, true)
 					await this.sessionService?.restoreSession(this.options.session)
 				} else if (this.options.fork) {
+					// Set flag BEFORE forking session (same race condition as restore)
+					this.store.set(taskResumedViaContinueOrSessionAtom, true)
 					logs.info("Forking session from share ID", "CLI", { shareId: this.options.fork })
 					await this.sessionService?.forkSession(this.options.fork)
 				}

cli/src/state/atoms/__tests__/effects-ci-completion.test.ts

@@ -0,0 +1,71 @@
+/**
+ * Tests for CI completion detection in effects.ts
+ */
+
+import { describe, it, expect, beforeEach, vi } from "vitest"
+import { createStore } from "jotai"
+import { messageHandlerEffectAtom } from "../effects.js"
+import { extensionServiceAtom } from "../service.js"
+import { ciCompletionDetectedAtom, ciCompletionIgnoreBeforeTimestampAtom } from "../ci.js"
+import { taskResumedViaContinueOrSessionAtom } from "../extension.js"
+import type { ExtensionMessage, ExtensionChatMessage } from "../../../types/messages.js"
+import type { ExtensionService } from "../../../services/extension.js"
+
+describe("CI completion detection in effects", () => {
+	let store: ReturnType<typeof createStore>
+
+	beforeEach(() => {
+		store = createStore()
+
+		const mockService: Partial<ExtensionService> = {
+			initialize: vi.fn(),
+			dispose: vi.fn(),
+			on: vi.fn(),
+			off: vi.fn(),
+		}
+		store.set(extensionServiceAtom, mockService as ExtensionService)
+	})
+
+	it("skips completion detection when session was resumed", () => {
+		const completionMessage: ExtensionChatMessage = {
+			ts: Date.now(),
+			type: "ask",
+			ask: "completion_result",
+			text: "Task completed",
+		}
+
+		const stateMessage: ExtensionMessage = {
+			type: "state",
+			state: {
+				chatMessages: [completionMessage],
+			} as ExtensionMessage["state"],
+		}
+
+		store.set(taskResumedViaContinueOrSessionAtom, true)
+		store.set(messageHandlerEffectAtom, stateMessage)
+
+		expect(store.get(ciCompletionDetectedAtom)).toBe(false)
+	})
+
+	it("uses the ignore timestamp to skip historical completion_result", () => {
+		const historicalTs = Date.now()
+		const completionMessage: ExtensionChatMessage = {
+			ts: historicalTs,
+			type: "ask",
+			ask: "completion_result",
+			text: "Task completed",
+		}
+
+		const stateMessage: ExtensionMessage = {
+			type: "state",
+			state: {
+				chatMessages: [completionMessage],
+			} as ExtensionMessage["state"],
+		}
+
+		store.set(ciCompletionIgnoreBeforeTimestampAtom, historicalTs)
+		store.set(messageHandlerEffectAtom, stateMessage)
+
+		expect(store.get(ciCompletionDetectedAtom)).toBe(false)
+	})
+})

cli/src/state/atoms/ci.ts

@@ -20,6 +20,11 @@ export const ciTimeoutAtom = atom<number | undefined>(undefined)
  */
 export const ciCompletionDetectedAtom = atom<boolean>(false)
 
+/**
+ * Ignore completion_result messages at or before this timestamp
+ */
+export const ciCompletionIgnoreBeforeTimestampAtom = atom<number>(0)
+
 /**
  * Atom to track if command/message execution has finished
  */
@@ -45,6 +50,7 @@ export const resetCIStateAtom = atom(null, (get, set) => {
 	set(ciModeAtom, false)
 	set(ciTimeoutAtom, undefined)
 	set(ciCompletionDetectedAtom, false)
+	set(ciCompletionIgnoreBeforeTimestampAtom, 0)
 	set(ciCommandFinishedAtom, false)
 	set(ciExitReasonAtom, null)
 })

cli/src/state/atoms/effects.ts

@@ -13,8 +13,9 @@ import {
 	updateRouterModelsAtom,
 	chatMessagesAtom,
 	updateChatMessagesAtom,
+	taskResumedViaContinueOrSessionAtom,
 } from "./extension.js"
-import { ciCompletionDetectedAtom } from "./ci.js"
+import { ciCompletionDetectedAtom, ciCompletionIgnoreBeforeTimestampAtom } from "./ci.js"
 import {
 	updateProfileDataAtom,
 	updateBalanceDataAtom,
@@ -611,11 +612,22 @@ export const messageHandlerEffectAtom = atom(null, (get, set, message: Extension
 		if (message.state?.chatMessages) {
 			const lastMessage = message.state.chatMessages[message.state.chatMessages.length - 1]
 			if (lastMessage?.type === "ask" && lastMessage?.ask === "completion_result") {
-				logs.info("Completion result detected in state update", "effects")
+				const completionIgnoreBeforeTimestamp = get(ciCompletionIgnoreBeforeTimestampAtom)
+				const taskResumedViaSession = get(taskResumedViaContinueOrSessionAtom)
+				const isHistoricalCompletion =
+					lastMessage.ts !== undefined && lastMessage.ts <= completionIgnoreBeforeTimestamp
+
+				// Skip completion detection if session was just restored via --session or --continue
+				// The historical completion_result from the previous task should not trigger CI exit
+				if (taskResumedViaSession || isHistoricalCompletion) {
+					logs.debug("Skipping completion_result detection - historical completion_result", "effects")
+				} else {
+					logs.info("Completion result detected in state update", "effects")
 
-				set(ciCompletionDetectedAtom, true)
+					set(ciCompletionDetectedAtom, true)
 
-				SessionManager.init()?.doSync(true)
+					SessionManager.init()?.doSync(true)
+				}
 			}
 		}
 	} catch (error) {

cli/src/state/hooks/__tests__/useCIMode.test.tsx

@@ -0,0 +1,126 @@
+/**
+ * Tests for useCIMode hook behavior
+ */
+
+import React from "react"
+import { describe, it, expect, beforeEach, vi } from "vitest"
+import { createStore } from "jotai"
+import { Provider } from "jotai"
+import { render } from "ink-testing-library"
+import { useCIMode } from "../useCIMode.js"
+import { chatMessagesAtom, taskResumedViaContinueOrSessionAtom } from "../../atoms/extension.js"
+import { ciCompletionDetectedAtom, ciCompletionIgnoreBeforeTimestampAtom, ciExitReasonAtom } from "../../atoms/ci.js"
+import type { ExtensionChatMessage } from "../../../types/messages.js"
+
+vi.mock("../../../services/logs.js", () => ({
+	logs: {
+		info: vi.fn(),
+		debug: vi.fn(),
+		warn: vi.fn(),
+		error: vi.fn(),
+	},
+}))
+
+const noop = () => {}
+
+const TestComponent = ({ enabled }: { enabled: boolean }) => {
+	useCIMode({ enabled, onExit: noop })
+	return null
+}
+
+describe("useCIMode", () => {
+	let store: ReturnType<typeof createStore>
+
+	beforeEach(() => {
+		store = createStore()
+	})
+
+	it("skips historical completion_result after session resume", async () => {
+		const completionMessage: ExtensionChatMessage = {
+			ts: Date.now(),
+			type: "ask",
+			ask: "completion_result",
+			text: "Completed",
+		}
+
+		store.set(taskResumedViaContinueOrSessionAtom, true)
+		store.set(ciCompletionIgnoreBeforeTimestampAtom, completionMessage.ts)
+		store.set(chatMessagesAtom, [completionMessage])
+
+		const { unmount } = render(
+			<Provider store={store}>
+				<TestComponent enabled={true} />
+			</Provider>,
+		)
+
+		await new Promise((resolve) => setTimeout(resolve, 0))
+
+		expect(store.get(ciExitReasonAtom)).toBeNull()
+
+		unmount()
+	})
+
+	it("exits on completion_result when not ignored", async () => {
+		const completionMessage: ExtensionChatMessage = {
+			ts: Date.now(),
+			type: "ask",
+			ask: "completion_result",
+			text: "Completed",
+		}
+
+		store.set(taskResumedViaContinueOrSessionAtom, false)
+		store.set(ciCompletionIgnoreBeforeTimestampAtom, 0)
+		store.set(chatMessagesAtom, [completionMessage])
+
+		const { unmount } = render(
+			<Provider store={store}>
+				<TestComponent enabled={true} />
+			</Provider>,
+		)
+
+		await new Promise((resolve) => setTimeout(resolve, 0))
+
+		expect(store.get(ciExitReasonAtom)).toBe("completion_result")
+
+		unmount()
+	})
+
+	it("exits when a new completion_result arrives after the ignore timestamp", async () => {
+		const historicalTs = Date.now()
+		const historicalMessage: ExtensionChatMessage = {
+			ts: historicalTs,
+			type: "ask",
+			ask: "completion_result",
+			text: "Completed",
+		}
+
+		store.set(taskResumedViaContinueOrSessionAtom, false)
+		store.set(ciCompletionIgnoreBeforeTimestampAtom, historicalTs)
+		store.set(chatMessagesAtom, [historicalMessage])
+
+		const { unmount } = render(
+			<Provider store={store}>
+				<TestComponent enabled={true} />
+			</Provider>,
+		)
+
+		await new Promise((resolve) => setTimeout(resolve, 0))
+		expect(store.get(ciExitReasonAtom)).toBeNull()
+
+		const newMessage: ExtensionChatMessage = {
+			ts: historicalTs + 1000,
+			type: "ask",
+			ask: "completion_result",
+			text: "Completed again",
+		}
+
+		store.set(chatMessagesAtom, [historicalMessage, newMessage])
+		await new Promise((resolve) => setTimeout(resolve, 0))
+		store.set(ciCompletionDetectedAtom, true)
+
+		await new Promise((resolve) => setTimeout(resolve, 0))
+		expect(store.get(ciExitReasonAtom)).toBe("completion_result")
+
+		unmount()
+	})
+})

cli/src/state/hooks/useCIMode.ts

@@ -5,7 +5,13 @@
 
 import { useAtomValue, useSetAtom } from "jotai"
 import { useEffect, useState, useCallback, useRef } from "react"
-import { ciCompletionDetectedAtom, ciCommandFinishedAtom, ciExitReasonAtom } from "../atoms/ci.js"
+import {
+	ciCompletionDetectedAtom,
+	ciCompletionIgnoreBeforeTimestampAtom,
+	ciCommandFinishedAtom,
+	ciExitReasonAtom,
+} from "../atoms/ci.js"
+import { taskResumedViaContinueOrSessionAtom } from "../atoms/extension.js"
 import { useExtensionMessage } from "./useExtensionMessage.js"
 import { logs } from "../../services/logs.js"
 
@@ -67,8 +73,9 @@ export function useCIMode(options: UseCIModeOptions): UseCIModeReturn {
 	const { enabled, timeout } = options
 
 	const completionDetected = useAtomValue(ciCompletionDetectedAtom)
+	const completionIgnoreBeforeTimestamp = useAtomValue(ciCompletionIgnoreBeforeTimestampAtom)
 	const commandFinished = useAtomValue(ciCommandFinishedAtom)
-
+	const taskResumedViaSession = useAtomValue(taskResumedViaContinueOrSessionAtom)
 	// Write atoms
 	const setCommandFinished = useSetAtom(ciCommandFinishedAtom)
 	const setExitReason = useSetAtom(ciExitReasonAtom)
@@ -86,10 +93,18 @@ export function useCIMode(options: UseCIModeOptions): UseCIModeReturn {
 	// Get extension messages to monitor for completion_result
 	const { lastMessage } = useExtensionMessage()
 
+	const isHistoricalCompletion = lastMessage?.ts !== undefined && lastMessage.ts <= completionIgnoreBeforeTimestamp
+
 	// Monitor for completion_result messages
 	useEffect(() => {
 		if (!enabled || !lastMessage || exitTriggeredRef.current) return
 
+		// Skip if session was just restored - the completion_result is historical
+		if (taskResumedViaSession || isHistoricalCompletion) {
+			logs.debug("CI mode: Skipping completion_result check - historical message", "useCIMode")
+			return
+		}
+
 		// Check if this is a completion_result message
 		if (lastMessage.type === "ask" && lastMessage.ask === "completion_result") {
 			logs.info("CI mode: completion_result message received", "useCIMode")
@@ -98,7 +113,7 @@ export function useCIMode(options: UseCIModeOptions): UseCIModeReturn {
 			setShouldExit(true)
 			exitTriggeredRef.current = true
 		}
-	}, [enabled, lastMessage])
+	}, [enabled, lastMessage, taskResumedViaSession, isHistoricalCompletion])
 
 	// Monitor for command finished
 	useEffect(() => {
@@ -115,12 +130,18 @@ export function useCIMode(options: UseCIModeOptions): UseCIModeReturn {
 	useEffect(() => {
 		if (!enabled || !completionDetected || exitTriggeredRef.current) return
 
+		// Skip if session was just restored - the completion_result is historical
+		if (taskResumedViaSession || isHistoricalCompletion) {
+			logs.debug("CI mode: Skipping completion detected atom - historical message", "useCIMode")
+			return
+		}
+
 		logs.info("CI mode: completion detected via atom", "useCIMode")
 		setLocalExitReason("completion_result")
 		setExitReason("completion_result")
 		setShouldExit(true)
 		exitTriggeredRef.current = true
-	}, [enabled, completionDetected])
+	}, [enabled, completionDetected, taskResumedViaSession, isHistoricalCompletion])
 
 	// Setup timeout if provided
 	useEffect(() => {