From 8dee60cfc2f12d337fd44828a772ced02995917f Mon Sep 17 00:00:00 2001
From: Tyler Rowsell <tylerfrowsell@gmail.com>
Date: Mon, 9 Jun 2025 13:32:20 -0400
Subject: [PATCH] Fix: Disallow $ref in tool input schemas

---
 lib/mcp/tool/input_schema.rb       | 14 ++++++++++++++
 test/mcp/tool/input_schema_test.rb | 12 ++++++++++++
 test/mcp/tool_test.rb              |  2 +-
 3 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/lib/mcp/tool/input_schema.rb b/lib/mcp/tool/input_schema.rb
index 4d1c46f..b56b640 100644
--- a/lib/mcp/tool/input_schema.rb
+++ b/lib/mcp/tool/input_schema.rb
@@ -37,6 +37,7 @@ module MCP
       private
 
       def validate_schema!
+        check_for_refs!
         schema = to_h
         schema_reader = JSON::Schema::Reader.new(
           accept_uri: false,
@@ -48,6 +49,19 @@ module MCP
           raise ArgumentError, "Invalid JSON Schema: #{errors.join(", ")}"
         end
       end
+
+      def check_for_refs!(obj = properties)
+        case obj
+        when Hash
+          if obj.key?("$ref") || obj.key?(:$ref)
+            raise ArgumentError, "Invalid JSON Schema: $ref is not allowed in tool input schemas"
+          end
+
+          obj.each_value { |value| check_for_refs!(value) }
+        when Array
+          obj.each { |item| check_for_refs!(item) }
+        end
+      end
     end
   end
 end
diff --git a/test/mcp/tool/input_schema_test.rb b/test/mcp/tool/input_schema_test.rb
index 75d1446..791f955 100644
--- a/test/mcp/tool/input_schema_test.rb
+++ b/test/mcp/tool/input_schema_test.rb
@@ -61,6 +61,18 @@ module MCP
           end
         end
       end
+
+      test "rejects schemas with $ref references" do
+        assert_raises(ArgumentError) do
+          InputSchema.new(properties: { foo: { "$ref" => "#/definitions/bar" } }, required: [:foo])
+        end
+      end
+
+      test "rejects schemas with symbol $ref references" do
+        assert_raises(ArgumentError) do
+          InputSchema.new(properties: { foo: { :$ref => "#/definitions/bar" } }, required: [:foo])
+        end
+      end
     end
   end
 end
diff --git a/test/mcp/tool_test.rb b/test/mcp/tool_test.rb
index f9551f1..e87c826 100644
--- a/test/mcp/tool_test.rb
+++ b/test/mcp/tool_test.rb
@@ -249,7 +249,7 @@ module MCP
         properties: {
           foo: { "$ref" => "#/definitions/bar" },
         },
-        required: [],
+        required: ["foo"],
         definitions: {
           bar: { type: "string" },
         },