1  import type {
  2    EnumSchema,
  3    MultiSelectEnumSchema,
  4    PrimitiveSchemaDefinition,
  5    StringSchema,
  6  } from '@modelcontextprotocol/sdk/types.js'
  7  import { z } from 'zod/v4'
  8  import { jsonStringify } from '../slowOperations.js'
  9  import { plural } from '../stringUtils.js'
 10  import {
 11    looksLikeISO8601,
 12    parseNaturalLanguageDateTime,
 13  } from './dateTimeParser.js'
 14  
 15  export type ValidationResult = {
 16    value?: string | number | boolean
 17    isValid: boolean
 18    error?: string
 19  }
 20  
 21  const STRING_FORMATS = {
 22    email: {
 23      description: 'email address',
 24      example: 'user@example.com',
 25    },
 26    uri: {
 27      description: 'URI',
 28      example: 'https://example.com',
 29    },
 30    date: {
 31      description: 'date',
 32      example: '2024-03-15',
 33    },
 34    'date-time': {
 35      description: 'date-time',
 36      example: '2024-03-15T14:30:00Z',
 37    },
 38  }
 39  
 40  /**
 41   * Check if schema is a single-select enum (either legacy `enum` format or new `oneOf` format)
 42   */
 43  export const isEnumSchema = (
 44    schema: PrimitiveSchemaDefinition,
 45  ): schema is EnumSchema => {
 46    return schema.type === 'string' && ('enum' in schema || 'oneOf' in schema)
 47  }
 48  
 49  /**
 50   * Check if schema is a multi-select enum (`type: "array"` with `items.enum` or `items.anyOf`)
 51   */
 52  export function isMultiSelectEnumSchema(
 53    schema: PrimitiveSchemaDefinition,
 54  ): schema is MultiSelectEnumSchema {
 55    return (
 56      schema.type === 'array' &&
 57      'items' in schema &&
 58      typeof schema.items === 'object' &&
 59      schema.items !== null &&
 60      ('enum' in schema.items || 'anyOf' in schema.items)
 61    )
 62  }
 63  
 64  /**
 65   * Get values from a multi-select enum schema
 66   */
 67  export function getMultiSelectValues(schema: MultiSelectEnumSchema): string[] {
 68    if ('anyOf' in schema.items) {
 69      return schema.items.anyOf.map(item => item.const)
 70    }
 71    if ('enum' in schema.items) {
 72      return schema.items.enum
 73    }
 74    return []
 75  }
 76  
 77  /**
 78   * Get display labels from a multi-select enum schema
 79   */
 80  export function getMultiSelectLabels(schema: MultiSelectEnumSchema): string[] {
 81    if ('anyOf' in schema.items) {
 82      return schema.items.anyOf.map(item => item.title)
 83    }
 84    if ('enum' in schema.items) {
 85      return schema.items.enum
 86    }
 87    return []
 88  }
 89  
 90  /**
 91   * Get label for a specific value in a multi-select enum
 92   */
 93  export function getMultiSelectLabel(
 94    schema: MultiSelectEnumSchema,
 95    value: string,
 96  ): string {
 97    const index = getMultiSelectValues(schema).indexOf(value)
 98    return index >= 0 ? (getMultiSelectLabels(schema)[index] ?? value) : value
 99  }
100  
101  /**
102   * Get enum values from EnumSchema (handles both legacy `enum` and new `oneOf` formats)
103   */
104  export function getEnumValues(schema: EnumSchema): string[] {
105    if ('oneOf' in schema) {
106      return schema.oneOf.map(item => item.const)
107    }
108    if ('enum' in schema) {
109      return schema.enum
110    }
111    return []
112  }
113  
114  /**
115   * Get enum display labels from EnumSchema
116   */
117  export function getEnumLabels(schema: EnumSchema): string[] {
118    if ('oneOf' in schema) {
119      return schema.oneOf.map(item => item.title)
120    }
121    if ('enum' in schema) {
122      return ('enumNames' in schema ? schema.enumNames : undefined) ?? schema.enum
123    }
124    return []
125  }
126  
127  /**
128   * Get label for a specific enum value
129   */
130  export function getEnumLabel(schema: EnumSchema, value: string): string {
131    const index = getEnumValues(schema).indexOf(value)
132    return index >= 0 ? (getEnumLabels(schema)[index] ?? value) : value
133  }
134  
135  function getZodSchema(schema: PrimitiveSchemaDefinition): z.ZodTypeAny {
136    if (isEnumSchema(schema)) {
137      const [first, ...rest] = getEnumValues(schema)
138      if (!first) {
139        return z.never()
140      }
141      return z.enum([first, ...rest])
142    }
143    if (schema.type === 'string') {
144      let stringSchema = z.string()
145      if (schema.minLength !== undefined) {
146        stringSchema = stringSchema.min(schema.minLength, {
147          message: `Must be at least ${schema.minLength} ${plural(schema.minLength, 'character')}`,
148        })
149      }
150      if (schema.maxLength !== undefined) {
151        stringSchema = stringSchema.max(schema.maxLength, {
152          message: `Must be at most ${schema.maxLength} ${plural(schema.maxLength, 'character')}`,
153        })
154      }
155      switch (schema.format) {
156        case 'email':
157          stringSchema = stringSchema.email({
158            message: 'Must be a valid email address, e.g. user@example.com',
159          })
160          break
161        case 'uri':
162          stringSchema = stringSchema.url({
163            message: 'Must be a valid URI, e.g. https://example.com',
164          })
165          break
166        case 'date':
167          stringSchema = stringSchema.date(
168            'Must be a valid date, e.g. 2024-03-15, today, next Monday',
169          )
170          break
171        case 'date-time':
172          stringSchema = stringSchema.datetime({
173            offset: true,
174            message:
175              'Must be a valid date-time, e.g. 2024-03-15T14:30:00Z, tomorrow at 3pm',
176          })
177          break
178        default:
179          // No specific format validation
180          break
181      }
182      return stringSchema
183    }
184    if (schema.type === 'number' || schema.type === 'integer') {
185      const typeLabel = schema.type === 'integer' ? 'an integer' : 'a number'
186      const isInteger = schema.type === 'integer'
187      const formatNum = (n: number) =>
188        Number.isInteger(n) && !isInteger ? `${n}.0` : String(n)
189  
190      // Build a single descriptive error message for range violations
191      const rangeMsg =
192        schema.minimum !== undefined && schema.maximum !== undefined
193          ? `Must be ${typeLabel} between ${formatNum(schema.minimum)} and ${formatNum(schema.maximum)}`
194          : schema.minimum !== undefined
195            ? `Must be ${typeLabel} >= ${formatNum(schema.minimum)}`
196            : schema.maximum !== undefined
197              ? `Must be ${typeLabel} <= ${formatNum(schema.maximum)}`
198              : `Must be ${typeLabel}`
199  
200      let numberSchema = z.coerce.number({
201        error: rangeMsg,
202      })
203      if (schema.type === 'integer') {
204        numberSchema = numberSchema.int({ message: rangeMsg })
205      }
206      if (schema.minimum !== undefined) {
207        numberSchema = numberSchema.min(schema.minimum, {
208          message: rangeMsg,
209        })
210      }
211      if (schema.maximum !== undefined) {
212        numberSchema = numberSchema.max(schema.maximum, {
213          message: rangeMsg,
214        })
215      }
216      return numberSchema
217    }
218    if (schema.type === 'boolean') {
219      return z.coerce.boolean()
220    }
221  
222    throw new Error(`Unsupported schema: ${jsonStringify(schema)}`)
223  }
224  
225  export function validateElicitationInput(
226    stringValue: string,
227    schema: PrimitiveSchemaDefinition,
228  ): ValidationResult {
229    const zodSchema = getZodSchema(schema)
230    const parseResult = zodSchema.safeParse(stringValue)
231  
232    if (parseResult.success) {
233      // zodSchema always produces primitive types for elicitation
234      return {
235        value: parseResult.data as string | number | boolean,
236        isValid: true,
237      }
238    }
239    return {
240      isValid: false,
241      error: parseResult.error.issues.map(e => e.message).join('; '),
242    }
243  }
244  
245  const hasStringFormat = (
246    schema: PrimitiveSchemaDefinition,
247  ): schema is StringSchema & { format: string } => {
248    return (
249      schema.type === 'string' &&
250      'format' in schema &&
251      typeof schema.format === 'string'
252    )
253  }
254  
255  /**
256   * Returns a helpful placeholder/hint for a given format
257   */
258  export function getFormatHint(
259    schema: PrimitiveSchemaDefinition,
260  ): string | undefined {
261    if (schema.type === 'string') {
262      if (!hasStringFormat(schema)) {
263        return undefined
264      }
265  
266      const { description, example } = STRING_FORMATS[schema.format] || {}
267      return `${description}, e.g. ${example}`
268    }
269  
270    if (schema.type === 'number' || schema.type === 'integer') {
271      const isInteger = schema.type === 'integer'
272      const formatNum = (n: number) =>
273        Number.isInteger(n) && !isInteger ? `${n}.0` : String(n)
274  
275      if (schema.minimum !== undefined && schema.maximum !== undefined) {
276        return `(${schema.type} between ${formatNum(schema.minimum!)} and ${formatNum(schema.maximum!)})`
277      } else if (schema.minimum !== undefined) {
278        return `(${schema.type} >= ${formatNum(schema.minimum!)})`
279      } else if (schema.maximum !== undefined) {
280        return `(${schema.type} <= ${formatNum(schema.maximum!)})`
281      } else {
282        const example = schema.type === 'integer' ? '42' : '3.14'
283        return `(${schema.type}, e.g. ${example})`
284      }
285    }
286  
287    return undefined
288  }
289  
290  /**
291   * Check if a schema is a date or date-time format that supports NL parsing
292   */
293  export function isDateTimeSchema(
294    schema: PrimitiveSchemaDefinition,
295  ): schema is StringSchema & { format: 'date' | 'date-time' } {
296    return (
297      schema.type === 'string' &&
298      'format' in schema &&
299      (schema.format === 'date' || schema.format === 'date-time')
300    )
301  }
302  
303  /**
304   * Async validation that attempts NL date/time parsing via Haiku
305   * when the input doesn't look like ISO 8601.
306   */
307  export async function validateElicitationInputAsync(
308    stringValue: string,
309    schema: PrimitiveSchemaDefinition,
310    signal: AbortSignal,
311  ): Promise<ValidationResult> {
312    const syncResult = validateElicitationInput(stringValue, schema)
313    if (syncResult.isValid) {
314      return syncResult
315    }
316  
317    if (isDateTimeSchema(schema) && !looksLikeISO8601(stringValue)) {
318      const parseResult = await parseNaturalLanguageDateTime(
319        stringValue,
320        schema.format,
321        signal,
322      )
323  
324      if (parseResult.success) {
325        const validatedParsed = validateElicitationInput(
326          parseResult.value,
327          schema,
328        )
329        if (validatedParsed.isValid) {
330          return validatedParsed
331        }
332      }
333    }
334  
335    return syncResult
336  }