review feedback 1

Author: tsck

src/tools/tool.ts

@@ -482,7 +482,7 @@ export abstract class ToolBase {
                 });
 
                 const result = await this.execute(args, { signal });
-                const finalResult = this.appendUIResource(result);
+                const finalResult = await this.appendUIResource(result);
 
                 this.emitToolEvent(args, { startTime, result: finalResult });
 
@@ -713,45 +713,24 @@ export abstract class ToolBase {
         return metadata;
     }
 
-    /**
-     * Get the UI HTML string for this tool from the registry.
-     * Returns the registered UI HTML, or undefined if no UI exists for this tool.
-     */
-    protected getUI(): string | undefined {
-        return this.uiRegistry?.get(this.name);
-    }
-
     /**
      * Appends a UIResource to the tool result.
      *
      * @param result - The result from the tool's `execute()` method
      * @returns The result with UIResource appended if conditions are met, otherwise unchanged
      */
-    private appendUIResource(result: CallToolResult): CallToolResult {
+    private async appendUIResource(result: CallToolResult): Promise<CallToolResult> {
         if (!this.isFeatureEnabled("mcpUI")) {
             return result;
         }
 
-        const uiHtml = this.getUI();
+        const uiHtml = await this.uiRegistry?.get(this.name);
         if (!uiHtml || !result.structuredContent) {
             return result;
         }
 
-        if (this.outputSchema) {
-            const schema = z.object(this.outputSchema);
-            const validation = schema.safeParse(result.structuredContent);
-            if (!validation.success) {
-                this.session.logger.warning({
-                    id: LogId.toolExecute,
-                    context: `tool - ${this.name}`,
-                    message: `structuredContent failed validation against outputSchema, skipping UI resource: ${validation.error.message}`,
-                });
-                return result;
-            }
-        }
-
         const uiResource = createUIResource({
-            uri: `ui://${this.name}/${Date.now()}`,
+            uri: `ui://${this.name}`,
             content: {
                 type: "rawHtml",
                 htmlString: uiHtml,

src/ui/registry/registry.ts

@@ -1,13 +1,9 @@
-import { uiHtml } from "../generated/uiHtml.js";
-
 /**
  * UI Registry that manages bundled UI HTML strings for tools.
- *
- * The default UIs are embedded at build time via the generated uiHtml module.
- * Custom UIs can be provided at runtime to override or extend the defaults.
  */
 export class UIRegistry {
     private customUIs: Map<string, string> = new Map();
+    private cache: Map<string, string> = new Map();
 
     constructor(options?: { customUIs?: Record<string, string> }) {
         if (options?.customUIs) {
@@ -18,30 +14,25 @@ export class UIRegistry {
     }
 
     /**
-     * Get the UI HTML string for a tool.
-     * @param toolName The name of the tool (kebab-case, e.g., "list-databases")
-     * @returns The HTML string, or undefined if no UI exists for this tool
+     * Gets the UI HTML string for a tool, or null if none exists.
      */
-    get(toolName: string): string | undefined {
-        return this.customUIs.get(toolName) ?? uiHtml[toolName];
-    }
+    async get(toolName: string): Promise<string | null> {
+        const customUI = this.customUIs.get(toolName);
+        if (customUI !== undefined) {
+            return customUI;
+        }
 
-    /**
-     * Check if a UI exists for a tool.
-     * @param toolName The name of the tool
-     * @returns True if a UI exists (custom or built-in)
-     */
-    has(toolName: string): boolean {
-        return this.customUIs.has(toolName) || toolName in uiHtml;
-    }
+        const cached = this.cache.get(toolName);
+        if (cached !== undefined) {
+            return cached;
+        }
 
-    /**
-     * Get all available tool names that have UIs.
-     * @returns Array of tool names
-     */
-    getAvailableTools(): string[] {
-        const builtIn = Object.keys(uiHtml);
-        const custom = Array.from(this.customUIs.keys());
-        return [...new Set([...builtIn, ...custom])];
+        try {
+            const module = (await import(`../generated/tools/${toolName}.js`)) as { default: string };
+            this.cache.set(toolName, module.default);
+            return module.default;
+        } catch {
+            return null;
+        }
     }
 }

vite.ui.config.ts

@@ -11,7 +11,7 @@ const componentsDir = resolve(__dirname, "src/ui/components");
 const entriesDir = resolve(__dirname, "node_modules/.cache/mongodb-mcp-server/ui-entries");
 const templatePath = resolve(__dirname, "src/ui/build/template.html");
 const mountPath = resolve(__dirname, "src/ui/build/mount.tsx");
-const generatedModulePath = resolve(__dirname, "src/ui/generated/uiHtml.ts");
+const generatedDir = resolve(__dirname, "src/ui/generated");
 const uiDistPath = resolve(__dirname, "dist/ui");
 
 // Unique component names from uiMap - only these will be built
@@ -56,11 +56,9 @@ function generateHtmlEntries(): Plugin {
 }
 
 /**
- * Vite plugin that generates the uiHtml.ts module after the build completes.
- * This embeds all built HTML strings into a TypeScript module so they can be
- * imported at runtime.
+ * Vite plugin that generates per-tool UI modules after the build completes.
  *
- * Uses the uiMap from src/ui/registry/uiMap.ts to map tool names to component HTML files.
+ * Uses the uiMap to map tool names to component HTML files.
  */
 function generateUIModule(): Plugin {
     return {
@@ -71,9 +69,16 @@ function generateUIModule(): Plugin {
                 return;
             }
 
-            const entries: Record<string, string> = {};
+            const toolsDir = join(generatedDir, "tools");
+            if (!existsSync(generatedDir)) {
+                mkdirSync(generatedDir, { recursive: true });
+            }
+            if (!existsSync(toolsDir)) {
+                mkdirSync(toolsDir, { recursive: true });
+            }
+
+            const generatedTools: string[] = [];
 
-            // Use uiMap to determine which tools get which UI
             for (const [toolName, componentName] of Object.entries(uiMap)) {
                 const htmlFile = join(uiDistPath, `${componentName}.html`);
                 if (!existsSync(htmlFile)) {
@@ -83,32 +88,21 @@ function generateUIModule(): Plugin {
                     continue;
                 }
                 const html = readFileSync(htmlFile, "utf-8");
-                entries[toolName] = html;
-            }
 
-            const moduleContent = `/**
+                const toolModuleContent = `/**
  * AUTO-GENERATED FILE - DO NOT EDIT MANUALLY
- *
- * This file is generated by the UI build process (vite build --config vite.ui.config.ts).
- * It contains the bundled HTML strings for each UI component, keyed by tool name.
- *
- * To add a new UI:
- * 1. Create a component in src/ui/components/YourComponent/
- * 2. Add the mapping in src/ui/registry/uiMap.ts
- * 3. Run \`pnpm build:ui\` to regenerate this file
+ * Generated by: vite build --config vite.ui.config.ts
+ * Tool: ${toolName}
+ * Component: ${componentName}
  */
-export const uiHtml: Record<string, string> = ${JSON.stringify(entries, null, 4)};
+export default ${JSON.stringify(html)};
 `;
-
-            // Ensure the generated directory exists
-            const generatedDir = dirname(generatedModulePath);
-            if (!existsSync(generatedDir)) {
-                mkdirSync(generatedDir, { recursive: true });
+                writeFileSync(join(toolsDir, `${toolName}.ts`), toolModuleContent);
+                generatedTools.push(toolName);
             }
 
-            writeFileSync(generatedModulePath, moduleContent);
             console.log(
-                `[generate-ui-module] Generated uiHtml.ts with ${Object.keys(entries).length} UI(s): ${Object.keys(entries).join(", ")}`
+                `[generate-ui-module] Generated ${generatedTools.length} lazy UI module(s): ${generatedTools.join(", ")}`
             );
         },
     };