diff --git a/lib/msf/core/mcp/server.rb b/lib/msf/core/mcp/server.rb index 864a3c45c9d9d..efebb875385c4 100644 --- a/lib/msf/core/mcp/server.rb +++ b/lib/msf/core/mcp/server.rb @@ -87,11 +87,13 @@ def start(transport: :stdio, host: 'localhost', port: 3000, auth_token: nil, min # Shutdown the MCP server and cleanup resources # def shutdown + @puma_server&.stop(true) @puma_launcher&.stop rescue StandardError => e elog({ message: 'Error stopping Puma', exception: e }, LOG_SOURCE, LOG_ERROR) ensure @puma_log_io&.close + @puma_server = nil @puma_launcher = nil @puma_log_io = nil @msf_client&.shutdown @@ -130,8 +132,7 @@ def start_stdio def start_http(host, port, auth_token, min_threads: PUMA_MIN_THREADS, max_threads: PUMA_MAX_THREADS, workers: PUMA_WORKERS) require 'rack' require 'puma' - require 'puma/configuration' - require 'puma/launcher' + require 'puma/server' require 'puma/log_writer' transport = ::MCP::Server::Transports::StreamableHTTPTransport.new(@mcp_server) @@ -148,24 +149,39 @@ def start_http(host, port, auth_token, min_threads: PUMA_MIN_THREADS, max_thread bind_host = host.include?(':') ? "[#{host}]" : host @puma_log_io = File.open(File::NULL, 'w') begin - puma_config = Puma::Configuration.new do |config| - config.bind "tcp://#{bind_host}:#{port}" - config.threads min_threads, max_threads - config.workers workers - config.log_requests false - config.app rack_app - end - - # Suppress Puma's startup banner by providing a silent log writer log_writer = Puma::LogWriter.new(@puma_log_io, @puma_log_io) - @puma_launcher = Puma::Launcher.new(puma_config, log_writer: log_writer) - @puma_launcher.run + + if workers.to_i > 0 + require 'puma/configuration' + require 'puma/launcher' + + puma_config = Puma::Configuration.new do |config| + config.bind "tcp://#{bind_host}:#{port}" + config.threads min_threads, max_threads + config.workers workers + config.log_requests false + config.app rack_app + end + + @puma_launcher = Puma::Launcher.new(puma_config, log_writer: log_writer) + @puma_launcher.run + else + @puma_server = Puma::Server.new(rack_app, nil, { + log_writer: log_writer, + min_threads: min_threads, + max_threads: max_threads + }) + @puma_server.add_tcp_listener(bind_host, port) + @puma_server.run.join + end rescue StandardError begin + @puma_server&.stop(true) @puma_launcher&.stop rescue StandardError nil end + @puma_server = nil @puma_launcher = nil @puma_log_io&.close @puma_log_io = nil