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

This commit is contained in:
William Valentin
2026-02-06 15:58:19 -08:00
parent ed1e290ddd
commit 1314ac0163
3 changed files with 93 additions and 1 deletions
+3
View File
@@ -0,0 +1,3 @@
export { DockerSandbox, type DockerSandboxConfig, type ExecOptions, type ExecResult } from './docker.js';
export { SandboxManager } from './manager.js';
export { createSandboxedShellTool, createSandboxedProcessStartTool } from './tools.js';
+70 -1
View File
@@ -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');
});
});
});
+20
View File
@@ -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);
}