feat(core): union-find context compaction for AgentHistoryProvider

packages/core/src/code_assist/experiments/flagNames.ts

@@ -20,6 +20,7 @@ export const ExperimentFlags = {
   PRO_MODEL_NO_ACCESS: 45768879,
   GEMINI_3_1_FLASH_LITE_LAUNCHED: 45771641,
   DEFAULT_REQUEST_TIMEOUT: 45773134,
+  COMPRESSION_STRATEGY: 45768880,
 } as const;
 
 export type ExperimentFlagName =

packages/core/src/config/config.ts

@@ -3073,6 +3073,17 @@ export class Config implements McpContext, AgentLoopContext {
     return remoteThreshold;
   }
 
+  async getCompressionStrategy(): Promise<'flat' | 'union-find'> {
+    await this.ensureExperimentsLoaded();
+    const remoteStrategy =
+      this.experiments?.flags[ExperimentFlags.COMPRESSION_STRATEGY]
+        ?.stringValue;
+    if (remoteStrategy === 'union-find' || remoteStrategy === 'flat') {
+      return remoteStrategy;
+    }
+    return 'flat';
+  }
+
   async getUserCaching(): Promise<boolean | undefined> {
     await this.ensureExperimentsLoaded();
 

packages/core/src/context/agentHistoryProvider.ts

@@ -19,6 +19,7 @@ import {
   truncateProportionally,
   normalizeFunctionResponse,
 } from './truncation.js';
+import { sanitizePromptValue } from '../utils/sanitizePromptInput.js';
 
 export class AgentHistoryProvider {
   // TODO(joshualitt): just pass the BaseLlmClient instead of the whole Config.
@@ -379,10 +380,10 @@ Distill these into a high-density Markdown block that orientates the agent on th
 - **Brevity:** Maximum 15 lines. No conversational preamble.
 
 ${hasPreviousSummary ? 'PREVIOUS SUMMARY AND TRUNCATED HISTORY:' : 'TRUNCATED HISTORY:'}
-${JSON.stringify(messagesToTruncate)}
+${sanitizePromptValue(messagesToTruncate)}
 
 ACTIVE BRIDGE (LOOKAHEAD):
-${JSON.stringify(bridge)}`;
+${sanitizePromptValue(bridge)}`;
 
     const summaryResponse = await this.config
       .getBaseLlmClient()

packages/core/src/context/chatCompressionService.test.ts

@@ -29,6 +29,13 @@ vi.mock('../telemetry/loggers.js');
 vi.mock('../utils/environmentContext.js');
 vi.mock('../core/tokenLimits.js');
 
+function makeHistory(n: number): Content[] {
+  return Array.from({ length: n }, (_, i) => ({
+    role: (i % 2 === 0 ? 'user' : 'model'),
+    parts: [{ text: `message ${i}` }],
+  }));
+}
+
 describe('findCompressSplitPoint', () => {
   it('should throw an error for non-positive numbers', () => {
     expect(() => findCompressSplitPoint([], 0)).toThrow(
@@ -193,6 +200,7 @@ describe('ChatCompressionService', () => {
         getProjectTempDir: vi.fn().mockReturnValue(testTempDir),
       },
       getApprovedPlanPath: vi.fn().mockReturnValue('/path/to/plan.md'),
+      getCompressionStrategy: vi.fn().mockReturnValue('flat'),
     } as unknown as Config;
 
     vi.mocked(getInitialChatHistory).mockImplementation(
@@ -897,4 +905,139 @@ describe('ChatCompressionService', () => {
       );
     });
   });
+
+  describe('Compression strategy dispatch', () => {
+    it('should route to flat compression when strategy is flat', async () => {
+      vi.mocked(
+        mockConfig as unknown as {
+          getCompressionStrategy: () => 'flat' | 'union-find';
+        },
+      ).getCompressionStrategy = vi.fn().mockReturnValue('flat');
+
+      const history: Content[] = [
+        { role: 'user', parts: [{ text: 'msg1' }] },
+        { role: 'model', parts: [{ text: 'msg2' }] },
+        { role: 'user', parts: [{ text: 'msg3' }] },
+        { role: 'model', parts: [{ text: 'msg4' }] },
+      ];
+      vi.mocked(mockChat.getHistory).mockReturnValue(history);
+      vi.mocked(mockChat.getLastPromptTokenCount).mockReturnValue(600000);
+
+      const result = await service.compress(
+        mockChat,
+        mockPromptId,
+        false,
+        mockModel,
+        mockConfig,
+        false,
+      );
+
+      expect(result.info.compressionStatus).toBe(CompressionStatus.COMPRESSED);
+      // Flat uses 2 LLM calls (generate + verify)
+      expect(
+        mockConfig.getBaseLlmClient().generateContent,
+      ).toHaveBeenCalledTimes(2);
+    });
+
+    it('should route to union-find compression when strategy is union-find', async () => {
+      vi.mocked(
+        mockConfig as unknown as {
+          getCompressionStrategy: () => 'flat' | 'union-find';
+        },
+      ).getCompressionStrategy = vi.fn().mockReturnValue('union-find');
+
+      // Mock for cluster summarization
+      const mockLlmClient = {
+        generateContent: vi.fn().mockResolvedValue({
+          candidates: [
+            {
+              content: {
+                parts: [{ text: 'Cluster summary' }],
+              },
+            },
+          ],
+        } as unknown as GenerateContentResponse),
+      };
+      vi.mocked(mockConfig.getBaseLlmClient).mockReturnValue(
+        mockLlmClient as unknown as BaseLlmClient,
+      );
+
+      const history = makeHistory(35);
+      vi.mocked(mockChat.getHistory).mockReturnValue(history);
+      vi.mocked(mockChat.getLastPromptTokenCount).mockReturnValue(600000);
+
+      const result = await service.compress(
+        mockChat,
+        mockPromptId,
+        false,
+        mockModel,
+        mockConfig,
+        false,
+      );
+
+      expect(result.info.compressionStatus).toBe(CompressionStatus.COMPRESSED);
+      expect(result.newHistory).not.toBeNull();
+    });
+
+    it('should return valid ChatCompressionInfo for union-find path', async () => {
+      vi.mocked(
+        mockConfig as unknown as {
+          getCompressionStrategy: () => 'flat' | 'union-find';
+        },
+      ).getCompressionStrategy = vi.fn().mockReturnValue('union-find');
+
+      const mockLlmClient = {
+        generateContent: vi.fn().mockResolvedValue({
+          candidates: [
+            {
+              content: {
+                parts: [{ text: 'Summary' }],
+              },
+            },
+          ],
+        } as unknown as GenerateContentResponse),
+      };
+      vi.mocked(mockConfig.getBaseLlmClient).mockReturnValue(
+        mockLlmClient as unknown as BaseLlmClient,
+      );
+
+      const history = makeHistory(35);
+      vi.mocked(mockChat.getHistory).mockReturnValue(history);
+      vi.mocked(mockChat.getLastPromptTokenCount).mockReturnValue(600000);
+
+      const result = await service.compress(
+        mockChat,
+        mockPromptId,
+        false,
+        mockModel,
+        mockConfig,
+        false,
+      );
+
+      expect(result.info.originalTokenCount).toBe(600000);
+      expect(result.info.newTokenCount).toBeDefined();
+      expect(typeof result.info.newTokenCount).toBe('number');
+    });
+
+    it('should return NOOP for union-find with empty history', async () => {
+      vi.mocked(
+        mockConfig as unknown as {
+          getCompressionStrategy: () => 'flat' | 'union-find';
+        },
+      ).getCompressionStrategy = vi.fn().mockReturnValue('union-find');
+      vi.mocked(mockChat.getHistory).mockReturnValue([]);
+
+      const result = await service.compress(
+        mockChat,
+        mockPromptId,
+        false,
+        mockModel,
+        mockConfig,
+        false,
+      );
+
+      expect(result.info.compressionStatus).toBe(CompressionStatus.NOOP);
+      expect(result.newHistory).toBeNull();
+    });
+  });
 });