Refresh hooks when refreshing extensions. (#14918)

packages/core/src/hooks/hookRegistry.test.ts

@@ -6,11 +6,7 @@
 
 import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
 import * as fs from 'node:fs';
-import {
-  HookRegistry,
-  ConfigSource,
-  HookRegistryNotInitializedError,
-} from './hookRegistry.js';
+import { HookRegistry, ConfigSource } from './hookRegistry.js';
 import type { Storage } from '../config/storage.js';
 import { HookEventName, HookType } from './types.js';
 import type { Config } from '../config/config.js';
@@ -258,14 +254,6 @@ describe('HookRegistry', () => {
       );
       expect(notificationHooks).toHaveLength(0);
     });
-
-    it('should throw error if not initialized', () => {
-      const uninitializedRegistry = new HookRegistry(mockConfig);
-
-      expect(() => {
-        uninitializedRegistry.getHooksForEvent(HookEventName.BeforeTool);
-      }).toThrow(HookRegistryNotInitializedError);
-    });
   });
 
   describe('setHookEnabled', () => {

packages/core/src/hooks/hookRegistry.ts

@@ -9,16 +9,6 @@ import type { HookDefinition, HookConfig } from './types.js';
 import { HookEventName } from './types.js';
 import { debugLogger } from '../utils/debugLogger.js';
 
-/**
- * Error thrown when attempting to use HookRegistry before initialization
- */
-export class HookRegistryNotInitializedError extends Error {
-  constructor(message = 'Hook registry not initialized') {
-    super(message);
-    this.name = 'HookRegistryNotInitializedError';
-  }
-}
-
 /**
  * Configuration source levels in precedence order (highest to lowest)
  */
@@ -47,7 +37,6 @@ export interface HookRegistryEntry {
 export class HookRegistry {
   private readonly config: Config;
   private entries: HookRegistryEntry[] = [];
-  private initialized = false;
 
   constructor(config: Config) {
     this.config = config;
@@ -57,13 +46,8 @@ export class HookRegistry {
    * Initialize the registry by processing hooks from config
    */
   async initialize(): Promise<void> {
-    if (this.initialized) {
-      return;
-    }
-
     this.entries = [];
     this.processHooksFromConfig();
-    this.initialized = true;
 
     debugLogger.log(
       `Hook registry initialized with ${this.entries.length} hook entries`,
@@ -74,10 +58,6 @@ export class HookRegistry {
    * Get all hook entries for a specific event
    */
   getHooksForEvent(eventName: HookEventName): HookRegistryEntry[] {
-    if (!this.initialized) {
-      throw new HookRegistryNotInitializedError();
-    }
-
     return this.entries
       .filter((entry) => entry.eventName === eventName && entry.enabled)
       .sort(
@@ -90,10 +70,6 @@ export class HookRegistry {
    * Get all registered hooks
    */
   getAllHooks(): HookRegistryEntry[] {
-    if (!this.initialized) {
-      throw new HookRegistryNotInitializedError();
-    }
-
     return [...this.entries];
   }
 

packages/core/src/hooks/hookSystem.test.ts

@@ -127,10 +127,7 @@ describe('HookSystem Integration', () => {
         'Hook system initialized successfully',
       );
 
-      // Verify system is initialized
-      const status = hookSystem.getStatus();
-      expect(status.initialized).toBe(true);
-      // Note: totalHooks might be 0 if hook validation rejects the test hooks
+      expect(hookSystem.getAllHooks().length).toBe(1);
     });
 
     it('should not initialize twice', async () => {
@@ -204,12 +201,6 @@ describe('HookSystem Integration', () => {
       });
       expect(result.success).toBe(true);
     });
-
-    it('should throw error when not initialized', () => {
-      expect(() => hookSystem.getEventHandler()).toThrow(
-        'Hook system not initialized',
-      );
-    });
   });
 
   describe('hook execution', () => {
@@ -261,24 +252,6 @@ describe('HookSystem Integration', () => {
     });
   });
 
-  describe('system management', () => {
-    it('should return correct status when initialized', async () => {
-      await hookSystem.initialize();
-
-      const status = hookSystem.getStatus();
-
-      expect(status.initialized).toBe(true);
-      // Note: totalHooks might be 0 if hook validation rejects the test hooks
-      expect(typeof status.totalHooks).toBe('number');
-    });
-
-    it('should return uninitialized status', () => {
-      const status = hookSystem.getStatus();
-
-      expect(status.initialized).toBe(false);
-    });
-  });
-
   describe('hook disabling via settings', () => {
     it('should not execute disabled hooks from settings', async () => {
       // Create config with two hooks, one enabled and one disabled via settings

packages/core/src/hooks/hookSystem.ts

@@ -24,7 +24,6 @@ export class HookSystem {
   private readonly hookAggregator: HookAggregator;
   private readonly hookPlanner: HookPlanner;
   private readonly hookEventHandler: HookEventHandler;
-  private initialized = false;
 
   constructor(config: Config) {
     const logger: Logger = logs.getLogger(SERVICE_NAME);
@@ -49,12 +48,7 @@ export class HookSystem {
    * Initialize the hook system
    */
   async initialize(): Promise<void> {
-    if (this.initialized) {
-      return;
-    }
-
     await this.hookRegistry.initialize();
-    this.initialized = true;
     debugLogger.debug('Hook system initialized successfully');
   }
 
@@ -62,9 +56,6 @@ export class HookSystem {
    * Get the hook event bus for firing events
    */
   getEventHandler(): HookEventHandler {
-    if (!this.initialized) {
-      throw new Error('Hook system not initialized');
-    }
     return this.hookEventHandler;
   }
 
@@ -88,19 +79,4 @@ export class HookSystem {
   getAllHooks(): HookRegistryEntry[] {
     return this.hookRegistry.getAllHooks();
   }
-
-  /**
-   * Get hook system status for debugging
-   */
-  getStatus(): {
-    initialized: boolean;
-    totalHooks: number;
-  } {
-    const allHooks = this.initialized ? this.hookRegistry.getAllHooks() : [];
-
-    return {
-      initialized: this.initialized,
-      totalHooks: allHooks.length,
-    };
-  }
 }

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> {