local.rb

# frozen_string_literal: true

# Custom code. Not generated by Stainless.
require "fileutils"
require "json"
require "net/http"
require "socket"
require "tmpdir"

module Stagehand
  module Local
    STAGEHAND_REPO = "browserbase/stagehand"
    DEFAULT_USER_AGENT = "stagehand-ruby/local"
    DEFAULT_HOST = "127.0.0.1"
    DEFAULT_READY_TIMEOUT_S = 30.0
    DOWNLOAD_TIMEOUT_S = 600

    @download_mutex = Mutex.new

    def self.download_mutex
      @download_mutex
    end

    def self.local_mode?(client)
      client.instance_variable_get(:@server_mode).to_s == "local"
    end

    def self.browser_type_from_params(params)
      return nil unless params.is_a?(Hash)

      browser = params[:browser] || params["browser"]
      return nil if browser.nil?

      type =
        case browser
        when Hash
          browser[:type] || browser["type"]
        else
          browser.respond_to?(:type) ? browser.type : nil
        end

      type&.to_sym
    end

    def self.ensure_browserbase_credentials!(client:, params:)
      return unless local_mode?(client)

      browser_type = browser_type_from_params(params)
      return if browser_type == :local

      missing = []
      missing << "browserbase_api_key" if client.browserbase_api_key.to_s.empty?
      return if missing.empty?

      message =
        "Browserbase credentials are required when launching a Browserbase browser: " \
        "missing #{missing.join(', ')}."
      raise ArgumentError, message
    end

    def self.windows?
      RbConfig::CONFIG["host_os"] =~ /mswin|mingw|cygwin/i
    end

    def self.macos?
      RbConfig::CONFIG["host_os"] =~ /darwin/i
    end

    module Binary
      module_function

      def platform_tag
        platform = if Stagehand::Local.macos?
          "darwin"
        elsif Stagehand::Local.windows?
          "win32"
        else
          "linux"
        end
        cpu = RbConfig::CONFIG["host_cpu"].to_s.downcase
        arch = cpu.include?("arm") || cpu.include?("aarch64") ? "arm64" : "x64"
        [platform, arch]
      end

      def binary_filename
        platform, arch = platform_tag
        name = "stagehand-server-v3-#{platform}-#{arch}"
        name += ".exe" if platform == "win32"
        name
      end

      def resolve_binary_path
        env = ENV["STAGEHAND_SEA_BINARY"].to_s
        return ensure_executable(env) unless env.empty?

        filename = binary_filename
        cache_path = File.join(cache_dir, filename)
        return ensure_executable(cache_path) if File.file?(cache_path)

        Stagehand::Local.download_mutex.synchronize do
          return ensure_executable(cache_path) if File.file?(cache_path)

          version = ENV["STAGEHAND_SERVER_VERSION"].to_s
          version = "latest" if version.empty?
          tag = resolve_version(version)
          download_binary(tag, cache_path)
        end

        ensure_executable(cache_path)
      end

      def cache_dir
        root =
          if Stagehand::Local.macos?
            File.join(home_dir || Dir.tmpdir, "Library", "Caches")
          elsif Stagehand::Local.windows?
            ENV["LOCALAPPDATA"] || File.join(home_dir || Dir.tmpdir, "AppData", "Local")
          else
            ENV["XDG_CACHE_HOME"] || File.join(home_dir || Dir.tmpdir, ".cache")
          end
        version = Stagehand::VERSION.to_s
        File.join(root, "stagehand", "lib", "ruby_#{version}")
      end

      def home_dir
        Dir.home
      rescue StandardError
        nil
      end

      def resolve_version(version)
        return fetch_latest_tag if version.empty? || version == "latest"
        return version if version.start_with?("stagehand-server-v3/")

        "stagehand-server-v3/#{version}"
      end

      def fetch_latest_tag
        url = URI("https://api.github.com/repos/#{STAGEHAND_REPO}/releases?per_page=15")
        request = Net::HTTP::Get.new(url)
        request["User-Agent"] = DEFAULT_USER_AGENT
        response = http_request(url, request)
        unless response.is_a?(Net::HTTPSuccess)
          raise "Failed to fetch releases: #{response.code} #{response.message}"
        end

        releases = JSON.parse(response.body.to_s)
        releases.each do |release|
          tag = release["tag_name"]
          return tag if tag.is_a?(String) && tag.start_with?("stagehand-server-v3/")
        end
        raise "Failed to find stagehand-server-v3 release tag"
      end

      def download_binary(tag, dest_path)
        filename = binary_filename
        url = URI("https://github.com/#{STAGEHAND_REPO}/releases/download/#{tag}/#{filename}")
        FileUtils.mkdir_p(File.dirname(dest_path))

        tmp_path = "#{dest_path}.tmp"
        download_with_redirects(url, tmp_path, limit: 3)
        FileUtils.mv(tmp_path, dest_path)
        ensure_executable(dest_path)
      rescue StandardError => e
        FileUtils.rm_f(tmp_path)
        hint = manual_download_hint(filename, dest_path)
        raise "Failed to download Stagehand driver binary: #{e.message}. #{hint}"
      end

      def manual_download_hint(filename, dest_path)
        "Download #{filename} from https://github.com/#{STAGEHAND_REPO}/releases " \
          "and save it to: #{dest_path}."
      end

      def download_with_redirects(url, dest_path, limit:)
        raise "Too many redirects while downloading Stagehand driver binary." if limit <= 0

        request = Net::HTTP::Get.new(url)
        request["User-Agent"] = DEFAULT_USER_AGENT
        response = http_request(url, request)

        case response
        when Net::HTTPRedirection
          location = response["location"]
          raise "Missing redirect location." if location.nil?

          download_with_redirects(URI(location), dest_path, limit: limit - 1)
        when Net::HTTPSuccess
          File.binwrite(dest_path, response.body)
          nil
        else
          raise "Failed to download binary: #{response.code} #{response.message}"
        end
      end

      def http_request(url, request)
        Net::HTTP.start(url.host, url.port, use_ssl: url.scheme == "https") do |http|
          http.open_timeout = 30
          http.read_timeout = DOWNLOAD_TIMEOUT_S
          http.request(request)
        end
      end

      def ensure_executable(path)
        return path if Stagehand::Local.windows?
        return path unless File.exist?(path)

        mode = File.stat(path).mode
        File.chmod(mode | 0o100, path)
        path
      end
    end

    class ServerManager
      def initialize(model_api_key:, browserbase_api_key:)
        @model_api_key = model_api_key
        @browserbase_api_key = browserbase_api_key
        @host = DEFAULT_HOST
        @port = 0
        @mutex = Mutex.new
        @pid = nil
        @base_url = nil
        @binary_path = nil
        @pgroup = !Stagehand::Local.windows?
        @at_exit_registered = false
      end

      def ensure_running
        @mutex.synchronize do
          return @base_url if running? && @base_url

          start
        end
      end

      def close
        @mutex.synchronize do
          return if @pid.nil?

          terminate(@pid)
          @pid = nil
          @base_url = nil
        end
      end

      private

      def start
        @binary_path ||= Stagehand::Local::Binary.resolve_binary_path
        port = @port.zero? ? pick_free_port(@host) : @port
        base_url = "http://#{@host}:#{port}"

        env = build_env(host: @host, port: port)
        spawn_opts = {out: $stdout, err: $stderr}
        spawn_opts[:pgroup] = true if @pgroup

        @pid = Process.spawn(env, @binary_path, **spawn_opts)
        register_at_exit_once

        begin
          wait_ready(base_url, timeout_s: DEFAULT_READY_TIMEOUT_S)
        rescue StandardError
          terminate(@pid)
          @pid = nil
          raise
        end

        @base_url = base_url
      end

      def running?
        return false if @pid.nil?

        Process.kill(0, @pid)
        true
      rescue Errno::ESRCH
        false
      rescue Errno::EPERM
        true
      end

      def register_at_exit_once
        return if @at_exit_registered

        at_exit { close }
        @at_exit_registered = true
      end

      def terminate(pid)
        return if pid.nil?

        target = @pgroup ? -pid : pid
        begin
          Process.kill("TERM", target)
        rescue StandardError
          nil
        end

        return if wait_for_exit(pid, timeout_s: 3.0)

        begin
          Process.kill("KILL", target)
        rescue StandardError
          nil
        end
        wait_for_exit(pid, timeout_s: 3.0)
      end

      def wait_for_exit(pid, timeout_s:)
        deadline = Process.clock_gettime(Process::CLOCK_MONOTONIC) + timeout_s
        loop do
          result = Process.waitpid(pid, Process::WNOHANG)
          return true if result
          return false if Process.clock_gettime(Process::CLOCK_MONOTONIC) >= deadline

          sleep(0.1)
        end
      rescue Errno::ECHILD
        true
      end

      def pick_free_port(host)
        TCPServer.open(host, 0) { |server| server.addr[1] }
      end

      def build_env(host:, port:)
        env = ENV.to_h
        env["NODE_ENV"] = "production"
        env["BB_ENV"] = "local"
        env["HOST"] = host
        env["PORT"] = port.to_s
        env["MODEL_API_KEY"] = @model_api_key if @model_api_key.to_s != ""
        if @browserbase_api_key.to_s != ""
          env["BROWSERBASE_API_KEY"] = @browserbase_api_key
        end
        env
      end

      def wait_ready(base_url, timeout_s:)
        deadline = Process.clock_gettime(Process::CLOCK_MONOTONIC) + timeout_s
        paths = ["/readyz", "/healthz", "/health"]

        while Process.clock_gettime(Process::CLOCK_MONOTONIC) < deadline
          raise "Stagehand local server exited unexpectedly" unless running?

          return if paths.any? { |path| ready_path?(base_url, path) }
          sleep(0.1)
        end

        raise "Stagehand local server not ready at #{base_url} after #{timeout_s}s"
      end

      def ready_path?(base_url, path)
        uri = URI.join(base_url, path)
        http = Net::HTTP.new(uri.host, uri.port)
        http.open_timeout = 1
        http.read_timeout = 1
        response = http.get(uri.request_uri)
        response.is_a?(Net::HTTPSuccess)
      rescue StandardError
        false
      end
    end

    module ClientPatch
      def initialize(server: "remote", **kwargs)
        @server_mode = server.to_s

        if @server_mode == "local"
          base_url = kwargs[:base_url]
          kwargs[:base_url] = base_url.nil? ? "http://#{DEFAULT_HOST}" : base_url
          kwargs[:browserbase_api_key] =
            kwargs[:browserbase_api_key] || ENV["BROWSERBASE_API_KEY"] || ""
        end

        super(**kwargs)

        return unless @server_mode == "local"

        @local_server_manager = Stagehand::Local::ServerManager.new(
          model_api_key: @model_api_key,
          browserbase_api_key: @browserbase_api_key
        )
      end

      def request(req)
        ensure_local_server!
        super
      end

      def close
        super if defined?(super)
      ensure
        @local_server_manager&.close
      end

      private

      def ensure_local_server!
        return unless @server_mode == "local"
        return if @local_server_manager.nil?

        base_url = @local_server_manager.ensure_running
        return if @base_url.to_s == base_url

        @base_url_components = Stagehand::Internal::Util.parse_uri(base_url)
        @base_url = Stagehand::Internal::Util.unparse_uri(@base_url_components)
      end

      def bb_api_key_auth
        return {} if @browserbase_api_key.to_s.empty?
        super
      end

      def bb_project_id_auth
        {}
      end
    end

    module SessionsPatch
      def start(params)
        Stagehand::Local.ensure_browserbase_credentials!(client: @client, params: params)
        super
      end
    end
  end
end

Stagehand::Client.prepend(Stagehand::Local::ClientPatch) if defined?(Stagehand::Client)
Stagehand::Resources::Sessions.prepend(Stagehand::Local::SessionsPatch) if defined?(Stagehand::Resources::Sessions)