diff --git a/packages/opencode/src/kilocode/plan-followup.ts b/packages/opencode/src/kilocode/plan-followup.ts index 0fb61f827d8..6aa0efa2f79 100644 --- a/packages/opencode/src/kilocode/plan-followup.ts +++ b/packages/opencode/src/kilocode/plan-followup.ts @@ -111,10 +111,24 @@ export namespace PlanFollowup { export const ANSWER_NEW_SESSION = "Start new session" export const ANSWER_CONTINUE = "Continue here" - async function resolvePlan(input: { assistant: MessageV2.WithParts; sessionID: string }) { - const text = toText(input.assistant) + async function resolvePlan(input: { assistant?: MessageV2.WithParts; messages: MessageV2.WithParts[]; sessionID: string }) { + // Fast path: check the last assistant message's text first (avoids array scanning) + if (input.assistant) { + const text = toText(input.assistant) + if (text) return text + } + + // Fallback: scan all assistant messages after the last user message (handles + // cases where plan text is on an earlier assistant and the last one is empty) + const lastUserIdx = input.messages.findLastIndex((m) => m.info.role === "user") + const assistantMessages = input.messages + .slice(lastUserIdx + 1) + .filter((m) => m.info.role === "assistant") + + const text = assistantMessages.map(toText).filter(Boolean).join("\n\n").trim() if (text) return text + // Fall back to plan file on disk const session = await Session.get(input.sessionID) const file = Bun.file(Session.plan(session)) const plan = await file.text().catch(() => "") @@ -241,7 +255,7 @@ export namespace PlanFollowup { const assistant = latest.find((msg) => msg.info.role === "assistant") if (!assistant) return "break" - const plan = await resolvePlan({ assistant, sessionID: input.sessionID }) + const plan = await resolvePlan({ assistant, messages: input.messages, sessionID: input.sessionID }) if (!plan) return "break" const user = latest.find((msg) => msg.info.role === "user")?.info diff --git a/packages/opencode/src/session/prompt.ts b/packages/opencode/src/session/prompt.ts index 85a0f376754..d8e6ea9052a 100644 --- a/packages/opencode/src/session/prompt.ts +++ b/packages/opencode/src/session/prompt.ts @@ -62,14 +62,15 @@ const STRUCTURED_OUTPUT_SYSTEM_PROMPT = `IMPORTANT: The user has requested struc export namespace SessionPrompt { // kilocode_change start - share follow-up trigger logic with tests - export function shouldAskPlanFollowup(input: { - assistant: MessageV2.WithParts | undefined - abort: AbortSignal - }) { + export function shouldAskPlanFollowup(input: { messages: MessageV2.WithParts[]; abort: AbortSignal }) { if (input.abort.aborted) return false - if (!input.assistant) return false if (!["cli", "vscode"].includes(Flag.KILO_CLIENT)) return false - return input.assistant.parts.some((p) => p.type === "tool" && p.tool === "plan_exit" && p.state.status === "completed") + const lastUserIdx = input.messages.findLastIndex((m) => m.info.role === "user") + return input.messages + .slice(lastUserIdx + 1) + .some((msg) => + msg.parts.some((p) => p.type === "tool" && p.tool === "plan_exit" && p.state.status === "completed"), + ) } // kilocode_change end @@ -337,16 +338,12 @@ export namespace SessionPrompt { let lastUser: MessageV2.User | undefined let lastAssistant: MessageV2.Assistant | undefined - let lastAssistantMsg: MessageV2.WithParts | undefined // kilocode_change - capture full msg for plan_exit detection let lastFinished: MessageV2.Assistant | undefined let tasks: (MessageV2.CompactionPart | MessageV2.SubtaskPart)[] = [] for (let i = msgs.length - 1; i >= 0; i--) { const msg = msgs[i] if (!lastUser && msg.info.role === "user") lastUser = msg.info as MessageV2.User - if (!lastAssistant && msg.info.role === "assistant") { - lastAssistant = msg.info as MessageV2.Assistant - lastAssistantMsg = msg // kilocode_change - } + if (!lastAssistant && msg.info.role === "assistant") lastAssistant = msg.info as MessageV2.Assistant if (!lastFinished && msg.info.role === "assistant" && msg.info.finish) lastFinished = msg.info as MessageV2.Assistant if (lastUser && lastFinished) break @@ -363,7 +360,7 @@ export namespace SessionPrompt { lastUser.id < lastAssistant.id ) { // kilocode_change start - ask follow-up when plan_exit tool was called - if (shouldAskPlanFollowup({ assistant: lastAssistantMsg, abort })) { + if (shouldAskPlanFollowup({ messages: msgs, abort })) { const action = await PlanFollowup.ask({ sessionID, messages: msgs, abort }) if (action === "continue") continue } diff --git a/packages/opencode/test/kilocode/plan-exit-detection.test.ts b/packages/opencode/test/kilocode/plan-exit-detection.test.ts index eb06137be23..722e9aea0f8 100644 --- a/packages/opencode/test/kilocode/plan-exit-detection.test.ts +++ b/packages/opencode/test/kilocode/plan-exit-detection.test.ts @@ -126,11 +126,7 @@ describe("plan_exit detection", () => { }, ], }) - const assistant = seeded.messages - .slice() - .reverse() - .find((msg) => msg.info.role === "assistant") - expect(SessionPrompt.shouldAskPlanFollowup({ assistant, abort: AbortSignal.any([]) })).toBe(true) + expect(SessionPrompt.shouldAskPlanFollowup({ messages: seeded.messages, abort: AbortSignal.any([]) })).toBe(true) const pending = PlanFollowup.ask({ sessionID: seeded.sessionID, @@ -190,11 +186,7 @@ describe("plan_exit detection", () => { const seeded = await seed({ text: "Here is a partial plan, I have questions", }) - const assistant = seeded.messages - .slice() - .reverse() - .find((msg) => msg.info.role === "assistant") - expect(SessionPrompt.shouldAskPlanFollowup({ assistant, abort: AbortSignal.any([]) })).toBe(false) + expect(SessionPrompt.shouldAskPlanFollowup({ messages: seeded.messages, abort: AbortSignal.any([]) })).toBe(false) const list = await Question.list() expect(list).toHaveLength(0) })) @@ -271,17 +263,97 @@ describe("plan_exit detection", () => { expect(toolPart!.type === "tool" && toolPart!.state.status).toBe("error") // Use the shared predicate — errored plan_exit should not trigger - const assistantMsg = messages - .slice() - .reverse() - .find((msg) => msg.info.role === "assistant") - expect(SessionPrompt.shouldAskPlanFollowup({ assistant: assistantMsg, abort: AbortSignal.any([]) })).toBe(false) + expect(SessionPrompt.shouldAskPlanFollowup({ messages, abort: AbortSignal.any([]) })).toBe(false) // Confirm no questions were posted const list = await Question.list() expect(list).toHaveLength(0) })) + test("plan_exit on earlier assistant message triggers when later message has text only", () => + withInstance(async () => { + const session = await Session.create({}) + // Use explicit timestamps to ensure deterministic message ordering + const now = Date.now() + const user = await Session.updateMessage({ + id: Identifier.ascending("message"), + role: "user", + sessionID: session.id, + time: { created: now }, + agent: "plan", + model, + }) + await Session.updatePart({ + id: Identifier.ascending("part"), + messageID: user.id, + sessionID: session.id, + type: "text", + text: "Create a plan", + }) + + // First assistant message: has plan_exit tool, finish = tool-calls + const assistant1: MessageV2.Assistant = { + id: Identifier.ascending("message"), + role: "assistant", + sessionID: session.id, + time: { created: now + 1 }, + parentID: user.id, + modelID: model.modelID, + providerID: model.providerID, + mode: "plan", + agent: "plan", + path: { cwd: Instance.directory, root: Instance.worktree }, + cost: 0, + tokens: { total: 0, input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } }, + finish: "tool-calls", + } + await Session.updateMessage(assistant1) + await Session.updatePart({ + id: Identifier.ascending("part"), + messageID: assistant1.id, + sessionID: session.id, + type: "tool", + callID: Identifier.ascending("tool"), + tool: "plan_exit", + state: { + status: "completed", + input: {}, + output: "Plan is ready. Ending planning turn.", + title: "plan_exit", + metadata: {}, + time: { start: now + 1, end: now + 1 }, + }, + } satisfies MessageV2.ToolPart) + + // Second assistant message: text only, finish = end_turn (this is what lastAssistantMsg would point to) + const assistant2: MessageV2.Assistant = { + id: Identifier.ascending("message"), + role: "assistant", + sessionID: session.id, + time: { created: now + 2 }, + parentID: user.id, + modelID: model.modelID, + providerID: model.providerID, + mode: "plan", + agent: "plan", + path: { cwd: Instance.directory, root: Instance.worktree }, + cost: 0, + tokens: { total: 0, input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } }, + finish: "end_turn", + } + await Session.updateMessage(assistant2) + await Session.updatePart({ + id: Identifier.ascending("part"), + messageID: assistant2.id, + sessionID: session.id, + type: "text", + text: "The plan is complete. I've called plan_exit.", + }) + + const messages = await Session.messages({ sessionID: session.id }) + expect(SessionPrompt.shouldAskPlanFollowup({ messages, abort: AbortSignal.any([]) })).toBe(true) + })) + test("PlanFollowup.ask falls back to plan file for tool-only plan_exit turns", () => withInstance(async () => { const seeded = await seed({ @@ -314,4 +386,107 @@ describe("plan_exit detection", () => { }) await expect(pending).resolves.toBe("continue") })) + + test("PlanFollowup.ask shows prompt when plan text is on earlier assistant and last assistant is empty", () => + withInstance(async () => { + const session = await Session.create({}) + // Use explicit timestamps to ensure deterministic message ordering + const now = Date.now() + const user = await Session.updateMessage({ + id: Identifier.ascending("message"), + role: "user", + sessionID: session.id, + time: { created: now }, + agent: "plan", + model, + }) + await Session.updatePart({ + id: Identifier.ascending("part"), + messageID: user.id, + sessionID: session.id, + type: "text", + text: "Create a plan", + }) + + // First assistant message: has plan text + plan_exit tool + const assistant1: MessageV2.Assistant = { + id: Identifier.ascending("message"), + role: "assistant", + sessionID: session.id, + time: { created: now + 1 }, + parentID: user.id, + modelID: model.modelID, + providerID: model.providerID, + mode: "plan", + agent: "plan", + path: { cwd: Instance.directory, root: Instance.worktree }, + cost: 0, + tokens: { total: 0, input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } }, + finish: "tool-calls", + } + await Session.updateMessage(assistant1) + await Session.updatePart({ + id: Identifier.ascending("part"), + messageID: assistant1.id, + sessionID: session.id, + type: "text", + text: "Here is the detailed plan:\n\n## Step 1\nDo something\n\n## Step 2\nDo something else", + }) + await Session.updatePart({ + id: Identifier.ascending("part"), + messageID: assistant1.id, + sessionID: session.id, + type: "tool", + callID: Identifier.ascending("tool"), + tool: "plan_exit", + state: { + status: "completed", + input: {}, + output: "Plan is ready. Ending planning turn.", + title: "plan_exit", + metadata: {}, + time: { start: now + 1, end: now + 1 }, + }, + } satisfies MessageV2.ToolPart) + + // Second assistant message: empty (LLM follow-up after tool result) + const assistant2: MessageV2.Assistant = { + id: Identifier.ascending("message"), + role: "assistant", + sessionID: session.id, + time: { created: now + 2 }, + parentID: user.id, + modelID: model.modelID, + providerID: model.providerID, + mode: "plan", + agent: "plan", + path: { cwd: Instance.directory, root: Instance.worktree }, + cost: 0, + tokens: { total: 0, input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } }, + finish: "end_turn", + } + await Session.updateMessage(assistant2) + + const messages = await Session.messages({ sessionID: session.id }) + + // shouldAskPlanFollowup should detect plan_exit on the earlier message + expect(SessionPrompt.shouldAskPlanFollowup({ messages, abort: AbortSignal.any([]) })).toBe(true) + + // PlanFollowup.ask should find plan text from the earlier assistant and show prompt + const pending = PlanFollowup.ask({ + sessionID: session.id, + messages, + abort: AbortSignal.any([]), + }) + + const question = await waitQuestion(session.id) + expect(question).toBeDefined() + if (!question) return + expect(question.questions[0].header).toBe("Implement") + await Question.reply({ + requestID: question.id, + answers: [[PlanFollowup.ANSWER_CONTINUE]], + }) + await expect(pending).resolves.toBe("continue") + })) })