headless_browser_tool 0.1.1

data/lib/headless_browser_tool/directory_setup.rb

@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+require "fileutils"
+
+module HeadlessBrowserTool
+  module DirectorySetup
+    HBT_DIR = "./.hbt"
+    SCREENSHOTS_DIR = File.join(HBT_DIR, "screenshots").freeze
+    SESSIONS_DIR = File.join(HBT_DIR, "sessions").freeze
+    LOGS_DIR = File.join(HBT_DIR, "logs").freeze
+
+    module_function
+
+    def setup_directories(include_logs: false)
+      # Create all necessary directories
+      FileUtils.mkdir_p(SCREENSHOTS_DIR)
+      FileUtils.mkdir_p(SESSIONS_DIR)
+      FileUtils.mkdir_p(LOGS_DIR) if include_logs
+
+      # Create .gitignore in .hbt directory
+      gitignore_path = File.join(HBT_DIR, ".gitignore")
+      File.write(gitignore_path, "*\n") unless File.exist?(gitignore_path)
+    end
+  end
+end

data/lib/headless_browser_tool/logger.rb

@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+require "logger"
+require "fileutils"
+require_relative "directory_setup"
+
+module HeadlessBrowserTool
+  class Logger
+    class << self
+      attr_accessor :instance
+
+      def initialize_logger(mode: :http)
+        @instance = if mode == :stdio
+                      # In stdio mode, write to log file
+                      log_file = File.join(DirectorySetup::LOGS_DIR, "#{Process.pid}.log")
+                      ::Logger.new(log_file, progname: "HBT")
+                    else
+                      # In HTTP mode, write to stdout
+                      ::Logger.new($stdout, progname: "HBT")
+                    end
+
+        @instance.level = ::Logger::INFO
+        @instance
+      end
+
+      def log
+        @instance ||= initialize_logger
+      end
+    end
+  end
+end

data/lib/headless_browser_tool/server.rb

@@ -0,0 +1,150 @@
+# frozen_string_literal: true
+
+require "sinatra/base"
+require "fast_mcp"
+require "rack"
+require "fileutils"
+require "securerandom"
+require "json"
+require_relative "browser"
+require_relative "tools"
+require_relative "session_manager"
+require_relative "session_middleware"
+require_relative "strict_session_middleware"
+require_relative "browser_adapter"
+require_relative "logger"
+require_relative "session_persistence"
+require_relative "directory_setup"
+
+module HeadlessBrowserTool
+  class Server < Sinatra::Base
+    class << self
+      attr_accessor :browser_instance, :session_manager, :single_session_mode, :show_headers, :session_id
+
+      def start_server(options = {})
+        # Initialize logger for HTTP mode
+        HeadlessBrowserTool::Logger.initialize_logger(mode: :http)
+
+        # Check if we should use single session mode
+        @single_session_mode = options[:single_session] || ENV["HBT_SINGLE_SESSION"] == "true"
+        @show_headers = options[:show_headers] || ENV["HBT_SHOW_HEADERS"] == "true"
+
+        # Validate session_id option
+        if options[:session_id] && !@single_session_mode
+          puts "Error: --session-id can only be used with --single-session"
+          exit 1
+        end
+
+        if @single_session_mode
+          puts "Running in single session mode"
+          if options[:session_id]
+            puts "Session ID: #{options[:session_id]}"
+            @session_id = options[:session_id]
+          end
+          @browser_instance = Browser.new(headless: options[:headless])
+
+          # Restore session if session_id provided
+          restore_single_session if @session_id
+        else
+          puts "Running in multi-session mode"
+          @session_manager = SessionManager.new(headless: options[:headless])
+        end
+
+        # Setup directory structure
+        DirectorySetup.setup_directories
+
+        puts "Starting HeadlessBrowserTool MCP server on port #{options[:port]}"
+        puts "Using fast-mcp for MCP protocol support"
+
+        # Register shutdown hook for single session persistence
+        at_exit { save_single_session } if @single_session_mode && @session_id
+
+        # Configure and run with Puma directly
+        require "puma"
+        require "puma/configuration"
+        require "puma/launcher"
+
+        puma_config = Puma::Configuration.new do |config|
+          config.bind "tcp://0.0.0.0:#{options[:port]}"
+          config.environment "production"
+          config.quiet false
+          config.app Server
+        end
+
+        launcher = Puma::Launcher.new(puma_config)
+        launcher.run
+      end
+
+      private
+
+      def restore_single_session
+        SessionPersistence.restore_session(@session_id, @browser_instance.session)
+      end
+
+      def save_single_session
+        return unless @session_id && @browser_instance
+
+        SessionPersistence.save_session(@session_id, @browser_instance.session)
+      end
+    end
+
+    # Use strict session middleware if in multi-session mode
+    use StrictSessionMiddleware unless Server.single_session_mode
+
+    # Create MCP server instance with custom context handler
+    mcp_server = FastMcp::Server.new(name: "headless-browser-tool", version: HeadlessBrowserTool::VERSION)
+
+    # Register all browser tools
+    HeadlessBrowserTool::Tools::ALL_TOOLS.each do |tool_class|
+      mcp_server.register_tool(tool_class)
+    end
+
+    # Create custom transport that passes session context
+    class SessionAwareTransport < FastMcp::Transports::RackTransport
+      def call(env)
+        # Get or create session ID for this connection
+        session_id = env["hbt.session_id"]
+
+        # Use the session ID from middleware (which handles persistence)
+        # Don't generate new ones here - let the middleware handle it
+
+        # Store session ID in thread local for tools to access
+        Thread.current[:hbt_session_id] = session_id
+
+        # Call parent
+        super
+      ensure
+        Thread.current[:hbt_session_id] = nil
+      end
+    end
+
+    # Use our custom transport
+    use SessionAwareTransport, mcp_server
+
+    # Session management endpoints
+    get "/sessions" do
+      if Server.single_session_mode
+        { mode: "single", message: "Server is running in single session mode" }.to_json
+      else
+        content_type :json
+        Server.session_manager.session_info.to_json
+      end
+    end
+
+    delete "/sessions/:id" do
+      if Server.single_session_mode
+        status 400
+        { error: "Cannot manage sessions in single session mode" }.to_json
+      else
+        Server.session_manager.close_session(params[:id])
+        { message: "Session closed", session_id: params[:id] }.to_json
+      end
+    end
+
+    # Default route
+    get "/" do
+      mode = Server.single_session_mode ? "single session" : "multi-session"
+      "HeadlessBrowserTool MCP Server is running in #{mode} mode"
+    end
+  end
+end

data/lib/headless_browser_tool/session_manager.rb

@@ -0,0 +1,199 @@
+# frozen_string_literal: true
+
+require "capybara"
+require "json"
+require "fileutils"
+require_relative "logger"
+require_relative "session_persistence"
+require_relative "directory_setup"
+
+module HeadlessBrowserTool
+  class SessionManager
+    SESSION_TIMEOUT = 30 * 60 # 30 minutes
+    CLEANUP_INTERVAL = 60 # 1 minute
+    MAX_SESSIONS = 10 # Maximum concurrent sessions
+
+    attr_reader :sessions_dir
+
+    def initialize(headless: true)
+      @sessions = {}
+      @session_data = {}
+      @mutex = Mutex.new
+      @headless = headless
+
+      # Enable Capybara threadsafe mode for per-session configuration
+      Capybara.threadsafe = true
+
+      # Use common sessions directory
+      @sessions_dir = DirectorySetup::SESSIONS_DIR
+
+      # Load existing session data
+      load_persisted_sessions
+
+      # Start cleanup thread
+      start_cleanup_thread
+
+      # Register shutdown hook
+      at_exit { shutdown_all_sessions }
+    end
+
+    def get_or_create_session(session_id)
+      @mutex.synchronize do
+        # Validate session_id
+        raise ArgumentError, "Invalid session ID: #{session_id}" unless valid_session_id?(session_id)
+
+        # Update last activity
+        session_data = @session_data[session_id] ||= {
+          created_at: Time.now,
+          last_activity: Time.now
+        }
+        session_data[:last_activity] = Time.now
+
+        # Check if we're at capacity
+        cleanup_least_recently_used if !@sessions[session_id] && @sessions.size >= MAX_SESSIONS
+
+        # Get or create the Capybara session
+        @sessions[session_id] ||= create_session(session_id)
+      end
+    end
+
+    def close_session(session_id)
+      @mutex.synchronize do
+        if (session = @sessions.delete(session_id))
+          begin
+            save_session_state(session_id, session)
+            session.quit
+          rescue StandardError => e
+            HeadlessBrowserTool::Logger.log.info "Error closing session #{session_id}: #{e.message}"
+          end
+        end
+        @session_data.delete(session_id)
+      end
+    end
+
+    def save_all_sessions
+      @mutex.synchronize do
+        @sessions.each do |session_id, session|
+          save_session_state(session_id, session)
+        end
+      end
+    end
+
+    def session_info
+      @mutex.synchronize do
+        {
+          active_sessions: @sessions.keys,
+          session_count: @sessions.size,
+          session_data: @session_data.transform_values do |data|
+            {
+              created_at: data[:created_at],
+              last_activity: data[:last_activity],
+              idle_time: Time.now - data[:last_activity]
+            }
+          end
+        }
+      end
+    end
+
+    private
+
+    def create_session(session_id)
+      HeadlessBrowserTool::Logger.log.info "Creating new Capybara session: #{session_id}"
+
+      # Create a new Capybara session
+      # With threadsafe mode enabled, each session is isolated
+      session = Capybara::Session.new(:selenium_chrome_headless)
+
+      # Try to restore previous state
+      restore_session_state(session_id, session)
+
+      session
+    rescue StandardError => e
+      HeadlessBrowserTool::Logger.log.info "Error creating session #{session_id}: #{e.message}"
+      raise
+    end
+
+    def save_session_state(session_id, session)
+      SessionPersistence.save_session(session_id, session)
+    end
+
+    def restore_session_state(session_id, session)
+      SessionPersistence.restore_session(session_id, session)
+    end
+
+    def load_persisted_sessions
+      Dir.glob(File.join(@sessions_dir, "*.json")).each do |file|
+        state = JSON.parse(File.read(file))
+        session_id = state["session_id"]
+
+        # Initialize session data for persisted sessions
+        @session_data[session_id] = {
+          created_at: Time.parse(state["saved_at"]),
+          last_activity: Time.now,
+          persisted: true
+        }
+
+        HeadlessBrowserTool::Logger.log.info "Found persisted session: #{session_id}"
+      rescue StandardError => e
+        HeadlessBrowserTool::Logger.log.info "Error loading persisted session from #{file}: #{e.message}"
+      end
+    end
+
+    def start_cleanup_thread
+      Thread.new do
+        loop do
+          sleep CLEANUP_INTERVAL
+          cleanup_idle_sessions
+        end
+      end
+    end
+
+    def cleanup_idle_sessions
+      @mutex.synchronize do
+        now = Time.now
+        sessions_to_close = []
+
+        @session_data.each do |session_id, data|
+          sessions_to_close << session_id if now - data[:last_activity] > SESSION_TIMEOUT
+        end
+
+        sessions_to_close.each do |session_id|
+          HeadlessBrowserTool::Logger.log.info "Cleaning up idle session: #{session_id} (idle for #{(now - @session_data[session_id][:last_activity]).to_i}s)" # rubocop:disable Layout/LineLength
+          close_session(session_id)
+        end
+      end
+    rescue StandardError => e
+      HeadlessBrowserTool::Logger.log.info "Error during cleanup: #{e.message}"
+    end
+
+    def cleanup_least_recently_used
+      # Find the least recently used session
+      lru_session = @session_data.min_by { |_, data| data[:last_activity] }&.first
+
+      return unless lru_session
+
+      HeadlessBrowserTool::Logger.log.info "Closing LRU session to make room: #{lru_session}"
+      close_session(lru_session)
+    end
+
+    def shutdown_all_sessions
+      HeadlessBrowserTool::Logger.log.info "Shutting down all sessions..."
+      @mutex.synchronize do
+        @sessions.each do |session_id, session|
+          save_session_state(session_id, session)
+          session.quit
+        rescue StandardError => e
+          HeadlessBrowserTool::Logger.log.info "Error shutting down session #{session_id}: #{e.message}"
+        end
+        @sessions.clear
+      end
+    end
+
+    def valid_session_id?(session_id)
+      session_id.is_a?(String) &&
+        session_id.length >= 1 &&
+        session_id.length <= 64 &&
+        session_id.match?(/\A[a-zA-Z0-9_\-]+\z/)
+    end
+  end
+end

data/lib/headless_browser_tool/session_middleware.rb

@@ -0,0 +1,158 @@
+# frozen_string_literal: true
+
+require "securerandom"
+require_relative "logger"
+
+module HeadlessBrowserTool
+  class SessionMiddleware
+    # Headers to check for session ID (in priority order)
+    SESSION_HEADERS = %w[
+      HTTP_X_SESSION_ID
+      HTTP_X_CLIENT_ID
+    ].freeze
+
+    DEFAULT_SESSION = "default"
+
+    # Class-level session storage for MCP connections
+    @mcp_sessions = {}
+    @session_mutex = Mutex.new
+
+    class << self
+      attr_accessor :mcp_sessions, :session_mutex
+    end
+
+    def initialize(app)
+      @app = app
+    end
+
+    def call(env)
+      # Log request headers if enabled
+      log_request_headers(env) if HeadlessBrowserTool::Server.show_headers
+
+      # Extract session ID from headers first
+      session_id = extract_session_id(env)
+
+      # For MCP connections, try to maintain session continuity
+      session_id = get_or_create_mcp_session(env) if session_id == DEFAULT_SESSION && mcp_request?(env)
+
+      # Store in environment for downstream use
+      env["hbt.session_id"] = session_id
+
+      # Call the app
+      status, headers, response = @app.call(env)
+
+      # Add session ID to response headers
+      headers["X-Session-ID"] = session_id
+
+      # For MCP SSE connections, send session info
+      if env["PATH_INFO"] == "/mcp/sse" && session_id != DEFAULT_SESSION
+        # Store for future requests from this client
+        store_mcp_session(env, session_id)
+      end
+
+      [status, headers, response]
+    end
+
+    private
+
+    def extract_session_id(env)
+      # Check for explicit session headers
+      SESSION_HEADERS.each do |header|
+        if (value = env[header])
+          sanitized = sanitize_session_id(value)
+          return sanitized if sanitized
+        end
+      end
+
+      # Default session if none provided
+      DEFAULT_SESSION
+    end
+
+    def sanitize_session_id(value)
+      return nil if value.nil? || value.empty?
+
+      # Ensure it's alphanumeric with underscores/hyphens only
+      sanitized = value.gsub(/[^a-zA-Z0-9_\-]/, "")
+
+      # Limit length
+      sanitized = sanitized[0..64]
+
+      # Return nil if invalid
+      return nil if sanitized.empty?
+
+      sanitized
+    end
+
+    def mcp_request?(env)
+      env["PATH_INFO"]&.start_with?("/mcp")
+    end
+
+    def get_or_create_mcp_session(env)
+      client_key = generate_client_key(env)
+
+      self.class.session_mutex.synchronize do
+        # Check if we have a session for this client
+        if self.class.mcp_sessions[client_key]
+          self.class.mcp_sessions[client_key][:last_seen] = Time.now
+          return self.class.mcp_sessions[client_key][:session_id]
+        end
+
+        # Create new session
+        session_id = "mcp_#{SecureRandom.hex(8)}"
+        self.class.mcp_sessions[client_key] = {
+          session_id: session_id,
+          created_at: Time.now,
+          last_seen: Time.now
+        }
+
+        # Clean old sessions
+        cleanup_old_sessions
+
+        session_id
+      end
+    end
+
+    def store_mcp_session(env, session_id)
+      client_key = generate_client_key(env)
+
+      self.class.session_mutex.synchronize do
+        self.class.mcp_sessions[client_key] = {
+          session_id: session_id,
+          created_at: Time.now,
+          last_seen: Time.now
+        }
+      end
+    end
+
+    def generate_client_key(env)
+      # Use IP + User-Agent as a simple client identifier
+      # This isn't perfect but works for most cases
+      ip = env["REMOTE_ADDR"] || "unknown"
+      user_agent = env["HTTP_USER_AGENT"] || "unknown"
+      "#{ip}:#{user_agent.hash}"
+    end
+
+    def cleanup_old_sessions
+      # Remove sessions older than 1 hour
+      cutoff = Time.now - 3600
+      self.class.mcp_sessions.delete_if { |_, data| data[:last_seen] < cutoff }
+    end
+
+    def log_request_headers(env)
+      HeadlessBrowserTool::Logger.log.info "\n=== REQUEST HEADERS ==="
+      HeadlessBrowserTool::Logger.log.info "Method: #{env["REQUEST_METHOD"]} #{env["PATH_INFO"]}"
+      HeadlessBrowserTool::Logger.log.info "Remote IP: #{env["REMOTE_ADDR"]}"
+
+      # Log all HTTP headers
+      env.select { |k, _v| k.start_with?("HTTP_") }.each do |key, value|
+        header_name = key.sub("HTTP_", "").split("_").map(&:capitalize).join("-")
+        HeadlessBrowserTool::Logger.log.info "#{header_name}: #{value}"
+      end
+
+      # Log specific session-related info
+      HeadlessBrowserTool::Logger.log.info "Session ID from headers: #{extract_session_id(env)}"
+      HeadlessBrowserTool::Logger.log.info "Client Key: #{generate_client_key(env)}"
+      HeadlessBrowserTool::Logger.log.info "=====================\n"
+    end
+  end
+end