fix(cli): preserve recall tail during queued turns

This commit is contained in:
marius-kilocode
2026-06-22 10:19:50 +02:00
parent 4c3aaf428a
commit aaa9d324d2
4 changed files with 145 additions and 5 deletions
@@ -57,7 +57,12 @@ export namespace RecallSearch {
WHERE p.id = ids.value
AND m.id = p.message_id
AND m.session_id = p.session_id
AND NOT (m.session_id = ? AND m.id >= ?)
AND NOT (
m.session_id = ? AND (
(json_extract(m.data, '$.role') = 'user' AND m.id >= ?)
OR (json_extract(m.data, '$.role') = 'assistant' AND json_extract(m.data, '$.parentID') >= ?)
)
)
AND (${FILTER_SQL})`
const PAGE_SQL = `
@@ -176,7 +181,7 @@ export namespace RecallSearch {
const sqlite = Database.Client().$client
const rowid = sqlite.prepare<{ rowid: number | null }, []>(END_ROWID_SQL).get()?.rowid ?? 0
const partID = sqlite.prepare<{ id: string | null }, []>(END_ID_SQL).get()?.id ?? ""
const statement = sqlite.prepare<Row, [string, string, string]>(SEARCH_SQL)
const statement = sqlite.prepare<Row, [string, string, string, string]>(SEARCH_SQL)
const page = sqlite.prepare<PageRow, [string, number, number, string]>(PAGE_SQL)
const excludeSessionID = input.excludeSessionID ?? ""
const excludeFromMessageID = input.excludeFromMessageID ?? ""
@@ -218,6 +223,7 @@ export namespace RecallSearch {
JSON.stringify(rows.map((entry) => entry.partID)),
excludeSessionID,
excludeFromMessageID,
excludeFromMessageID,
)) {
consume(row)
}
@@ -265,6 +271,15 @@ export namespace RecallSearch {
return user?.info.id ?? messageID
}
export function visible(messages: MessageV2.WithParts[], messageID: MessageID) {
return messages.filter((message) => before(message.info, messageID))
}
function before(info: MessageV2.Info, messageID: MessageID) {
if (info.role === "user") return info.id < messageID
return info.parentID < messageID
}
function family(id: string) {
const row = Database.use((db) =>
db
+1 -1
View File
@@ -151,7 +151,7 @@ async function read(
const msgs = await bridge.promise(sessions.messages({ sessionID: session.id }))
const boundary = KiloSessionPromptQueue.active(ctx.sessionID) ?? RecallSearch.active(ctx.messages, ctx.messageID)
const visible = session.id === ctx.sessionID ? msgs.filter((message) => message.info.id < boundary) : msgs
const visible = session.id === ctx.sessionID ? RecallSearch.visible(msgs, boundary) : msgs
const lines: string[] = [
`# Session: ${session.title}`,
`Directory: ${session.directory}`,
@@ -15,7 +15,12 @@ type Stored<T> = T extends unknown ? Omit<T, "id" | "sessionID" | "messageID"> :
afterEach(resetDatabase)
function add(sessionID: SessionID, role: "user" | "assistant", data: Stored<MessageV2.Part>) {
function add(
sessionID: SessionID,
role: "user" | "assistant",
data: Stored<MessageV2.Part>,
opts?: { parentID?: MessageID },
) {
const messageID = MessageID.ascending()
const message: Stored<MessageV2.Info> =
role === "user"
@@ -28,7 +33,7 @@ function add(sessionID: SessionID, role: "user" | "assistant", data: Stored<Mess
: {
role,
time: { created: Date.now(), completed: Date.now() },
parentID: MessageID.ascending(),
parentID: opts?.parentID ?? MessageID.ascending(),
modelID: ModelID.make("test"),
providerID: ProviderID.make("test"),
mode: "code",
@@ -114,6 +119,55 @@ describe("RecallSearch", () => {
})
})
test("keeps prior assistant tail written after an active queued prompt", async () => {
await using tmp = await tmpdir({ git: true })
await provideTestInstance({
directory: tmp.path,
fn: async () => {
const sessions = await Effect.runPromise(Session.Service.pipe(Effect.provide(Session.defaultLayer)))
const session = await Effect.runPromise(sessions.create({ title: "Queued turn" }))
const previous = add(session.id, "user", { type: "text", text: "previous request" })
const active = add(session.id, "user", { type: "text", text: "queued prompt current-turn-needle" })
const tail = add(
session.id,
"assistant",
{ type: "text", text: "prior assistant tail tail-turn-needle" },
{ parentID: previous.messageID },
)
add(
session.id,
"assistant",
{ type: "text", text: "current assistant current-turn-needle" },
{ parentID: active.messageID },
)
const messages = await Effect.runPromise(sessions.messages({ sessionID: session.id }))
expect(RecallSearch.visible(messages, active.messageID).map((message) => message.info.id)).toEqual([
previous.messageID,
tail.messageID,
])
const result = await RecallSearch.search({
query: "tail-turn-needle",
projectID: Instance.project.id,
directories: [Instance.worktree],
excludeSessionID: session.id,
excludeFromMessageID: active.messageID,
})
expect(result.results.map((item) => item.id)).toEqual([session.id])
const current = await RecallSearch.search({
query: "current-turn-needle",
projectID: Instance.project.id,
directories: [Instance.worktree],
excludeSessionID: session.id,
excludeFromMessageID: active.messageID,
})
expect(current.results).toEqual([])
},
})
})
test("searches references and errors while excluding noisy content", async () => {
await using tmp = await tmpdir({ git: true })
await provideTestInstance({
@@ -55,6 +55,49 @@ const create = (title: string, text?: string | string[]) =>
),
)
const add = (input: { sessionID: SessionID; role: "user" | "assistant"; text: string; parentID?: MessageID }) =>
AppRuntime.runPromise(
Session.Service.use((svc) =>
Effect.gen(function* () {
const messageID = MessageID.ascending()
if (input.role === "user") {
yield* svc.updateMessage({
id: messageID,
sessionID: input.sessionID,
role: "user",
time: { created: Date.now() },
agent: "code",
model: { providerID: ProviderID.make("test"), modelID: ModelID.make("test") },
})
} else {
yield* svc.updateMessage({
id: messageID,
sessionID: input.sessionID,
role: "assistant",
time: { created: Date.now(), completed: Date.now() },
parentID: input.parentID ?? MessageID.ascending(),
modelID: ModelID.make("test"),
providerID: ProviderID.make("test"),
mode: "code",
agent: "code",
path: { cwd: "/tmp", root: "/tmp" },
cost: 0,
tokens: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } },
finish: "stop",
})
}
yield* svc.updatePart({
id: PartID.ascending(),
messageID,
sessionID: input.sessionID,
type: "text",
text: input.text,
})
return messageID
}),
),
)
describe("tool.recall", () => {
test("search is limited to the current project worktrees", async () => {
await using first = await tmpdir({ git: true })
@@ -130,6 +173,34 @@ describe("tool.recall", () => {
}
})
test("read includes prior assistant tail written after an active queued prompt", async () => {
await using tmp = await tmpdir({ git: true })
const result = await provideTestInstance({
directory: tmp.path,
fn: async () => {
const session = await create("queued read")
const previous = await add({ sessionID: session.id, role: "user", text: "previous request" })
const active = await add({ sessionID: session.id, role: "user", text: "active queued prompt" })
await add({ sessionID: session.id, role: "assistant", text: "visible prior assistant tail", parentID: previous })
await add({ sessionID: session.id, role: "assistant", text: "hidden current assistant", parentID: active })
const info = await AppRuntime.runPromise(RecallTool)
const tool = await AppRuntime.runPromise(info.init())
return AppRuntime.runPromise(
Effect.gen(function* () {
const sessions = yield* Session.Service
const messages = yield* sessions.messages({ sessionID: session.id })
return yield* tool.execute({ mode: "read", sessionID: session.id }, { ...ctx, sessionID: session.id, messages })
}),
)
},
})
expect(result.output).toContain("visible prior assistant tail")
expect(result.output).not.toContain("active queued prompt")
expect(result.output).not.toContain("hidden current assistant")
})
test("read rejects sessions from another project", async () => {
await using first = await tmpdir({ git: true })
await using second = await tmpdir({ git: true })