feat: add ToolRegistry.clone() and replace() for per-session registries

src/tools/registry.test.ts

@@ -1,4 +1,4 @@
-import { describe, it, expect } from 'vitest';
+import { describe, it, expect, vi } from 'vitest';
 import { ToolRegistry } from './registry.js';
 import type { Tool } from './types.js';
 
@@ -92,4 +92,73 @@ describe('ToolRegistry', () => {
       },
     }]);
   });
+
+  describe('ToolRegistry — clone()', () => {
+    function makeTool(name: string): Tool {
+      return {
+        name,
+        description: `Mock ${name}`,
+        inputSchema: { type: 'object', properties: {} },
+        execute: async () => ({ success: true, output: '' }),
+      };
+    }
+
+    it('creates a copy with all tools', () => {
+      const reg = new ToolRegistry();
+      reg.register(makeTool('tool.a'));
+      reg.register(makeTool('tool.b'));
+
+      const cloned = reg.clone();
+      expect(cloned.list().map(t => t.name).sort()).toEqual(['tool.a', 'tool.b']);
+    });
+
+    it('inherits the policy from original', () => {
+      const reg = new ToolRegistry();
+      const mockPolicy = { filterTools: vi.fn(), isAllowed: vi.fn(), resolveAllowedNames: vi.fn(), getEffectiveProfile: vi.fn() };
+      reg.setPolicy(mockPolicy as any);
+
+      const cloned = reg.clone();
+      expect(cloned.getPolicy()).toBe(mockPolicy);
+    });
+
+    it('allows replacing tools in clone without affecting original', () => {
+      const reg = new ToolRegistry();
+      const originalTool = makeTool('shell.exec');
+      reg.register(originalTool);
+
+      const cloned = reg.clone();
+      const replacementTool = makeTool('shell.exec');
+      replacementTool.description = 'Sandboxed version';
+
+      cloned.replace(replacementTool);
+      expect(cloned.get('shell.exec')!.description).toBe('Sandboxed version');
+      expect(reg.get('shell.exec')!.description).toBe('Mock shell.exec');
+    });
+  });
+
+  describe('ToolRegistry — replace()', () => {
+    function makeTool(name: string): Tool {
+      return {
+        name,
+        description: `Mock ${name}`,
+        inputSchema: { type: 'object', properties: {} },
+        execute: async () => ({ success: true, output: '' }),
+      };
+    }
+
+    it('replaces an existing tool', () => {
+      const reg = new ToolRegistry();
+      reg.register(makeTool('tool.a'));
+      const replacement = makeTool('tool.a');
+      replacement.description = 'New description';
+
+      reg.replace(replacement);
+      expect(reg.get('tool.a')!.description).toBe('New description');
+    });
+
+    it('throws if tool does not exist', () => {
+      const reg = new ToolRegistry();
+      expect(() => reg.replace(makeTool('nonexistent'))).toThrow('not registered');
+    });
+  });
 });

src/tools/registry.ts

@@ -31,6 +31,26 @@ export class ToolRegistry {
     return this.tools.delete(name);
   }
 
+  /** Replace an existing tool with a new implementation. Throws if not registered. */
+  replace(tool: Tool): void {
+    if (!this.tools.has(tool.name)) {
+      throw new Error(`Tool '${tool.name}' is not registered — cannot replace`);
+    }
+    this.tools.set(tool.name, tool);
+  }
+
+  /** Create a shallow clone of this registry (new Map, same Tool objects + policy). */
+  clone(): ToolRegistry {
+    const cloned = new ToolRegistry();
+    for (const tool of this.tools.values()) {
+      cloned.register(tool);
+    }
+    if (this._policy) {
+      cloned.setPolicy(this._policy);
+    }
+    return cloned;
+  }
+
   get(name: string): Tool | undefined {
     return this.tools.get(name);
   }