Merge pull request #597 from objectstack-ai/copilot/add-query-profiling-hook

Author: hotlong

packages/core/src/kernel-base.ts

@@ -85,6 +85,21 @@ export abstract class ObjectKernelBase {
                     return this.services.get<T>(name);
                 }
             },
+            replaceService: <T>(name: string, implementation: T): void => {
+                if (this.services instanceof Map) {
+                    if (!this.services.has(name)) {
+                        throw new Error(`[Kernel] Service '${name}' not found. Use registerService() to add new services.`);
+                    }
+                    this.services.set(name, implementation);
+                } else {
+                    // IServiceRegistry implementation
+                    if (!this.services.has(name)) {
+                        throw new Error(`[Kernel] Service '${name}' not found. Use registerService() to add new services.`);
+                    }
+                    this.services.register(name, implementation);
+                }
+                this.logger.info(`Service '${name}' replaced`, { service: name });
+            },
             hook: (name, handler) => {
                 if (!this.hooks.has(name)) {
                     this.hooks.set(name, []);

packages/core/src/kernel.test.ts

@@ -534,4 +534,91 @@ describe('ObjectKernel', () => {
             }).rejects.toThrow('not running');
         });
     });
+
+    describe('Service Replacement', () => {
+        it('should replace an existing service via replaceService', async () => {
+            const originalService = { value: 'original' };
+            const replacementService = { value: 'replaced' };
+
+            const plugin: Plugin = {
+                name: 'register-plugin',
+                version: '1.0.0',
+                init: async (ctx) => {
+                    ctx.registerService('metadata', originalService);
+                },
+            };
+
+            const optimizationPlugin: Plugin = {
+                name: 'optimization-plugin',
+                version: '1.0.0',
+                dependencies: ['register-plugin'],
+                init: async (ctx) => {
+                    const existing = ctx.getService('metadata');
+                    expect(existing).toBe(originalService);
+                    ctx.replaceService('metadata', replacementService);
+                },
+            };
+
+            await kernel.use(plugin);
+            await kernel.use(optimizationPlugin);
+            await kernel.bootstrap();
+
+            const result = kernel.getService('metadata');
+            expect(result).toBe(replacementService);
+
+            await kernel.shutdown();
+        });
+
+        it('should throw when replacing a non-existent service', async () => {
+            const plugin: Plugin = {
+                name: 'bad-replace-plugin',
+                version: '1.0.0',
+                init: async (ctx) => {
+                    expect(() => {
+                        ctx.replaceService('nonexistent', { value: 'test' });
+                    }).toThrow("Service 'nonexistent' not found");
+                },
+            };
+
+            await kernel.use(plugin);
+            await kernel.bootstrap();
+            await kernel.shutdown();
+        });
+
+        it('should allow decorator pattern via replaceService', async () => {
+            const original = {
+                getData: () => 'raw-data',
+            };
+
+            const plugin: Plugin = {
+                name: 'data-plugin',
+                version: '1.0.0',
+                init: async (ctx) => {
+                    ctx.registerService('data', original);
+                },
+            };
+
+            const wrapperPlugin: Plugin = {
+                name: 'wrapper-plugin',
+                version: '1.0.0',
+                dependencies: ['data-plugin'],
+                init: async (ctx) => {
+                    const existing = ctx.getService<typeof original>('data');
+                    const decorated = {
+                        getData: () => `cached(${existing.getData()})`,
+                    };
+                    ctx.replaceService('data', decorated);
+                },
+            };
+
+            await kernel.use(plugin);
+            await kernel.use(wrapperPlugin);
+            await kernel.bootstrap();
+
+            const result = kernel.getService<typeof original>('data');
+            expect(result.getData()).toBe('cached(raw-data)');
+
+            await kernel.shutdown();
+        });
+    });
 });

packages/core/src/kernel.ts

@@ -119,6 +119,15 @@ export class ObjectKernel {
                     throw new Error(`[Kernel] Service '${name}' not found`);
                 }
             },
+            replaceService: <T>(name: string, implementation: T): void => {
+                const hasService = this.services.has(name) || this.pluginLoader.hasService(name);
+                if (!hasService) {
+                    throw new Error(`[Kernel] Service '${name}' not found. Use registerService() to add new services.`);
+                }
+                this.services.set(name, implementation);
+                this.pluginLoader.replaceService(name, implementation);
+                this.logger.info(`Service '${name}' replaced`, { service: name });
+            },
             hook: (name, handler) => {
                 if (!this.hooks.has(name)) {
                     this.hooks.set(name, []);

packages/core/src/lite-kernel.test.ts

@@ -197,4 +197,52 @@ describe('LiteKernel with Configurable Logger', () => {
             await browserKernel.shutdown();
         });
     });
+
+    describe('Service Replacement', () => {
+        it('should replace an existing service via replaceService', async () => {
+            const originalService = { value: 'original' };
+            const replacementService = { value: 'replaced' };
+
+            const plugin: Plugin = {
+                name: 'register-plugin',
+                init: async (ctx) => {
+                    ctx.registerService('metadata', originalService);
+                },
+            };
+
+            const optimizationPlugin: Plugin = {
+                name: 'optimization-plugin',
+                dependencies: ['register-plugin'],
+                init: async (ctx) => {
+                    const existing = ctx.getService('metadata');
+                    expect(existing).toBe(originalService);
+                    ctx.replaceService('metadata', replacementService);
+                },
+            };
+
+            kernel.use(plugin);
+            kernel.use(optimizationPlugin);
+            await kernel.bootstrap();
+
+            const result = kernel.getService('metadata');
+            expect(result).toBe(replacementService);
+
+            await kernel.shutdown();
+        });
+
+        it('should throw when replacing a non-existent service', async () => {
+            const plugin: Plugin = {
+                name: 'bad-replace-plugin',
+                init: async (ctx) => {
+                    expect(() => {
+                        ctx.replaceService('nonexistent', { value: 'test' });
+                    }).toThrow("Service 'nonexistent' not found");
+                },
+            };
+
+            kernel.use(plugin);
+            await kernel.bootstrap();
+            await kernel.shutdown();
+        });
+    });
 });

packages/core/src/plugin-loader.ts

@@ -248,6 +248,18 @@ export class PluginLoader {
         this.serviceInstances.set(name, service);
     }
 
+    /**
+     * Replace an existing service instance.
+     * Used by optimization plugins to swap kernel internals.
+     * @throws Error if service does not exist
+     */
+    replaceService(name: string, service: any): void {
+        if (!this.hasService(name)) {
+            throw new Error(`Service '${name}' not found`);
+        }
+        this.serviceInstances.set(name, service);
+    }
+
     /**
      * Check if a service is registered (either as instance or factory)
      */

packages/core/src/security/plugin-permission-enforcer.ts

@@ -394,6 +394,12 @@ export class SecurePluginContext implements PluginContext {
     return this.baseContext.getService<T>(name);
   }
 
+  replaceService<T>(name: string, implementation: T): void {
+    // Check permission before replacing service
+    this.permissionEnforcer.enforceServiceAccess(this.pluginName, name);
+    this.baseContext.replaceService(name, implementation);
+  }
+  
   getServices(): Map<string, any> {
     // Return all services (no permission check for listing)
     return this.baseContext.getServices();

packages/core/src/types.ts

@@ -28,6 +28,17 @@ export interface PluginContext {
      */
     getService<T>(name: string): T;
 
+    /**
+     * Replace an existing service with a new implementation.
+     * Useful for optimization plugins that wrap or swap kernel internals
+     * (e.g., metadata registry, connection pooling).
+     * 
+     * @param name - Service name to replace
+     * @param implementation - New service implementation
+     * @throws Error if the service does not exist
+     */
+    replaceService<T>(name: string, implementation: T): void;
+
     /**
      * Get all registered services
      */