1 files changed, 30 insertions, 34 deletions

diff --git a/packages/cli/src/ui/commands/ideCommand.test.ts b/packages/cli/src/ui/commands/ideCommand.test.ts
index 51322843..d1d72466 100644
--- a/packages/cli/src/ui/commands/ideCommand.test.ts
+++ b/packages/cli/src/ui/commands/ideCommand.test.ts
@@ -19,25 +19,10 @@ import { type Config } from '@google/gemini-cli-core';
 import * as child_process from 'child_process';
 import { glob } from 'glob';
 
-import {
-  getMCPDiscoveryState,
-  getMCPServerStatus,
-  IDE_SERVER_NAME,
-  MCPDiscoveryState,
-  MCPServerStatus,
-} from '@google/gemini-cli-core';
+import { IDEConnectionStatus } from '@google/gemini-cli-core/index.js';
 
 vi.mock('child_process');
 vi.mock('glob');
-vi.mock('@google/gemini-cli-core', async (importOriginal) => {
-  const original =
-    await importOriginal<typeof import('@google/gemini-cli-core')>();
-  return {
-    ...original,
-    getMCPServerStatus: vi.fn(),
-    getMCPDiscoveryState: vi.fn(),
-  };
-});
 
 function regexEscape(value: string) {
   return value.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
@@ -49,8 +34,6 @@ describe('ideCommand', () => {
   let execSyncSpy: MockInstance;
   let globSyncSpy: MockInstance;
   let platformSpy: MockInstance;
-  let getMCPServerStatusSpy: MockInstance;
-  let getMCPDiscoveryStateSpy: MockInstance;
 
   beforeEach(() => {
     mockContext = {
@@ -61,13 +44,12 @@ describe('ideCommand', () => {
 
     mockConfig = {
       getIdeMode: vi.fn(),
+      getIdeClient: vi.fn(),
     } as unknown as Config;
 
     execSyncSpy = vi.spyOn(child_process, 'execSync');
     globSyncSpy = vi.spyOn(glob, 'sync');
     platformSpy = vi.spyOn(process, 'platform', 'get');
-    getMCPServerStatusSpy = vi.mocked(getMCPServerStatus);
-    getMCPDiscoveryStateSpy = vi.mocked(getMCPDiscoveryState);
   });
 
   afterEach(() => {
@@ -91,15 +73,21 @@ describe('ideCommand', () => {
   });
 
   describe('status subcommand', () => {
+    const mockGetConnectionStatus = vi.fn();
     beforeEach(() => {
       vi.mocked(mockConfig.getIdeMode).mockReturnValue(true);
+      vi.mocked(mockConfig.getIdeClient).mockReturnValue({
+        getConnectionStatus: mockGetConnectionStatus,
+      } as ReturnType<Config['getIdeClient']>);
     });
 
     it('should show connected status', () => {
-      getMCPServerStatusSpy.mockReturnValue(MCPServerStatus.CONNECTED);
+      mockGetConnectionStatus.mockReturnValue({
+        status: IDEConnectionStatus.Connected,
+      });
       const command = ideCommand(mockConfig);
       const result = command!.subCommands![0].action!(mockContext, '');
-      expect(getMCPServerStatusSpy).toHaveBeenCalledWith(IDE_SERVER_NAME);
+      expect(mockGetConnectionStatus).toHaveBeenCalled();
       expect(result).toEqual({
         type: 'message',
         messageType: 'info',
@@ -108,37 +96,45 @@ describe('ideCommand', () => {
     });
 
     it('should show connecting status', () => {
-      getMCPServerStatusSpy.mockReturnValue(MCPServerStatus.CONNECTING);
+      mockGetConnectionStatus.mockReturnValue({
+        status: IDEConnectionStatus.Connecting,
+      });
       const command = ideCommand(mockConfig);
       const result = command!.subCommands![0].action!(mockContext, '');
+      expect(mockGetConnectionStatus).toHaveBeenCalled();
       expect(result).toEqual({
         type: 'message',
         messageType: 'info',
-        content: '🔄 Initializing...',
+        content: `🟡 Connecting...`,
       });
     });
-
-    it('should show discovery in progress status', () => {
-      getMCPServerStatusSpy.mockReturnValue(MCPServerStatus.DISCONNECTED);
-      getMCPDiscoveryStateSpy.mockReturnValue(MCPDiscoveryState.IN_PROGRESS);
+    it('should show disconnected status', () => {
+      mockGetConnectionStatus.mockReturnValue({
+        status: IDEConnectionStatus.Disconnected,
+      });
       const command = ideCommand(mockConfig);
       const result = command!.subCommands![0].action!(mockContext, '');
+      expect(mockGetConnectionStatus).toHaveBeenCalled();
       expect(result).toEqual({
         type: 'message',
-        messageType: 'info',
-        content: '🔄 Initializing...',
+        messageType: 'error',
+        content: `🔴 Disconnected`,
       });
     });
 
-    it('should show disconnected status', () => {
-      getMCPServerStatusSpy.mockReturnValue(MCPServerStatus.DISCONNECTED);
-      getMCPDiscoveryStateSpy.mockReturnValue(MCPDiscoveryState.COMPLETED);
+    it('should show disconnected status with details', () => {
+      const details = 'Something went wrong';
+      mockGetConnectionStatus.mockReturnValue({
+        status: IDEConnectionStatus.Disconnected,
+        details,
+      });
       const command = ideCommand(mockConfig);
       const result = command!.subCommands![0].action!(mockContext, '');
+      expect(mockGetConnectionStatus).toHaveBeenCalled();
       expect(result).toEqual({
         type: 'message',
         messageType: 'error',
-        content: '🔴 Disconnected',
+        content: `🔴 Disconnected: ${details}`,
       });
     });
   });