CR changes

Author: tsck

src/server.ts

@@ -63,20 +63,25 @@ export interface ServerOptions {
      */
     tools?: ToolClass[];
     /**
-     * Custom UIs for tools. Maps tool names to their HTML strings.
+     * Custom UIs for tools. Function that returns HTML strings for tool names.
      * Use this to add UIs to tools or replace the default bundled UIs.
+     * The function is called lazily when a UI is requested, allowing you to
+     * defer loading large HTML files until needed.
      *
      * ```ts
      * import { readFileSync } from 'fs';
      * const server = new Server({
      *     // ... other options
-     *     customUIs: {
-     *         'list-databases': readFileSync('./my-custom-ui.html', 'utf-8'),
+     *     customUIs: (toolName) => {
+     *         if (toolName === 'list-databases') {
+     *             return readFileSync('./my-custom-ui.html', 'utf-8');
+     *         }
+     *         return null;
      *     }
      * });
      * ```
      */
-    customUIs?: Record<string, string>;
+    customUIs?: (toolName: string) => string | null | Promise<string | null>;
 }
 
 export class Server {

src/tools/tool.ts

@@ -22,7 +22,7 @@ export type ToolExecutionContext = {
 };
 
 /**
- * The type of operation the tool performs. This is used when evaluating if a tool is allowed to run based on
+ * The ~`type of operation the tool performs. This is used when evaluating if a tool is allowed to run based on
  * the config's `disabledTools` and `readOnly` settings.
  * - `metadata` is used for tools that read but do not access potentially user-generated
  *   data, such as listing databases, collections, or indexes, or inferring collection schema.
@@ -94,7 +94,7 @@ export type ToolConstructorParams = {
      * See `src/elicitation.ts` for further reference.
      */
     elicitation: Elicitation;
-    uiRegistry?: UIRegistry;
+    uiRegistry: UIRegistry;
 };
 
 /**
@@ -432,7 +432,7 @@ export abstract class ToolBase {
      * or inputs during tool execution.
      */
     protected readonly elicitation: Elicitation;
-    protected readonly uiRegistry?: UIRegistry;
+    private readonly uiRegistry: UIRegistry;
     constructor({
         category,
         operationType,
@@ -728,7 +728,7 @@ export abstract class ToolBase {
             return result;
         }
 
-        const uiHtml = await this.uiRegistry?.get(this.name);
+        const uiHtml = await this.uiRegistry.get(this.name);
         if (!uiHtml || !result.structuredContent) {
             return result;
         }

src/ui/registry/registry.ts

@@ -10,24 +10,22 @@ function toPascalCase(kebabCase: string): string {
  * UI Registry that manages bundled UI HTML strings for tools.
  */
 export class UIRegistry {
-    private customUIs: Map<string, string> = new Map();
+    private customUIs?: (toolName: string) => string | null | Promise<string | null>;
     private cache: Map<string, string> = new Map();
 
-    constructor(options?: { customUIs?: Record<string, string> }) {
-        if (options?.customUIs) {
-            for (const [toolName, html] of Object.entries(options.customUIs)) {
-                this.customUIs.set(toolName, html);
-            }
-        }
+    constructor(options?: { customUIs?: (toolName: string) => string | null | Promise<string | null> }) {
+        this.customUIs = options?.customUIs;
     }
 
     /**
      * Gets the UI HTML string for a tool, or null if none exists.
      */
     async get(toolName: string): Promise<string | null> {
-        const customUI = this.customUIs.get(toolName);
-        if (customUI !== undefined) {
-            return customUI;
+        if (this.customUIs) {
+            const customUI = await this.customUIs(toolName);
+            if (customUI !== null && customUI !== undefined) {
+                return customUI;
+            }
         }
 
         const cached = this.cache.get(toolName);
@@ -38,7 +36,7 @@ export class UIRegistry {
         try {
             const module = (await import(`../lib/tools/${toolName}.js`)) as Record<string, string>;
             const exportName = `${toPascalCase(toolName)}Html`;
-            const html = module[exportName];
+            const html = module[exportName]; // HTML generated at build time
             if (html === undefined) {
                 return null;
             }