conformance/src/validation/wire-schema.ts

// Validates wire JSON-RPC messages against the vendored per-version spec schema
// (src/spec-types/*.schema.json): envelope, then typed defs via method/error.code consts and
// request→result pairs. Known gap: the client-auth scenarios' express mock is not instrumented.

import { AsyncLocalStorage } from 'node:async_hooks';
import { Ajv, type ValidateFunction, type ErrorObject } from 'ajv';
import { Ajv2020 } from 'ajv/dist/2020.js';
import { default as addFormats } from 'ajv-formats';
import {
  DRAFT_PROTOCOL_VERSION,
  type ConformanceCheck,
  type SpecVersion
} from '../types';
import schema2025_03_26 from '../spec-types/2025-03-26.schema.json';
import schema2025_06_18 from '../spec-types/2025-06-18.schema.json';
import schema2025_11_25 from '../spec-types/2025-11-25.schema.json';
import schemaDraft from '../spec-types/draft.schema.json';

export type WireOrigin = 'harness' | 'implementation';

export interface WireMessageInfo {
  /** Who authored the message: the harness or the implementation-under-test. */
  origin: WireOrigin;
  /** Human-readable location, e.g. `stateless request 'tools/call'`. */
  context: string;
  /** For responses: the answered request's method, so the result is validated
   * against its typed result definition. */
  requestMethod?: string;
}

export interface WireSchemaViolation {
  origin: WireOrigin;
  specVersion: SpecVersion;
  context: string;
  errors: string[];
  message: unknown;
}

const SCHEMAS: Record<SpecVersion, Record<string, unknown>> = {
  '2025-03-26': schema2025_03_26,
  '2025-06-18': schema2025_06_18,
  '2025-11-25': schema2025_11_25,
  [DRAFT_PROTOCOL_VERSION]: schemaDraft
};

/** Spec-repo directory name for a version (the draft is pinned unversioned). */
export function schemaDirFor(specVersion: SpecVersion): string {
  return specVersion === DRAFT_PROTOCOL_VERSION ? 'draft' : specVersion;
}

/** Union definitions that alias a single concrete type (and so carry a
 * `method` const) but are not the canonical definition for that method. */
const NON_CANONICAL_DEFS = new Set([
  'ClientRequest',
  'ClientNotification',
  'ClientResult',
  'ClientMessage',
  'ServerRequest',
  'ServerNotification',
  'ServerResult',
  'ServerMessage'
]);

interface CompiledSpec {
  defsKey: '$defs' | 'definitions';
  defs: Record<string, Record<string, unknown>>;
  /** JSON-RPC method → canonical request/notification definition name. */
  methodDefs: Map<string, string>;
  /** `error.code` const → typed error-response definition name. */
  errorDefs: Map<number, string>;
  /** JSON-RPC method → typed result definition name. */
  resultDefs: Map<string, string>;
  validatorFor(defName: string): ValidateFunction;
}

const compiledSpecs = new Map<SpecVersion, CompiledSpec>();

/** Patch known generator bugs in released spec schema.json files at load time
 * so validation matches the schema.ts source of truth; delete each branch (a
 * test trips) once the dated schema is fixed upstream and re-vendored. */
function applySchemaErrata(
  specVersion: SpecVersion,
  schema: Record<string, unknown>
): Record<string, unknown> {
  if (specVersion !== '2025-11-25' && specVersion !== '2025-06-18') {
    return schema;
  }
  // NumberSchema minimum/maximum (plus default at 2025-11-25) are generated
  // as `integer`, contradicting schema.ts (`number`). Fixed for draft in
  // modelcontextprotocol#2710; dated fixes proposed in modelcontextprotocol#3139.
  const patched = structuredClone(schema);
  const defs = (patched.$defs ?? patched.definitions) as Record<
    string,
    { properties?: Record<string, { type?: string }> }
  >;
  for (const prop of ['minimum', 'maximum', 'default']) {
    const p = defs.NumberSchema?.properties?.[prop];
    if (p?.type === 'integer') p.type = 'number';
  }
  return patched;
}

function compileSpec(specVersion: SpecVersion): CompiledSpec {
  let compiled = compiledSpecs.get(specVersion);
  if (compiled) return compiled;

  const schema = applySchemaErrata(specVersion, SCHEMAS[specVersion]);
  const is2020 =
    typeof schema.$schema === 'string' && schema.$schema.includes('2020-12');
  // The spec schemas are not authored for ajv strict mode; validate them
  // as-is rather than editing vendored files.
  const options = { strict: false, allErrors: true } as const;
  const ajv = is2020 ? new Ajv2020(options) : new Ajv(options);
  addFormats(ajv);
  // "byte" (base64) is an OpenAPI format ajv-formats doesn't ship; the spec
  // uses it on blob contents. Accept any string rather than reimplementing it.
  ajv.addFormat('byte', true);

  const schemaId = `mcp://${schemaDirFor(specVersion)}/schema.json`;
  ajv.addSchema({ ...schema, $id: schemaId });

  const defsKey: '$defs' | 'definitions' =
    '$defs' in schema ? '$defs' : 'definitions';
  const defs = schema[defsKey] as Record<string, Record<string, unknown>>;

  const methodDefs = new Map<string, string>();
  const errorDefs = new Map<number, string>();
  const resultDefs = new Map<string, string>();
  for (const [name, def] of Object.entries(defs)) {
    if (NON_CANONICAL_DEFS.has(name)) continue;
    const props = (def.properties ?? {}) as Record<
      string,
      Record<string, unknown>
    >;
    const method = props.method?.const;
    if (typeof method === 'string') {
      methodDefs.set(method, name);
      if (name.endsWith('Request')) {
        const resultName = `${name.slice(0, -'Request'.length)}Result`;
        if (resultName in defs) resultDefs.set(method, resultName);
      }
    }
    const errorProp = props.error as
      | { allOf?: { properties?: { code?: { const?: unknown } } }[] }
      | undefined;
    for (const part of errorProp?.allOf ?? []) {
      const code = part.properties?.code?.const;
      if (typeof code === 'number') errorDefs.set(code, name);
    }
  }

  const validators = new Map<string, ValidateFunction>();
  compiled = {
    defsKey,
    defs,
    methodDefs,
    errorDefs,
    resultDefs,
    validatorFor(defName: string): ValidateFunction {
      let v = validators.get(defName);
      if (!v) {
        v = ajv.compile({ $ref: `${schemaId}#/${defsKey}/${defName}` });
        validators.set(defName, v);
      }
      return v;
    }
  };
  compiledSpecs.set(specVersion, compiled);
  return compiled;
}

/** Dispatch maps extracted from a version's schema (method → request/notification def,
 * `error.code` const → error def). Exposed so a unit test can pin them: if a schema sync
 * restructures what the extraction walks, validation silently degrades to envelope-only. */
export function specDispatchMaps(specVersion: SpecVersion): {
  methodDefs: ReadonlyMap<string, string>;
  errorDefs: ReadonlyMap<number, string>;
  resultDefs: ReadonlyMap<string, string>;
} {
  const spec = compileSpec(specVersion);
  return {
    methodDefs: spec.methodDefs,
    errorDefs: spec.errorDefs,
    resultDefs: spec.resultDefs
  };
}

function formatErrors(
  defName: string,
  errors: ErrorObject[] | null | undefined
): string[] {
  const MAX = 6;
  const formatted = (errors ?? []).map(
    (e) =>
      `${defName}${e.instancePath || ''}: ${e.message ?? 'invalid'}` +
      (e.keyword === 'const' || e.keyword === 'enum'
        ? ` (${JSON.stringify(e.params)})`
        : '')
  );
  if (formatted.length > MAX) {
    return [
      ...formatted.slice(0, MAX),
      `... ${formatted.length - MAX} more error(s)`
    ];
  }
  return formatted.length > 0 ? formatted : [`${defName}: invalid`];
}

/** Validate a single JSON-RPC message against the given spec version's schema.
 * Returns `[]` when the message is valid. Pure — does not touch the recorder. */
export function wireSchemaErrors(
  specVersion: SpecVersion,
  message: unknown,
  requestMethod?: string
): string[] {
  const spec = compileSpec(specVersion);
  const validateAgainst = (defName: string, value: unknown): string[] => {
    const validate = spec.validatorFor(defName);
    return validate(value) ? [] : formatErrors(defName, validate.errors);
  };
  const firstDef = (...names: string[]): string =>
    names.find((n) => n in spec.defs) ?? names[names.length - 1];

  if (Array.isArray(message)) {
    // A batch: only legal where the envelope union admits arrays (2025-03-26).
    // Limitation: `requestMethod` is forwarded to every element, so batched responses all
    // validate against one result type — no caller sends batched requests today.
    const elementErrors = message.flatMap((m, i) =>
      wireSchemaErrors(specVersion, m, requestMethod).map((e) => `[${i}] ${e}`)
    );
    if (elementErrors.length > 0) return elementErrors;
    if (validateAgainst('JSONRPCMessage', message).length > 0) {
      return [
        'JSONRPCMessage: batch arrays are not valid in this spec version'
      ];
    }
    return [];
  }

  const msg = (
    typeof message === 'object' && message !== null ? message : {}
  ) as Record<string, unknown>;

  // Classify the message and validate against the most specific definition
  // the schema has for it. The typed definitions include the envelope
  // requirements (jsonrpc, id, method, ...), so they subsume the union check.
  if (typeof msg.method === 'string') {
    const defName =
      spec.methodDefs.get(msg.method) ??
      ('id' in msg ? 'JSONRPCRequest' : 'JSONRPCNotification');
    return validateAgainst(defName, message);
  }

  if (msg.error !== undefined && msg.error !== null) {
    const code = (msg.error as Record<string, unknown>).code;
    const defName =
      (typeof code === 'number' ? spec.errorDefs.get(code) : undefined) ??
      firstDef('JSONRPCErrorResponse', 'JSONRPCError');
    return validateAgainst(defName, message);
  }

  if (msg.result !== undefined) {
    // SEP-2322 (MRTR): any request may be answered with an InputRequiredResult
    // instead of its method's result type; discriminate on resultType.
    const inputRequired =
      (msg.result as Record<string, unknown> | null)?.resultType ===
        'input_required' && 'InputRequiredResult' in spec.defs;
    const resultDefName = inputRequired
      ? 'InputRequiredResult'
      : requestMethod !== undefined
        ? spec.resultDefs.get(requestMethod)
        : undefined;
    if (resultDefName) {
      const typed = validateAgainst(resultDefName, msg.result).map(
        (e) => `${e} (result of '${requestMethod}')`
      );
      if (typed.length > 0) return typed;
    }
    return validateAgainst(
      firstDef('JSONRPCResultResponse', 'JSONRPCResponse'),
      message
    );
  }

  return [
    'JSONRPCMessage: not a valid JSON-RPC request, notification, or response'
  ];
}

// Recorder: choke points record every message; runners (and the vitest hook)
// drain the accumulated violations per scenario / test. Suite runs execute
// scenarios concurrently, so each scenario gets its own recorder via
// AsyncLocalStorage; code outside any scope falls back to the global recorder.

interface RecorderState {
  violations: WireSchemaViolation[];
  observed: number;
}

const recorderStorage = new AsyncLocalStorage<RecorderState>();
const globalRecorder: RecorderState = { violations: [], observed: 0 };

function recorder(): RecorderState {
  return recorderStorage.getStore() ?? globalRecorder;
}

/** Run `fn` with its own isolated wire recorder. Async work started inside
 * `fn` (mock servers, transport hooks) records into that recorder, so
 * concurrent scenarios in a suite run cannot steal or wipe each other's
 * violations. */
export function withWireRecorder<T>(fn: () => Promise<T>): Promise<T> {
  return recorderStorage.run({ violations: [], observed: 0 }, fn);
}

/** Validate a wire message and record any violation. Called by choke points. */
export function validateWireMessage(
  specVersion: SpecVersion,
  message: unknown,
  info: WireMessageInfo
): void {
  const state = recorder();
  state.observed++;
  const errors = wireSchemaErrors(specVersion, message, info.requestMethod);
  if (errors.length > 0) {
    state.violations.push({
      origin: info.origin,
      specVersion,
      context: info.context,
      errors,
      message
    });
  }
}

export function resetWireValidation(): void {
  const state = recorder();
  state.violations = [];
  state.observed = 0;
}

/** Return everything recorded since the last reset, and reset. */
export function takeWireViolations(): {
  violations: WireSchemaViolation[];
  observed: number;
} {
  const state = recorder();
  const result = { violations: state.violations, observed: state.observed };
  resetWireValidation();
  return result;
}

function violationDetails(v: WireSchemaViolation): Record<string, unknown> {
  return {
    origin: v.origin,
    specVersion: v.specVersion,
    context: v.context,
    errors: v.errors,
    message: v.message
  };
}

export function formatWireViolation(v: WireSchemaViolation): string {
  return (
    `[${v.origin}] ${v.context} (spec ${v.specVersion}): ` +
    `${v.errors.join('; ')} — message: ${JSON.stringify(v.message)}`
  );
}

/** Drain the recorder and synthesize the per-scenario conformance checks. Returns `[]`
 * when no wire traffic was observed (e.g. a scenario asserting only on raw HTTP). */
export function wireSchemaChecks(specVersion: SpecVersion): ConformanceCheck[] {
  const { violations: all, observed: count } = takeWireViolations();
  if (count === 0 && all.length === 0) return [];

  const timestamp = new Date().toISOString();
  const specReferences = [
    {
      id: 'MCP-Schema',
      url: `https://github.com/modelcontextprotocol/modelcontextprotocol/blob/main/schema/${schemaDirFor(specVersion)}/schema.json`
    }
  ];
  const checks: ConformanceCheck[] = [];

  const implViolations = all.filter((v) => v.origin === 'implementation');
  checks.push({
    id: 'wire-schema-valid',
    name: 'WireSchemaValid',
    description:
      'Every JSON-RPC message the implementation sent is valid per the spec JSON schema for the negotiated spec version',
    status: implViolations.length === 0 ? 'SUCCESS' : 'FAILURE',
    timestamp,
    specReferences,
    details: {
      messagesValidated: count,
      violations: implViolations.map(violationDetails)
    },
    errorMessage:
      implViolations.length > 0
        ? implViolations.map(formatWireViolation).join('\n')
        : undefined
  });

  const harnessViolations = all.filter((v) => v.origin === 'harness');
  if (harnessViolations.length > 0) {
    checks.push({
      id: 'wire-schema-harness-error',
      name: 'WireSchemaHarnessError',
      description:
        'HARNESS ERROR: the conformance harness itself sent a JSON-RPC message the spec JSON schema forbids. This is a bug in the harness or a scenario fixture, not in the implementation under test.',
      status: 'FAILURE',
      timestamp,
      specReferences,
      details: { violations: harnessViolations.map(violationDetails) },
      errorMessage:
        'HARNESS ERROR: ' +
        harnessViolations.map(formatWireViolation).join('\n')
    });
  }

  return checks;
}