Refresh hooks when refreshing extensions. (#14918)

packages/core/src/utils/extensionLoader.test.ts

@@ -35,6 +35,7 @@ describe('SimpleExtensionLoader', () => {
   let mockGeminiClientSetTools: MockInstance<
     typeof GeminiClient.prototype.setTools
   >;
+  let mockHookSystemInit: MockInstance;
 
   const activeExtension: GeminiCLIExtension = {
     name: 'test-extension',
@@ -61,6 +62,7 @@ describe('SimpleExtensionLoader', () => {
     } as unknown as McpClientManager;
     extensionReloadingEnabled = false;
     mockGeminiClientSetTools = vi.fn();
+    mockHookSystemInit = vi.fn();
     mockConfig = {
       getMcpClientManager: () => mockMcpClientManager,
       getEnableExtensionReloading: () => extensionReloadingEnabled,
@@ -68,6 +70,9 @@ describe('SimpleExtensionLoader', () => {
         isInitialized: () => true,
         setTools: mockGeminiClientSetTools,
       })),
+      getHookSystem: () => ({
+        initialize: mockHookSystemInit,
+      }),
     } as unknown as Config;
   });
 
@@ -125,13 +130,16 @@ describe('SimpleExtensionLoader', () => {
               mockMcpClientManager.startExtension,
             ).toHaveBeenCalledExactlyOnceWith(activeExtension);
             expect(mockRefreshServerHierarchicalMemory).toHaveBeenCalledOnce();
+            expect(mockHookSystemInit).toHaveBeenCalledOnce();
             expect(mockGeminiClientSetTools).toHaveBeenCalledOnce();
           } else {
             expect(mockMcpClientManager.startExtension).not.toHaveBeenCalled();
             expect(mockRefreshServerHierarchicalMemory).not.toHaveBeenCalled();
+            expect(mockHookSystemInit).not.toHaveBeenCalled();
             expect(mockGeminiClientSetTools).not.toHaveBeenCalledOnce();
           }
           mockRefreshServerHierarchicalMemory.mockClear();
+          mockHookSystemInit.mockClear();
           mockGeminiClientSetTools.mockClear();
 
           await loader.unloadExtension(activeExtension);
@@ -140,10 +148,12 @@ describe('SimpleExtensionLoader', () => {
               mockMcpClientManager.stopExtension,
             ).toHaveBeenCalledExactlyOnceWith(activeExtension);
             expect(mockRefreshServerHierarchicalMemory).toHaveBeenCalledOnce();
+            expect(mockHookSystemInit).toHaveBeenCalledOnce();
             expect(mockGeminiClientSetTools).toHaveBeenCalledOnce();
           } else {
             expect(mockMcpClientManager.stopExtension).not.toHaveBeenCalled();
             expect(mockRefreshServerHierarchicalMemory).not.toHaveBeenCalled();
+            expect(mockHookSystemInit).not.toHaveBeenCalled();
             expect(mockGeminiClientSetTools).not.toHaveBeenCalledOnce();
           }
         });
@@ -164,6 +174,7 @@ describe('SimpleExtensionLoader', () => {
               loader.loadExtension(anotherExtension),
             ]);
             expect(mockRefreshServerHierarchicalMemory).toHaveBeenCalledOnce();
+            expect(mockHookSystemInit).toHaveBeenCalledOnce();
           },
         );
       },

packages/core/src/utils/extensionLoader.ts

@@ -111,6 +111,7 @@ export abstract class ExtensionLoader {
       // reload memory, this is somewhat expensive and also busts the context
       // cache, we want to only do it once.
       await refreshServerHierarchicalMemory(this.config);
+      await this.config.getHookSystem()?.initialize();
     }
   }
 
@@ -134,13 +135,12 @@ export abstract class ExtensionLoader {
    * then calls `startExtension` to include all extension features into the
    * program.
    */
-  protected maybeStartExtension(
+  protected async maybeStartExtension(
     extension: GeminiCLIExtension,
-  ): Promise<void> | undefined {
+  ): Promise<void> {
     if (this.config && this.config.getEnableExtensionReloading()) {
-      return this.startExtension(extension);
+      await this.startExtension(extension);
     }
-    return;
   }
 
   /**
@@ -192,13 +192,12 @@ export abstract class ExtensionLoader {
    * then this also performs all necessary steps to remove all extension
    * features from the rest of the system.
    */
-  protected maybeStopExtension(
+  protected async maybeStopExtension(
     extension: GeminiCLIExtension,
-  ): Promise<void> | undefined {
+  ): Promise<void> {
     if (this.config && this.config.getEnableExtensionReloading()) {
-      return this.stopExtension(extension);
+      await this.stopExtension(extension);
     }
-    return;
   }
 
   async restartExtension(extension: GeminiCLIExtension): Promise<void> {