🍃 fix: Strip $-Prefixed Schema Keywords Before Persisting MCP Tool Params (#14464)

* 🐛 fix: strip $-prefixed schema keywords from MCP tool params before storage

MCP tools whose inputSchema carries a spec-compliant $schema keyword (or any
other $-prefixed JSON Schema keyword) failed to register: MongoDB rejects field
names beginning with $, so persisting the tool's parameters blob threw
"The dollar ($) prefixed field '...$schema' is not valid for storage".

Normalize the schema when building stored toolFunctions (resolve $refs and drop
$-prefixed keywords via the existing resolveJsonSchemaRefs + normalizeJsonSchema
pipeline). normalizeJsonSchema now strips every $-prefixed keyword, not just
$defs, while preserving property names that happen to start with $.

* fix: recurse through every schema-valued keyword when stripping $ keys

MongoDB rejects $-prefixed field names at any depth, but the normalizer only
recursed through properties, items, additionalProperties and unions, so a
$schema or $comment nested under not, if/then/else, contains, propertyNames,
patternProperties, dependentSchemas or prefixItems survived into the persisted
tool parameters and still failed registration.

The keyword sets are now explicit, covering the single-subschema, map-of-schema
and list-of-schema forms.

A $-prefixed property name is deliberately left alone: it is an argument the
tool actually accepts, so dropping it would silently remove the parameter from
the schema the model sees.

* fix: recurse into draft-07 dependencies and 2020-12 contentSchema

Both are schema-bearing and were absent from the traversal sets, so a nested
annotation survived into the persisted tool parameters and still hit the
MongoDB dollar-prefixed-field failure. dependencies is polymorphic - a value
may be an array of property names rather than a subschema - and that form
round-trips unchanged.

* fix: keep __proto__ entries when normalizing schema maps

Schema-map keys name instance properties, so __proto__ is a legal entry and
arrives as a real own property via JSON.parse. Plain assignment invoked the
prototype setter instead, silently dropping the constraint; entries are now
defined rather than assigned.

* fix: bound MCP schema reference expansion and keep __proto__ arguments

A remote MCP server controls the schema fetched at registration, and sibling
references to the same definition each re-expand because visited is cleared
after resolving - so an acyclic graph where each Dn holds two refs to Dn-1
expands 2^n. At depth 24 that is over 16 million nodes, enough to block the
event loop or exhaust memory before registration finishes.

Resolution now carries a node budget and leaves a reference unexpanded once it
is spent, and assignments use defineProperty so an argument legitimately named
__proto__ is not swallowed by the inherited setter during resolution.

---------

Co-authored-by: Arham Wani <arhamwani765@gmail.com>

packages/api/src/mcp/__tests__/zod.spec.ts

@@ -2298,6 +2298,70 @@ describe('normalizeJsonSchema', () => {
     expect(result.properties.name).toEqual({ type: 'string' });
   });
 
+  it('should strip the spec-compliant $schema keyword (MongoDB rejects $-prefixed keys)', () => {
+    const schema = {
+      $schema: 'https://json-schema.org/draft/2020-12/schema',
+      type: 'object',
+      properties: {
+        title: { type: 'string' },
+      },
+      required: ['title'],
+    } as any;
+
+    const result = normalizeJsonSchema(schema);
+    expect(result).not.toHaveProperty('$schema');
+    expect(result.type).toBe('object');
+    expect(result.properties.title).toEqual({ type: 'string' });
+    expect(result.required).toEqual(['title']);
+    // Nothing left behind that MongoDB would reject.
+    expect(Object.keys(result).some((k) => k.startsWith('$'))).toBe(false);
+  });
+
+  it('should strip $-prefixed annotation keywords at all nesting levels', () => {
+    const schema = {
+      $schema: 'https://json-schema.org/draft/2020-12/schema',
+      $id: 'urn:tool:create_citation',
+      $comment: 'root comment',
+      type: 'object',
+      properties: {
+        note: {
+          type: 'string',
+          $comment: 'nested comment',
+          $anchor: 'note',
+        },
+        items: {
+          type: 'array',
+          items: { type: 'string', $id: 'urn:item' },
+        },
+      },
+    } as any;
+
+    const result = normalizeJsonSchema(schema);
+    expect(result).not.toHaveProperty('$schema');
+    expect(result).not.toHaveProperty('$id');
+    expect(result).not.toHaveProperty('$comment');
+    expect(result.properties.note).toEqual({ type: 'string' });
+    expect(result.properties.items.items).toEqual({ type: 'string' });
+  });
+
+  it('should preserve property names that begin with $', () => {
+    const schema = {
+      $schema: 'https://json-schema.org/draft/2020-12/schema',
+      type: 'object',
+      properties: {
+        $ref: { type: 'string', description: 'a property literally named $ref' },
+      },
+    } as any;
+
+    const result = normalizeJsonSchema(schema);
+    expect(result).not.toHaveProperty('$schema');
+    // A `$`-prefixed name under `properties` is a data field, not a keyword.
+    expect(result.properties.$ref).toEqual({
+      type: 'string',
+      description: 'a property literally named $ref',
+    });
+  });
+
   it('should strip x-* fields inside oneOf/anyOf/allOf', () => {
     const schema = {
       type: 'object',
@@ -3016,3 +3080,159 @@ describe('resolveJsonSchemaRefs local pointer refs', () => {
     expect(resolved.properties?.alias).toEqual({ type: 'string' });
   });
 });
+
+describe('normalizeJsonSchema $-key recursion', () => {
+  /** Mongo rejects `$`-prefixed field names at any depth, so a `$` keyword nested
+   *  under a schema-valued container has to be stripped too or the stored
+   *  `parameters` blob still fails to persist. */
+  const nested = (container: Record<string, unknown>) =>
+    normalizeJsonSchema({
+      type: 'object',
+      properties: { q: { type: 'string', ...container } },
+    } as Record<string, unknown>) as {
+      properties: { q: Record<string, unknown> };
+    };
+
+  it.each([
+    ['not', { not: { $schema: 'https://json-schema.org/draft/2020-12/schema', type: 'null' } }],
+    ['if', { if: { $comment: 'x', type: 'string' } }],
+    ['then', { then: { $comment: 'x', type: 'string' } }],
+    ['else', { else: { $comment: 'x', type: 'string' } }],
+    ['contains', { contains: { $id: 'x', type: 'string' } }],
+    ['propertyNames', { propertyNames: { $comment: 'x', type: 'string' } }],
+  ])('strips a $ keyword nested under %s', (name, container) => {
+    const result = nested(container);
+    expect(JSON.stringify(result)).not.toContain('"$');
+    expect(result.properties.q[name]).toBeDefined();
+  });
+
+  it('strips $ keywords under patternProperties and dependentSchemas', () => {
+    const result = normalizeJsonSchema({
+      type: 'object',
+      patternProperties: { '^a': { $schema: 'x', type: 'string' } },
+      dependentSchemas: { a: { $comment: 'x', type: 'object' } },
+    } as Record<string, unknown>);
+
+    expect(JSON.stringify(result)).not.toContain('"$');
+  });
+
+  it('strips $ keywords under prefixItems', () => {
+    const result = normalizeJsonSchema({
+      type: 'array',
+      prefixItems: [{ $schema: 'x', type: 'string' }],
+    } as Record<string, unknown>);
+
+    expect(JSON.stringify(result)).not.toContain('"$');
+  });
+
+  it('preserves a $-prefixed property name, which is data rather than a keyword', () => {
+    /** `$filter` is a real argument the tool accepts; dropping it would silently
+     *  remove the parameter from the schema the model sees. Modern MongoDB accepts
+     *  `$`-prefixed field names, so this is left intact deliberately. */
+    const result = normalizeJsonSchema({
+      type: 'object',
+      properties: { $filter: { type: 'string', $comment: 'odata' } },
+    } as Record<string, unknown>) as { properties: Record<string, unknown> };
+
+    expect(result.properties.$filter).toEqual({ type: 'string' });
+  });
+});
+
+describe('normalizeJsonSchema draft-07 and 2020-12 containers', () => {
+  it('strips a $ keyword under draft-07 dependencies', () => {
+    const result = normalizeJsonSchema({
+      type: 'object',
+      dependencies: { foo: { $comment: 'x', type: 'object' } },
+    } as Record<string, unknown>);
+
+    expect(JSON.stringify(result)).not.toContain('"$');
+  });
+
+  it('leaves a draft-07 dependencies property-name array intact', () => {
+    /** `dependencies` is polymorphic: an array of required property names is data,
+     *  not a subschema, and must round-trip unchanged. */
+    const result = normalizeJsonSchema({
+      type: 'object',
+      dependencies: { foo: ['bar', 'baz'] },
+    } as Record<string, unknown>) as { dependencies: Record<string, unknown> };
+
+    expect(result.dependencies.foo).toEqual(['bar', 'baz']);
+  });
+
+  it('strips a $ keyword under contentSchema', () => {
+    const result = normalizeJsonSchema({
+      type: 'string',
+      contentMediaType: 'application/json',
+      contentSchema: { $schema: 'x', type: 'object' },
+    } as Record<string, unknown>);
+
+    expect(JSON.stringify(result)).not.toContain('"$');
+  });
+});
+
+describe('normalizeJsonSchema prototype-polluting map keys', () => {
+  /** An MCP server's schema arrives as JSON, and `JSON.parse` creates a real own
+   *  `__proto__` property — unlike an object literal, where it sets the prototype. */
+  const parse = (json: string) => JSON.parse(json) as Record<string, unknown>;
+
+  it('keeps a __proto__ entry in a schema map as an own property', () => {
+    const result = normalizeJsonSchema(
+      parse('{"type":"object","properties":{"__proto__":{"type":"string","$comment":"x"}}}'),
+    ) as { properties: Record<string, unknown> };
+
+    expect(Object.prototype.hasOwnProperty.call(result.properties, '__proto__')).toBe(true);
+    expect(JSON.parse(JSON.stringify(result)).properties.__proto__).toEqual({ type: 'string' });
+  });
+
+  it('keeps a __proto__ entry under dependentSchemas', () => {
+    const result = normalizeJsonSchema(
+      parse('{"type":"object","dependentSchemas":{"__proto__":{"type":"object"}}}'),
+    ) as { dependentSchemas: Record<string, unknown> };
+
+    expect(Object.prototype.hasOwnProperty.call(result.dependentSchemas, '__proto__')).toBe(true);
+  });
+});
+
+describe('resolveJsonSchemaRefs expansion safety', () => {
+  /** A remote MCP server controls this schema. Each `Dn` holding two refs to
+   *  `Dn-1` is a compact acyclic graph that expands 2^n, so registration must
+   *  stay bounded rather than exhaust memory. */
+  const fanOutSchema = (depth: number) => {
+    const $defs: Record<string, unknown> = { D0: { type: 'string' } };
+    for (let i = 1; i <= depth; i++) {
+      $defs[`D${i}`] = {
+        type: 'object',
+        properties: { a: { $ref: `#/$defs/D${i - 1}` }, b: { $ref: `#/$defs/D${i - 1}` } },
+      };
+    }
+    return { $defs, $ref: `#/$defs/D${depth}` } as Record<string, unknown>;
+  };
+
+  it('stays bounded on an exponentially-expanding reference graph', () => {
+    const start = Date.now();
+    const result = resolveJsonSchemaRefs(fanOutSchema(30));
+    const nodes = JSON.stringify(result).length;
+
+    expect(Date.now() - start).toBeLessThan(10_000);
+    expect(nodes).toBeLessThan(50_000_000);
+  });
+
+  it('still resolves an ordinary reference graph fully', () => {
+    const result = resolveJsonSchemaRefs({
+      $defs: { Name: { type: 'string' } },
+      type: 'object',
+      properties: { first: { $ref: '#/$defs/Name' } },
+    } as Record<string, unknown>) as { properties: { first: Record<string, unknown> } };
+
+    expect(result.properties.first).toEqual({ type: 'string' });
+  });
+
+  it('keeps a __proto__ argument through reference resolution', () => {
+    const parsed = JSON.parse(
+      '{"type":"object","properties":{"__proto__":{"type":"string"}}}',
+    ) as Record<string, unknown>;
+    const result = resolveJsonSchemaRefs(parsed) as { properties: Record<string, unknown> };
+
+    expect(Object.prototype.hasOwnProperty.call(result.properties, '__proto__')).toBe(true);
+  });
+});

packages/api/src/mcp/registry/MCPServerInspector.ts

@@ -10,6 +10,7 @@ import {
   isUserSourced,
 } from '~/mcp/utils';
 import { isMCPDomainAllowed, extractMCPServerDomain } from '~/auth/domain';
+import { normalizeJsonSchema, resolveJsonSchemaRefs } from '~/mcp/zod';
 import { MCPConnectionFactory } from '~/mcp/MCPConnectionFactory';
 import { MCPDomainNotAllowedError } from '~/mcp/errors';
 import { detectOAuthRequirement } from '~/mcp/oauth';
@@ -187,7 +188,13 @@ export class MCPServerInspector {
         ['function']: {
           name,
           description: tool.description,
-          parameters: tool.inputSchema as JsonSchemaType,
+          // Normalize before persisting: resolves `$ref`s and strips
+          // `$`-prefixed keywords (e.g. a spec-compliant `$schema`), which
+          // MongoDB rejects as field names and would otherwise crash storage
+          // of this `parameters` blob during server registration.
+          parameters: normalizeJsonSchema(
+            resolveJsonSchemaRefs(tool.inputSchema as Record<string, unknown>),
+          ) as JsonSchemaType,
         },
       };
     });

packages/api/src/mcp/zod.ts

@@ -211,11 +211,36 @@ function resolveLocalPointer(
  * @param root - The root schema local pointers resolve against (defaults to `schema`)
  * @returns The resolved schema
  */
+/**
+ * Caps how many nodes a single resolution may emit. A remote MCP server controls
+ * this schema, and sibling references to the same definition each re-expand, so a
+ * compact acyclic graph can blow up exponentially (`Dn` holding two refs to
+ * `Dn-1` is 2^n). Past the cap the reference is left unexpanded rather than
+ * exhausting memory during registration.
+ */
+export const MAX_RESOLVED_SCHEMA_NODES = 50_000;
+
+interface ResolveBudget {
+  remaining: number;
+}
+
+/** Assigns without invoking the inherited `__proto__` setter, which would drop a
+ *  legitimately-named argument instead of creating an own property. */
+function setOwn(target: Record<string, unknown>, key: string, value: unknown): void {
+  Object.defineProperty(target, key, {
+    value,
+    enumerable: true,
+    writable: true,
+    configurable: true,
+  });
+}
+
 export function resolveJsonSchemaRefs<T extends Record<string, unknown>>(
   schema: T,
   definitions?: Record<string, unknown>,
   visited: Set<string> = new Set<string>(),
   root?: Record<string, unknown>,
+  budget: ResolveBudget = { remaining: MAX_RESOLVED_SCHEMA_NODES },
 ): T {
   // Handle null, undefined, or non-object values first
   if (!schema || typeof schema !== 'object') {
@@ -232,10 +257,12 @@ export function resolveJsonSchemaRefs<T extends Record<string, unknown>>(
   // Handle arrays
   if (Array.isArray(schema)) {
     return schema.map((item) =>
-      resolveJsonSchemaRefs(item, definitions, visited, rootSchema),
+      resolveJsonSchemaRefs(item, definitions, visited, rootSchema, budget),
     ) as unknown as T;
   }
 
+  budget.remaining -= 1;
+
   // Handle objects
   const result: Record<string, unknown> = {};
 
@@ -263,33 +290,39 @@ export function resolveJsonSchemaRefs<T extends Record<string, unknown>>(
         resolved = resolveLocalPointer(rootSchema, value);
       }
 
-      if (resolved) {
+      if (resolved && budget.remaining > 0) {
         visited.add(value);
         const resolvedSchema = resolveJsonSchemaRefs(
           resolved as Record<string, unknown>,
           definitions,
           visited,
           rootSchema,
+          budget,
         );
         visited.delete(value);
 
         // Merge the resolved schema into the result
         Object.assign(result, resolvedSchema);
       } else {
-        // If we can't resolve the reference, keep it as is
-        result[key] = value;
+        /** Unresolvable, or the expansion budget is spent: leave the reference. */
+        setOwn(result, key, value);
       }
     } else if (value && typeof value === 'object') {
       // Recursively resolve nested objects/arrays
-      result[key] = resolveJsonSchemaRefs(
-        value as Record<string, unknown>,
-        definitions,
-        visited,
-        rootSchema,
+      setOwn(
+        result,
+        key,
+        resolveJsonSchemaRefs(
+          value as Record<string, unknown>,
+          definitions,
+          visited,
+          rootSchema,
+          budget,
+        ),
       );
     } else {
       // Copy primitive values as is
-      result[key] = value;
+      setOwn(result, key, value);
     }
   }
 
@@ -302,11 +335,46 @@ export function resolveJsonSchemaRefs<T extends Record<string, unknown>>(
  * Transformations applied:
  * - Converts `const` values to `enum` arrays (Gemini/Vertex AI rejects `const`)
  * - Strips vendor extension fields (`x-*` prefixed keys, e.g. `x-google-enum-descriptions`)
- * - Strips leftover `$defs`/`definitions` blocks that may survive ref resolution
+ * - Strips `definitions` and `$`-prefixed schema keywords (`$defs`, `$schema`,
+ *   `$id`, `$comment`, ...) that may survive ref resolution
+ *
+ * Beyond LLM compatibility, dropping every `$`-prefixed keyword also makes the
+ * output safe to persist: MongoDB rejects field names beginning with `$`, so a
+ * standard, spec-compliant `$schema` keyword in an MCP tool's `inputSchema`
+ * would otherwise crash storage of the tool's `parameters` blob.
  *
  * @param schema - The JSON schema to normalize
  * @returns The normalized schema
  */
+/** Keywords whose value is a single subschema. */
+const SCHEMA_KEYWORDS = new Set([
+  'items',
+  'additionalItems',
+  'unevaluatedItems',
+  'additionalProperties',
+  'unevaluatedProperties',
+  'propertyNames',
+  'contains',
+  'contentSchema',
+  'not',
+  'if',
+  'then',
+  'else',
+]);
+
+/** Keywords whose value maps names to subschemas. */
+const SCHEMA_MAP_KEYWORDS = new Set([
+  'properties',
+  'patternProperties',
+  'dependentSchemas',
+  /** draft-07, where a value is either a subschema or an array of property
+   *  names; an array round-trips unchanged through the recursion. */
+  'dependencies',
+]);
+
+/** Keywords whose value is an array of subschemas. */
+const SCHEMA_LIST_KEYWORDS = new Set(['oneOf', 'anyOf', 'allOf', 'prefixItems']);
+
 export function normalizeJsonSchema<T extends Record<string, unknown>>(schema: T): T {
   if (!schema || typeof schema !== 'object') {
     return schema;
@@ -327,9 +395,14 @@ export function normalizeJsonSchema<T extends Record<string, unknown>>(schema: T
       continue;
     }
 
-    // Strip leftover $defs/definitions (should already be resolved by resolveJsonSchemaRefs,
-    // but strip as a safety net for schemas that bypass ref resolution).
-    if (key === '$defs' || key === 'definitions') {
+    // Strip `definitions` and any `$`-prefixed JSON Schema keyword (`$defs`,
+    // `$schema`, `$id`, `$comment`, ...). `$defs`/`$ref` should already be
+    // resolved away by resolveJsonSchemaRefs; the remaining `$`-prefixed keys
+    // are informational annotations the LLM function schema doesn't need — and
+    // MongoDB rejects `$`-prefixed field names, so leaving them in a stored MCP
+    // tool `parameters` blob breaks persistence. Property names (which live
+    // under `properties` and are handled below) are never reached here.
+    if (key === 'definitions' || key.startsWith('$')) {
       continue;
     }
 
@@ -343,22 +416,31 @@ export function normalizeJsonSchema<T extends Record<string, unknown>>(schema: T
       continue;
     }
 
-    if (key === 'properties' && value && typeof value === 'object' && !Array.isArray(value)) {
+    if (
+      SCHEMA_MAP_KEYWORDS.has(key) &&
+      value &&
+      typeof value === 'object' &&
+      !Array.isArray(value)
+    ) {
       const newProps: Record<string, unknown> = {};
       for (const [propKey, propValue] of Object.entries(value as Record<string, unknown>)) {
-        newProps[propKey] =
+        const normalized =
           propValue && typeof propValue === 'object'
             ? normalizeJsonSchema(propValue as Record<string, unknown>)
             : propValue;
+        /** These keys name instance properties, so `__proto__` is legal here.
+         *  Plain assignment would hit the prototype setter and drop the entry. */
+        Object.defineProperty(newProps, propKey, {
+          value: normalized,
+          enumerable: true,
+          writable: true,
+          configurable: true,
+        });
       }
       result[key] = newProps;
-    } else if (
-      (key === 'items' || key === 'additionalProperties') &&
-      value &&
-      typeof value === 'object'
-    ) {
+    } else if (SCHEMA_KEYWORDS.has(key) && value && typeof value === 'object') {
       result[key] = normalizeJsonSchema(value as Record<string, unknown>);
-    } else if ((key === 'oneOf' || key === 'anyOf' || key === 'allOf') && Array.isArray(value)) {
+    } else if (SCHEMA_LIST_KEYWORDS.has(key) && Array.isArray(value)) {
       result[key] = value.map((item) =>
         item && typeof item === 'object' ? normalizeJsonSchema(item) : item,
       );