Merge remote-tracking branch 'upstream/main' into make_server_context_optional_and_sorbet_friendly

test/mcp/configuration_test.rb

@@ -61,5 +61,45 @@ module MCP
       merged = config3.merge(config1)
       assert_equal "2025-03-27", merged.protocol_version
     end
+
+    test "defaults validate_tool_call_arguments to true" do
+      config = Configuration.new
+      assert config.validate_tool_call_arguments
+    end
+
+    test "can set validate_tool_call_arguments to false" do
+      config = Configuration.new(validate_tool_call_arguments: false)
+      refute config.validate_tool_call_arguments
+    end
+
+    test "validate_tool_call_arguments? returns false when set" do
+      config = Configuration.new(validate_tool_call_arguments: false)
+      refute config.validate_tool_call_arguments?
+    end
+
+    test "validate_tool_call_arguments? returns true when not set" do
+      config = Configuration.new
+      assert config.validate_tool_call_arguments?
+    end
+
+    test "merge preserves validate_tool_call_arguments from other config" do
+      config1 = Configuration.new(validate_tool_call_arguments: false)
+      config2 = Configuration.new
+      merged = config1.merge(config2)
+      assert merged.validate_tool_call_arguments?
+    end
+
+    test "merge preserves validate_tool_call_arguments from self when other not set" do
+      config1 = Configuration.new(validate_tool_call_arguments: false)
+      config2 = Configuration.new
+      merged = config2.merge(config1)
+      refute merged.validate_tool_call_arguments
+    end
+
+    test "raises ArgumentError when validate_tool_call_arguments is not a boolean" do
+      assert_raises(ArgumentError) do
+        Configuration.new(validate_tool_call_arguments: "true")
+      end
+    end
   end
 end

test/mcp/server/transports/stdio_notification_integration_test.rb

@@ -0,0 +1,249 @@
+# frozen_string_literal: true
+
+require "test_helper"
+require "json"
+require "mcp/server"
+require "mcp/server/transports/stdio_transport"
+
+module MCP
+  class Server
+    module Transports
+      class StdioNotificationIntegrationTest < ActiveSupport::TestCase
+        class MockIO
+          attr_reader :output
+
+          def initialize
+            @output = []
+            @closed = false
+          end
+
+          def puts(message)
+            @output << message
+          end
+
+          def write(message)
+            @output << message
+            message.length
+          end
+
+          def gets
+            nil # Simulate end of input
+          end
+
+          def set_encoding(encoding) # rubocop:disable Naming/AccessorMethodName
+            # Mock implementation
+          end
+
+          def flush
+            # Mock implementation
+          end
+
+          def close
+            @closed = true
+          end
+
+          def closed?
+            @closed
+          end
+        end
+
+        setup do
+          @original_stdout = $stdout
+          @original_stdin = $stdin
+
+          @mock_stdout = MockIO.new
+          @mock_stdin = MockIO.new
+
+          $stdout = @mock_stdout
+          $stdin = @mock_stdin
+
+          @server = Server.new(
+            name: "test_server",
+            tools: [],
+            prompts: [],
+            resources: [],
+          )
+          @transport = StdioTransport.new(@server)
+          @server.transport = @transport
+        end
+
+        teardown do
+          $stdout = @original_stdout
+          $stdin = @original_stdin
+        end
+
+        test "server notification methods send JSON-RPC notifications through StdioTransport" do
+          # Test tools notification
+          @server.notify_tools_list_changed
+
+          # Test prompts notification
+          @server.notify_prompts_list_changed
+
+          # Test resources notification
+          @server.notify_resources_list_changed
+
+          # Check the notifications were sent
+          assert_equal 3, @mock_stdout.output.size
+
+          # Parse and verify each notification
+          notifications = @mock_stdout.output.map { |msg| JSON.parse(msg) }
+
+          assert_equal "2.0", notifications[0]["jsonrpc"]
+          assert_equal Methods::NOTIFICATIONS_TOOLS_LIST_CHANGED, notifications[0]["method"]
+          assert_nil notifications[0]["params"]
+
+          assert_equal "2.0", notifications[1]["jsonrpc"]
+          assert_equal Methods::NOTIFICATIONS_PROMPTS_LIST_CHANGED, notifications[1]["method"]
+          assert_nil notifications[1]["params"]
+
+          assert_equal "2.0", notifications[2]["jsonrpc"]
+          assert_equal Methods::NOTIFICATIONS_RESOURCES_LIST_CHANGED, notifications[2]["method"]
+          assert_nil notifications[2]["params"]
+        end
+
+        test "notifications include params when provided" do
+          # Test the transport's send_notification directly with params
+          result = @transport.send_notification("test/notification", { data: "test_value" })
+
+          assert result
+          assert_equal 1, @mock_stdout.output.size
+
+          notification = JSON.parse(@mock_stdout.output.first)
+          assert_equal "2.0", notification["jsonrpc"]
+          assert_equal "test/notification", notification["method"]
+          assert_equal({ "data" => "test_value" }, notification["params"])
+        end
+
+        test "server continues to work when stdout is closed" do
+          # Close stdout
+          @mock_stdout.close
+
+          # Server notifications should not raise errors
+          assert_nothing_raised do
+            @server.notify_tools_list_changed
+            @server.notify_prompts_list_changed
+            @server.notify_resources_list_changed
+          end
+        end
+
+        test "notifications work with dynamic tool additions" do
+          # Define a new tool
+          @server.define_tool(
+            name: "dynamic_tool",
+            description: "A dynamically added tool",
+          ) do |**_args|
+            { result: "success" }
+          end
+
+          # Clear previous output
+          @mock_stdout.output.clear
+
+          # Manually trigger notification
+          @server.notify_tools_list_changed
+
+          # Check the notification was sent
+          assert_equal 1, @mock_stdout.output.size
+
+          notification = JSON.parse(@mock_stdout.output.first)
+          assert_equal Methods::NOTIFICATIONS_TOOLS_LIST_CHANGED, notification["method"]
+
+          # Verify the tool was added to the server
+          assert @server.tools.key?("dynamic_tool")
+        end
+
+        test "notifications are properly formatted JSON-RPC 2.0 messages" do
+          # Send a notification
+          @server.notify_prompts_list_changed
+
+          # Verify format
+          assert_equal 1, @mock_stdout.output.size
+          output = @mock_stdout.output.first
+
+          # Should be valid JSON
+          notification = JSON.parse(output)
+
+          # Should have required JSON-RPC 2.0 fields
+          assert_equal "2.0", notification["jsonrpc"]
+          assert notification.key?("method")
+          refute notification.key?("id") # Notifications should not have an id
+
+          # Method should be the expected notification type
+          assert_equal Methods::NOTIFICATIONS_PROMPTS_LIST_CHANGED, notification["method"]
+        end
+
+        test "multiple notifications are sent as separate JSON messages" do
+          # Send multiple notifications rapidly
+          5.times do
+            @server.notify_tools_list_changed
+          end
+
+          # Each should be a separate JSON message
+          assert_equal 5, @mock_stdout.output.size
+
+          # All should be parseable as JSON
+          @mock_stdout.output.each do |msg|
+            notification = JSON.parse(msg)
+            assert_equal "2.0", notification["jsonrpc"]
+            assert_equal Methods::NOTIFICATIONS_TOOLS_LIST_CHANGED, notification["method"]
+          end
+        end
+
+        test "transport handles errors gracefully" do
+          # Create a stdout that raises errors
+          error_stdout = Class.new(MockIO) do
+            def puts(message)
+              raise IOError, "Simulated IO error"
+            end
+          end.new
+
+          $stdout = error_stdout
+
+          # Notification should return false but not raise
+          result = @transport.send_notification("test/notification")
+          refute result
+        end
+
+        test "server notification flow works end-to-end with StdioTransport" do
+          # This test verifies the complete integration from server to transport
+
+          # Start with no output
+          assert_empty @mock_stdout.output
+
+          # Add a prompt and notify
+          @server.define_prompt(
+            name: "test_prompt",
+            description: "Test prompt",
+          ) do |_args, _server_context:|
+            MCP::PromptResponse.new(messages: [{ role: "user", content: "Test" }])
+          end
+
+          # Manually trigger notification
+          @server.notify_prompts_list_changed
+
+          # Verify notification was sent
+          assert_equal 1, @mock_stdout.output.size
+          notification = JSON.parse(@mock_stdout.output.first)
+          assert_equal Methods::NOTIFICATIONS_PROMPTS_LIST_CHANGED, notification["method"]
+
+          # Add a resource and notify
+          @server.resources = [
+            MCP::Resource.new(
+              uri: "test://resource",
+              name: "Test Resource",
+              description: "A test resource",
+              mime_type: "text/plain",
+            ),
+          ]
+
+          # Manually trigger notification
+          @server.notify_resources_list_changed
+
+          # Verify both notifications were sent
+          assert_equal 2, @mock_stdout.output.size
+          second_notification = JSON.parse(@mock_stdout.output.last)
+          assert_equal Methods::NOTIFICATIONS_RESOURCES_LIST_CHANGED, second_notification["method"]
+        end
+      end
+    end
+  end
+end

test/mcp/server/transports/streamable_http_notification_integration_test.rb

@@ -0,0 +1,249 @@
+# frozen_string_literal: true
+
+require "test_helper"
+require "json"
+require "rack"
+require "mcp/server"
+require "mcp/server/transports/streamable_http_transport"
+
+module MCP
+  class Server
+    module Transports
+      class StreamableHTTPNotificationIntegrationTest < ActiveSupport::TestCase
+        setup do
+          @server = Server.new(
+            name: "test_server",
+            tools: [],
+            prompts: [],
+            resources: [],
+          )
+          @transport = StreamableHTTPTransport.new(@server)
+          @server.transport = @transport
+        end
+
+        test "server notification methods send SSE notifications through HTTP transport" do
+          # Initialize a session
+          init_request = create_rack_request(
+            "POST",
+            "/",
+            { "CONTENT_TYPE" => "application/json" },
+            { jsonrpc: "2.0", method: "initialize", id: "init" }.to_json,
+          )
+          init_response = @transport.handle_request(init_request)
+          session_id = init_response[1]["Mcp-Session-Id"]
+
+          # Connect with SSE
+          io = StringIO.new
+          get_request = create_rack_request(
+            "GET",
+            "/",
+            { "HTTP_MCP_SESSION_ID" => session_id },
+          )
+          response = @transport.handle_request(get_request)
+          response[2].call(io) if response[2].is_a?(Proc)
+
+          # Give the stream time to set up
+          sleep(0.1)
+
+          # Test tools notification
+          @server.notify_tools_list_changed
+
+          # Test prompts notification
+          @server.notify_prompts_list_changed
+
+          # Test resources notification
+          @server.notify_resources_list_changed
+
+          # Check the notifications were received
+          io.rewind
+          output = io.read
+
+          assert_includes output, "data: #{Methods::NOTIFICATIONS_TOOLS_LIST_CHANGED}"
+          assert_includes output, "data: #{Methods::NOTIFICATIONS_PROMPTS_LIST_CHANGED}"
+          assert_includes output, "data: #{Methods::NOTIFICATIONS_RESOURCES_LIST_CHANGED}"
+        end
+
+        test "notifications are broadcast to all connected sessions" do
+          # Create two sessions
+          init_request1 = create_rack_request(
+            "POST",
+            "/",
+            { "CONTENT_TYPE" => "application/json" },
+            { jsonrpc: "2.0", method: "initialize", id: "123" }.to_json,
+          )
+          init_response1 = @transport.handle_request(init_request1)
+          session_id1 = init_response1[1]["Mcp-Session-Id"]
+
+          init_request2 = create_rack_request(
+            "POST",
+            "/",
+            { "CONTENT_TYPE" => "application/json" },
+            { jsonrpc: "2.0", method: "initialize", id: "456" }.to_json,
+          )
+          init_response2 = @transport.handle_request(init_request2)
+          session_id2 = init_response2[1]["Mcp-Session-Id"]
+
+          # Connect both sessions with SSE
+          io1 = StringIO.new
+          get_request1 = create_rack_request(
+            "GET",
+            "/",
+            { "HTTP_MCP_SESSION_ID" => session_id1 },
+          )
+          response1 = @transport.handle_request(get_request1)
+          response1[2].call(io1) if response1[2].is_a?(Proc)
+
+          io2 = StringIO.new
+          get_request2 = create_rack_request(
+            "GET",
+            "/",
+            { "HTTP_MCP_SESSION_ID" => session_id2 },
+          )
+          response2 = @transport.handle_request(get_request2)
+          response2[2].call(io2) if response2[2].is_a?(Proc)
+
+          # Give the streams time to set up
+          sleep(0.1)
+
+          # Send notification through server
+          @server.notify_tools_list_changed
+
+          # Check both sessions received the notification
+          io1.rewind
+          output1 = io1.read
+          assert_includes output1, "data: #{Methods::NOTIFICATIONS_TOOLS_LIST_CHANGED}"
+
+          io2.rewind
+          output2 = io2.read
+          assert_includes output2, "data: #{Methods::NOTIFICATIONS_TOOLS_LIST_CHANGED}"
+        end
+
+        test "server continues to work when SSE connection is closed" do
+          # Initialize a session
+          init_request = create_rack_request(
+            "POST",
+            "/",
+            { "CONTENT_TYPE" => "application/json" },
+            { jsonrpc: "2.0", method: "initialize", id: "init" }.to_json,
+          )
+          init_response = @transport.handle_request(init_request)
+          session_id = init_response[1]["Mcp-Session-Id"]
+
+          # Connect with SSE
+          io = StringIO.new
+          get_request = create_rack_request(
+            "GET",
+            "/",
+            { "HTTP_MCP_SESSION_ID" => session_id },
+          )
+          response = @transport.handle_request(get_request)
+          response[2].call(io) if response[2].is_a?(Proc)
+
+          # Give the stream time to set up
+          sleep(0.1)
+
+          # Close the stream
+          io.close
+
+          # Server notifications should not raise errors
+          assert_nothing_raised do
+            @server.notify_tools_list_changed
+            @server.notify_prompts_list_changed
+            @server.notify_resources_list_changed
+          end
+        end
+
+        test "notifications work with dynamic tool additions" do
+          # Initialize a session
+          init_request = create_rack_request(
+            "POST",
+            "/",
+            { "CONTENT_TYPE" => "application/json" },
+            { jsonrpc: "2.0", method: "initialize", id: "init" }.to_json,
+          )
+          init_response = @transport.handle_request(init_request)
+          session_id = init_response[1]["Mcp-Session-Id"]
+
+          # Connect with SSE
+          io = StringIO.new
+          get_request = create_rack_request(
+            "GET",
+            "/",
+            { "HTTP_MCP_SESSION_ID" => session_id },
+          )
+          response = @transport.handle_request(get_request)
+          response[2].call(io) if response[2].is_a?(Proc)
+
+          # Give the stream time to set up
+          sleep(0.1)
+
+          # Define a new tool (simulating dynamic tool addition)
+          @server.define_tool(
+            name: "dynamic_tool",
+            description: "A dynamically added tool",
+          ) do |**_args|
+            { result: "success" }
+          end
+
+          # Manually trigger notification (since we removed the automatic triggers)
+          @server.notify_tools_list_changed
+
+          # Check the notification was received
+          io.rewind
+          output = io.read
+          assert_includes output, "data: #{Methods::NOTIFICATIONS_TOOLS_LIST_CHANGED}"
+
+          # Verify the tool was added to the server
+          assert @server.tools.key?("dynamic_tool")
+        end
+
+        test "SSE format is correct for notifications" do
+          # Initialize a session
+          init_request = create_rack_request(
+            "POST",
+            "/",
+            { "CONTENT_TYPE" => "application/json" },
+            { jsonrpc: "2.0", method: "initialize", id: "init" }.to_json,
+          )
+          init_response = @transport.handle_request(init_request)
+          session_id = init_response[1]["Mcp-Session-Id"]
+
+          # Connect with SSE
+          io = StringIO.new
+          get_request = create_rack_request(
+            "GET",
+            "/",
+            { "HTTP_MCP_SESSION_ID" => session_id },
+          )
+          response = @transport.handle_request(get_request)
+          response[2].call(io) if response[2].is_a?(Proc)
+
+          # Give the stream time to set up
+          sleep(0.1)
+
+          # Send a notification
+          @server.notify_tools_list_changed
+
+          # Check SSE format
+          io.rewind
+          output = io.read
+
+          # SSE format should be "data: <message>\n\n"
+          assert_match(/data: #{Regexp.escape(Methods::NOTIFICATIONS_TOOLS_LIST_CHANGED)}\n/, output)
+        end
+
+        private
+
+        def create_rack_request(method, path, headers, body = nil)
+          env = {
+            "REQUEST_METHOD" => method,
+            "PATH_INFO" => path,
+            "rack.input" => StringIO.new(body.to_s),
+          }.merge(headers)
+
+          Rack::Request.new(env)
+        end
+      end
+    end
+  end
+end

test/mcp/server/transports/streamable_http_transport_test.rb

@@ -0,0 +1,583 @@
+# frozen_string_literal: true
+
+require "test_helper"
+require "json"
+require "rack"
+require "mcp/server"
+require "mcp/server/transports/streamable_http_transport"
+
+module MCP
+  class Server
+    module Transports
+      class StreamableHTTPTransportTest < ActiveSupport::TestCase
+        setup do
+          @server = Server.new(
+            name: "test_server",
+            tools: [],
+            prompts: [],
+            resources: [],
+          )
+          @transport = StreamableHTTPTransport.new(@server)
+        end
+
+        test "handles POST request with valid JSON-RPC message" do
+          # First create a session
+          init_request = create_rack_request(
+            "POST",
+            "/",
+            { "CONTENT_TYPE" => "application/json" },
+            { jsonrpc: "2.0", method: "initialize", id: "init" }.to_json,
+          )
+          init_response = @transport.handle_request(init_request)
+          session_id = init_response[1]["Mcp-Session-Id"]
+
+          # Now make the ping request with the session ID
+          request = create_rack_request(
+            "POST",
+            "/",
+            {
+              "CONTENT_TYPE" => "application/json",
+              "HTTP_MCP_SESSION_ID" => session_id,
+            },
+            { jsonrpc: "2.0", method: "ping", id: "123" }.to_json,
+          )
+
+          response = @transport.handle_request(request)
+          assert_equal 200, response[0]
+          assert_equal({ "Content-Type" => "application/json" }, response[1])
+
+          body = JSON.parse(response[2][0])
+          assert_equal "2.0", body["jsonrpc"]
+          assert_equal "123", body["id"]
+          assert_empty(body["result"])
+        end
+
+        test "handles POST request with invalid JSON" do
+          request = create_rack_request(
+            "POST",
+            "/",
+            { "CONTENT_TYPE" => "application/json" },
+            "invalid json",
+          )
+
+          response = @transport.handle_request(request)
+          assert_equal 400, response[0]
+          assert_equal({ "Content-Type" => "application/json" }, response[1])
+
+          body = JSON.parse(response[2][0])
+          assert_equal "Invalid JSON", body["error"]
+        end
+
+        test "handles POST request with initialize method" do
+          request = create_rack_request(
+            "POST",
+            "/",
+            { "CONTENT_TYPE" => "application/json" },
+            { jsonrpc: "2.0", method: "initialize", id: "123" }.to_json,
+          )
+
+          response = @transport.handle_request(request)
+          assert_equal 200, response[0]
+          assert_equal "application/json", response[1]["Content-Type"]
+          assert response[1]["Mcp-Session-Id"]
+
+          body = JSON.parse(response[2][0])
+          assert_equal "2.0", body["jsonrpc"]
+          assert_equal "123", body["id"]
+          assert_equal "2024-11-05", body["result"]["protocolVersion"]
+        end
+
+        test "handles GET request with valid session ID" do
+          # First create a session with initialize
+          init_request = create_rack_request(
+            "POST",
+            "/",
+            { "CONTENT_TYPE" => "application/json" },
+            { jsonrpc: "2.0", method: "initialize", id: "123" }.to_json,
+          )
+          init_response = @transport.handle_request(init_request)
+          session_id = init_response[1]["Mcp-Session-Id"]
+
+          # Then try to connect with GET
+          request = create_rack_request(
+            "GET",
+            "/",
+            {
+              "HTTP_MCP_SESSION_ID" => session_id,
+            },
+          )
+
+          response = @transport.handle_request(request)
+          assert_equal 200, response[0]
+          assert_equal "text/event-stream", response[1]["Content-Type"]
+          assert response[2].is_a?(Proc) # The body should be a Proc for streaming
+        end
+
+        test "handles POST request when IOError raised" do
+          # Create and initialize a session
+          init_request = create_rack_request(
+            "POST",
+            "/",
+            { "CONTENT_TYPE" => "application/json" },
+            { jsonrpc: "2.0", method: "initialize", id: "123" }.to_json,
+          )
+          init_response = @transport.handle_request(init_request)
+          session_id = init_response[1]["Mcp-Session-Id"]
+
+          # Connect with SSE
+          io = StringIO.new
+          get_request = create_rack_request(
+            "GET",
+            "/",
+            { "HTTP_MCP_SESSION_ID" => session_id },
+          )
+          response = @transport.handle_request(get_request)
+          response[2].call(io) if response[2].is_a?(Proc)
+
+          # Give the stream time to set up
+          sleep(0.1)
+
+          # Close the stream
+          io.close
+
+          request = create_rack_request(
+            "POST",
+            "/",
+            {
+              "CONTENT_TYPE" => "application/json",
+              "HTTP_MCP_SESSION_ID" => session_id,
+            },
+            { jsonrpc: "2.0", method: "ping", id: "456" }.to_json,
+          )
+
+          # This should handle IOError and return the original response
+          response = @transport.handle_request(request)
+          assert_equal 200, response[0]
+          assert_equal({ "Content-Type" => "application/json" }, response[1])
+
+          # Verify session was cleaned up
+          assert_not @transport.instance_variable_get(:@sessions).key?(session_id)
+        end
+
+        test "handles POST request when Errno::EPIPE raised" do
+          # Create and initialize a session
+          init_request = create_rack_request(
+            "POST",
+            "/",
+            { "CONTENT_TYPE" => "application/json" },
+            { jsonrpc: "2.0", method: "initialize", id: "123" }.to_json,
+          )
+          init_response = @transport.handle_request(init_request)
+          session_id = init_response[1]["Mcp-Session-Id"]
+
+          # Create a pipe to simulate EPIPE condition
+          reader, writer = IO.pipe
+
+          # Connect with SSE using the writer end of the pipe
+          get_request = create_rack_request(
+            "GET",
+            "/",
+            { "HTTP_MCP_SESSION_ID" => session_id },
+          )
+          response = @transport.handle_request(get_request)
+          response[2].call(writer) if response[2].is_a?(Proc)
+
+          # Give the stream time to set up
+          sleep(0.1)
+
+          # Close the reader end to break the pipe - this will cause EPIPE on write
+          reader.close
+
+          request = create_rack_request(
+            "POST",
+            "/",
+            {
+              "CONTENT_TYPE" => "application/json",
+              "HTTP_MCP_SESSION_ID" => session_id,
+            },
+            { jsonrpc: "2.0", method: "ping", id: "789" }.to_json,
+          )
+
+          # This should handle Errno::EPIPE and return the original response
+          response = @transport.handle_request(request)
+          assert_equal 200, response[0]
+          assert_equal({ "Content-Type" => "application/json" }, response[1])
+
+          # Verify session was cleaned up
+          assert_not @transport.instance_variable_get(:@sessions).key?(session_id)
+
+          begin
+            writer.close
+          rescue
+            nil
+          end
+        end
+
+        test "handles GET request with missing session ID" do
+          request = create_rack_request(
+            "GET",
+            "/",
+            {},
+          )
+
+          response = @transport.handle_request(request)
+          assert_equal 400, response[0]
+          assert_equal({ "Content-Type" => "application/json" }, response[1])
+
+          body = JSON.parse(response[2][0])
+          assert_equal "Missing session ID", body["error"]
+        end
+
+        test "handles GET request with invalid session ID" do
+          request = create_rack_request(
+            "GET",
+            "/",
+            { "HTTP_MCP_SESSION_ID" => "invalid_id" },
+          )
+
+          response = @transport.handle_request(request)
+          assert_equal 404, response[0]
+          assert_equal({ "Content-Type" => "application/json" }, response[1])
+
+          body = JSON.parse(response[2][0])
+          assert_equal "Session not found", body["error"]
+        end
+
+        test "handles DELETE request with valid session ID" do
+          # First create a session with initialize
+          init_request = create_rack_request(
+            "POST",
+            "/",
+            { "CONTENT_TYPE" => "application/json" },
+            { jsonrpc: "2.0", method: "initialize", id: "123" }.to_json,
+          )
+          init_response = @transport.handle_request(init_request)
+          session_id = init_response[1]["Mcp-Session-Id"]
+
+          # Then try to delete it
+          request = create_rack_request(
+            "DELETE",
+            "/",
+            { "HTTP_MCP_SESSION_ID" => session_id },
+          )
+
+          response = @transport.handle_request(request)
+          assert_equal 200, response[0]
+          assert_equal({ "Content-Type" => "application/json" }, response[1])
+
+          body = JSON.parse(response[2][0])
+          assert body["success"]
+        end
+
+        test "handles DELETE request with missing session ID" do
+          request = create_rack_request(
+            "DELETE",
+            "/",
+            {},
+          )
+
+          response = @transport.handle_request(request)
+          assert_equal 400, response[0]
+          assert_equal({ "Content-Type" => "application/json" }, response[1])
+
+          body = JSON.parse(response[2][0])
+          assert_equal "Missing session ID", body["error"]
+        end
+
+        test "closes transport and cleans up session" do
+          # First create a session with initialize
+          init_request = create_rack_request(
+            "POST",
+            "/",
+            { "CONTENT_TYPE" => "application/json" },
+            { jsonrpc: "2.0", method: "initialize", id: "123" }.to_json,
+          )
+          init_response = @transport.handle_request(init_request)
+          session_id = init_response[1]["Mcp-Session-Id"]
+
+          # Then connect with GET
+          io = StringIO.new
+          request = create_rack_request(
+            "GET",
+            "/",
+            {
+              "HTTP_MCP_SESSION_ID" => session_id,
+            },
+          )
+          response = @transport.handle_request(request)
+          # Call the body proc with our StringIO
+          response[2].call(io) if response[2].is_a?(Proc)
+
+          # Give the background thread a moment to set up
+          sleep(0.01)
+
+          # Verify session exists before closing
+          assert @transport.instance_variable_get(:@sessions).key?(session_id)
+
+          # Close the transport without session context (closes all sessions)
+          @transport.close
+
+          # Verify session was cleaned up
+          assert_empty @transport.instance_variable_get(:@sessions)
+        end
+
+        test "sends notification to correct session with multiple active sessions" do
+          # Create first session
+          init_request1 = create_rack_request(
+            "POST",
+            "/",
+            { "CONTENT_TYPE" => "application/json" },
+            { jsonrpc: "2.0", method: "initialize", id: "123" }.to_json,
+          )
+          init_response1 = @transport.handle_request(init_request1)
+          session_id1 = init_response1[1]["Mcp-Session-Id"]
+
+          # Create second session
+          init_request2 = create_rack_request(
+            "POST",
+            "/",
+            { "CONTENT_TYPE" => "application/json" },
+            { jsonrpc: "2.0", method: "initialize", id: "456" }.to_json,
+          )
+          init_response2 = @transport.handle_request(init_request2)
+          session_id2 = init_response2[1]["Mcp-Session-Id"]
+
+          # Connect first session with GET
+          io1 = StringIO.new
+          get_request1 = create_rack_request(
+            "GET",
+            "/",
+            {
+              "HTTP_MCP_SESSION_ID" => session_id1,
+            },
+          )
+          response1 = @transport.handle_request(get_request1)
+          response1[2].call(io1) if response1[2].is_a?(Proc)
+
+          # Connect second session with GET
+          io2 = StringIO.new
+          get_request2 = create_rack_request(
+            "GET",
+            "/",
+            {
+              "HTTP_MCP_SESSION_ID" => session_id2,
+            },
+          )
+          response2 = @transport.handle_request(get_request2)
+          response2[2].call(io2) if response2[2].is_a?(Proc)
+
+          # Give the streams time to be fully set up
+          sleep(0.2)
+
+          # Verify sessions are set up
+          assert @transport.instance_variable_get(:@sessions).key?(session_id1), "Session 1 not found in @sessions"
+          assert @transport.instance_variable_get(:@sessions).key?(session_id2), "Session 2 not found in @sessions"
+
+          # Test that notifications go to the correct session based on the request context
+          # First, make a request as session 1
+          request_as_session1 = create_rack_request(
+            "POST",
+            "/",
+            {
+              "CONTENT_TYPE" => "application/json",
+              "HTTP_MCP_SESSION_ID" => session_id1,
+            },
+            { jsonrpc: "2.0", method: "ping", id: "789" }.to_json,
+          )
+
+          # Monkey-patch handle_json on the server to send a notification when called
+          original_handle_json = @server.method(:handle_json)
+          transport = @transport # Capture the transport in a local variable
+          @server.define_singleton_method(:handle_json) do |request|
+            result = original_handle_json.call(request)
+            # Send notification while still in request context - broadcast to all sessions
+            notification = { jsonrpc: "2.0", method: "test_notification", params: { session: "current" } }
+            transport.send_notification(notification)
+            result
+          end
+
+          # Handle request from session 1
+          @transport.handle_request(request_as_session1)
+
+          # Make a request as session 2
+          request_as_session2 = create_rack_request(
+            "POST",
+            "/",
+            {
+              "CONTENT_TYPE" => "application/json",
+              "HTTP_MCP_SESSION_ID" => session_id2,
+            },
+            { jsonrpc: "2.0", method: "ping", id: "890" }.to_json,
+          )
+
+          # Handle request from session 2
+          @transport.handle_request(request_as_session2)
+
+          # Check that each session received one notification
+          io1.rewind
+          output1 = io1.read
+          # Session 1 should have received two notifications (one from each request since we broadcast)
+          assert_equal 2,
+            output1.scan(/data: {"jsonrpc":"2.0","method":"test_notification","params":{"session":"current"}}/).count
+
+          io2.rewind
+          output2 = io2.read
+          # Session 2 should have received two notifications (one from each request since we broadcast)
+          assert_equal 2,
+            output2.scan(/data: {"jsonrpc":"2.0","method":"test_notification","params":{"session":"current"}}/).count
+        end
+
+        test "send_notification to specific session" do
+          # Create and initialize a session
+          init_request = create_rack_request(
+            "POST",
+            "/",
+            { "CONTENT_TYPE" => "application/json" },
+            { jsonrpc: "2.0", method: "initialize", id: "123" }.to_json,
+          )
+          init_response = @transport.handle_request(init_request)
+          session_id = init_response[1]["Mcp-Session-Id"]
+
+          # Connect with SSE
+          io = StringIO.new
+          get_request = create_rack_request(
+            "GET",
+            "/",
+            { "HTTP_MCP_SESSION_ID" => session_id },
+          )
+          response = @transport.handle_request(get_request)
+          response[2].call(io) if response[2].is_a?(Proc)
+
+          # Give the stream time to set up
+          sleep(0.1)
+
+          # Send notification to specific session
+          notification = { jsonrpc: "2.0", method: "test_notification", params: { message: "Hello" } }
+          result = @transport.send_notification(notification, session_id: session_id)
+
+          assert result
+
+          # Check the notification was received
+          io.rewind
+          output = io.read
+          assert_includes output,
+            "data: {\"jsonrpc\":\"2.0\",\"method\":\"test_notification\",\"params\":{\"message\":\"Hello\"}}"
+        end
+
+        test "send_notification broadcasts to all sessions when no session_id" do
+          # Create two sessions
+          init_request1 = create_rack_request(
+            "POST",
+            "/",
+            { "CONTENT_TYPE" => "application/json" },
+            { jsonrpc: "2.0", method: "initialize", id: "123" }.to_json,
+          )
+          init_response1 = @transport.handle_request(init_request1)
+          session_id1 = init_response1[1]["Mcp-Session-Id"]
+
+          init_request2 = create_rack_request(
+            "POST",
+            "/",
+            { "CONTENT_TYPE" => "application/json" },
+            { jsonrpc: "2.0", method: "initialize", id: "456" }.to_json,
+          )
+          init_response2 = @transport.handle_request(init_request2)
+          session_id2 = init_response2[1]["Mcp-Session-Id"]
+
+          # Connect both sessions with SSE
+          io1 = StringIO.new
+          get_request1 = create_rack_request(
+            "GET",
+            "/",
+            { "HTTP_MCP_SESSION_ID" => session_id1 },
+          )
+          response1 = @transport.handle_request(get_request1)
+          response1[2].call(io1) if response1[2].is_a?(Proc)
+
+          io2 = StringIO.new
+          get_request2 = create_rack_request(
+            "GET",
+            "/",
+            { "HTTP_MCP_SESSION_ID" => session_id2 },
+          )
+          response2 = @transport.handle_request(get_request2)
+          response2[2].call(io2) if response2[2].is_a?(Proc)
+
+          # Give the streams time to set up
+          sleep(0.1)
+
+          # Broadcast notification to all sessions
+          notification = { jsonrpc: "2.0", method: "broadcast", params: { message: "Hello everyone" } }
+          sent_count = @transport.send_notification(notification)
+
+          assert_equal 2, sent_count
+
+          # Check both sessions received the notification
+          io1.rewind
+          output1 = io1.read
+          assert_includes output1,
+            "data: {\"jsonrpc\":\"2.0\",\"method\":\"broadcast\",\"params\":{\"message\":\"Hello everyone\"}}"
+
+          io2.rewind
+          output2 = io2.read
+          assert_includes output2,
+            "data: {\"jsonrpc\":\"2.0\",\"method\":\"broadcast\",\"params\":{\"message\":\"Hello everyone\"}}"
+        end
+
+        test "send_notification returns false for non-existent session" do
+          result = @transport.send_notification({ message: "test" }, session_id: "non_existent")
+          refute result
+        end
+
+        test "send_notification handles closed streams gracefully" do
+          # Create and initialize a session
+          init_request = create_rack_request(
+            "POST",
+            "/",
+            { "CONTENT_TYPE" => "application/json" },
+            { jsonrpc: "2.0", method: "initialize", id: "123" }.to_json,
+          )
+          init_response = @transport.handle_request(init_request)
+          session_id = init_response[1]["Mcp-Session-Id"]
+
+          # Connect with SSE
+          io = StringIO.new
+          get_request = create_rack_request(
+            "GET",
+            "/",
+            { "HTTP_MCP_SESSION_ID" => session_id },
+          )
+          response = @transport.handle_request(get_request)
+          response[2].call(io) if response[2].is_a?(Proc)
+
+          # Give the stream time to set up
+          sleep(0.1)
+
+          # Close the stream
+          io.close
+
+          # Try to send notification
+          result = @transport.send_notification({ message: "test" }, session_id: session_id)
+
+          # Should return false and clean up the session
+          refute result
+
+          # Verify session was cleaned up
+          assert_not @transport.instance_variable_get(:@sessions).key?(session_id)
+        end
+
+        private
+
+        def create_rack_request(method, path, headers, body = nil)
+          env = {
+            "REQUEST_METHOD" => method,
+            "PATH_INFO" => path,
+            "rack.input" => StringIO.new(body.to_s),
+          }.merge(headers)
+
+          Rack::Request.new(env)
+        end
+      end
+    end
+  end
+end

test/mcp/server_notification_test.rb

@@ -0,0 +1,130 @@
+# typed: true
+# frozen_string_literal: true
+
+require "test_helper"
+
+module MCP
+  class ServerNotificationTest < ActiveSupport::TestCase
+    include InstrumentationTestHelper
+
+    class MockTransport < Transport
+      attr_reader :notifications
+
+      def initialize(server)
+        super
+        @notifications = []
+      end
+
+      def send_notification(method, params = nil)
+        @notifications << { method: method, params: params }
+        true
+      end
+
+      def send_response(response); end
+      def open; end
+      def close; end
+      def handle_request(request); end
+    end
+
+    setup do
+      configuration = MCP::Configuration.new
+      configuration.instrumentation_callback = instrumentation_helper.callback
+
+      @server = Server.new(
+        name: "test_server",
+        version: "1.0.0",
+        configuration: configuration,
+      )
+
+      @mock_transport = MockTransport.new(@server)
+      @server.transport = @mock_transport
+    end
+
+    test "#notify_tools_list_changed sends notification through transport" do
+      @server.notify_tools_list_changed
+
+      assert_equal 1, @mock_transport.notifications.size
+      notification = @mock_transport.notifications.first
+      assert_equal Methods::NOTIFICATIONS_TOOLS_LIST_CHANGED, notification[:method]
+      assert_nil notification[:params]
+    end
+
+    test "#notify_prompts_list_changed sends notification through transport" do
+      @server.notify_prompts_list_changed
+
+      assert_equal 1, @mock_transport.notifications.size
+      notification = @mock_transport.notifications.first
+      assert_equal Methods::NOTIFICATIONS_PROMPTS_LIST_CHANGED, notification[:method]
+      assert_nil notification[:params]
+    end
+
+    test "#notify_resources_list_changed sends notification through transport" do
+      @server.notify_resources_list_changed
+
+      assert_equal 1, @mock_transport.notifications.size
+      notification = @mock_transport.notifications.first
+      assert_equal Methods::NOTIFICATIONS_RESOURCES_LIST_CHANGED, notification[:method]
+      assert_nil notification[:params]
+    end
+
+    test "notification methods work without transport" do
+      server_without_transport = Server.new(name: "test_server")
+
+      # Should not raise any errors
+      assert_nothing_raised do
+        server_without_transport.notify_tools_list_changed
+        server_without_transport.notify_prompts_list_changed
+        server_without_transport.notify_resources_list_changed
+      end
+    end
+
+    test "notification methods handle transport errors gracefully" do
+      # Create a transport that raises errors
+      error_transport = Class.new(MockTransport) do
+        def send_notification(method, params = nil)
+          raise StandardError, "Transport error"
+        end
+      end.new(@server)
+
+      @server.transport = error_transport
+
+      # Mock the exception reporter
+      expected_contexts = [
+        { notification: "tools_list_changed" },
+        { notification: "prompts_list_changed" },
+        { notification: "resources_list_changed" },
+      ]
+
+      call_count = 0
+      @server.configuration.exception_reporter.expects(:call).times(3).with do |exception, context|
+        assert_kind_of StandardError, exception
+        assert_equal "Transport error", exception.message
+        assert_includes expected_contexts, context
+        call_count += 1
+        true
+      end
+
+      # Should not raise errors to the caller
+      assert_nothing_raised do
+        @server.notify_tools_list_changed
+        @server.notify_prompts_list_changed
+        @server.notify_resources_list_changed
+      end
+
+      assert_equal 3, call_count
+    end
+
+    test "multiple notification methods can be called in sequence" do
+      @server.notify_tools_list_changed
+      @server.notify_prompts_list_changed
+      @server.notify_resources_list_changed
+
+      assert_equal 3, @mock_transport.notifications.size
+
+      notifications = @mock_transport.notifications
+      assert_equal Methods::NOTIFICATIONS_TOOLS_LIST_CHANGED, notifications[0][:method]
+      assert_equal Methods::NOTIFICATIONS_PROMPTS_LIST_CHANGED, notifications[1][:method]
+      assert_equal Methods::NOTIFICATIONS_RESOURCES_LIST_CHANGED, notifications[2][:method]
+    end
+  end
+end

test/mcp/server_test.rb

@@ -839,5 +839,123 @@ module MCP
 
       refute_includes server_without_resources.capabilities, :resources
     end
+
+    test "tools/call validates arguments against input schema when validate_tool_call_arguments is true" do
+      server = Server.new(
+        tools: [TestTool],
+        configuration: Configuration.new(validate_tool_call_arguments: true),
+      )
+
+      response = server.handle(
+        {
+          jsonrpc: "2.0",
+          id: 1,
+          method: "tools/call",
+          params: {
+            name: "test_tool",
+            arguments: { message: 123 },
+          },
+        },
+      )
+
+      assert_equal "2.0", response[:jsonrpc]
+      assert_equal 1, response[:id]
+      assert_equal(-32603, response[:error][:code])
+      assert_includes response[:error][:data], "Invalid arguments"
+    end
+
+    test "tools/call skips argument validation when validate_tool_call_arguments is false" do
+      server = Server.new(
+        tools: [TestTool],
+        configuration: Configuration.new(validate_tool_call_arguments: false),
+      )
+
+      response = server.handle(
+        {
+          jsonrpc: "2.0",
+          id: 1,
+          method: "tools/call",
+          params: {
+            name: "test_tool",
+            arguments: { message: 123 },
+          },
+        },
+      )
+
+      assert_equal "2.0", response[:jsonrpc]
+      assert_equal 1, response[:id]
+      assert response[:result], "Expected result key in response"
+      assert_equal "text", response[:result][:content][0][:type]
+      assert_equal "OK", response[:result][:content][0][:content]
+    end
+
+    test "tools/call validates arguments with complex types" do
+      server = Server.new(
+        tools: [ComplexTypesTool],
+        configuration: Configuration.new(validate_tool_call_arguments: true),
+      )
+
+      response = server.handle(
+        {
+          jsonrpc: "2.0",
+          id: 1,
+          method: "tools/call",
+          params: {
+            name: "complex_types_tool",
+            arguments: {
+              numbers: [1, 2, 3],
+              strings: ["a", "b", "c"],
+              objects: [{ name: "test" }],
+            },
+          },
+        },
+      )
+
+      assert_equal "2.0", response[:jsonrpc]
+      assert_equal 1, response[:id]
+      assert response[:result], "Expected result key in response"
+      assert_equal "text", response[:result][:content][0][:type]
+      assert_equal "OK", response[:result][:content][0][:content]
+    end
+
+    class TestTool < Tool
+      tool_name "test_tool"
+      description "a test tool for testing"
+      input_schema({ properties: { message: { type: "string" } }, required: ["message"] })
+
+      class << self
+        def call(message:, server_context: nil)
+          Tool::Response.new([{ type: "text", content: "OK" }])
+        end
+      end
+    end
+
+    class ComplexTypesTool < Tool
+      tool_name "complex_types_tool"
+      description "a test tool with complex types"
+      input_schema({
+        properties: {
+          numbers: { type: "array", items: { type: "number" } },
+          strings: { type: "array", items: { type: "string" } },
+          objects: {
+            type: "array",
+            items: {
+              type: "object",
+              properties: {
+                name: { type: "string" },
+              },
+              required: ["name"],
+            },
+          },
+        },
+        required: ["numbers", "strings", "objects"],
+      })
+
+      class << self
+        def call(numbers:, strings:, objects:, server_context: nil)
+          Tool::Response.new([{ type: "text", content: "OK" }])
+        end
+      end
+    end
   end
 end

test/mcp/tool/input_schema_test.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require "test_helper"
+require "mcp/tool/input_schema"
 
 module MCP
   class Tool
@@ -27,6 +28,51 @@ module MCP
         input_schema = InputSchema.new(properties: { message: { type: "string" } }, required: [:message])
         assert_empty input_schema.missing_required_arguments({ message: "Hello, world!" })
       end
+
+      test "valid schema initialization" do
+        schema = InputSchema.new(properties: { foo: { type: "string" } }, required: [:foo])
+        assert_equal({ type: "object", properties: { foo: { type: "string" } }, required: [:foo] }, schema.to_h)
+      end
+
+      test "invalid schema raises argument error" do
+        assert_raises(ArgumentError) do
+          InputSchema.new(properties: { foo: { type: "invalid_type" } }, required: [:foo])
+        end
+      end
+
+      test "validate arguments with valid data" do
+        schema = InputSchema.new(properties: { foo: { type: "string" } }, required: [:foo])
+        assert_nil(schema.validate_arguments({ foo: "bar" }))
+      end
+
+      test "validate arguments with invalid data" do
+        schema = InputSchema.new(properties: { foo: { type: "string" } }, required: [:foo])
+        assert_raises(InputSchema::ValidationError) do
+          schema.validate_arguments({ foo: 123 })
+        end
+      end
+
+      test "unexpected errors bubble up from validate_arguments" do
+        schema = InputSchema.new(properties: { foo: { type: "string" } }, required: [:foo])
+
+        JSON::Validator.stub(:fully_validate, ->(*) { raise "unexpected error" }) do
+          assert_raises(RuntimeError) do
+            schema.validate_arguments({ foo: "bar" })
+          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

test/mcp/tool_test.rb

@@ -82,6 +82,23 @@ module MCP
       assert_equal expected, tool.input_schema.to_h
     end
 
+    test "raises detailed error message for invalid schema" do
+      error = assert_raises(ArgumentError) do
+        Class.new(MCP::Tool) do
+          input_schema(
+            properties: {
+              count: { type: "integer", minimum: "not a number" },
+            },
+            required: [:count],
+          )
+        end
+      end
+
+      assert_includes error.message, "Invalid JSON Schema"
+      assert_includes error.message, "#/properties/count/minimum"
+      assert_includes error.message, "string did not match the following type: number"
+    end
+
     test ".define allows definition of simple tools with a block" do
       tool = Tool.define(name: "mock_tool", description: "a mock tool for testing") do |_|
         Tool::Response.new([{ type: "text", content: "OK" }])
@@ -261,5 +278,23 @@ module MCP
       response = tool.call("test", server_context: { foo: "bar" })
       assert_equal response.content, [{ type: "text", content: "OK" }]
     end
+
+    test "input_schema rejects any $ref in schema" do
+      schema_with_ref = {
+        properties: {
+          foo: { "$ref" => "#/definitions/bar" },
+        },
+        required: ["foo"],
+        definitions: {
+          bar: { type: "string" },
+        },
+      }
+      error = assert_raises(ArgumentError) do
+        Class.new(MCP::Tool) do
+          input_schema schema_with_ref
+        end
+      end
+      assert_match(/Invalid JSON Schema/, error.message)
+    end
   end
 end

test/mcp/transport_test.rb

@@ -0,0 +1,50 @@
+# frozen_string_literal: true
+
+require "test_helper"
+require "mcp/server"
+require "mcp/transport"
+
+module MCP
+  class TransportTest < ActiveSupport::TestCase
+    class TestTransport < Transport
+      def handle_request(request)
+        [200, {}, ["OK"]]
+      end
+
+      def send_request(method, params = nil)
+        true
+      end
+
+      def close
+        true
+      end
+    end
+
+    setup do
+      @server = Server.new(
+        name: "test_server",
+        tools: [],
+        prompts: [],
+        resources: [],
+      )
+      @transport = TestTransport.new(@server)
+    end
+
+    test "initializes with server instance" do
+      assert_equal @server, @transport.instance_variable_get(:@server)
+    end
+
+    test "handles request" do
+      response = @transport.handle_request(nil)
+      assert_equal [200, {}, ["OK"]], response
+    end
+
+    test "sends request" do
+      assert @transport.send_request("test_method", { foo: "bar" })
+    end
+
+    test "closes connection" do
+      assert @transport.close
+    end
+  end
+end