diff --git a/src/server/index.ts b/src/server/index.ts index 720cf658..df7a36e5 100644 --- a/src/server/index.ts +++ b/src/server/index.ts @@ -11,7 +11,9 @@ import { McpUiToolMeta, } from "../app.js"; import type { + BaseToolCallback, McpServer, + RegisteredTool, ResourceMetadata, ToolCallback, ReadResourceCallback, @@ -33,6 +35,7 @@ export interface ToolConfig { title?: string; description?: string; inputSchema?: ZodRawShapeCompat | AnySchema; + outputSchema?: ZodRawShapeCompat | AnySchema; annotations?: ToolAnnotations; _meta?: Record; } @@ -100,15 +103,17 @@ export interface McpUiAppResourceConfig extends ResourceMetadata { * ``` */ export function registerAppTool< - TInputSchema extends ZodRawShapeCompat | AnySchema | undefined = undefined, + OutputArgs extends ZodRawShapeCompat | AnySchema, + InputArgs extends undefined | ZodRawShapeCompat | AnySchema = undefined, >( server: Pick, name: string, - config: Omit & { - inputSchema?: TInputSchema; + config: McpUiAppToolConfig & { + inputSchema?: InputArgs; + outputSchema?: OutputArgs; }, - handler: ToolCallback, -): void { + cb: ToolCallback, +): RegisteredTool { // Normalize metadata for backward compatibility: // - If _meta.ui.resourceUri is set, also set the legacy flat key // - If the legacy flat key is set, also set _meta.ui.resourceUri @@ -125,7 +130,7 @@ export function registerAppTool< normalizedMeta = { ...meta, ui: { ...uiMeta, resourceUri: legacyUri } }; } - server.registerTool(name, { ...config, _meta: normalizedMeta }, handler); + return server.registerTool(name, { ...config, _meta: normalizedMeta }, cb); } /**