Cache Tool::Schema Validation to Avoid Re-validating Identical Schemas

lib/mcp/tool/schema.rb

@@ -1,10 +1,41 @@
 # frozen_string_literal: true
 
+require "digest"
 require "json-schema"
 
 module MCP
   class Tool
     class Schema
+      # Metaschema validation depends only on schema content, so a given schema
+      # never needs to be validated more than once. Caching the result lets repeated
+      # (e.g. dynamically rebuilt) schemas skip the costly traversal.
+      class ValidationCache
+        DEFAULT_MAX_SIZE = 1000
+
+        def initialize(max_size: DEFAULT_MAX_SIZE)
+          @max_size = max_size
+          @entries = {}
+          @mutex = Mutex.new
+        end
+
+        def validated?(key)
+          @mutex.synchronize { @entries.key?(key) }
+        end
+
+        def store(key)
+          @mutex.synchronize do
+            @entries.delete(key)
+            @entries[key] = true
+            @entries.shift while @entries.size > @max_size
+          end
+        end
+
+        def clear
+          @mutex.synchronize { @entries.clear }
+        end
+      end
+      VALIDATION_CACHE = ValidationCache.new
+
       # JSON Schema 2020-12 is the default dialect for MCP schema definitions
       # per MCP 2025-11-25 (SEP-1613). Note: emission only — runtime validation
       # is still performed against the JSON Schema draft-04 metaschema because
@@ -36,6 +67,14 @@ def fully_validate(data)
       end
 
       def validate_schema!
+        target = schema_for_validation
+
+        # `max_nesting: false` because normalization uses `JSON.dump` (no nesting limit),
+        # so the default `JSON.generate` limit would raise on a deeply nested schema that
+        # the initializer already accepted.
+        key = Digest::SHA256.hexdigest(JSON.generate(target, max_nesting: false))
+        return if VALIDATION_CACHE.validated?(key)
+
         gem_path = File.realpath(Gem.loaded_specs["json-schema"].full_gem_path)
         schema_reader = JSON::Schema::Reader.new(
           accept_uri: false,
@@ -45,10 +84,12 @@ def validate_schema!
         # Converts metaschema to a file URI for cross-platform compatibility
         metaschema_uri = JSON::Util::URI.file_uri(metaschema_path.expand_path.cleanpath.to_s.tr("\\", "/"))
         metaschema = metaschema_uri.to_s
-        errors = JSON::Validator.fully_validate(metaschema, schema_for_validation, schema_reader: schema_reader)
+        errors = JSON::Validator.fully_validate(metaschema, target, schema_reader: schema_reader)
         if errors.any?
           raise ArgumentError, "Invalid JSON Schema: #{errors.join(", ")}"
         end
+
+        VALIDATION_CACHE.store(key)
       end
 
       # The `json-schema` gem's draft-04 validator cannot resolve newer or unknown `$schema`

test/mcp/tool/schema_test.rb

@@ -0,0 +1,95 @@
+# frozen_string_literal: true
+
+require "test_helper"
+
+module MCP
+  class Tool
+    class SchemaTest < ActiveSupport::TestCase
+      setup do
+        Schema::VALIDATION_CACHE.clear
+      end
+
+      test "validates a schema once and reuses the result for identical schemas" do
+        JSON::Validator.expects(:fully_validate).once.returns([])
+
+        schema = { properties: { validates_once: { type: "string" } } }
+        InputSchema.new(schema)
+        InputSchema.new(schema)
+      end
+
+      test "validates distinct schemas separately" do
+        JSON::Validator.expects(:fully_validate).twice.returns([])
+
+        InputSchema.new(properties: { distinct_a: { type: "string" } })
+        InputSchema.new(properties: { distinct_b: { type: "string" } })
+      end
+
+      test "a cache hit still yields a usable, validated schema" do
+        schema = { properties: { cache_hit: { type: "string" } }, required: ["cache_hit"] }
+        InputSchema.new(schema)
+        cached = InputSchema.new(schema)
+
+        assert_equal(
+          {
+            "$schema": "https://json-schema.org/draft/2020-12/schema",
+            type: "object",
+            properties: { cache_hit: { type: "string" } },
+            required: ["cache_hit"],
+          },
+          cached.to_h,
+        )
+        assert_nil(cached.validate_arguments(cache_hit: "value"))
+        assert_raises(InputSchema::ValidationError) do
+          cached.validate_arguments(cache_hit: 123)
+        end
+      end
+
+      test "an invalid schema raises every time and is not cached" do
+        invalid = { properties: { not_cached: { type: "invalid_type" } } }
+
+        assert_raises(ArgumentError) { InputSchema.new(invalid) }
+        assert_raises(ArgumentError) { InputSchema.new(invalid) }
+      end
+
+      test "a schema at the normalization depth limit is cached without a nesting error" do
+        # The deepest schema the initializer can still normalize via JSON.dump/parse.
+        # The cache key must tolerate the same depth; the default JSON.generate
+        # nesting limit (100) is stricter than normalization and would raise here.
+        schema = { properties: { leaf: { type: "string" } } }
+        loop do
+          candidate = { properties: { child: schema } }
+          JSON.parse(JSON.dump(candidate))
+          schema = candidate
+        rescue JSON::NestingError
+          break
+        end
+
+        JSON::Validator.stub(:fully_validate, []) do
+          assert_nothing_raised do
+            InputSchema.new(schema)
+            InputSchema.new(schema)
+          end
+        end
+      end
+
+      test "ValidationCache evicts the oldest entry beyond its max size" do
+        cache = Schema::ValidationCache.new(max_size: 2)
+        cache.store("a")
+        cache.store("b")
+        cache.store("c")
+
+        refute cache.validated?("a")
+        assert cache.validated?("b")
+        assert cache.validated?("c")
+      end
+
+      test "ValidationCache#clear empties the cache" do
+        cache = Schema::ValidationCache.new
+        cache.store("a")
+        cache.clear
+
+        refute cache.validated?("a")
+      end
+    end
+  end
+end