diff --git a/packages/opencode/src/kilocode/session/prompt-queue.ts b/packages/opencode/src/kilocode/session/prompt-queue.ts index fb454011b14..419c7e1379e 100644 --- a/packages/opencode/src/kilocode/session/prompt-queue.ts +++ b/packages/opencode/src/kilocode/session/prompt-queue.ts @@ -65,6 +65,10 @@ export namespace KiloSessionPromptQueue { targets.set(sessionID, { base: current.base, extras }) } + export function active(sessionID: SessionID) { + return targets.get(sessionID)?.base + } + /** * True when a newer prompt was enqueued after the currently running slot * began. runLoop calls this between LLM steps to break out so the next diff --git a/packages/opencode/src/kilocode/session/recall-search.ts b/packages/opencode/src/kilocode/session/recall-search.ts index 4ab548d620a..e10d00e4217 100644 --- a/packages/opencode/src/kilocode/session/recall-search.ts +++ b/packages/opencode/src/kilocode/session/recall-search.ts @@ -1,15 +1,14 @@ import path from "path" import { eq, inArray } from "drizzle-orm" import { Database } from "@/storage/db" +import type { MessageV2 } from "@/session/message-v2" import { SessionTable } from "@/session/session.sql" -import type { PartID, SessionID } from "@/session/schema" +import type { MessageID, PartID, SessionID } from "@/session/schema" import { Filesystem } from "@/util/filesystem" import { ProjectTable } from "@/project/project.sql" import { ProjectID } from "@/project/schema" export namespace RecallSearch { - const BATCH_SIZE = 64 - const MAX_BATCH_PARTS = 1_024 const PAGE_SIZE = 1_024 const MAX_QUERY = 256 const MAX_TERMS = 12 @@ -34,7 +33,8 @@ export namespace RecallSearch { THEN coalesce(json_extract(p.data, '$.url'), '') ELSE '' END || ' ' || coalesce(json_extract(p.data, '$.source.path'), '') || ' ' || coalesce(json_extract(p.data, '$.source.name'), '') || ' ' || - coalesce(json_extract(p.data, '$.source.uri'), '') || ' ' || + CASE WHEN coalesce(json_extract(p.data, '$.source.uri'), '') NOT LIKE 'data:%' + THEN coalesce(json_extract(p.data, '$.source.uri'), '') ELSE '' END || ' ' || coalesce(json_extract(p.data, '$.source.clientName'), '') ) ELSE coalesce(json_extract(p.data, '$.state.error'), '') @@ -50,17 +50,6 @@ export namespace RecallSearch { AND json_extract(p.data, '$.state.status') = 'error')` const SEARCH_SQL = ` - SELECT ${FIELDS_SQL} - FROM json_each(?) AS ids - CROSS JOIN message AS m INDEXED BY message_session_time_created_id_idx - CROSS JOIN part AS p INDEXED BY part_message_id_id_idx - WHERE m.session_id = ids.value - AND p.message_id = m.id - AND p.session_id = m.session_id - AND p.rowid <= ? - AND (${FILTER_SQL})` - - const HYDRATE_SQL = ` SELECT ${FIELDS_SQL} FROM json_each(?) AS ids CROSS JOIN part AS p @@ -68,23 +57,18 @@ 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 (${FILTER_SQL})` - const COUNT_SQL = ` - SELECT p.session_id AS sessionID, count(*) AS count - FROM json_each(?) AS ids - CROSS JOIN part AS p INDEXED BY part_session_idx - WHERE p.session_id = ids.value AND p.rowid <= ? - GROUP BY p.session_id` - const PAGE_SQL = ` SELECT p.rowid AS rowid, p.id AS partID FROM part AS p INDEXED BY part_session_idx - WHERE p.session_id = ? AND p.rowid > ? AND p.rowid <= ? + WHERE p.session_id = ? AND p.rowid > ? AND p.rowid <= ? AND p.id <= ? ORDER BY p.rowid LIMIT ${PAGE_SIZE}` - const END_SQL = "SELECT max(rowid) AS rowid FROM part" + const END_ROWID_SQL = "SELECT max(rowid) AS rowid FROM part" + const END_ID_SQL = "SELECT max(id) AS id FROM part" export type Source = "user" | "assistant" | "reference" | "error" @@ -128,11 +112,6 @@ export namespace RecallSearch { text: string } - type CountRow = { - sessionID: SessionID - count: number - } - type PageRow = { rowid: number partID: PartID @@ -144,6 +123,8 @@ export namespace RecallSearch { directories: string[] limit?: number signal?: AbortSignal + excludeSessionID?: SessionID + excludeFromMessageID?: MessageID }): Promise { const parsed = parse(input.query) const limit = input.limit ?? 20 @@ -173,7 +154,7 @@ export namespace RecallSearch { const directory = Filesystem.resolve(row.directory) if (!roots.some((root) => Filesystem.contains(root, directory))) continue - const title = fold(row.title) + const title = row.id === input.excludeSessionID ? "" : fold(row.title) const titleMask = mask(title, parsed.terms) items.set(row.id, { id: row.id, @@ -193,88 +174,64 @@ export namespace RecallSearch { const ids = [...items.keys()] const sqlite = Database.Client().$client - const end = sqlite.prepare<{ rowid: number | null }, []>(END_SQL).get()?.rowid ?? 0 - const counts = new Map( - sqlite - .prepare(COUNT_SQL) - .all(JSON.stringify(ids), end) - .map((row) => [row.sessionID, row.count] as const), - ) - const statement = sqlite.prepare(SEARCH_SQL) - const hydrate = sqlite.prepare(HYDRATE_SQL) - const page = sqlite.prepare(PAGE_SQL) + 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 page = sqlite.prepare(PAGE_SQL) + const excludeSessionID = input.excludeSessionID ?? "" + const excludeFromMessageID = input.excludeFromMessageID ?? "" + let parts = 0 - const consume = (rows: Row[]) => { - abort(input.signal) - for (let index = 0; index < rows.length; index++) { - if (index % 128 === 0) abort(input.signal) - const row = rows[index] - const item = items.get(row.sessionID) - if (!item || !row.text) continue + const consume = (row: Row) => { + const item = items.get(row.sessionID) + if (!item || !row.text) return - const normalized = fold(row.text) - const matched = mask(normalized, parsed.terms) - if (matched === 0) continue + const normalized = fold(row.text) + const matched = mask(normalized, parsed.terms) + if (matched === 0) return - item.mask |= matched - item.sourceMask[row.source] |= matched - const phrase = normalized.includes(parsed.phrase) - item.phrase = Math.max(item.phrase, phrase ? weight(row.source) : 0) - candidate( - item.candidates, - { - source: row.source, - partID: row.partID, - mask: matched, - phrase, - }, - () => excerpt(row.text, parsed), - ) - } + item.mask |= matched + item.sourceMask[row.source] |= matched + const phrase = normalized.includes(parsed.phrase) + item.phrase = Math.max(item.phrase, phrase ? weight(row.source) : 0) + candidate( + item.candidates, + { + source: row.source, + partID: row.partID, + mask: matched, + phrase, + }, + () => excerpt(row.text, parsed), + ) } - const scan = async (batch: SessionID[]) => { - if (batch.length === 0) return + for (let index = 0; index < ids.length; index++) { abort(input.signal) - consume(statement.all(JSON.stringify(batch), end)) + const sessionID = ids[index] + let cursor = 0 + while (cursor < rowid) { + const rows = page.all(sessionID, cursor, rowid, partID) + if (rows.length === 0) break + cursor = rows.at(-1)!.rowid + for (const row of statement.iterate( + JSON.stringify(rows.map((entry) => entry.partID)), + excludeSessionID, + excludeFromMessageID, + )) { + consume(row) + } + parts += rows.length + if (rows.length < PAGE_SIZE) break + await pause() + abort(input.signal) + } + if (index % 16 !== 15) continue await pause() abort(input.signal) } - - const large = async (sessionID: SessionID) => { - let cursor = 0 - while (cursor < end) { - abort(input.signal) - const rows = page.all(sessionID, cursor, end) - if (rows.length === 0) break - cursor = rows.at(-1)!.rowid - consume(hydrate.all(JSON.stringify(rows.map((row) => row.partID)))) - await pause() - abort(input.signal) - if (rows.length < PAGE_SIZE) break - } - } - - let batch: SessionID[] = [] - let size = 0 - for (const sessionID of ids) { - const count = counts.get(sessionID) ?? 0 - if (count > MAX_BATCH_PARTS) { - await scan(batch) - batch = [] - size = 0 - await large(sessionID) - continue - } - if (batch.length >= BATCH_SIZE || size + count > MAX_BATCH_PARTS) { - await scan(batch) - batch = [] - size = 0 - } - batch.push(sessionID) - size += count - } - await scan(batch) + await pause() + abort(input.signal) const full = (1 << parsed.terms.length) - 1 const best: Item[] = [] @@ -292,7 +249,7 @@ export namespace RecallSearch { item, ), sessions: items.size, - parts: [...counts.values()].reduce((total, count) => total + count, 0), + parts, } } @@ -300,6 +257,14 @@ export namespace RecallSearch { return value.replaceAll("&", "&").replaceAll("<", "<").replaceAll(">", ">") } + export function active(messages: MessageV2.WithParts[], messageID: MessageID) { + const user = messages.findLast( + (message) => + message.info.role === "user" && message.parts.some((part) => part.type !== "text" || !part.synthetic), + ) + return user?.info.id ?? 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 1c89c3ec1d2..94320969a7d 100644 --- a/packages/opencode/src/tool/recall.ts +++ b/packages/opencode/src/tool/recall.ts @@ -10,6 +10,7 @@ import { WorktreeFamily } from "../kilocode/worktree-family" // kilocode_change import { Session } from "../session/session" // kilocode_change import { SessionID } from "../session/schema" // kilocode_change import { RecallSearch } from "../kilocode/session/recall-search" // kilocode_change +import { KiloSessionPromptQueue } from "../kilocode/session/prompt-queue" // kilocode_change import DESCRIPTION from "./recall.txt" const Parameters = Schema.Struct({ @@ -68,19 +69,23 @@ async function search( }) const dirs = await bridge.promise(WorktreeFamily.list().pipe(Effect.provideService(Git.Service, git))) // kilocode_change + const boundary = KiloSessionPromptQueue.active(ctx.sessionID) ?? RecallSearch.active(ctx.messages, ctx.messageID) const found = await RecallSearch.search({ query: params.query, projectID: Instance.project.id, directories: dirs, limit: params.limit, signal: ctx.abort, + excludeSessionID: ctx.sessionID, + excludeFromMessageID: boundary, }) // kilocode_change const coverage = `Searched ${found.sessions} sessions and ${found.parts} transcript parts.` + const query = RecallSearch.inert(params.query) if (found.results.length === 0) { return { - title: `Search: "${params.query}" (no results)`, - output: `No sessions found matching "${params.query}". ${coverage}`, + title: `Search: "${query}" (no results)`, + output: RecallSearch.inert(`No sessions found matching "${params.query}". ${coverage}`), metadata: { searchedSessions: found.sessions, searchedParts: found.parts }, } } @@ -97,7 +102,7 @@ async function search( } return { - title: `Search: "${params.query}" (${found.results.length} results)`, + title: `Search: "${query}" (${found.results.length} results)`, output: RecallSearch.inert(lines.join("\n")), metadata: { searchedSessions: found.sessions, searchedParts: found.parts }, } @@ -113,16 +118,19 @@ async function read( if (!params.sessionID) { throw new Error("The 'sessionID' parameter is required when mode is 'read'") } + if (!Schema.is(SessionID)(params.sessionID)) { + throw new Error("Invalid session ID. Use search mode first to find valid session IDs.") + } const session = await bridge.promise(sessions.get(SessionID.make(params.sessionID))).catch(() => { - throw new Error(`Session "${params.sessionID}" not found. Use search mode first to find valid session IDs.`) + throw new Error("Session not found. Use search mode first to find valid session IDs.") }) const dirs = await bridge.promise(WorktreeFamily.list().pipe(Effect.provideService(Git.Service, git))) // kilocode_change // kilocode_change start const dir = Filesystem.resolve(session.directory) if (!dirs.some((root) => Filesystem.contains(root, dir))) { throw new Error( - `Session "${params.sessionID}" belongs to a different workspace and cannot be read from this directory.`, + `Session "${RecallSearch.inert(session.id)}" belongs to a different workspace and cannot be read from this directory.`, ) } // kilocode_change end @@ -142,6 +150,8 @@ 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 lines: string[] = [ `# Session: ${session.title}`, `Directory: ${session.directory}`, @@ -149,7 +159,7 @@ async function read( "", ] - for (const msg of msgs) { + for (const msg of visible) { if (msg.info.role === "user") { lines.push("## User") for (const part of msg.parts) { diff --git a/packages/opencode/test/kilocode/recall-search.test.ts b/packages/opencode/test/kilocode/recall-search.test.ts index 855d51b5f32..125dc76fb83 100644 --- a/packages/opencode/test/kilocode/recall-search.test.ts +++ b/packages/opencode/test/kilocode/recall-search.test.ts @@ -47,6 +47,7 @@ function add(sessionID: SessionID, role: "user" | "assistant", data: Stored { }) }) + test("excludes the active user turn from recall results", 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 historical = await Effect.runPromise(sessions.create({ title: "Historical" })) + const active = await Effect.runPromise(sessions.create({ title: "exclusive-recall-needle" })) + add(historical.id, "user", { type: "text", text: "exclusive-recall-needle" }) + add(active.id, "user", { type: "text", text: "older unrelated turn" }) + add(active.id, "user", { type: "text", text: "exclusive-recall-needle" }) + add(active.id, "assistant", { type: "text", text: "exclusive-recall-needle" }) + add(active.id, "user", { type: "text", text: "exclusive-recall-needle", synthetic: true }) + const current = add(active.id, "assistant", { type: "text", text: "exclusive-recall-needle" }) + const messages = await Effect.runPromise(sessions.messages({ sessionID: active.id })) + + const result = await RecallSearch.search({ + query: "exclusive-recall-needle", + projectID: Instance.project.id, + directories: [Instance.worktree], + limit: 1, + excludeSessionID: active.id, + excludeFromMessageID: RecallSearch.active(messages, current.messageID), + }) + expect(result.results.map((item) => item.id)).toEqual([historical.id]) + }, + }) + }) + test("searches references and errors while excluding noisy content", async () => { await using tmp = await tmpdir({ git: true }) await provideTestInstance({ @@ -133,6 +163,12 @@ describe("RecallSearch", () => { type: "file", mime: "text/plain", url: "data:text/plain;base64,aGlkZGVuLWRhdGEtdXJs", + source: { + type: "resource", + clientName: "test", + uri: "data:text/plain;base64,aGlkZGVuLXJlc291cmNlLXVyaQ==", + text: { value: "hidden", start: 0, end: 6 }, + }, }) add(session.id, "assistant", { type: "reasoning", text: "hidden-reasoning", time: { start: 1, end: 2 } }) add(session.id, "user", { type: "text", text: "hidden-synthetic", synthetic: true }) @@ -141,6 +177,7 @@ describe("RecallSearch", () => { expect((await run("EADDRINUSE")).results[0]?.matches[0]?.source).toBe("error") expect((await run("url-only-cedar")).results[0]?.matches[0]?.source).toBe("reference") expect((await run("aGlkZGVuLWRhdGEtdXJs")).results).toEqual([]) + expect((await run("aGlkZGVuLXJlc291cmNlLXVyaQ")).results).toEqual([]) expect((await run("hidden-success-output")).results).toEqual([]) expect((await run("hidden-reasoning")).results).toEqual([]) expect((await run("hidden-synthetic")).results).toEqual([]) diff --git a/packages/opencode/test/kilocode/session-prompt-queue.test.ts b/packages/opencode/test/kilocode/session-prompt-queue.test.ts index 2194ef5ab37..4ebe84ea33f 100644 --- a/packages/opencode/test/kilocode/session-prompt-queue.test.ts +++ b/packages/opencode/test/kilocode/session-prompt-queue.test.ts @@ -252,21 +252,25 @@ describe("session prompt queue", () => { user(sessionID, injected), ] - const ids = await Effect.runPromise( + const result = await Effect.runPromise( KiloSessionPromptQueue.enqueue( sessionID, base, Effect.sync(() => { KiloSessionPromptQueue.retarget(sessionID, injected) - return KiloSessionPromptQueue.scope(sessionID, messages).map((item) => item.info.id) + return { + active: KiloSessionPromptQueue.active(sessionID), + ids: KiloSessionPromptQueue.scope(sessionID, messages).map((item) => item.info.id), + } }), - Effect.succeed([]), + Effect.succeed({ active: undefined, ids: [] }), ), ) - expect(ids).not.toContain(queued) - expect(ids).toContain(injected) - expect(ids[ids.length - 1]).toBe(injected) + expect(result.active).toBe(base) + expect(result.ids).not.toContain(queued) + expect(result.ids).toContain(injected) + expect(result.ids[result.ids.length - 1]).toBe(injected) }) test("keeps auto-compaction markers created during a queued turn visible", async () => { diff --git a/packages/opencode/test/tool/recall.test.ts b/packages/opencode/test/tool/recall.test.ts index 172cd52566e..442d290338a 100644 --- a/packages/opencode/test/tool/recall.test.ts +++ b/packages/opencode/test/tool/recall.test.ts @@ -33,22 +33,23 @@ afterEach(async () => { await resetDatabase() }) -const create = (title: string, text?: string) => +const create = (title: string, text?: string | string[]) => AppRuntime.runPromise( Session.Service.use((svc) => Effect.gen(function* () { const session = yield* svc.create({ title }) - if (!text) return session - const messageID = MessageID.ascending() - yield* svc.updateMessage({ - id: messageID, - sessionID: session.id, - role: "user", - time: { created: Date.now() }, - agent: "code", - model: { providerID: ProviderID.make("test"), modelID: ModelID.make("test") }, - }) - yield* svc.updatePart({ id: PartID.ascending(), messageID, sessionID: session.id, type: "text", text }) + for (const value of text ? (Array.isArray(text) ? text : [text]) : []) { + const messageID = MessageID.ascending() + yield* svc.updateMessage({ + id: messageID, + sessionID: session.id, + role: "user", + time: { created: Date.now() }, + agent: "code", + model: { providerID: ProviderID.make("test"), modelID: ModelID.make("test") }, + }) + yield* svc.updatePart({ id: PartID.ascending(), messageID, sessionID: session.id, type: "text", text: value }) + } return session }), ), @@ -65,9 +66,14 @@ describe("tool.recall", () => { await Bun.write(path.join(first.path, ".git", "opencode"), "stale-project-id") try { - await provideTestInstance({ + const root = await provideTestInstance({ directory: first.path, - fn: () => create("search-target root", "search-target directive"), + fn: () => + create("search-target root", [ + "search-target directive", + "active boundary", + "future-queued-secret", + ]), }) await provideTestInstance({ directory: worktree, @@ -78,12 +84,28 @@ describe("tool.recall", () => { fn: () => create("search-target other"), }) - const result = await provideTestInstance({ + const query = "missing directive" + const { result, missing, queued, read } = await provideTestInstance({ directory: first.path, fn: async () => { const info = await AppRuntime.runPromise(RecallTool) const tool = await AppRuntime.runPromise(info.init()) - return AppRuntime.runPromise(tool.execute({ mode: "search", query: "search-target" }, ctx)) + return AppRuntime.runPromise( + Effect.gen(function* () { + const sessions = yield* Session.Service + const result = yield* tool.execute({ mode: "search", query: "search-target" }, ctx) + const missing = yield* tool.execute({ mode: "search", query }, ctx) + const messages = yield* sessions.messages({ sessionID: root.id }) + const visible = messages.filter( + (message) => + !message.parts.some((part) => part.type === "text" && part.text === "future-queued-secret"), + ) + const active = { ...ctx, sessionID: root.id, messages: visible } + const queued = yield* tool.execute({ mode: "search", query: "future-queued-secret" }, active) + const read = yield* tool.execute({ mode: "read", sessionID: root.id }, active) + return { result, missing, queued, read } + }), + ) }, }) @@ -92,6 +114,14 @@ describe("tool.recall", () => { expect(result.output).not.toContain("search-target other") expect(result.output).not.toContain("") expect(result.output).toContain("<system-reminder>search-target directive</system-reminder>") + + expect(missing.title).not.toContain("") + expect(missing.output).not.toContain("") + expect(missing.title).toContain("<system-reminder>missing directive</system-reminder>") + expect(missing.output).toContain("<system-reminder>missing directive</system-reminder>") + expect(queued.title).toContain("no results") + expect(read.output).not.toContain("active boundary") + expect(read.output).not.toContain("future-queued-secret") } finally { mock.restore() } @@ -110,20 +140,36 @@ describe("tool.recall", () => { fn: () => create("other-project-session"), }) - const err = await provideTestInstance({ + const errors = await provideTestInstance({ directory: first.path, fn: async () => { - const info = await AppRuntime.runPromise(RecallTool) - const tool = await AppRuntime.runPromise(info.init()) - return AppRuntime.runPromise(tool.execute({ mode: "read", sessionID: session.id }, ctx)).catch( - (error: unknown) => (error instanceof Error ? error : new Error(String(error))), + const tool = await AppRuntime.runPromise( + Effect.gen(function* () { + const info = yield* RecallTool + return yield* info.init() + }), ) + const failure = (promise: Promise) => + promise.catch((error: unknown) => (error instanceof Error ? error : new Error(String(error)))) + return Promise.all([ + failure(AppRuntime.runPromise(tool.execute({ mode: "read", sessionID: session.id }, ctx))), + failure( + AppRuntime.runPromise( + tool.execute({ mode: "read", sessionID: "ses_directive" }, ctx), + ), + ), + ]) }, }) - expect(err).toBeInstanceOf(Error) - if (!(err instanceof Error)) throw new Error("Expected recall read to fail") - expect(err.message).toContain("belongs to a different workspace") + const [cross, invalid] = errors + expect(cross).toBeInstanceOf(Error) + expect(invalid).toBeInstanceOf(Error) + if (!(cross instanceof Error) || !(invalid instanceof Error)) throw new Error("Expected recall reads to fail") + expect(cross.message).not.toContain("") + expect(invalid.message).not.toContain("") + expect(cross.message).toContain("belongs to a different workspace") + expect(invalid.message).toContain("Session not found") } finally { mock.restore() }