refactor: remove UI business logic from subclass using structured content

tsck · tsck · commit 0f9a81229ca3 · 2025-12-05T09:44:23.000-05:00
diff --git a/src/tools/mongodb/metadata/listDatabases.tsx b/src/tools/mongodb/metadata/listDatabases.tsx
@@ -3,54 +3,35 @@ import { MongoDBToolBase } from "../mongodbTool.js";
 import type * as bson from "bson";
 import type { OperationType } from "../../tool.js";
 import { formatUntrustedData } from "../../tool.js";
-import { createUIResource } from "@mcp-ui/server";
+import { ListDatabasesOutputSchema, type ListDatabasesOutput } from "../../../ui/components/ListDatabases/schema.js";
+
+// Re-export for consumers who need the schema/type
+export { ListDatabasesOutputSchema, type ListDatabasesOutput };
 
 export class ListDatabasesTool extends MongoDBToolBase {
     public name = "list-databases";
     protected description = "List all databases for a MongoDB connection";
     protected argsShape = {};
+    protected override outputSchema = ListDatabasesOutputSchema;
     static operationType: OperationType = "metadata";
 
     protected async execute(): Promise<CallToolResult> {
         const provider = await this.ensureConnected();
         const dbs = (await provider.listDatabases("")).databases as { name: string; sizeOnDisk: bson.Long }[];
+        const databases = dbs.map((db) => ({
+            name: db.name,
+            size: Number(db.sizeOnDisk),
+        }));
 
-        const toolResult = {
+        return {
             content: formatUntrustedData(
                 `Found ${dbs.length} databases`,
                 ...dbs.map((db) => `Name: ${db.name}, Size: ${db.sizeOnDisk.toString()} bytes`)
             ),
-        };
-
-        const ui = this.getUI();
-
-        if (!ui) {
-            return toolResult;
-        }
-
-        const uiDatabases = dbs.map((db) => ({
-            name: db.name,
-            size: Number(db.sizeOnDisk),
-        }));
-
-        const uiResource = createUIResource({
-            uri: `ui://list-databases/${Date.now()}`,
-            content: {
-                type: "rawHtml",
-                htmlString: this.getUI() || "",
-            },
-            encoding: "text",
-            uiMetadata: {
-                "initial-render-data": {
-                    databases: uiDatabases,
-                    totalCount: uiDatabases.length,
-                },
+            structuredContent: {
+                databases,
+                totalCount: databases.length,
             },
-        });
-
-        return {
-            ...toolResult,
-            content: [...(toolResult.content || []), uiResource],
         };
     }
 }
diff --git a/src/tools/tool.ts b/src/tools/tool.ts
@@ -1,4 +1,4 @@
-import type { z, ZodRawShape } from "zod";
+import { z, type ZodRawShape } from "zod";
 import type { RegisteredTool } from "@modelcontextprotocol/sdk/server/mcp.js";
 import type { CallToolResult, ToolAnnotations } from "@modelcontextprotocol/sdk/types.js";
 import type { Session } from "../common/session.js";
@@ -10,6 +10,7 @@ import type { Server } from "../server.js";
 import type { Elicitation } from "../elicitation.js";
 import type { PreviewFeature } from "../common/schemas.js";
 import type { UIRegistry } from "../ui/registry/index.js";
+import { createUIResource } from "@mcp-ui/server";
 
 export type ToolArgs<T extends ZodRawShape> = {
     [K in keyof T]: z.infer<T[K]>;
@@ -278,6 +279,37 @@ export abstract class ToolBase {
      */
     protected abstract argsShape: ZodRawShape;
 
+    /**
+     * Optional Zod schema defining the tool's structured output.
+     *
+     * When defined:
+     * 1. The MCP SDK will validate `structuredContent` against this schema
+     * 2. If the tool has a registered UI component, the base class will validate
+     *    `structuredContent` before creating a UIResource. If validation fails,
+     *    the UI is skipped and only the text result is returned.
+     *
+     * This enables a clean separation: tools define their output schema and return
+     * structured data, while the base class handles validation and UI integration
+     * automatically.
+     *
+     * @example
+     * ```typescript
+     * protected outputSchema = {
+     *   items: z.array(z.object({ name: z.string(), count: z.number() })),
+     *   totalCount: z.number(),
+     * };
+     *
+     * protected async execute(): Promise<CallToolResult> {
+     *   const items = await this.fetchItems();
+     *   return {
+     *     content: [{ type: "text", text: `Found ${items.length} items` }],
+     *     structuredContent: { items, totalCount: items.length },
+     *   };
+     * }
+     * ```
+     */
+    protected outputSchema?: ZodRawShape;
+
     private registeredTool: RegisteredTool | undefined;
 
     protected get annotations(): ToolAnnotations {
@@ -460,15 +492,17 @@ export abstract class ToolBase {
                 });
 
                 const result = await this.execute(args, { signal });
-                this.emitToolEvent(args, { startTime, result });
+                const finalResult = this.appendUIResourceIfAvailable(result);
+
+                this.emitToolEvent(args, { startTime, result: finalResult });
 
                 this.session.logger.debug({
                     id: LogId.toolExecute,
                     context: "tool",
                     message: `Executed tool ${this.name}`,
                     noRedaction: true,
                 });
-                return result;
+                return finalResult;
             } catch (error: unknown) {
                 this.session.logger.error({
                     id: LogId.toolExecuteFailure,
@@ -491,6 +525,7 @@ export abstract class ToolBase {
                     config: {
                         description?: string;
                         inputSchema?: ZodRawShape;
+                        outputSchema?: ZodRawShape;
                         annotations?: ToolAnnotations;
                     },
                     cb: (args: ToolArgs<ZodRawShape>, extra: ToolExecutionContext) => Promise<CallToolResult>
@@ -500,6 +535,7 @@ export abstract class ToolBase {
                 {
                     description: this.description,
                     inputSchema: this.argsShape,
+                    outputSchema: this.outputSchema,
                     annotations: this.annotations,
                 },
                 callback
@@ -694,6 +730,56 @@ export abstract class ToolBase {
     protected getUI(): string | undefined {
         return this.uiRegistry?.get(this.name);
     }
+
+    /**
+     * Automatically appends a UIResource to the tool result if:
+     * 1. The tool has a registered UI component
+     * 2. The result contains `structuredContent`
+     * 3. If `outputSchema` is defined, `structuredContent` must validate against it
+     *
+     * This method is called automatically after `execute()` in the `register()` callback.
+     * Tools don't need to call this directly - they just need to return `structuredContent`
+     * in their result and the base class handles the rest.
+     *
+     * @param result - The result from the tool's `execute()` method
+     * @returns The result with UIResource appended if conditions are met, otherwise unchanged
+     */
+    private appendUIResourceIfAvailable(result: CallToolResult): CallToolResult {
+        const uiHtml = this.getUI();
+        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()}`,
+            content: {
+                type: "rawHtml",
+                htmlString: uiHtml,
+            },
+            encoding: "text",
+            uiMetadata: {
+                "initial-render-data": result.structuredContent,
+            },
+        });
+
+        return {
+            ...result,
+            content: [...(result.content || []), uiResource],
+        };
+    }
 }
 
 /**
diff --git a/src/ui/components/ListDatabases/ListDatabases.tsx b/src/ui/components/ListDatabases/ListDatabases.tsx
@@ -11,22 +11,13 @@ import {
     TableHead,
 } from "@leafygreen-ui/table";
 import { tableStyles } from "./ListDatabases.styles.js";
+import { ListDatabasesOutputSchema, type ListDatabasesOutput } from "./schema.js";
 
 const HeaderCell = LGHeaderCell as React.FC<React.ComponentPropsWithoutRef<"th">>;
 const Cell = LGCell as React.FC<React.ComponentPropsWithoutRef<"td">>;
 const Row = LGRow as React.FC<React.ComponentPropsWithoutRef<"tr">>;
 
-const DatabaseInfoSchema = z.object({
-    name: z.string(),
-    size: z.number(),
-});
-
-const ListDatabasesDataSchema = z.object({
-    databases: z.array(DatabaseInfoSchema),
-    totalCount: z.number(),
-});
-
-type ListDatabasesRenderData = z.infer<typeof ListDatabasesDataSchema>;
+const ListDatabasesDataSchema = z.object(ListDatabasesOutputSchema);
 
 function formatBytes(bytes: number): string {
     if (bytes === 0) return "0 Bytes";
@@ -39,7 +30,7 @@ function formatBytes(bytes: number): string {
 }
 
 export const ListDatabases = () => {
-    const { data, isLoading, error } = useRenderData<ListDatabasesRenderData>();
+    const { data, isLoading, error } = useRenderData<ListDatabasesOutput>();
 
     if (isLoading) {
         return <div>Loading...</div>;
@@ -49,7 +40,6 @@ export const ListDatabases = () => {
         return <div>Error: {error}</div>;
     }
 
-    // Validate props against schema
     const validationResult = ListDatabasesDataSchema.safeParse(data);
 
     if (!validationResult.success) {
diff --git a/src/ui/components/ListDatabases/index.ts b/src/ui/components/ListDatabases/index.ts
@@ -1 +1,2 @@
 export { ListDatabases } from "./ListDatabases.js";
+export { ListDatabasesOutputSchema, type ListDatabasesOutput } from "./schema.js";
diff --git a/src/ui/components/ListDatabases/schema.ts b/src/ui/components/ListDatabases/schema.ts
@@ -0,0 +1,21 @@
+import { z } from "zod";
+
+/**
+ * Shared schema for the list-databases tool output.
+ *
+ * This schema is the single source of truth for the data contract between:
+ * - The ListDatabasesTool (which returns structuredContent matching this schema)
+ * - The ListDatabases UI component (which renders this data)
+ */
+export const ListDatabasesOutputSchema = {
+    databases: z.array(
+        z.object({
+            name: z.string(),
+            size: z.number(),
+        })
+    ),
+    totalCount: z.number(),
+};
+
+/** Type derived from the output schema */
+export type ListDatabasesOutput = z.infer<z.ZodObject<typeof ListDatabasesOutputSchema>>;

Original file line number	Diff line number	Diff line change
`@@ -1 +1,2 @@`
`1`	`1`	`export { ListDatabases } from "./ListDatabases.js";`
	`2`	`+export { ListDatabasesOutputSchema, type ListDatabasesOutput } from "./schema.js";`