fix(vscode): filter next edit history through access policy

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<MercuryEditRequestContext> {
     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(),
     }
   }
 

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()
+  })
 })

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<boolean>
+}
+
+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<string, string>()
   private readonly pendingTimers = new Map<string, NodeJS.Timeout>()
-  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<void> {
     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<string[]> {
+    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<void> {
     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<void> {
+    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<void> {
     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<boolean> {
+    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)
+  }
 }