model-context-protocol-rb 0.6.0 → 0.7.0

data/tasks/templates/dev-http-puma.erb

@@ -0,0 +1,251 @@
+#!/usr/bin/env <%= @ruby_path %>
+
+require "bundler/setup"
+require "fileutils"
+require "rack"
+require "puma"
+require "puma/configuration"
+require "puma/launcher"
+require "securerandom"
+require "redis"
+require "logger"
+require "json"
+require "stringio"
+
+require_relative "../lib/model_context_protocol"
+require_relative "../lib/puma/plugin/mcp"
+
+# Only require test handler files (prompts, resources, tools, etc.), not RSpec helpers
+%w[prompts resources resource_templates tools completions].each do |subdir|
+  Dir[File.join(__dir__, "../spec/support/#{subdir}/**/*.rb")].each { |file| require file }
+end
+
+# Flag files for dynamic handler registration (checked at startup)
+FLAGS_DIR = File.join(__dir__, '..', 'tmp', 'flags')
+FileUtils.mkdir_p(FLAGS_DIR) unless Dir.exist?(FLAGS_DIR)
+
+def flag_enabled?(flag_name)
+  File.exist?(File.join(FLAGS_DIR, flag_name))
+end
+
+# Configure server logging
+ModelContextProtocol::Server.configure_server_logging do |config|
+  config.logdev = $stdout
+  config.level = Logger::DEBUG
+  config.progname = "MCP-Dev-Server"
+end
+
+# Configure MCP server (safe before fork - the Puma plugin handles starting after fork)
+ModelContextProtocol::Server.with_streamable_http_transport do |config|
+  config.name = "MCP Development Server (Puma)"
+  config.version = "1.0.0"
+
+  config.redis_url = "redis://localhost:6379/0"
+  config.redis_pool_size = 10
+  config.redis_enable_reaper = true
+  config.redis_reaper_interval = 10
+  config.redis_idle_timeout = 15
+
+  config.pagination = {
+    default_page_size: 2,
+    max_page_size: 3,
+    cursor_ttl: 1800
+  }
+
+  config.require_sessions = true
+  config.session_ttl = 3600
+  config.allowed_origins = ["*"]
+
+  config.registry do
+    prompts do
+      register TestPrompt
+      register TestPromptWithCompletionClass if flag_enabled?('extra_prompts')
+    end
+
+    resources subscribe: true do
+      register TestResource
+      register TestAnnotatedResource
+      register TestProgressiveResource
+      register TestBinaryResource if flag_enabled?('extra_resources')
+    end
+
+    resource_templates do
+      register TestResourceTemplate
+    end
+
+    tools do
+      register TestToolWithStructuredContentResponse
+      register TestToolWithTextResponse
+      register TestToolWithImageResponse
+      register TestToolWithMixedContentResponse
+      register TestToolWithResourceResponse
+      register TestToolWithToolErrorResponse
+      register TestToolWithCancellableSleep
+      register TestToolWithProgressableAndCancellable
+      register TestToolWithAnnotations
+      register TestToolWithSecuritySchemes
+      register TestToolWithResourceLinkResponse
+      register TestToolWithAudioResponse if flag_enabled?('extra_tools')
+    end
+  end
+end
+
+# Rack application
+class MCPHttpApp
+  def initialize
+    @logger = Logger.new(STDOUT)
+    @logger.level = Logger::INFO
+    @logger.formatter = proc do |severity, datetime, progname, msg|
+      request_id = Thread.current[:request_id] || "----"
+      "[#{datetime.strftime('%Y-%m-%d %H:%M:%S')}] #{severity} [#{request_id}]: #{msg}\n"
+    end
+  end
+
+  def call(env)
+    request_id = SecureRandom.hex(4)
+    Thread.current[:request_id] = request_id
+
+    request = Rack::Request.new(env)
+    body_content = request.body.read
+
+    log_request(env, body_content)
+
+    env['rack.input'] = StringIO.new(body_content)
+
+    unless request.path == "/mcp"
+      return [404, {"Content-Type" => "application/json"}, ['{"error": "Not found"}']]
+    end
+
+    if request.request_method == "OPTIONS"
+      return [200, cors_headers, [""]]
+    end
+
+    begin
+      result = ModelContextProtocol::Server.serve(
+        env: env,
+        session_context: {
+          user_id: "dev-user-123",
+          request_id: request_id
+        }
+      )
+
+      handle_result(result, body_content)
+    rescue => e
+      @logger.error("Error handling request: #{e.message}")
+      @logger.debug("Full backtrace:\n#{e.backtrace.join("\n")}")
+      [500, {"Content-Type" => "application/json"}, [%Q({"error": "Internal server error: #{e.message}"})]]
+    ensure
+      Thread.current[:request_id] = nil
+    end
+  end
+
+  private
+
+  def log_request(env, body_content)
+    case env['REQUEST_METHOD']
+    when 'POST'
+      begin
+        request_json = JSON.parse(body_content)
+        method = request_json['method']
+        id = request_json['id']
+
+        if method&.start_with?('notifications/') || id.nil?
+          @logger.info("→ #{method} [NOTIFICATION]")
+        else
+          @logger.info("→ #{method} (id: #{id}) [REQUEST]")
+        end
+        @logger.info("  Request: #{body_content}")
+      rescue JSON::ParserError
+        @logger.info("→ POST #{env['PATH_INFO']} [INVALID JSON]")
+      end
+    when 'GET'
+      accept_header = env['HTTP_ACCEPT'] || ''
+      if accept_header.include?('text/event-stream')
+        @logger.info("→ GET #{env['PATH_INFO']} [SSE STREAM REQUEST]")
+      else
+        @logger.info("→ GET #{env['PATH_INFO']}")
+      end
+    when 'DELETE'
+      session_id = env['HTTP_MCP_SESSION_ID']
+      @logger.info("→ DELETE #{env['PATH_INFO']} [SESSION CLEANUP#{session_id ? ": #{session_id}" : ""}]")
+    else
+      @logger.info("→ #{env['REQUEST_METHOD']} #{env['PATH_INFO']}")
+    end
+  end
+
+  def handle_result(result, body_content)
+    case result
+    when Hash
+      if result[:stream]
+        @logger.info("← SSE STREAM OPENED [PERSISTENT CONNECTION]")
+        headers = (result[:headers] || {}).merge(cors_headers)
+        [200, headers, result[:stream_proc]]
+      elsif result[:json]
+        response_body = result[:json].to_json
+        status = result[:status] || 200
+        @logger.info("← RESPONSE (status: #{status})")
+        @logger.info("  Response: #{response_body}") unless status == 202
+        headers = (result[:headers] || {}).merge(cors_headers).merge("Content-Type" => "application/json")
+        [status, headers, [response_body]]
+      else
+        @logger.error("← Invalid transport response")
+        [500, {"Content-Type" => "application/json"}, ['{"error": "Invalid transport response"}']]
+      end
+    else
+      @logger.error("← Unexpected response format")
+      [500, {"Content-Type" => "application/json"}, ['{"error": "Unexpected response format"}']]
+    end
+  end
+
+  def cors_headers
+    {
+      "Access-Control-Allow-Origin" => "*",
+      "Access-Control-Allow-Methods" => "GET, POST, DELETE, OPTIONS",
+      "Access-Control-Allow-Headers" => "Content-Type, Accept, Mcp-Session-Id, MCP-Protocol-Version, Origin",
+      "Access-Control-Max-Age" => "86400"
+    }
+  end
+end
+
+# Parse command line options
+workers = ENV.fetch("WEB_CONCURRENCY", "0").to_i
+threads = ENV.fetch("PUMA_THREADS", "5").to_i
+port = ENV.fetch("PORT", "9292").to_i
+
+puts "=" * 60
+puts "MCP Development Server (Puma)"
+puts "=" * 60
+puts "Port:    #{port}"
+puts "Workers: #{workers} (set WEB_CONCURRENCY to change)"
+puts "Threads: #{threads} (set PUMA_THREADS to change)"
+puts "Mode:    #{workers > 0 ? "Clustered (forking)" : "Single (no forking)"}"
+puts "=" * 60
+puts ""
+puts "Flag files (checked at startup only):"
+puts "  touch tmp/flags/extra_tools     # Add TestToolWithAudioResponse"
+puts "  touch tmp/flags/extra_resources # Add TestBinaryResource"
+puts "  touch tmp/flags/extra_prompts   # Add TestPromptWithCompletionClass"
+puts ""
+
+# Build the Rack app
+app = Rack::Builder.new do
+  map '/mcp' do
+    run MCPHttpApp.new
+  end
+end
+
+# Configure Puma
+puma_config = Puma::Configuration.new do |config|
+  config.bind "tcp://0.0.0.0:#{port}"
+  config.workers workers
+  config.threads 1, threads
+
+  # Use the MCP plugin to handle server lifecycle (start after fork, shutdown on exit)
+  config.plugin :mcp
+
+  config.app app
+end
+
+# Launch Puma
+launcher = Puma::Launcher.new(puma_config)
+launcher.run

data/tasks/templates/dev-http.erb

@@ -1,6 +1,7 @@
 #!/usr/bin/env <%= @ruby_path %>
 
 require "bundler/setup"
+require "fileutils"
 require "rack"
 require 'rackup/handler/webrick'
 require "webrick"
@@ -10,31 +11,94 @@ require "securerandom"
 require "redis"
 require "logger"
 require "json"
-require 'stringio'
+require "stringio"
 
 require_relative "../lib/model_context_protocol"
 
-ModelContextProtocol::Server.configure_redis do |config|
-  config.redis_url = "redis://localhost:6379/0"
-  config.pool_size = 10
-  config.enable_reaper = true
-  config.reaper_interval = 10
-  config.idle_timeout = 15
+# Only require test handler files (prompts, resources, tools, etc.), not RSpec helpers
+%w[prompts resources resource_templates tools completions].each do |subdir|
+  Dir[File.join(__dir__, "../spec/support/#{subdir}/**/*.rb")].each { |file| require file }
 end
 
-Dir[File.join(__dir__, "../spec/support/**/*.rb")].each { |file| require file }
+# Flag files for dynamic handler registration (checked at startup)
+FLAGS_DIR = File.join(__dir__, '..', 'tmp', 'flags')
+FileUtils.mkdir_p(FLAGS_DIR) unless Dir.exist?(FLAGS_DIR)
+
+def flag_enabled?(flag_name)
+  File.exist?(File.join(FLAGS_DIR, flag_name))
+end
 
-logger = Logger.new(STDOUT)
-logger.level = Logger::INFO
-logger.formatter = proc do |severity, datetime, progname, msg|
-  request_id = Thread.current[:request_id] || "----"
-  "[#{datetime.strftime('%Y-%m-%d %H:%M:%S')}] #{severity} [#{request_id}]: #{msg}\n"
+# Configure server logging
+ModelContextProtocol::Server.configure_server_logging do |config|
+  config.logdev = $stdout
+  config.level = Logger::DEBUG
+  config.progname = "MCP-Dev-Server"
 end
 
+# Configure MCP server
+ModelContextProtocol::Server.with_streamable_http_transport do |config|
+  config.name = "MCP Development Server"
+  config.version = "1.0.0"
+
+  config.redis_url = "redis://localhost:6379/0"
+  config.redis_pool_size = 10
+  config.redis_enable_reaper = true
+  config.redis_reaper_interval = 10
+  config.redis_idle_timeout = 15
+
+  config.pagination = {
+    default_page_size: 2,
+    max_page_size: 3,
+    cursor_ttl: 1800
+  }
+
+  config.require_sessions = true
+  config.session_ttl = 3600
+  config.allowed_origins = ["*"]
+
+  config.registry do
+    prompts do
+      register TestPrompt
+      register TestPromptWithCompletionClass if flag_enabled?('extra_prompts')
+    end
+
+    resources subscribe: true do
+      register TestResource
+      register TestAnnotatedResource
+      register TestProgressiveResource
+      register TestBinaryResource if flag_enabled?('extra_resources')
+    end
+
+    resource_templates do
+      register TestResourceTemplate
+    end
+
+    tools do
+      register TestToolWithStructuredContentResponse
+      register TestToolWithTextResponse
+      register TestToolWithImageResponse
+      register TestToolWithMixedContentResponse
+      register TestToolWithResourceResponse
+      register TestToolWithToolErrorResponse
+      register TestToolWithCancellableSleep
+      register TestToolWithProgressableAndCancellable
+      register TestToolWithAnnotations
+      register TestToolWithSecuritySchemes
+      register TestToolWithResourceLinkResponse
+      register TestToolWithAudioResponse if flag_enabled?('extra_tools')
+    end
+  end
+end
 
+# Rack application
 class MCPHttpApp
-  def initialize(logger)
-    @logger = logger
+  def initialize
+    @logger = Logger.new(STDOUT)
+    @logger.level = Logger::INFO
+    @logger.formatter = proc do |severity, datetime, progname, msg|
+      request_id = Thread.current[:request_id] || "----"
+      "[#{datetime.strftime('%Y-%m-%d %H:%M:%S')}] #{severity} [#{request_id}]: #{msg}\n"
+    end
   end
 
   def call(env)
@@ -44,6 +108,40 @@ class MCPHttpApp
     request = Rack::Request.new(env)
     body_content = request.body.read
 
+    log_request(env, body_content)
+
+    env['rack.input'] = StringIO.new(body_content)
+
+    unless request.path == "/mcp"
+      return [404, {"Content-Type" => "application/json"}, ['{"error": "Not found"}']]
+    end
+
+    if request.request_method == "OPTIONS"
+      return [200, cors_headers, [""]]
+    end
+
+    begin
+      result = ModelContextProtocol::Server.serve(
+        env: env,
+        session_context: {
+          user_id: "dev-user-123",
+          request_id: request_id
+        }
+      )
+
+      handle_result(result, body_content)
+    rescue => e
+      @logger.error("Error handling request: #{e.message}")
+      @logger.debug("Full backtrace:\n#{e.backtrace.join("\n")}")
+      [500, {"Content-Type" => "application/json"}, [%Q({"error": "Internal server error: #{e.message}"})]]
+    ensure
+      Thread.current[:request_id] = nil
+    end
+  end
+
+  private
+
+  def log_request(env, body_content)
     case env['REQUEST_METHOD']
     when 'POST'
       begin
@@ -51,9 +149,7 @@ class MCPHttpApp
         method = request_json['method']
         id = request_json['id']
 
-        if method&.start_with?('notifications/')
-          @logger.info("→ #{method} [NOTIFICATION]")
-        elsif id.nil?
+        if method&.start_with?('notifications/') || id.nil?
           @logger.info("→ #{method} [NOTIFICATION]")
         else
           @logger.info("→ #{method} (id: #{id}) [REQUEST]")
@@ -61,7 +157,6 @@ class MCPHttpApp
         @logger.info("  Request: #{body_content}")
       rescue JSON::ParserError
         @logger.info("→ POST #{env['PATH_INFO']} [INVALID JSON]")
-        @logger.info("  Request: #{body_content}")
       end
     when 'GET'
       accept_header = env['HTTP_ACCEPT'] || ''
@@ -70,173 +165,45 @@ class MCPHttpApp
       else
         @logger.info("→ GET #{env['PATH_INFO']}")
       end
-      @logger.info("  Headers: Accept=#{accept_header}") unless accept_header.empty?
     when 'DELETE'
       session_id = env['HTTP_MCP_SESSION_ID']
-      if session_id
-        @logger.info("→ DELETE #{env['PATH_INFO']} [SESSION CLEANUP: #{session_id}]")
-      else
-        @logger.info("→ DELETE #{env['PATH_INFO']} [SESSION CLEANUP]")
-      end
+      @logger.info("→ DELETE #{env['PATH_INFO']} [SESSION CLEANUP#{session_id ? ": #{session_id}" : ""}]")
     else
       @logger.info("→ #{env['REQUEST_METHOD']} #{env['PATH_INFO']}")
-      @logger.info("  Request: #{body_content}") unless body_content.empty?
-    end
-
-    if ModelContextProtocol::Server::RedisConfig.configured?
-      pool_stats = ModelContextProtocol::Server::RedisConfig.stats
-      @logger.info("  Redis Pool: #{pool_stats}")
-    end
-
-    env['rack.input'] = StringIO.new(body_content)
-    request = Rack::Request.new(env)
-
-    unless request.path == "/mcp"
-      return [404, {"Content-Type" => "application/json"}, ['{"error": "Not found"}']]
     end
+  end
 
-    if request.request_method == "OPTIONS"
-      return [200, {
-        "Access-Control-Allow-Origin" => "*",
-        "Access-Control-Allow-Methods" => "GET, POST, DELETE, OPTIONS",
-        "Access-Control-Allow-Headers" => "Content-Type, Accept, Mcp-Session-Id, MCP-Protocol-Version, Origin",
-        "Access-Control-Max-Age" => "86400"
-      }, [""]]
-    end
-
-    transport_config = {
-      type: :streamable_http,
-      env:,
-      require_sessions: false,
-      session_ttl: 3600,
-      allowed_origins: ["*"]
-    }
-
-    @logger.debug("Creating MCP server with transport config")
-    server = create_mcp_server(transport_config)
-    transport = nil
-
-    begin
-      @logger.debug("Starting MCP server")
-      result = server.start
-
-      if server.respond_to?(:transport)
-        transport = server.transport
-      end
-
-      case result
-      when Hash
-        if result[:stream]
-          @logger.info("← SSE STREAM OPENED [PERSISTENT CONNECTION]")
-          @logger.info("  Connection will remain open for real-time notifications")
-          headers = result[:headers] || {}
-          headers["Access-Control-Allow-Origin"] = "*"
-          headers["Access-Control-Allow-Methods"] = "GET, POST, DELETE, OPTIONS"
-          headers["Access-Control-Allow-Headers"] = "Content-Type, Accept, Mcp-Session-Id, MCP-Protocol-Version, Origin"
-
-          return [200, headers, result[:stream_proc]]
-        elsif result[:json]
-          response_body = result[:json].to_json
-          status = result[:status] || 200
-
-          begin
-            response_json = result[:json]
-            if response_json[:error]
-              @logger.info("← ERROR RESPONSE (code: #{response_json[:error][:code]})")
-            elsif status == 202
-              @logger.info("← NOTIFICATION ACCEPTED [NO RESPONSE REQUIRED]")
-            elsif response_json[:accepted] == true && status == 200
-              method = request_json['method'] rescue 'unknown'
-              id = request_json['id'] rescue 'unknown'
-              @logger.info("← #{method} RESPONSE (id: #{id}) [DELIVERED VIA SSE STREAM]")
-            elsif response_json[:result]
-              method = request_json['method'] rescue 'unknown'
-              @logger.info("← #{method} RESPONSE (id: #{response_json[:id]})")
-            else
-              @logger.info("← RESPONSE (status: #{status})")
-            end
-            @logger.info("  Response: #{response_body}") unless status == 202 && response_body == '{}'
-          rescue
-            @logger.info("← RESPONSE (status: #{status})")
-            @logger.info("  Response: #{response_body}") unless response_body.empty?
-          end
-
-          headers = result[:headers] || {}
-          headers["Content-Type"] = "application/json"
-          headers["Access-Control-Allow-Origin"] = "*"
-          headers["Access-Control-Allow-Methods"] = "GET, POST, DELETE, OPTIONS" 
-          headers["Access-Control-Allow-Headers"] = "Content-Type, Accept, Mcp-Session-Id, MCP-Protocol-Version, Origin"
-
-          [result[:status] || 200, headers, [response_body]]
-        else
-          # Fallback
-          @logger.error("← Invalid transport response")
-          [500, {"Content-Type" => "application/json"}, ['{"error": "Invalid transport response"}']]
-        end
+  def handle_result(result, body_content)
+    case result
+    when Hash
+      if result[:stream]
+        @logger.info("← SSE STREAM OPENED [PERSISTENT CONNECTION]")
+        headers = (result[:headers] || {}).merge(cors_headers)
+        [200, headers, result[:stream_proc]]
+      elsif result[:json]
+        response_body = result[:json].to_json
+        status = result[:status] || 200
+        @logger.info("← RESPONSE (status: #{status})")
+        @logger.info("  Response: #{response_body}") unless status == 202
+        headers = (result[:headers] || {}).merge(cors_headers).merge("Content-Type" => "application/json")
+        [status, headers, [response_body]]
       else
-        @logger.error("← Unexpected response format")
-        [500, {"Content-Type" => "application/json"}, ['{"error": "Unexpected response format"}']]
+        @logger.error("← Invalid transport response")
+        [500, {"Content-Type" => "application/json"}, ['{"error": "Invalid transport response"}']]
       end
-    rescue => e
-      @logger.error("Error handling request: #{e.message}")
-      @logger.debug("Full backtrace:\n#{e.backtrace.join("\n")}")
-      [500, {"Content-Type" => "application/json"}, [%Q({"error": "Internal server error: #{e.message}"})]]
-    ensure
-      transport&.cleanup if transport&.respond_to?(:cleanup)
-      Thread.current[:request_id] = nil
+    else
+      @logger.error("← Unexpected response format")
+      [500, {"Content-Type" => "application/json"}, ['{"error": "Unexpected response format"}']]
     end
   end
 
-  private
-
-  def create_mcp_server(transport_config)
-    ModelContextProtocol::Server.new do |config|
-      config.name = "MCP Development Server"
-      config.version = "1.0.0"
-      config.logging_enabled = true
-
-      config.pagination = {
-        default_page_size: 2,
-        max_page_size: 3,
-        cursor_ttl: 1800
-      }
-
-      config.set_environment_variable("MCP_ENV", "development")
-
-      config.context = {
-        user_id: "123456",
-        request_id: Thread.current[:request_id]
-      }
-
-      config.transport = transport_config
-
-      config.registry = ModelContextProtocol::Server::Registry.new do
-        prompts list_changed: true do
-          register TestPrompt
-          register TestPromptWithCompletionClass
-        end
-
-        resources list_changed: true, subscribe: true do
-          register TestResource
-          register TestAnnotatedResource
-          register TestBinaryResource
-        end
-
-        resource_templates do
-          register TestResourceTemplate
-        end
-
-        tools list_changed: true do
-          register TestToolWithStructuredContentResponse
-          register TestToolWithTextResponse
-          register TestToolWithImageResponse
-          register TestToolWithMixedContentResponse
-          register TestToolWithResourceResponse
-          register TestToolWithToolErrorResponse
-          register TestToolWithCancellableSleep
-        end
-      end
-    end
+  def cors_headers
+    {
+      "Access-Control-Allow-Origin" => "*",
+      "Access-Control-Allow-Methods" => "GET, POST, DELETE, OPTIONS",
+      "Access-Control-Allow-Headers" => "Content-Type, Accept, Mcp-Session-Id, MCP-Protocol-Version, Origin",
+      "Access-Control-Max-Age" => "86400"
+    }
   end
 end
 
@@ -244,11 +211,23 @@ use_ssl = ENV['SSL'] == 'true'
 port = use_ssl ? 9293 : 9292
 protocol = use_ssl ? 'https' : 'http'
 
-logger.info("Starting MCP #{protocol.upcase} Development Server on #{protocol}://localhost:#{port}/mcp")
+puts "=" * 60
+puts "MCP Development Server (WEBrick)"
+puts "=" * 60
+puts "URL:      #{protocol}://localhost:#{port}/mcp"
+puts "Transport: Streamable HTTP"
+puts "=" * 60
+puts ""
+puts "Flag files (checked at startup only - restart server to apply changes):"
+puts "  touch tmp/flags/extra_tools     # Add TestToolWithAudioResponse"
+puts "  touch tmp/flags/extra_resources # Add TestBinaryResource"
+puts "  touch tmp/flags/extra_prompts   # Add TestPromptWithCompletionClass"
+puts "  rm tmp/flags/<flag>             # Remove the handler"
+puts ""
 
 app = Rack::Builder.new do
   map '/mcp' do
-    run MCPHttpApp.new(logger)
+    run MCPHttpApp.new
   end
 end
 
@@ -262,8 +241,8 @@ if use_ssl
   key_path = File.join(__dir__, '..', 'tmp', 'ssl', 'server.key')
 
   unless File.exist?(cert_path) && File.exist?(key_path)
-    logger.error("SSL certificates not found at tmp/ssl/")
-    logger.error("Generate them with: openssl req -x509 -newkey rsa:4096 -keyout tmp/ssl/server.key -out tmp/ssl/server.crt -days 365 -nodes -subj \"/C=US/ST=Dev/L=Dev/O=Dev/CN=localhost\"")
+    $stderr.puts "SSL certificates not found at tmp/ssl/"
+    $stderr.puts "Generate them with: openssl req -x509 -newkey rsa:4096 -keyout tmp/ssl/server.key -out tmp/ssl/server.crt -days 365 -nodes -subj \"/C=US/ST=Dev/L=Dev/O=Dev/CN=localhost\""
     exit(1)
   end
 
@@ -275,14 +254,17 @@ if use_ssl
   )
 end
 
-server = WEBrick::HTTPServer.new(server_options)
-server.mount '/', Rackup::Handler::WEBrick, app
+# Start the MCP server (lifecycle hook)
+ModelContextProtocol::Server.start
+
+webrick_server = WEBrick::HTTPServer.new(server_options)
+webrick_server.mount '/', Rackup::Handler::WEBrick, app
 
 ['INT', 'TERM'].each do |signal|
   Signal.trap(signal) do
-    server.shutdown
-    exit(0)
+    ModelContextProtocol::Server.shutdown
+    webrick_server.shutdown
   end
 end
 
-server.start
+webrick_server.start