1 files changed, 9 insertions, 0 deletions
diff --git a/packages/cli/src/ui/commands/mcpCommand.test.ts b/packages/cli/src/ui/commands/mcpCommand.test.ts
index ad04cb69..0f339665 100644
--- a/packages/cli/src/ui/commands/mcpCommand.test.ts
+++ b/packages/cli/src/ui/commands/mcpCommand.test.ts
@@ -881,9 +881,14 @@ describe('mcpCommand', () => {
             }),
             getToolRegistry: vi.fn().mockResolvedValue(mockToolRegistry),
             getGeminiClient: vi.fn().mockReturnValue(mockGeminiClient),
+            getPromptRegistry: vi.fn().mockResolvedValue({
+              removePromptsByServer: vi.fn(),
+            }),
           },
         },
       });
+      // Mock the reloadCommands function
+      context.ui.reloadCommands = vi.fn();
 
       const { MCPOAuthProvider } = await import('@google/gemini-cli-core');
 
@@ -901,6 +906,7 @@ describe('mcpCommand', () => {
         'test-server',
       );
       expect(mockGeminiClient.setTools).toHaveBeenCalled();
+      expect(context.ui.reloadCommands).toHaveBeenCalledTimes(1);
 
       expect(isMessageAction(result)).toBe(true);
       if (isMessageAction(result)) {
@@ -985,6 +991,8 @@ describe('mcpCommand', () => {
           },
         },
       });
+      // Mock the reloadCommands function, which is new logic.
+      context.ui.reloadCommands = vi.fn();
 
       const refreshCommand = mcpCommand.subCommands?.find(
         (cmd) => cmd.name === 'refresh',
@@ -1002,6 +1010,7 @@ describe('mcpCommand', () => {
       );
       expect(mockToolRegistry.discoverMcpTools).toHaveBeenCalled();
       expect(mockGeminiClient.setTools).toHaveBeenCalled();
+      expect(context.ui.reloadCommands).toHaveBeenCalledTimes(1);
 
       expect(isMessageAction(result)).toBe(true);
       if (isMessageAction(result)) {