deivid-rodriguez/pry-byebug

# frozen_string_literal: true

require "pry-remote"

module PryRemote
  #
  # Overrides PryRemote::Server
  #
  class Server
    #
    # Override the call to Pry.start to save off current Server, and not
    # teardown the server right after Pry.start finishes.
    #
    def run
      raise("Already running a pry-remote session!") if
        PryByebug.current_remote_server

      PryByebug.current_remote_server = self

      setup
      Pry.start @object, input: client.input_proxy, output: client.output
    end

    #
    # Override to reset our saved global current server session.
    #
    alias teardown_without_pry_byebug teardown
    def teardown_with_pry_byebug
      return if @torn

      teardown_without_pry_byebug
      PryByebug.current_remote_server = nil
      @torn = true
    end
    alias teardown teardown_with_pry_byebug
  end
end

# Ensure cleanup when a program finishes without another break. For example,
# 'next' on the last line of a program won't hit Byebug::PryProcessor#run,
# which normally handles cleanup.
at_exit do
  PryByebug.current_remote_server&.teardown
end