daytona 0.197.0 → 0.199.0

data/lib/daytona/sandbox.rb

@@ -17,6 +17,29 @@ module Daytona
     include Instrumentation
 
     DEFAULT_TIMEOUT = 60
+    def self.with_events(*method_names)
+      method_names.each do |method_name|
+        original = instance_method(method_name)
+
+        visibility = if private_method_defined?(method_name, false)
+                       :private
+                     elsif protected_method_defined?(method_name, false)
+                       :protected
+                     else
+                       :public
+                     end
+
+        define_method(method_name) do |*args, **kwargs, &blk|
+          ensure_subscribed
+          original.bind_call(self, *args, **kwargs, &blk)
+        end
+
+        case visibility
+        when :private then private method_name
+        when :protected then protected method_name
+        end
+      end
+    end
 
     SANDBOX_METRIC_FIELD_BY_NAME = {
       'daytona.sandbox.cpu.utilization' => :cpu_used_pct,
@@ -117,6 +140,9 @@ module Daytona
     #   Not returned by list results; call #refresh on each item to populate.
     attr_reader :build_info
 
+    # @return [String, nil] When the sandbox will be automatically destroyed (nil if no TTL is set)
+    attr_reader :auto_destroy_at
+
     # @return [String] The creation timestamp of the sandbox
     attr_reader :created_at
 
@@ -129,6 +155,9 @@ module Daytona
     # @return [String] The version of the daemon running in the sandbox
     attr_reader :daemon_version
 
+    # @return [String] The toolbox proxy URL used to reach the sandbox's toolbox API
+    attr_reader :toolbox_proxy_url
+
     # @return [Daytona::Config]
     attr_reader :config
 
@@ -150,19 +179,27 @@ module Daytona
     # @return [Daytona::CodeInterpreter]
     attr_reader :code_interpreter
 
+    # Internal — obtain sandboxes via Daytona::Daytona#create, Daytona::Daytona#get,
+    # or Daytona::Daytona#list rather than constructing directly.
+    #
     # @params config [Daytona::Config]
     # @params sandbox_api [DaytonaApiClient::SandboxApi]
     # @params sandbox_dto [DaytonaApiClient::Sandbox, DaytonaApiClient::SandboxListItem]
     # @params otel_state [Daytona::OtelState, nil]
-    def initialize(sandbox_dto:, config:, sandbox_api:, otel_state: nil, analytics_api_url_provider: nil) # rubocop:disable Metrics/MethodLength
+    def initialize(sandbox_dto:, config:, sandbox_api:, subscription_manager:, otel_state: nil, analytics_api_url_provider: nil) # rubocop:disable Metrics/MethodLength
+      @state_waiters = []
+      @waiter_mutex = Mutex.new
       process_response(sandbox_dto)
       @config = config
       @sandbox_api = sandbox_api
       @otel_state = otel_state
       @analytics_api_url_provider = analytics_api_url_provider
+      @subscription_manager = subscription_manager
+      @sub_id = nil
 
       # Create toolbox API clients with dynamic configuration
-      toolbox_api_config = build_toolbox_api_config
+      @toolbox_api_config = build_toolbox_api_config
+      toolbox_api_config = @toolbox_api_config
 
       # Helper to create API client with authentication header
       create_authenticated_client = lambda do
@@ -205,6 +242,8 @@ module Daytona
       @info_api = info_api
       @server_api = server_api
       @system_api = system_api
+
+      ensure_subscribed
     end
 
     # Archives the sandbox, making it inactive and preserving its state. When sandboxes are
@@ -318,20 +357,73 @@ module Daytona
       @auto_pause_interval = interval
     end
 
+    # Sets the TTL (time to live) for the Sandbox.
+    # The TTL is re-anchored from the current time. When it elapses the Sandbox is destroyed,
+    # regardless of its current state. Use 0 to disable.
+    #
+    # @param minutes [Integer]
+    # @return [void]
+    # @raise [Daytona::Sdk::Error]
+    def ttl_minutes=(minutes)
+      raise Sdk::Error, 'TTL must be a non-negative integer' if minutes.negative?
+
+      sandbox_api.set_ttl(id, minutes)
+    end
+
     # Creates an SSH access token for the sandbox.
     #
     # @param expires_in_minutes [Integer] TThe number of minutes the SSH access token will be valid for
     # @return [DaytonaApiClient::SshAccessDto]
-    def create_ssh_access(expires_in_minutes) = sandbox_api.create_ssh_access(id, { expires_in_minutes: })
+    def create_ssh_access(expires_in_minutes)
+      sandbox_api.create_ssh_access(id, { expires_in_minutes: })
+    end
 
+    # Deletes the Sandbox.
+    #
+    # By default returns as soon as the deletion request is accepted (matching
+    # origin/main behavior). Pass +wait: true+ to block until the Sandbox
+    # reaches the +destroyed+ state.
+    #
+    # @param timeout [Numeric] Maximum wait time in seconds (defaults to 60 s).
+    #   Only meaningful when +wait+ is true.
+    # @param wait [Boolean] When +true+, block until the Sandbox is destroyed.
     # @return [void]
-    def delete
-      sandbox_api.delete_sandbox(id)
-      refresh
-    rescue DaytonaApiClient::ApiError => e
-      raise unless e.code == 404
+    def delete(timeout = DEFAULT_TIMEOUT, wait: false) # rubocop:disable Metrics/AbcSize,Metrics/MethodLength
+      begin
+        process_response(sandbox_api.delete_sandbox(id))
+      rescue DaytonaApiClient::ApiError => e
+        raise unless e.code == 404
+
+        apply_state(DaytonaApiClient::SandboxState::DESTROYED)
+      end
+
+      if wait
+        refresh_data_safe
+      else
+        begin
+          refresh
+        rescue DaytonaApiClient::ApiError => e
+          raise unless e.code == 404
+
+          apply_state(DaytonaApiClient::SandboxState::DESTROYED)
+        end
+        return
+      end
+      return if state.to_s == DaytonaApiClient::SandboxState::DESTROYED.to_s
 
-      @state = 'destroyed'
+      with_timeout(
+        timeout:,
+        message: "Sandbox #{id} failed to be destroyed within the #{timeout} seconds timeout period",
+        setup: nil
+      ) do
+        wait_for_state(
+          target_states: [DaytonaApiClient::SandboxState::DESTROYED],
+          error_states: [DaytonaApiClient::SandboxState::ERROR, DaytonaApiClient::SandboxState::BUILD_FAILED],
+          safe_refresh: true
+        )
+      end
+    ensure
+      unsubscribe_from_events
     end
 
     # Gets the user's home directory path for the logged in user inside the Sandbox.
@@ -428,7 +520,9 @@ module Daytona
     #
     # @param port [Integer]
     # @return [DaytonaApiClient::PortPreviewUrl]
-    def preview_url(port) = sandbox_api.get_port_preview_url(id, port)
+    def preview_url(port)
+      sandbox_api.get_port_preview_url(id, port)
+    end
 
     # Creates a signed preview URL for the sandbox at the specified port.
     #
@@ -460,7 +554,9 @@ module Daytona
     # Refresh the Sandbox data from the API.
     #
     # @return [void]
-    def refresh = process_response(sandbox_api.get_sandbox(id))
+    def refresh
+      process_response(sandbox_api.get_sandbox(id))
+    end
 
     # Refreshes the sandbox activity to reset the timer for automated lifecycle management actions.
     #
@@ -482,7 +578,9 @@ module Daytona
     #
     # @param token [String]
     # @return [void]
-    def revoke_ssh_access(token) = sandbox_api.revoke_ssh_access(id, token:)
+    def revoke_ssh_access(token)
+      sandbox_api.revoke_ssh_access(id, token:)
+    end
 
     # Starts the Sandbox and waits for it to be ready.
     #
@@ -493,7 +591,12 @@ module Daytona
         timeout:,
         message: "Sandbox #{id} failed to become ready within the #{timeout} seconds timeout period",
         setup: proc { process_response(sandbox_api.start_sandbox(id)) }
-      ) { wait_for_states(operation: OPERATION_START, target_states: [DaytonaApiClient::SandboxState::STARTED]) }
+      ) do
+        wait_for_state(
+          target_states: [DaytonaApiClient::SandboxState::STARTED],
+          error_states: [DaytonaApiClient::SandboxState::ERROR, DaytonaApiClient::SandboxState::BUILD_FAILED]
+        )
+      end
     end
 
     # Recovers the Sandbox from a recoverable error and waits for it to be ready.
@@ -510,7 +613,12 @@ module Daytona
         timeout:,
         message: "Sandbox #{id} failed to recover within the #{timeout} seconds timeout period",
         setup: proc { process_response(sandbox_api.recover_sandbox(id)) }
-      ) { wait_for_states(operation: OPERATION_START, target_states: [DaytonaApiClient::SandboxState::STARTED]) }
+      ) do
+        wait_for_state(
+          target_states: [DaytonaApiClient::SandboxState::STARTED],
+          error_states: [DaytonaApiClient::SandboxState::ERROR, DaytonaApiClient::SandboxState::BUILD_FAILED]
+        )
+      end
     rescue StandardError => e
       raise Sdk::Error, "Failed to recover sandbox: #{e.message}"
     end
@@ -529,9 +637,9 @@ module Daytona
           refresh
         }
       ) do
-        wait_for_states(
-          operation: OPERATION_STOP,
-          target_states: [DaytonaApiClient::SandboxState::STOPPED, DaytonaApiClient::SandboxState::DESTROYED]
+        wait_for_state(
+          target_states: [DaytonaApiClient::SandboxState::STOPPED, DaytonaApiClient::SandboxState::DESTROYED],
+          error_states: [DaytonaApiClient::SandboxState::ERROR, DaytonaApiClient::SandboxState::BUILD_FAILED]
         )
       end
     end
@@ -572,7 +680,7 @@ module Daytona
           resize_request = DaytonaApiClient::ResizeSandbox.new(resize_attrs)
           process_response(sandbox_api.resize_sandbox(id, resize_request))
         }
-      ) { wait_for_resize_complete }
+      ) { wait_for_resize_complete(timeout) }
     end
 
     # Waits for the Sandbox resize operation to complete.
@@ -580,9 +688,18 @@ module Daytona
     #
     # @param timeout [Numeric] Maximum wait time in seconds (defaults to 60 s)
     # @return [void]
-    def wait_for_resize_complete(_timeout = DEFAULT_TIMEOUT)
-      wait_for_states(operation: OPERATION_RESIZE, target_states: [DaytonaApiClient::SandboxState::STARTED,
-                                                                   DaytonaApiClient::SandboxState::STOPPED])
+    def wait_for_resize_complete(timeout = DEFAULT_TIMEOUT)
+      with_timeout(
+        timeout:,
+        message: "Sandbox #{id} resize did not complete within the #{timeout} seconds timeout period",
+        setup: nil
+      ) do
+        error_states = [DaytonaApiClient::SandboxState::ERROR, DaytonaApiClient::SandboxState::BUILD_FAILED]
+        exclude = [DaytonaApiClient::SandboxState::RESIZING] + error_states
+        target_states = DaytonaApiClient::SandboxState.all_vars.reject { |s| exclude.include?(s) }
+
+        wait_for_state(target_states:, error_states:)
+      end
     end
 
     # Creates a new Language Server Protocol (LSP) server instance.
@@ -601,15 +718,26 @@ module Daytona
     #
     # @param token [String]
     # @return [DaytonaApiClient::SshAccessValidationDto]
-    def validate_ssh_access(token) = sandbox_api.validate_ssh_access(token)
+    def validate_ssh_access(token)
+      sandbox_api.validate_ssh_access(token)
+    end
 
     # Waits for the Sandbox to reach the 'started' state. Polls the Sandbox status until it
     # reaches the 'started' state or encounters an error.
     #
     # @param timeout [Numeric] Maximum wait time in seconds (defaults to 60 s).
     # @return [void]
-    def wait_for_sandbox_start(_timeout = DEFAULT_TIMEOUT)
-      wait_for_states(operation: OPERATION_START, target_states: [DaytonaApiClient::SandboxState::STARTED])
+    def wait_for_sandbox_start(timeout = DEFAULT_TIMEOUT)
+      with_timeout(
+        timeout:,
+        message: "Sandbox #{id} failed to start within the #{timeout} seconds timeout period",
+        setup: nil
+      ) do
+        wait_for_state(
+          target_states: [DaytonaApiClient::SandboxState::STARTED],
+          error_states: [DaytonaApiClient::SandboxState::ERROR, DaytonaApiClient::SandboxState::BUILD_FAILED]
+        )
+      end
     end
 
     # Waits for the Sandbox to reach the 'stopped' state. Polls the Sandbox status until it
@@ -618,9 +746,17 @@ module Daytona
     #
     # @param timeout [Numeric] Maximum wait time in seconds (defaults to 60 s).
     # @return [void]
-    def wait_for_sandbox_stop(_timeout = DEFAULT_TIMEOUT)
-      wait_for_states(operation: OPERATION_STOP, target_states: [DaytonaApiClient::SandboxState::STOPPED,
-                                                                 DaytonaApiClient::SandboxState::DESTROYED])
+    def wait_for_sandbox_stop(timeout = DEFAULT_TIMEOUT)
+      with_timeout(
+        timeout:,
+        message: "Sandbox #{id} failed to stop within the #{timeout} seconds timeout period",
+        setup: nil
+      ) do
+        wait_for_state(
+          target_states: [DaytonaApiClient::SandboxState::STOPPED, DaytonaApiClient::SandboxState::DESTROYED],
+          error_states: [DaytonaApiClient::SandboxState::ERROR, DaytonaApiClient::SandboxState::BUILD_FAILED]
+        )
+      end
     end
 
     # Forks the Sandbox, creating a new Sandbox with an identical filesystem.
@@ -643,12 +779,11 @@ module Daytona
           sandbox_dto: forked_dto,
           config:,
           sandbox_api:,
-          code_toolbox:,
+          subscription_manager: @subscription_manager,
           otel_state:,
           analytics_api_url_provider: @analytics_api_url_provider
         )
-        forked.send(:wait_for_states, operation: OPERATION_START,
-                                      target_states: [DaytonaApiClient::SandboxState::STARTED])
+        forked.wait_for_sandbox_start(timeout)
         return forked
       end
     end
@@ -664,14 +799,15 @@ module Daytona
         timeout:,
         message: "Sandbox #{id} snapshot failed within the #{timeout} seconds timeout period",
         setup: proc {
-          sandbox_api.create_sandbox_snapshot(id, DaytonaApiClient::CreateSandboxSnapshot.new(name:))
-          refresh
+          response = sandbox_api.create_sandbox_snapshot(id, DaytonaApiClient::CreateSandboxSnapshot.new(name:))
+          process_response(response)
         }
       ) { wait_for_snapshot_complete }
     end
 
     # Pauses the Sandbox, freezing all running processes.
-    # The Sandbox will enter a 'pausing' state and transition to 'paused' when complete.
+    # Completes when the Sandbox has left the +pausing+ state (paused, stopped,
+    # archived, etc.); error states still raise.
     #
     # @param timeout [Numeric] Maximum wait time in seconds (defaults to 60 s)
     # @return [void]
@@ -683,10 +819,25 @@ module Daytona
           sandbox_api.pause_sandbox(id)
           refresh
         }
-      ) { wait_for_pause_complete }
+      ) do
+        error_states = [DaytonaApiClient::SandboxState::ERROR, DaytonaApiClient::SandboxState::BUILD_FAILED]
+        exclude = [DaytonaApiClient::SandboxState::PAUSING] + error_states
+        target_states = DaytonaApiClient::SandboxState.all_vars.reject { |s| exclude.include?(s) }
+
+        wait_for_state(target_states:, error_states:)
+      end
     end
 
+    with_events :archive, :auto_archive_interval=, :auto_delete_interval=, :auto_stop_interval=,
+                :create_ssh_access, :delete, :get_user_home_dir, :get_work_dir, :labels=,
+                :preview_url, :create_signed_preview_url, :expire_signed_preview_url,
+                :refresh, :refresh_activity, :revoke_ssh_access, :start, :recover, :stop,
+                :create_lsp_server, :validate_ssh_access, :wait_for_sandbox_start,
+                :wait_for_sandbox_stop, :resize, :wait_for_resize_complete,
+                :experimental_fork, :experimental_create_snapshot, :pause
+
     instrument :archive, :auto_archive_interval=, :auto_delete_interval=, :auto_pause_interval=, :auto_stop_interval=,
+               :ttl_minutes=,
                :update_network_settings, :update_secrets, :update_env,
                :create_ssh_access, :delete, :get_user_home_dir, :get_work_dir, :get_metrics, :get_metrics_latest,
                :labels=,
@@ -782,6 +933,11 @@ module Daytona
     # @return [DaytonaToolboxApiClient::Configuration]
     def build_toolbox_api_config
       DaytonaToolboxApiClient::Configuration.new.configure do |cfg|
+        if @toolbox_proxy_url.nil? || @toolbox_proxy_url.empty?
+          proxy_response = @sandbox_api.get_toolbox_proxy_url(id)
+          @toolbox_proxy_url = proxy_response&.url || ''
+        end
+
         proxy_url = @toolbox_proxy_url
         proxy_url += '/' unless proxy_url.end_with?('/')
         full_url = "#{proxy_url}#{id}"
@@ -795,9 +951,26 @@ module Daytona
       end
     end
 
+    def update_toolbox_base_url
+      return unless @toolbox_api_config && @toolbox_proxy_url
+
+      proxy_url = @toolbox_proxy_url
+      proxy_url += '/' unless proxy_url.end_with?('/')
+      full_url = "#{proxy_url}#{id}"
+      uri = URI(full_url)
+
+      @toolbox_api_config.scheme = uri.scheme
+      @toolbox_api_config.host = uri.authority
+      @toolbox_api_config.base_path = uri.path.empty? ? '/' : uri.path
+    end
+
     # @params sandbox_dto [DaytonaApiClient::Sandbox, DaytonaApiClient::SandboxListItem]
     # @return [void]
     def process_response(sandbox_dto) # rubocop:disable Metrics/MethodLength, Metrics/AbcSize
+      return if sandbox_dto.nil?
+
+      new_state = sandbox_dto.state
+
       # Fields shared by both DaytonaApiClient::Sandbox and DaytonaApiClient::SandboxListItem.
       @id = sandbox_dto.id
       @organization_id = sandbox_dto.organization_id
@@ -810,7 +983,6 @@ module Daytona
       @gpu = sandbox_dto.gpu
       @memory = sandbox_dto.memory
       @disk = sandbox_dto.disk
-      @state = sandbox_dto.state
       @desired_state = sandbox_dto.desired_state
       @error_reason = sandbox_dto.error_reason
       @backup_state = sandbox_dto.backup_state
@@ -818,23 +990,43 @@ module Daytona
       @auto_pause_interval = sandbox_dto.auto_pause_interval
       @auto_archive_interval = sandbox_dto.auto_archive_interval
       @auto_delete_interval = sandbox_dto.auto_delete_interval
+      @auto_destroy_at = sandbox_dto.auto_destroy_at
       @created_at = sandbox_dto.created_at
       @updated_at = sandbox_dto.updated_at
       @last_activity_at = sandbox_dto.last_activity_at
       @daemon_version = sandbox_dto.daemon_version
-      @toolbox_proxy_url = sandbox_dto.toolbox_proxy_url
+
+      new_proxy_url = sandbox_dto.toolbox_proxy_url
+      if new_proxy_url && !new_proxy_url.empty? && new_proxy_url != @toolbox_proxy_url
+        @toolbox_proxy_url = new_proxy_url
+        update_toolbox_base_url
+      end
 
       # Fields only present on the full DaytonaApiClient::Sandbox DTO (not returned by list
       # results; call #refresh on each item to populate them).
-      return unless sandbox_dto.is_a?(DaytonaApiClient::Sandbox)
+      if sandbox_dto.is_a?(DaytonaApiClient::Sandbox)
+        @env = sandbox_dto.env
+        @network_block_all = sandbox_dto.network_block_all
+        @network_allow_list = sandbox_dto.network_allow_list
+        @domain_allow_list = sandbox_dto.domain_allow_list
+        @volumes = sandbox_dto.volumes
+        @build_info = sandbox_dto.build_info
+        @backup_created_at = sandbox_dto.backup_created_at
+      end
 
-      @env = sandbox_dto.env
-      @network_block_all = sandbox_dto.network_block_all
-      @network_allow_list = sandbox_dto.network_allow_list
-      @domain_allow_list = sandbox_dto.domain_allow_list
-      @volumes = sandbox_dto.volumes
-      @build_info = sandbox_dto.build_info
-      @backup_created_at = sandbox_dto.backup_created_at
+      apply_state(new_state)
+    end
+
+    def apply_state(new_state)
+      return if new_state.to_s == @state.to_s
+
+      @state = new_state
+
+      @waiter_mutex.synchronize do
+        @state_waiters.each do |waiter|
+          waiter[:notify].call(new_state.to_s) if waiter[:matches].call(new_state.to_s)
+        end
+      end
     end
 
     # Monitors block not to exceed max execution time.
@@ -848,99 +1040,200 @@ module Daytona
       start_at = Time.now
       setup&.call
 
-      Timeout.timeout(
-        setup ? [NO_TIMEOUT, timeout - (Time.now - start_at)].max : timeout,
-        Sdk::Error,
-        message,
-        &
-      )
+      effective = if setup
+                    remaining = timeout - (Time.now - start_at)
+                    raise Sdk::TimeoutError, message if timeout.positive? && remaining <= 0
+
+                    [NO_TIMEOUT, remaining].max
+                  else
+                    timeout
+                  end
+
+      Timeout.timeout(effective, Sdk::TimeoutError, message, &)
     end
 
-    # Waits for the Sandbox to reach the one of the target states. Polls the Sandbox status until it
-    # reaches the one of the target states or encounters an error. It will wait up to 60 seconds
-    # for the Sandbox to reach one of the target states.
+    # Waits for the Sandbox to reach one of the target states via WebSocket events
+    # with periodic polling safety net.
     #
-    # @param operation [#to_s] Operation name for error message
-    # @param target_states [Array<DaytonaApiClient::SandboxState>] List of the target states
+    # @param target_states [Array<DaytonaApiClient::SandboxState>] States that indicate success.
+    # @param error_states [Array<DaytonaApiClient::SandboxState>] States that indicate failure.
+    # @param safe_refresh [Boolean] If true, use refresh_data_safe for polling (swallows transient errors).
     # @return [void]
     # @raise [Daytona::Sdk::Error]
-    def wait_for_states(operation:, target_states:)
-      interval = INITIAL_POLL_INTERVAL
-      start_time = ::Process.clock_gettime(::Process::CLOCK_MONOTONIC)
-      loop do
-        case state
-        when *target_states then return
-        when DaytonaApiClient::SandboxState::ERROR, DaytonaApiClient::SandboxState::BUILD_FAILED
-          raise Sdk::Error, "Sandbox #{id} failed to #{operation} with state: #{state}, error reason: #{error_reason}"
+    def wait_for_state(target_states:, error_states:, safe_refresh: false) # rubocop:disable Metrics/MethodLength,Metrics/CyclomaticComplexity,Metrics/PerceivedComplexity
+      ensure_subscribed
+      subscribed = !@sub_id.nil?
+
+      target_strings = target_states.map(&:to_s)
+      error_strings = error_states.map(&:to_s)
+      started_at = ::Process.clock_gettime(::Process::CLOCK_MONOTONIC)
+      poll_interval = subscribed ? MAX_POLL_INTERVAL : INITIAL_POLL_INTERVAL
+
+      mutex = Mutex.new
+      state_changed = ConditionVariable.new
+      result_state = nil
+
+      waiter = {
+        matches: lambda do |new_state|
+          target_strings.include?(new_state) || error_strings.include?(new_state)
+        end,
+        notify: lambda do |new_state|
+          mutex.synchronize do
+            next if result_state
+
+            result_state = new_state
+            state_changed.signal
+          end
+        end
+      }
+
+      @waiter_mutex.synchronize do
+        @state_waiters << waiter
+
+        current_state = state.to_s
+        result_state = current_state if target_strings.include?(current_state)
+      end
+
+      begin
+        first_poll = true
+        until mutex.synchronize { result_state }
+          unless first_poll
+            should_refresh = mutex.synchronize do
+              next false if result_state
+
+              state_changed.wait(mutex, poll_interval)
+              result_state.nil?
+            end
+            next unless should_refresh
+          end
+          first_poll = false
+
+          safe_refresh ? refresh_data_safe : refresh
+
+          mutex.synchronize do
+            next if result_state
+
+            current = state.to_s
+            result_state = current if target_strings.include?(current) || error_strings.include?(current)
+          end
+
+          poll_interval = next_wait_for_state_poll_interval(
+            current_interval: poll_interval,
+            started_at: started_at,
+            subscribed: subscribed
+          )
         end
 
-        sleep(interval)
-        if ::Process.clock_gettime(::Process::CLOCK_MONOTONIC) - start_time > 5
-          interval = [interval * BACKOFF_MULTIPLIER, MAX_POLL_INTERVAL].min
+        if result_state && error_strings.include?(result_state)
+          raise Sdk::Error,
+                "Sandbox #{id} entered error state: #{result_state}, error reason: #{error_reason}"
+        end
+      rescue Sdk::TimeoutError => e
+        begin
+          safe_refresh ? refresh_data_safe : refresh
+        rescue StandardError
+          raise e
+        end
+        current = state.to_s
+        return if target_strings.include?(current)
+
+        if error_strings.include?(current)
+          raise Sdk::Error,
+                "Sandbox #{id} entered error state: #{current}, error reason: #{error_reason}"
+        end
+
+        raise e
+      ensure
+        @waiter_mutex.synchronize do
+          @state_waiters.delete(waiter)
         end
-        refresh
       end
     end
 
-    INITIAL_POLL_INTERVAL = 0.1
-    private_constant :INITIAL_POLL_INTERVAL
+    def refresh_data_safe
+      refresh
+    rescue DaytonaApiClient::ApiError => e
+      apply_state(DaytonaApiClient::SandboxState::DESTROYED) if e.code == 404
+    rescue StandardError # rubocop:disable Lint/SuppressedException
+    end
 
-    MAX_POLL_INTERVAL = 1.0
-    private_constant :MAX_POLL_INTERVAL
+    def ensure_subscribed
+      @waiter_mutex.synchronize do
+        return unless @subscription_manager
 
-    BACKOFF_MULTIPLIER = 1.1
-    private_constant :BACKOFF_MULTIPLIER
+        if @sub_id
+          return if @subscription_manager.refresh(@sub_id)
 
-    NO_TIMEOUT = 0
-    private_constant :NO_TIMEOUT
-
-    OPERATION_START = :start
-    private_constant :OPERATION_START
+          @sub_id = nil
+        end
 
-    OPERATION_STOP = :stop
-    private_constant :OPERATION_STOP
+        subscribe_to_events
+      end
+    end
 
-    OPERATION_RESIZE = :resize
-    private_constant :OPERATION_RESIZE
+    def subscribe_to_events
+      return if @sub_id || !@subscription_manager
 
-    def wait_for_snapshot_complete
-      interval = INITIAL_POLL_INTERVAL
-      start_time = ::Process.clock_gettime(::Process::CLOCK_MONOTONIC)
-      while state == DaytonaApiClient::SandboxState::SNAPSHOTTING
-        refresh
+      @sub_id = @subscription_manager.subscribe(
+        resource_id: id,
+        handler: method(:handle_event),
+        events: ['sandbox.state.updated', 'sandbox.created']
+      )
+    end
 
-        if [DaytonaApiClient::SandboxState::ERROR, DaytonaApiClient::SandboxState::BUILD_FAILED].include?(state)
-          raise Sdk::Error,
-                "Sandbox #{id} snapshot failed with state: #{state}, error reason: #{error_reason}"
-        end
+    def handle_event(event_name, data)
+      return unless data.is_a?(Hash)
 
-        break if state != DaytonaApiClient::SandboxState::SNAPSHOTTING
+      raw = data['sandbox'] || data
+      return unless raw.is_a?(Hash)
 
-        sleep(interval)
-        if ::Process.clock_gettime(::Process::CLOCK_MONOTONIC) - start_time > 5
-          interval = [interval * BACKOFF_MULTIPLIER, MAX_POLL_INTERVAL].min
-        end
+      if event_name == 'sandbox.created'
+        process_response(DaytonaApiClient::Sandbox.build_from_hash(raw))
+      else
+        new_state = raw['state'] || data['newState']
+        apply_state(new_state) if new_state
       end
+    rescue StandardError
+      nil
     end
 
-    def wait_for_pause_complete
-      interval = INITIAL_POLL_INTERVAL
-      start_time = ::Process.clock_gettime(::Process::CLOCK_MONOTONIC)
-      while state == DaytonaApiClient::SandboxState::PAUSING
-        refresh
+    def unsubscribe_from_events
+      return unless @sub_id
 
-        if [DaytonaApiClient::SandboxState::ERROR, DaytonaApiClient::SandboxState::BUILD_FAILED].include?(state)
-          raise Sdk::Error,
-                "Sandbox #{id} pause failed with state: #{state}, error reason: #{error_reason}"
-        end
+      @subscription_manager&.unsubscribe(@sub_id)
+      @sub_id = nil
+    end
 
-        break if state != DaytonaApiClient::SandboxState::PAUSING
+    def next_wait_for_state_poll_interval(current_interval:, started_at:, subscribed:)
+      return [current_interval * BACKOFF_MULTIPLIER, MAX_POLL_INTERVAL].min if subscribed
 
-        sleep(interval)
-        if ::Process.clock_gettime(::Process::CLOCK_MONOTONIC) - start_time > 5
-          interval = [interval * BACKOFF_MULTIPLIER, MAX_POLL_INTERVAL].min
-        end
-      end
+      elapsed = ::Process.clock_gettime(::Process::CLOCK_MONOTONIC) - started_at
+      return current_interval if elapsed <= POLL_BACKOFF_DELAY
+
+      [current_interval * BACKOFF_MULTIPLIER, MAX_POLL_INTERVAL].min
+    end
+
+    INITIAL_POLL_INTERVAL = 0.1
+    private_constant :INITIAL_POLL_INTERVAL
+
+    MAX_POLL_INTERVAL = 1.0
+    private_constant :MAX_POLL_INTERVAL
+
+    BACKOFF_MULTIPLIER = 1.1
+    private_constant :BACKOFF_MULTIPLIER
+
+    POLL_BACKOFF_DELAY = 5.0
+    private_constant :POLL_BACKOFF_DELAY
+
+    NO_TIMEOUT = 0
+    private_constant :NO_TIMEOUT
+
+    def wait_for_snapshot_complete
+      error_states = [DaytonaApiClient::SandboxState::ERROR, DaytonaApiClient::SandboxState::BUILD_FAILED]
+      exclude = [DaytonaApiClient::SandboxState::SNAPSHOTTING] + error_states
+      target_states = DaytonaApiClient::SandboxState.all_vars.reject { |s| exclude.include?(s) }
+
+      wait_for_state(target_states:, error_states:)
     end
   end
 end