mirror of
https://github.com/Kilo-Org/kilocode.git
synced 2026-10-29 18:52:45 +08:00
fix(cli): preserve recall tail during queued turns
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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 })
|
||||
|
||||
Reference in New Issue
Block a user