@@ -38,10 +38,54 @@ import {
3838 SUPPORTED_PROTOCOL_VERSIONS ,
3939 type SubscribeRequest ,
4040 type Tool ,
41- type UnsubscribeRequest
41+ type UnsubscribeRequest ,
42+ ElicitResultSchema ,
43+ ElicitRequestSchema
4244} from '../types.js' ;
4345import { AjvJsonSchemaValidator } from '../validation/ajv-provider.js' ;
4446import type { JsonSchemaType , JsonSchemaValidator , jsonSchemaValidator } from '../validation/types.js' ;
47+ import { ZodLiteral , ZodObject , z } from 'zod' ;
48+ import type { RequestHandlerExtra } from '../shared/protocol.js' ;
49+
50+ /**
51+ * Elicitation default application helper. Applies defaults to the data based on the schema.
52+ *
53+ * @param schema - The schema to apply defaults to.
54+ * @param data - The data to apply defaults to.
55+ */
56+ function applyElicitationDefaults ( schema : JsonSchemaType | undefined , data : unknown ) : void {
57+ if ( ! schema || data === null || typeof data !== 'object' ) return ;
58+
59+ // Handle object properties
60+ if ( schema . type === 'object' && schema . properties && typeof schema . properties === 'object' ) {
61+ const obj = data as Record < string , unknown > ;
62+ const props = schema . properties as Record < string , JsonSchemaType & { default ?: unknown } > ;
63+ for ( const key of Object . keys ( props ) ) {
64+ const propSchema = props [ key ] ;
65+ // If missing or explicitly undefined, apply default if present
66+ if ( obj [ key ] === undefined && Object . prototype . hasOwnProperty . call ( propSchema , 'default' ) ) {
67+ obj [ key ] = propSchema . default ;
68+ }
69+ // Recurse into existing nested objects/arrays
70+ if ( obj [ key ] !== undefined ) {
71+ applyElicitationDefaults ( propSchema , obj [ key ] ) ;
72+ }
73+ }
74+ }
75+
76+ if ( Array . isArray ( schema . anyOf ) ) {
77+ for ( const sub of schema . anyOf ) {
78+ applyElicitationDefaults ( sub , data ) ;
79+ }
80+ }
81+
82+ // Combine schemas
83+ if ( Array . isArray ( schema . oneOf ) ) {
84+ for ( const sub of schema . oneOf ) {
85+ applyElicitationDefaults ( sub , data ) ;
86+ }
87+ }
88+ }
4589
4690export type ClientOptions = ProtocolOptions & {
4791 /**
@@ -143,6 +187,64 @@ export class Client<
143187 this . _capabilities = mergeCapabilities ( this . _capabilities , capabilities ) ;
144188 }
145189
190+ /**
191+ * Override request handler registration to enforce client-side validation for elicitation.
192+ */
193+ public override setRequestHandler <
194+ T extends ZodObject < {
195+ method : ZodLiteral < string > ;
196+ } >
197+ > (
198+ requestSchema : T ,
199+ handler : (
200+ request : z . infer < T > ,
201+ extra : RequestHandlerExtra < ClientRequest | RequestT , ClientNotification | NotificationT >
202+ ) => ClientResult | ResultT | Promise < ClientResult | ResultT >
203+ ) : void {
204+ const method = requestSchema . shape . method . value ;
205+ if ( method === 'elicitation/create' ) {
206+ const wrappedHandler = async (
207+ request : z . infer < T > ,
208+ extra : RequestHandlerExtra < ClientRequest | RequestT , ClientNotification | NotificationT >
209+ ) : Promise < ClientResult | ResultT > => {
210+ const validatedRequest = ElicitRequestSchema . safeParse ( request ) ;
211+ if ( ! validatedRequest . success ) {
212+ throw new McpError ( ErrorCode . InvalidParams , `Invalid elicitation request: ${ validatedRequest . error . message } ` ) ;
213+ }
214+
215+ const result = await Promise . resolve ( handler ( request , extra ) ) ;
216+
217+ const validationResult = ElicitResultSchema . safeParse ( result ) ;
218+ if ( ! validationResult . success ) {
219+ throw new McpError ( ErrorCode . InvalidParams , `Invalid elicitation result: ${ validationResult . error . message } ` ) ;
220+ }
221+
222+ const validatedResult = validationResult . data ;
223+
224+ if (
225+ this . _capabilities . elicitation ?. applyDefaults &&
226+ validatedResult . action === 'accept' &&
227+ validatedResult . content &&
228+ validatedRequest . data . params . requestedSchema
229+ ) {
230+ try {
231+ applyElicitationDefaults ( validatedRequest . data . params . requestedSchema , validatedResult . content ) ;
232+ } catch {
233+ // gracefully ignore errors in default application
234+ }
235+ }
236+
237+ return validatedResult ;
238+ } ;
239+
240+ // Install the wrapped handler
241+ return super . setRequestHandler ( requestSchema , wrappedHandler as unknown as typeof handler ) ;
242+ }
243+
244+ // Non-elicitation handlers use default behavior
245+ return super . setRequestHandler ( requestSchema , handler ) ;
246+ }
247+
146248 protected assertCapability ( capability : keyof ServerCapabilities , method : string ) : void {
147249 if ( ! this . _serverCapabilities ?. [ capability ] ) {
148250 throw new Error ( `Server does not support ${ capability } (required for ${ method } )` ) ;
0 commit comments