feat(core): Implement JIT context memory loading and UI sync (#14469)

packages/cli/src/config/config.ts

@@ -435,9 +435,15 @@ export async function loadCliConfig(
   });
   await extensionManager.loadExtensions();
 
-  // Call the (now wrapper) loadHierarchicalGeminiMemory which calls the server's version
-  const { memoryContent, fileCount, filePaths } =
-    await loadServerHierarchicalMemory(
+  const experimentalJitContext = settings.experimental?.jitContext ?? false;
+
+  let memoryContent = '';
+  let fileCount = 0;
+  let filePaths: string[] = [];
+
+  if (!experimentalJitContext) {
+    // Call the (now wrapper) loadHierarchicalGeminiMemory which calls the server's version
+    const result = await loadServerHierarchicalMemory(
       cwd,
       [],
       debugMode,
@@ -448,6 +454,10 @@ export async function loadCliConfig(
       memoryFileFiltering,
       settings.context?.discoveryMaxDirs,
     );
+    memoryContent = result.memoryContent;
+    fileCount = result.fileCount;
+    filePaths = result.filePaths;
+  }
 
   const question = argv.promptInteractive || argv.prompt || '';
 

packages/cli/src/ui/commands/memoryCommand.test.ts

@@ -157,12 +157,14 @@ describe('memoryCommand', () => {
     let mockSetUserMemory: Mock;
     let mockSetGeminiMdFileCount: Mock;
     let mockSetGeminiMdFilePaths: Mock;
+    let mockContextManagerRefresh: Mock;
 
     beforeEach(() => {
       refreshCommand = getSubCommand('refresh');
       mockSetUserMemory = vi.fn();
       mockSetGeminiMdFileCount = vi.fn();
       mockSetGeminiMdFilePaths = vi.fn();
+      mockContextManagerRefresh = vi.fn().mockResolvedValue(undefined);
 
       const mockConfig = {
         setUserMemory: mockSetUserMemory,
@@ -185,6 +187,12 @@ describe('memoryCommand', () => {
         updateSystemInstructionIfInitialized: vi
           .fn()
           .mockResolvedValue(undefined),
+        isJitContextEnabled: vi.fn().mockReturnValue(false),
+        getContextManager: vi.fn().mockReturnValue({
+          refresh: mockContextManagerRefresh,
+        }),
+        getUserMemory: vi.fn().mockReturnValue(''),
+        getGeminiMdFileCount: vi.fn().mockReturnValue(0),
       };
 
       mockContext = createMockCommandContext({
@@ -203,7 +211,32 @@ describe('memoryCommand', () => {
       mockRefreshServerHierarchicalMemory.mockClear();
     });
 
-    it('should display success message when memory is refreshed with content', async () => {
+    it('should use ContextManager.refresh when JIT is enabled', async () => {
+      if (!refreshCommand.action) throw new Error('Command has no action');
+
+      // Enable JIT in mock config
+      const config = mockContext.services.config;
+      if (!config) throw new Error('Config is undefined');
+
+      vi.mocked(config.isJitContextEnabled).mockReturnValue(true);
+      vi.mocked(config.getUserMemory).mockReturnValue('JIT Memory Content');
+      vi.mocked(config.getGeminiMdFileCount).mockReturnValue(3);
+
+      await refreshCommand.action(mockContext, '');
+
+      expect(mockContextManagerRefresh).toHaveBeenCalledOnce();
+      expect(mockRefreshServerHierarchicalMemory).not.toHaveBeenCalled();
+
+      expect(mockContext.ui.addItem).toHaveBeenCalledWith(
+        {
+          type: MessageType.INFO,
+          text: 'Memory refreshed successfully. Loaded 18 characters from 3 file(s).',
+        },
+        expect.any(Number),
+      );
+    });
+
+    it('should display success message when memory is refreshed with content (Legacy)', async () => {
       if (!refreshCommand.action) throw new Error('Command has no action');
 
       const refreshResult: LoadServerHierarchicalMemoryResponse = {

packages/cli/src/ui/commands/memoryCommand.ts

@@ -87,8 +87,18 @@ export const memoryCommand: SlashCommand = {
         try {
           const config = context.services.config;
           if (config) {
-            const { memoryContent, fileCount } =
-              await refreshServerHierarchicalMemory(config);
+            let memoryContent = '';
+            let fileCount = 0;
+
+            if (config.isJitContextEnabled()) {
+              await config.getContextManager()?.refresh();
+              memoryContent = config.getUserMemory();
+              fileCount = config.getGeminiMdFileCount();
+            } else {
+              const result = await refreshServerHierarchicalMemory(config);
+              memoryContent = result.memoryContent;
+              fileCount = result.fileCount;
+            }
 
             await config.updateSystemInstructionIfInitialized();