From 57e359a30218034ec45a5ab35be315b9dbed076f Mon Sep 17 00:00:00 2001 From: "kiloconnect[bot]" <240665456+kiloconnect[bot]@users.noreply.github.com> Date: Wed, 27 May 2026 10:59:11 +0000 Subject: [PATCH] fix(vscode): filter next edit history through access policy --- .../NextEditInlineCompletionProvider.ts | 13 +-- .../__tests__/editHistoryTracker.spec.ts | 35 +++++++- .../next-edit/editHistoryTracker.ts | 86 +++++++++++++++---- 3 files changed, 109 insertions(+), 25 deletions(-) diff --git a/packages/kilo-vscode/src/services/autocomplete/next-edit/NextEditInlineCompletionProvider.ts b/packages/kilo-vscode/src/services/autocomplete/next-edit/NextEditInlineCompletionProvider.ts index 34f5e5762dc..0c04e285129 100644 --- a/packages/kilo-vscode/src/services/autocomplete/next-edit/NextEditInlineCompletionProvider.ts +++ b/packages/kilo-vscode/src/services/autocomplete/next-edit/NextEditInlineCompletionProvider.ts @@ -49,7 +49,7 @@ export class NextEditInlineCompletionProvider implements vscode.InlineCompletion private currentAbort: AbortController | null = null constructor(private readonly deps: NextEditProviderDeps) { - this.editHistoryTracker = new EditHistoryTracker() + this.editHistoryTracker = new EditHistoryTracker({ isFileAllowed: deps.isFileAllowed }) } dispose(): void { @@ -77,7 +77,7 @@ export class NextEditInlineCompletionProvider implements vscode.InlineCompletion } const abort = this.swapAbortController(token) - const ctx = this.buildRequestContext(document, position) + const ctx = await this.buildRequestContext(document, position) const provider = new MercuryEditProvider({ connectionService: this.deps.connectionService, signal: abort.signal, @@ -103,12 +103,15 @@ export class NextEditInlineCompletionProvider implements vscode.InlineCompletion return abort } - private buildRequestContext(document: vscode.TextDocument, position: vscode.Position): MercuryEditRequestContext { + private async buildRequestContext( + document: vscode.TextDocument, + position: vscode.Position, + ): Promise { const { startLine, endLine } = computeEditableRegion({ cursorLine: position.line, totalLines: document.lineCount, }) - this.editHistoryTracker.flush(document) + await this.editHistoryTracker.flush(document) return { currentFilePath: document.uri.fsPath, currentFileContent: document.getText(), @@ -117,7 +120,7 @@ export class NextEditInlineCompletionProvider implements vscode.InlineCompletion editableRegionStartLine: startLine, editableRegionEndLine: endLine, recentlyViewedSnippets: this.deps.getRecentlyViewedSnippets?.(document) ?? [], - editDiffHistory: this.editHistoryTracker.getRecentDiffs(), + editDiffHistory: await this.editHistoryTracker.getRecentDiffs(), } } diff --git a/packages/kilo-vscode/src/services/autocomplete/next-edit/__tests__/editHistoryTracker.spec.ts b/packages/kilo-vscode/src/services/autocomplete/next-edit/__tests__/editHistoryTracker.spec.ts index 739b66ca0c5..0b5f0406d3d 100644 --- a/packages/kilo-vscode/src/services/autocomplete/next-edit/__tests__/editHistoryTracker.spec.ts +++ b/packages/kilo-vscode/src/services/autocomplete/next-edit/__tests__/editHistoryTracker.spec.ts @@ -33,7 +33,7 @@ function doc(path: string, initial: string): Doc { } describe("EditHistoryTracker", () => { - it("retains chronological edits across files for Mercury context", () => { + it("retains chronological edits across files for Mercury context", async () => { const tracker = new EditHistoryTracker() const a = doc("/workspace/a.ts", "const a = 1\n") const b = doc("/workspace/b.ts", "const b = 1\n") @@ -42,11 +42,11 @@ describe("EditHistoryTracker", () => { open(a) open(b) a.setText("const a = 2\n") - tracker.flush(a) + await tracker.flush(a) b.setText("const b = 2\n") - tracker.flush(b) + await tracker.flush(b) - const diffs = tracker.getRecentDiffs() + const diffs = await tracker.getRecentDiffs() expect(diffs).toHaveLength(2) expect(diffs[0]).toContain("a.ts") expect(diffs[0]).toContain("+const a = 2") @@ -55,4 +55,31 @@ describe("EditHistoryTracker", () => { tracker.dispose() }) + + it("never returns edits from denied documents", async () => { + const denied = new Set(["/workspace/.env"]) + const tracker = new EditHistoryTracker({ isFileAllowed: async (path) => !denied.has(path) }) + const safe = doc("/workspace/app.ts", "const safe = 1\n") + const secret = doc("/workspace/.env", "TOKEN=old\n") + const open = (vscode.workspace as unknown as { open(doc: vscode.TextDocument): void }).open + + open(safe) + open(secret) + await Promise.resolve() + await Promise.resolve() + secret.setText("TOKEN=secret\n") + await tracker.flush(secret) + safe.setText("const safe = 2\n") + await tracker.flush(safe) + + const diffs = await tracker.getRecentDiffs() + expect(diffs).toHaveLength(1) + expect(diffs[0]).toContain("app.ts") + expect(diffs[0]).not.toContain("TOKEN=secret") + + denied.add("/workspace/app.ts") + expect(await tracker.getRecentDiffs()).toEqual([]) + + tracker.dispose() + }) }) diff --git a/packages/kilo-vscode/src/services/autocomplete/next-edit/editHistoryTracker.ts b/packages/kilo-vscode/src/services/autocomplete/next-edit/editHistoryTracker.ts index ce9a00af02e..215fefc626d 100644 --- a/packages/kilo-vscode/src/services/autocomplete/next-edit/editHistoryTracker.ts +++ b/packages/kilo-vscode/src/services/autocomplete/next-edit/editHistoryTracker.ts @@ -4,6 +4,17 @@ import * as vscode from "vscode" const DEFAULT_DEBOUNCE_MS = 1500 const DEFAULT_MAX_DIFFS = 5 +type Options = { + debounceMs?: number + maxDiffs?: number + isFileAllowed?: (fsPath: string) => Promise +} + +type Diff = { + key: string + patch: string +} + /** * Tracks per-file snapshots and emits a workspace-wide chronological stream * of range-based unidiffs after a short idle window. Cross-file history is @@ -16,31 +27,30 @@ const DEFAULT_MAX_DIFFS = 5 export class EditHistoryTracker implements vscode.Disposable { private readonly snapshots = new Map() private readonly pendingTimers = new Map() - private readonly diffs: string[] = [] + private readonly diffs: Diff[] = [] private readonly subscriptions: vscode.Disposable[] = [] - constructor( - private readonly options: { debounceMs?: number; maxDiffs?: number } = {}, - ) { + constructor(private readonly options: Options = {}) { const debounceMs = options.debounceMs ?? DEFAULT_DEBOUNCE_MS // Seed snapshots on open so the FIRST edit in a freshly-opened file is // captured in the diff history (otherwise the common "open, type, trigger" - // flow ships an empty edit-history block). + // flow ships an empty edit-history block). Access checks happen before + // reading text so ignored documents are never retained as edit context. this.subscriptions.push( vscode.workspace.onDidOpenTextDocument((doc) => { if (doc.uri.scheme !== "file") return - if (!this.snapshots.has(doc.uri.fsPath)) this.snapshots.set(doc.uri.fsPath, doc.getText()) + void this.seed(doc) }), ) for (const doc of vscode.workspace.textDocuments) { - if (doc.uri.scheme === "file") this.snapshots.set(doc.uri.fsPath, doc.getText()) + if (doc.uri.scheme === "file") void this.seed(doc) } this.subscriptions.push( vscode.workspace.onDidChangeTextDocument((event) => { if (event.document.uri.scheme !== "file") return if (event.contentChanges.length === 0) return - this.scheduleSnapshotDiff(event.document, debounceMs) + void this.scheduleSnapshotDiff(event.document, debounceMs) }), ) this.subscriptions.push( @@ -59,17 +69,25 @@ export class EditHistoryTracker implements vscode.Disposable { * this immediately before building a request so the freshest user edit * makes it into the prompt. */ - public flush(document: vscode.TextDocument): void { + public async flush(document: vscode.TextDocument): Promise { const key = document.uri.fsPath + if (!(await this.allowed(key))) { + this.reject(key) + return + } const t = this.pendingTimers.get(key) if (t) clearTimeout(t) this.pendingTimers.delete(key) - this.emitDiffNow(document) + await this.emitDiffNow(document) } /** Workspace-wide oldest to newest, matching the Mercury prompt-history convention. */ - public getRecentDiffs(): string[] { - return [...this.diffs] + public async getRecentDiffs(): Promise { + const kept = ( + await Promise.all(this.diffs.map(async (diff) => ((await this.allowed(diff.key)) ? diff : undefined))) + ).filter((diff): diff is Diff => diff !== undefined) + this.diffs.splice(0, this.diffs.length, ...kept) + return kept.map((diff) => diff.patch) } public dispose(): void { @@ -79,8 +97,26 @@ export class EditHistoryTracker implements vscode.Disposable { this.subscriptions.length = 0 } - private scheduleSnapshotDiff(document: vscode.TextDocument, debounceMs: number): void { + private async seed(document: vscode.TextDocument): Promise { const key = document.uri.fsPath + if (this.snapshots.has(key)) return + if (!this.options.isFileAllowed) { + this.snapshots.set(key, document.getText()) + return + } + if (!(await this.allowed(key))) { + this.reject(key) + return + } + if (!this.snapshots.has(key)) this.snapshots.set(key, document.getText()) + } + + private async scheduleSnapshotDiff(document: vscode.TextDocument, debounceMs: number): Promise { + const key = document.uri.fsPath + if (!(await this.allowed(key))) { + this.reject(key) + return + } if (!this.snapshots.has(key)) { // Fallback seed for documents we never saw open (e.g. opened before the // tracker existed). The triggering change is lost, but subsequent edits @@ -92,13 +128,17 @@ export class EditHistoryTracker implements vscode.Disposable { if (existing) clearTimeout(existing) const timer = setTimeout(() => { this.pendingTimers.delete(key) - this.emitDiffNow(document) + void this.emitDiffNow(document) }, debounceMs) this.pendingTimers.set(key, timer) } - private emitDiffNow(document: vscode.TextDocument): void { + private async emitDiffNow(document: vscode.TextDocument): Promise { const key = document.uri.fsPath + if (!(await this.allowed(key))) { + this.reject(key) + return + } const previous = this.snapshots.get(key) if (previous === undefined) return const current = document.getText() @@ -108,10 +148,24 @@ export class EditHistoryTracker implements vscode.Disposable { const patch = createPatch(filename, previous, current, undefined, undefined, { context: 1 }) // `createPatch` returns "" for identical inputs; guard anyway. if (patch && patch.trim().length > 0) { - this.diffs.push(patch) + this.diffs.push({ key, patch }) const maxDiffs = this.options.maxDiffs ?? DEFAULT_MAX_DIFFS if (this.diffs.length > maxDiffs) this.diffs.shift() } this.snapshots.set(key, current) } + + private async allowed(key: string): Promise { + if (!this.options.isFileAllowed) return true + return this.options.isFileAllowed(key).catch(() => false) + } + + private reject(key: string): void { + const timer = this.pendingTimers.get(key) + if (timer) clearTimeout(timer) + this.pendingTimers.delete(key) + this.snapshots.delete(key) + const kept = this.diffs.filter((diff) => diff.key !== key) + this.diffs.splice(0, this.diffs.length, ...kept) + } }