From aaa9d324d2761e568392e7f940651f2185864a4e Mon Sep 17 00:00:00 2001 From: marius-kilocode Date: Mon, 22 Jun 2026 10:19:50 +0200 Subject: [PATCH] fix(cli): preserve recall tail during queued turns --- .../src/kilocode/session/recall-search.ts | 19 ++++- packages/opencode/src/tool/recall.ts | 2 +- .../test/kilocode/recall-search.test.ts | 58 ++++++++++++++- packages/opencode/test/tool/recall.test.ts | 71 +++++++++++++++++++ 4 files changed, 145 insertions(+), 5 deletions(-) diff --git a/packages/opencode/src/kilocode/session/recall-search.ts b/packages/opencode/src/kilocode/session/recall-search.ts index e10d00e4217..b401a1ca7a8 100644 --- a/packages/opencode/src/kilocode/session/recall-search.ts +++ b/packages/opencode/src/kilocode/session/recall-search.ts @@ -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(SEARCH_SQL) + const statement = sqlite.prepare(SEARCH_SQL) const page = sqlite.prepare(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 diff --git a/packages/opencode/src/tool/recall.ts b/packages/opencode/src/tool/recall.ts index 94320969a7d..e5aeaafe736 100644 --- a/packages/opencode/src/tool/recall.ts +++ b/packages/opencode/src/tool/recall.ts @@ -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}`, diff --git a/packages/opencode/test/kilocode/recall-search.test.ts b/packages/opencode/test/kilocode/recall-search.test.ts index 125dc76fb83..e9f803702ab 100644 --- a/packages/opencode/test/kilocode/recall-search.test.ts +++ b/packages/opencode/test/kilocode/recall-search.test.ts @@ -15,7 +15,12 @@ type Stored = T extends unknown ? Omit : afterEach(resetDatabase) -function add(sessionID: SessionID, role: "user" | "assistant", data: Stored) { +function add( + sessionID: SessionID, + role: "user" | "assistant", + data: Stored, + opts?: { parentID?: MessageID }, +) { const messageID = MessageID.ascending() const message: Stored = role === "user" @@ -28,7 +33,7 @@ function add(sessionID: SessionID, role: "user" | "assistant", data: Stored { }) }) + 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({ diff --git a/packages/opencode/test/tool/recall.test.ts b/packages/opencode/test/tool/recall.test.ts index 442d290338a..2e666607a08 100644 --- a/packages/opencode/test/tool/recall.test.ts +++ b/packages/opencode/test/tool/recall.test.ts @@ -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 })