fix: Fix plan_exit call

This commit is contained in:
Alex Alecu
2026-02-25 09:46:17 +02:00
parent 0886ef8fa0
commit 7143683e20
3 changed files with 216 additions and 30 deletions
@@ -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
+9 -12
View File
@@ -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
}
@@ -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")
}))
})