From 8ceefacfac67fb8f57bccee72a3f0b6077430efd Mon Sep 17 00:00:00 2001
From: Shivam Mishra <scm.mymail@gmail.com>
Date: Fri, 3 Oct 2025 16:28:37 +0530
Subject: [PATCH] feat: add schema validation

---
 enterprise/app/models/captain/custom_tool.rb  | 18 ++++++
 .../models/captain/custom_tool_spec.rb        | 58 +++++++++++++++++++
 2 files changed, 76 insertions(+)

diff --git a/enterprise/app/models/captain/custom_tool.rb b/enterprise/app/models/captain/custom_tool.rb
index 3cb5c8726..6d39c4e1a 100644
--- a/enterprise/app/models/captain/custom_tool.rb
+++ b/enterprise/app/models/captain/custom_tool.rb
@@ -28,6 +28,21 @@ class Captain::CustomTool < ApplicationRecord
 
   self.table_name = 'captain_custom_tools'
 
+  PARAM_SCHEMA_VALIDATION = {
+    'type': 'array',
+    'items': {
+      'type': 'object',
+      'properties': {
+        'name': { 'type': 'string' },
+        'type': { 'type': 'string' },
+        'description': { 'type': 'string' },
+        'required': { 'type': 'boolean' }
+      },
+      'required': %w[name type description],
+      'additionalProperties': false
+    }
+  }.to_json.freeze
+
   belongs_to :account
 
   enum :http_method, %w[GET POST].index_by(&:itself), validate: true
@@ -38,6 +53,9 @@ class Captain::CustomTool < ApplicationRecord
   validates :slug, presence: true, uniqueness: { scope: :account_id }
   validates :title, presence: true
   validates :endpoint_url, presence: true
+  validates_with JsonSchemaValidator,
+                 schema: PARAM_SCHEMA_VALIDATION,
+                 attribute_resolver: ->(record) { record.param_schema }
 
   scope :enabled, -> { where(enabled: true) }
 
diff --git a/spec/enterprise/models/captain/custom_tool_spec.rb b/spec/enterprise/models/captain/custom_tool_spec.rb
index 23f7cbd2c..c7c0451b1 100644
--- a/spec/enterprise/models/captain/custom_tool_spec.rb
+++ b/spec/enterprise/models/captain/custom_tool_spec.rb
@@ -34,6 +34,64 @@ RSpec.describe Captain::CustomTool, type: :model do
         expect(different_account_tool).to be_valid
       end
     end
+
+    describe 'param_schema validation' do
+      let(:account) { create(:account) }
+
+      it 'is valid with proper param_schema' do
+        tool = build(:captain_custom_tool, account: account, param_schema: [
+                       { 'name' => 'order_id', 'type' => 'string', 'description' => 'Order ID', 'required' => true }
+                     ])
+
+        expect(tool).to be_valid
+      end
+
+      it 'is valid with empty param_schema' do
+        tool = build(:captain_custom_tool, account: account, param_schema: [])
+
+        expect(tool).to be_valid
+      end
+
+      it 'is invalid when param_schema is missing name' do
+        tool = build(:captain_custom_tool, account: account, param_schema: [
+                       { 'type' => 'string', 'description' => 'Order ID' }
+                     ])
+
+        expect(tool).not_to be_valid
+      end
+
+      it 'is invalid when param_schema is missing type' do
+        tool = build(:captain_custom_tool, account: account, param_schema: [
+                       { 'name' => 'order_id', 'description' => 'Order ID' }
+                     ])
+
+        expect(tool).not_to be_valid
+      end
+
+      it 'is invalid when param_schema is missing description' do
+        tool = build(:captain_custom_tool, account: account, param_schema: [
+                       { 'name' => 'order_id', 'type' => 'string' }
+                     ])
+
+        expect(tool).not_to be_valid
+      end
+
+      it 'is invalid with additional properties in param_schema' do
+        tool = build(:captain_custom_tool, account: account, param_schema: [
+                       { 'name' => 'order_id', 'type' => 'string', 'description' => 'Order ID', 'extra_field' => 'value' }
+                     ])
+
+        expect(tool).not_to be_valid
+      end
+
+      it 'is valid when required field is omitted (defaults to optional param)' do
+        tool = build(:captain_custom_tool, account: account, param_schema: [
+                       { 'name' => 'order_id', 'type' => 'string', 'description' => 'Order ID' }
+                     ])
+
+        expect(tool).to be_valid
+      end
+    end
   end
 
   describe 'scopes' do