fix(core): harden union-find compaction after code review

Five bugs fixed, found via GPT-5.4 code review:

1. cosineSimilarity: handle mismatched vector dimensions (TF-IDF vocab
   grows over time, making newer vectors longer). Treats missing dims
   as zero; trailing dims still contribute to norm.

2. Forest.resolveDirty: replace _dirtyInputs.clear() with per-entry
   deletion and reference-equality guard. If union() replaces the
   inputs array during an async summarization, the stale summary is
   discarded. Combined inputs resolve in a future call.

3. TFIDFEmbedder.embedQuery: non-mutating embed for retrieval paths.
   render(query) was contaminating the corpus by calling embed(),
   which updates vocabulary and IDF. embedQuery uses current vocab
   without modifying state.

4. Forest.union centroid merging: handle mismatched embedding
   dimensions with Math.max and nullish coalescing.

5. ContextWindow constructor: reject evictAt < graduateAt
   (would corrupt _graduatedIndex tracking).

62/62 tests pass (11 new tests for the above).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: kimjune01

packages/core/src/services/contextWindow.test.ts

@@ -52,6 +52,25 @@ describe('cosineSimilarity', () => {
   it('should handle negative values', () => {
     expect(cosineSimilarity([1, 0], [-1, 0])).toBeCloseTo(-1.0);
   });
+
+  it('should handle mismatched dimensions without NaN', () => {
+    const short = [1, 0];
+    const long = [1, 0, 0.5, 0.3];
+    const sim = cosineSimilarity(long, short);
+    expect(Number.isNaN(sim)).toBe(false);
+    expect(sim).toBeGreaterThan(0);
+    // Symmetric
+    expect(cosineSimilarity(short, long)).toBeCloseTo(sim);
+  });
+
+  it('should return 0 for mismatched zero-overlap vectors', () => {
+    // short has values only in dims 0-1, long only in dims 2-3
+    const a = [1, 0];
+    const b = [0, 0, 1, 0];
+    const sim = cosineSimilarity(a, b);
+    expect(Number.isNaN(sim)).toBe(false);
+    expect(sim).toBeCloseTo(0.0);
+  });
 });
 
 // -- Forest --
@@ -281,6 +300,27 @@ describe('Forest', () => {
     expect(roots).toContain(root);
   });
 
+  it('should handle centroid merging with mismatched embedding dimensions', () => {
+    const embedder: Embedder = {
+      embed(text: string): number[] {
+        // Simulate growing vocab: earlier messages have shorter embeddings
+        if (text === 'early') return [1, 0];
+        return [0.5, 0.5, 0.3]; // later messages have longer embeddings
+      },
+    };
+    const forest = new Forest(embedder, stubSummarizer);
+    forest.insert(0, 'early');
+    forest.insert(1, 'later');
+    forest.union(0, 1);
+
+    const root = forest.find(0);
+    const centroid = forest.getCentroid(root);
+    expect(centroid).toBeDefined();
+    expect(centroid!.every((v) => !Number.isNaN(v))).toBe(true);
+    // Merged centroid should have max dimension length
+    expect(centroid!.length).toBe(3);
+  });
+
   it('should return no-op for union of same cluster', () => {
     const forest = new Forest(stubEmbedder, stubSummarizer);
     forest.insert(0, 'a');
@@ -400,6 +440,72 @@ describe('Forest', () => {
     expect(members).toContain(1);
   });
 
+  it('should not drop dirty state added by union during resolveDirty', async () => {
+    const slowSummarizer: Summarizer = {
+      async summarize(messages: string[]): Promise<string> {
+        // Simulate slow LLM call
+        await new Promise((resolve) => setTimeout(resolve, 10));
+        return messages.join('; ');
+      },
+    };
+    const forest = new Forest(stubEmbedder, slowSummarizer);
+    forest.insert(0, 'a');
+    forest.insert(1, 'b');
+    forest.union(0, 1); // dirty cluster {0,1}
+
+    // Start resolving — the await inside gives us a window
+    const resolvePromise = forest.resolveDirty();
+
+    // While resolve is in flight, add new dirty state
+    forest.insert(2, 'c');
+    forest.insert(3, 'd');
+    forest.union(2, 3); // new dirty cluster {2,3}
+
+    await resolvePromise;
+
+    // The new dirty cluster should NOT have been wiped
+    expect(forest.isDirty(forest.find(2))).toBe(true);
+
+    // Resolve it now
+    await forest.resolveDirty();
+    expect(forest.isDirty(forest.find(2))).toBe(false);
+  });
+
+  it('should not overwrite merged cluster dirty state when in-flight root is merged', async () => {
+    const slowSummarizer: Summarizer = {
+      async summarize(messages: string[]): Promise<string> {
+        await new Promise((resolve) => setTimeout(resolve, 10));
+        return messages.join('; ');
+      },
+    };
+    const forest = new Forest(stubEmbedder, slowSummarizer);
+    forest.insert(0, 'a');
+    forest.insert(1, 'b');
+    forest.union(0, 1); // dirty cluster {0,1}
+    const originalRoot = forest.find(0);
+
+    // Start resolving {0,1}
+    const resolvePromise = forest.resolveDirty();
+
+    // While {0,1} is being summarized, merge it into a new cluster
+    forest.insert(2, 'c');
+    forest.union(originalRoot, 2); // now {0,1,2} is dirty with combined inputs
+
+    await resolvePromise;
+
+    // The merged cluster should still be dirty — the stale summary
+    // from the in-flight call should NOT have resolved it
+    const mergedRoot = forest.find(0);
+    expect(forest.isDirty(mergedRoot)).toBe(true);
+
+    // Resolve it properly now
+    await forest.resolveDirty();
+    expect(forest.isDirty(forest.find(0))).toBe(false);
+    // Summary should include all three messages
+    const summary = forest.summary(forest.find(0))!;
+    expect(summary).toBeDefined();
+  });
+
   it('should list all roots', () => {
     const forest = new Forest(stubEmbedder, stubSummarizer);
     forest.insert(0, 'a');
@@ -688,6 +794,40 @@ describe('ContextWindow', () => {
     expect(cw.totalMessages).toBe(2);
   });
 
+  it('render(query) should not mutate the embedder corpus', () => {
+    const embedder = {
+      embed(text: string): number[] {
+        if (text.includes('cat')) return [1, 0, 0];
+        return [0, 0, 1];
+      },
+      embedQuery: vi.fn().mockReturnValue([1, 0, 0]),
+    };
+
+    const cw = new ContextWindow(embedder, stubSummarizer, {
+      graduateAt: 2,
+      evictAt: 4,
+      maxColdClusters: 10,
+      mergeThreshold: 0.0,
+    });
+
+    cw.append('cat info');
+    cw.append('dog info');
+    cw.append('hot1');
+
+    // render with query should call embedQuery, not embed
+    cw.render('cat question', 1, 0.0);
+    expect(embedder.embedQuery).toHaveBeenCalledWith('cat question');
+  });
+
+  it('should throw if evictAt < graduateAt', () => {
+    expect(() => {
+      new ContextWindow(stubEmbedder, stubSummarizer, {
+        graduateAt: 5,
+        evictAt: 3,
+      });
+    }).toThrow('evictAt (3) must be >= graduateAt (5)');
+  });
+
   it('should expose forest for direct access', () => {
     const cw = new ContextWindow(stubEmbedder, stubSummarizer);
     expect(cw.forest).toBeInstanceOf(Forest);

packages/core/src/services/contextWindow.ts

@@ -7,6 +7,12 @@
 /**
  * Union-find context compaction with overlap window and deferred summarization.
  *
+ * Reading order for reviewers:
+ *   1. cosineSimilarity() — handles mismatched vector dimensions safely
+ *   2. Forest class — union-find with path compression, deferred summarization
+ *   3. ContextWindow class — overlap window, graduation/eviction
+ *   4. Integration: chatCompressionService.ts compactWithUnionFind()
+ *
  * v2 architecture:
  * - append() is synchronous — no LLM calls. Graduation triggers structural
  *   union() only.
@@ -15,13 +21,17 @@
  *   in background during main LLM call wait.
  * - Overlap window (graduateAt/evictAt): graduated messages stay in hot zone
  *   for ~2 turns. By the time they evict, background resolveDirty() has
- *   resolved their cluster summaries. Zero blocking, zero staleness.
+ *   resolved their cluster summaries.
+ *
+ * Design doc: https://github.com/kimjune01/union-find-compaction-for-gemini-cli/blob/main/transformation-design.md
  */
 
 // -- Interfaces --
 
 export interface Embedder {
   embed(text: string): number[];
+  /** Embed without mutating internal state. Used for queries/retrieval. */
+  embedQuery?(text: string): number[];
 }
 
 export interface Summarizer {
@@ -41,15 +51,27 @@ export interface Message {
 
 // -- Helpers --
 
+// TF-IDF vocabulary grows over time, so newer vectors are longer than older
+// ones. We handle mismatched dimensions by treating missing entries as zero:
+// only shared dimensions contribute to the dot product, but trailing dimensions
+// still contribute to the norm (lowering similarity, as expected).
 export function cosineSimilarity(a: number[], b: number[]): number {
+  const len = Math.min(a.length, b.length);
   let dot = 0;
   let normA = 0;
   let normB = 0;
-  for (let i = 0; i < a.length; i++) {
+  for (let i = 0; i < len; i++) {
     dot += a[i] * b[i];
     normA += a[i] * a[i];
     normB += b[i] * b[i];
   }
+  // Include trailing dimensions from the longer vector in its norm
+  for (let i = len; i < a.length; i++) {
+    normA += a[i] * a[i];
+  }
+  for (let i = len; i < b.length; i++) {
+    normB += b[i] * b[i];
+  }
   normA = Math.sqrt(normA);
   normB = Math.sqrt(normB);
   if (normA === 0 || normB === 0) return 0.0;
@@ -167,7 +189,12 @@ export class Forest {
     if (ca && cb) {
       const na = membersA.length - membersB.length;
       const nb = membersB.length;
-      const merged = ca.map((v, i) => (v * na + cb[i] * nb) / (na + nb));
+      const total = na + nb;
+      const maxLen = Math.max(ca.length, cb.length);
+      const merged = new Array<number>(maxLen);
+      for (let i = 0; i < maxLen; i++) {
+        merged[i] = ((ca[i] ?? 0) * na + (cb[i] ?? 0) * nb) / total;
+      }
       this._centroids.set(rootA, merged);
     }
 
@@ -200,14 +227,26 @@ export class Forest {
   /**
    * Batch-summarize all dirty clusters. One LLM call per dirty root.
    * Called as fire-and-forget after render(), runs during main LLM call wait.
+   *
+   * Concurrency safety: union() can run between awaits (JS is single-threaded
+   * but yields at each await). When union() merges into a dirty root, it
+   * replaces _dirtyInputs with a new array containing combined content.
+   * We detect this via reference equality (=== check on the inputs array).
+   * If the array changed, we skip — the combined entry resolves next call.
    */
   async resolveDirty(): Promise<void> {
     const entries = [...this._dirtyInputs.entries()];
     for (const [root, inputs] of entries) {
+      if (!this._dirtyInputs.has(root)) continue;
       const summary = await this._summarizer.summarize(inputs);
-      this._summaries.set(root, summary);
+      if (this._dirtyInputs.get(root) === inputs) {
+        this._summaries.set(root, summary);
+        this._dirtyInputs.delete(root);
+      }
+      // If union() replaced the inputs (merged new content into this root)
+      // or merged this root away, skip — the combined dirty entry will be
+      // resolved in a future resolveDirty() call.
     }
-    this._dirtyInputs.clear();
   }
 
   /** Whether this cluster has unsummarized content. */
@@ -329,6 +368,11 @@ export class ContextWindow {
     this._evictAt = options.evictAt ?? 30;
     this._maxColdClusters = options.maxColdClusters ?? 10;
     this._mergeThreshold = options.mergeThreshold ?? 0.15;
+    if (this._evictAt < this._graduateAt) {
+      throw new Error(
+        `evictAt (${this._evictAt}) must be >= graduateAt (${this._graduateAt})`,
+      );
+    }
   }
 
   /**
@@ -420,7 +464,10 @@ export class ContextWindow {
     let cold: string[];
 
     if (query != null && this._forest.clusterCount() > 0) {
-      const queryEmb = this._embedder.embed(query);
+      // Use embedQuery (non-mutating) to avoid contaminating the TF-IDF corpus.
+      // embed() would add query terms to the vocabulary, changing future embeddings.
+      const embedFn = this._embedder.embedQuery ?? this._embedder.embed;
+      const queryEmb = embedFn.call(this._embedder, query);
       const topRoots = this._forest.nearest(queryEmb, k, minSim);
       cold = topRoots.map((r) => this._forest.compact(r));
     } else {

packages/core/src/services/embeddingService.test.ts

@@ -86,6 +86,32 @@ describe('TFIDFEmbedder', () => {
     }
   });
 
+  it('embedQuery should not mutate vocabulary or doc count', () => {
+    const embedder = new TFIDFEmbedder();
+    embedder.embed('alpha beta');
+    embedder.embed('gamma delta');
+    const vocabBefore = embedder.getVocabulary().length;
+
+    // embedQuery with new terms should not grow vocab
+    const qvec = embedder.embedQuery('epsilon zeta');
+    const vocabAfter = embedder.getVocabulary().length;
+
+    expect(vocabAfter).toBe(vocabBefore);
+    // Unknown terms should produce a zero vector (no known terms matched)
+    expect(qvec.every((v) => v === 0)).toBe(true);
+  });
+
+  it('embedQuery should use existing vocabulary for known terms', () => {
+    const embedder = new TFIDFEmbedder();
+    embedder.embed('cat dog fish');
+
+    const qvec = embedder.embedQuery('cat');
+    // Should produce non-zero vector since 'cat' is in vocab
+    expect(qvec.some((v) => v !== 0)).toBe(true);
+    // Same dimension as current vocab
+    expect(qvec.length).toBe(embedder.getVocabulary().length);
+  });
+
   it('should normalize vectors', () => {
     const embedder = new TFIDFEmbedder();
     const vec = embedder.embed('test normalization vector');

packages/core/src/services/embeddingService.ts

@@ -68,6 +68,46 @@ export class TFIDFEmbedder implements Embedder {
     return vec;
   }
 
+  /**
+   * Embed without mutating vocabulary, docCount, or termDocFreq.
+   * Used for queries/retrieval so searching doesn't contaminate the corpus.
+   */
+  embedQuery(text: string): number[] {
+    const tokens = this._tokenize(text);
+    if (tokens.length === 0) {
+      return new Array<number>(Math.max(this._vocab.size, 1)).fill(0);
+    }
+
+    const vec = new Array<number>(this._vocab.size).fill(0);
+    const termFreq = new Map<string, number>();
+
+    for (const token of tokens) {
+      termFreq.set(token, (termFreq.get(token) ?? 0) + 1);
+    }
+
+    for (const [term, count] of termFreq) {
+      const idx = this._vocab.get(term);
+      if (idx === undefined) continue; // unknown terms ignored
+
+      const tf = count / tokens.length;
+      const df = this._termDocFreq.get(term) ?? 1;
+      const idf = Math.log(1 + this._docCount / df);
+      vec[idx] = tf * idf;
+    }
+
+    // L2 normalize
+    const norm = Math.sqrt(
+      vec.reduce((sum: number, v: number) => sum + v * v, 0),
+    );
+    if (norm > 0) {
+      for (let i = 0; i < vec.length; i++) {
+        vec[i] /= norm;
+      }
+    }
+
+    return vec;
+  }
+
   getVocabulary(): string[] {
     const vocab = new Array<string>(this._vocab.size);
     for (const [term, idx] of this._vocab) {