antigravity-sdk 0.3.0 → 0.4.2

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 665ab001fa484b0bc5ba97c2e7feb960712b0f8e5070ed521f0caf5ad0f16a58
-  data.tar.gz: 3f556321b1083daadea10792632ac216ba168abb26708fc5cc3a81824301d460
+  metadata.gz: b2a5cad1edda5eb5efaa58d7d5104ba433f8469a923bb5f862b9f1bcb4219021
+  data.tar.gz: 75c9924910ef8aab407444abfecfa0f7e989473d8be251691e4e07c9d04499aa
 SHA512:
-  metadata.gz: d60ff58a66fec4ed17a85ad51015568e743363f313eb1a27b9e6512a767ae5a7abca062796ae0bc19edbe59545202828415b693f6551b77a22d8a1ff5cca99ea
-  data.tar.gz: 9e09ed8b23e59d29f644d10f421b73d70f357c7d9fd601ba67ba429df8c2814768d0264084e5e6439cd9161d2f8c13313f6890bcd11faf6f66f7e43e3ea2578f
+  metadata.gz: ac637d9792051869f94b4a8ebc716fd604b2db18886e19b105190a9636c7e7bce418aeece52a360b139b1c534e560605b0dc578fab10b3e06ac913dc454663eb
+  data.tar.gz: c8b6e4bee541853abd1f993d5423cc3426c830a24aba58f6790ec30e093279163b83e258a824811380ac5a603cba749dd7a099dee5ad4a2160fe25a7373c6e51

data/lib/antigravity/agent.rb

@@ -4,27 +4,122 @@ module Antigravity
   class Agent < Base
 
     attr_accessor :model, :system_instruction, :api_key
-    attr_reader :tools, :skills, :hooks, :sidecars, :client, :logger_guard
+    attr_reader :tools, :skills, :hooks, :sidecars, :client, :logger_guard,
+                :workspace, :connection, :conversation
 
-    def initialize(model: nil, auto_logger: true, &block)
+    def initialize(model: nil, system_instruction: nil, tools: [],
+                   skills: [], workspace: nil, auto_logger: true, log_file: nil, &block)
       @model = model || Antigravity.config.default_model
       @api_key = Antigravity.config.api_key
-      @system_instruction = nil
-      @tools = []
+      @system_instruction = system_instruction
+      @workspace = workspace ? File.expand_path(workspace) : nil
+      @tools = tools.dup
       @skills = []
       @sidecars = []
       @hooks = Hooks.new
       @client = Client.new
       @logger_guard = nil
+      @connection = nil
+      @conversation = nil
+      @connected = false
+
+      # Register pre-provided tools into the tool runner
+      @tool_runner = ToolRunner.new
+      @tools.each { |t| @tool_runner.register(t) }
+
+      # Load skills provided at construction (local paths or GitHub URLs)
+      add_skills(skills) unless Array(skills).empty?
 
       # Automagic Logger attachment unless disabled via ENV["ANTIGRAVITY_LOGGER"]=false or auto_logger: false
       if auto_logger && logger_enabled?
-        attach_logger
+        attach_logger(log_file)
       end
 
       yield(self) if block_given?
     end
 
+    # --- Class Methods ---
+
+    # Block form: opens connection, yields agent, auto-closes.
+    def self.open(**kwargs, &block)
+      agent = new(**kwargs)
+      agent.connect!
+      begin
+        block.call(agent)
+      ensure
+        agent.close!
+      end
+    end
+
+    # --- Connection Lifecycle ---
+
+    def connect!
+      return self if @connected
+
+      @connection = Connection::LocalConnection.new
+      @connection.connect!
+
+      @conversation = Conversation.new(
+        ws_client: @connection.ws_client,
+        tool_runner: @tool_runner,
+        hooks: @hooks
+      )
+
+      harness_config = build_harness_config
+      @conversation.initialize_session!(harness_config: harness_config)
+      @connected = true
+      self
+    end
+
+    def connected?
+      @connected && @connection&.connected?
+    end
+
+    def close!
+      @connected = false
+      @connection&.disconnect!
+      @connection = nil
+      @conversation = nil
+    end
+
+    # --- Chat ---
+
+    def prompt(message, timeout: Antigravity.config.timeout_llm, &block)
+      emit_sidecar_event(:prompt_started, prompt: message)
+      hooks.run_pre_prompt(message)
+
+      if @connected && @conversation
+        response = @conversation.chat(message, timeout: timeout, &block)
+      else
+        # Legacy mock-client path (unit tests, pre-connection)
+        response = client.send_turn(self, message, &block)
+      end
+
+      hooks.run_post_response(response)
+      emit_sidecar_event(:turn_completed, response: response.content, model: model)
+
+      response
+    end
+    alias ask prompt
+
+    # --- Metadata Accessors (mirrors Python SDK) ---
+
+    def conversation_id
+      @conversation&.conversation_id
+    end
+
+    def turn_count
+      @conversation&.turn_count || 0
+    end
+
+    def session_summary
+      return {} unless @conversation
+
+      @conversation.session_summary(model: @model)
+    end
+
+    # --- Tool Registration ---
+
     def register_tool(tool_or_name = nil, description: "", &block)
       if block_given? && tool_or_name
         tool = Tool::Dynamic.new(tool_or_name, description: description, &block)
@@ -34,6 +129,7 @@ module Antigravity
         raise ArgumentError, "Invalid tool definition"
       end
       @tools << tool
+      @tool_runner.register(tool) if @tool_runner
       tool
     end
 
@@ -42,18 +138,61 @@ module Antigravity
       sidecar
     end
 
-    def attach_logger(log_target = nil, level: ::Logger::INFO, silent_notice: false)
+    def attach_logger(log_target = nil, level: :info, silent_notice: false)
       @logger_guard = Guards::AgentLogger.new(log_target, level: level, silent_notice: silent_notice)
       @logger_guard.attach_to(self)
       @logger_guard
     end
 
-    def load_skill(skill_path)
-      skill = Skill.load(skill_path)
-      @skills << skill
+    # Add a single skill by path or GitHub URL.
+    # Raises if the path resolves to multiple skills (use add_skills instead).
+    # @param path_or_url [String] local path or GitHub URL
+    # @param skill_name [String, nil] optional specific skill name within a repo
+    # @return [Skill] the loaded skill
+    def add_skill(path_or_url, skill_name: nil)
+      target = skill_name ? "#{path_or_url.to_s.chomp('/')}/#{skill_name}" : path_or_url.to_s
+      paths = SkillResolver.resolve(target)
+      if paths.size > 1
+        raise ArgumentError,
+              "add_skill resolved to #{paths.size} skills. Use add_skills instead, " \
+              "or specify skill_name: to pick one."
+      end
+      raise ArgumentError, "No skill found at #{target}" if paths.empty?
+
+      load_single_skill(paths.first)
+    end
+
+    # Add one or more skills by path or GitHub URL.
+    # Accepts a single string or an array. Each entry is resolved (may expand to multiple).
+    # @param paths_or_urls [String, Array<String>] local paths or GitHub URLs
+    # @return [Array<Skill>] all loaded skills
+    def add_skills(paths_or_urls)
+      Array(paths_or_urls).flat_map do |p|
+        SkillResolver.resolve(p).map { |skill_path| load_single_skill(skill_path) }
+      end
+    end
+
+    # Create and add an inline skill (no file needed).
+    # @param name [String] skill name
+    # @param description [String] what the skill does
+    # @param instructions [String] the skill body (markdown)
+    # @return [Skill] the inline skill
+    def add_inline_skill(name:, description:, instructions:)
+      skill = Skill.inline(name: name, description: description, instructions: instructions)
+      @skills << skill unless @skills.any? { |s| s.name == skill.name }
       skill
     end
 
+    # List discovered skills in a container path (without loading them).
+    # @param path_or_url [String] local path or GitHub URL
+    # @return [Array<String>] skill directory paths
+    def self.list_skills(path_or_url)
+      SkillResolver.resolve(path_or_url)
+    end
+
+    # Legacy alias
+    alias_method :load_skill, :add_skill
+
     def before_prompt(&block)
       hooks.before_prompt(&block)
     end
@@ -75,19 +214,6 @@ module Antigravity
       @sidecars.each { |sidecar| sidecar.emit(event_type, payload) }
     end
 
-    def prompt(message, &block)
-      emit_sidecar_event(:prompt_started, prompt: message)
-      hooks.run_pre_prompt(message)
-
-      response = client.send_turn(self, message, &block)
-
-      hooks.run_post_response(response)
-      emit_sidecar_event(:turn_completed, response: response.content, model: model)
-
-      response
-    end
-    alias ask prompt
-
     private
 
     def logger_enabled?
@@ -96,5 +222,91 @@ module Antigravity
 
       true
     end
+
+    def build_harness_config
+      api_key = ENV.fetch('GEMINI_API_KEY') {
+        raise ConfigError, 'GEMINI_API_KEY environment variable is required'
+      }
+
+      config = {
+        config: {
+          models: [
+            {
+              name: @model,
+              geminiApiEndpoint: {
+                apiKey: api_key
+              }
+            }
+          ]
+        }
+      }
+
+      # Add workspaces if specified
+      if @workspace
+        expanded = File.expand_path(@workspace)
+        $stderr.puts "\u231B Indexing workspace: #{expanded} — this may take a moment..."
+        config[:config][:workspaces] = [
+          {
+            filesystemWorkspace: {
+              directory: expanded
+            }
+          }
+        ]
+      end
+
+      # Add system instructions if specified (protobuf: SystemInstructions.custom.part[])
+      effective_instructions = @system_instruction || ''
+
+      # Auto-append workspace tool hints — models (esp. flash-lite) won't use tools unless told
+      if @workspace && !effective_instructions.match?(/list_dir|view_file|tools/i)
+        tool_hint = 'You have access to the workspace filesystem. Use the available tools (list_dir, view_file, grep_search) to explore it.'
+        effective_instructions = effective_instructions.empty? ? tool_hint : "#{effective_instructions}\n#{tool_hint}"
+      end
+
+      unless effective_instructions.empty?
+        config[:config][:systemInstructions] = {
+          custom: {
+            part: [{ text: effective_instructions }]
+          }
+        }
+      end
+
+      # Add custom tools
+      if @tool_runner && !@tool_runner.empty?
+        config[:config][:tools] = @tool_runner.to_harness_tools
+      end
+
+      # Enable harness-side built-in tools by default (list_dir, view_file, grep_search, etc.)
+      config[:config][:harnessSideTools] = {
+        listDir: { enabled: true },
+        viewFile: { enabled: true },
+        grepSearch: { enabled: true },
+        find: { enabled: true },
+        writeToFile: { enabled: true },
+        fileEdit: { enabled: true },
+        readUrlContent: { enabled: true },
+        searchWeb: { enabled: true }
+      }
+
+      # Wire skills paths for harness (proto field: skills_paths)
+      unless @skills.empty?
+        skill_paths = @skills.select(&:path).map(&:path)
+        config[:config][:skillsPaths] = skill_paths unless skill_paths.empty?
+      end
+
+      # App data dir
+      config[:config][:appDataDir] = File.expand_path('~/.gemini/antigravity')
+
+      config
+    end
+
+    def load_single_skill(skill_path)
+      # Dedup by path
+      return @skills.find { |s| s.path == skill_path } if @skills.any? { |s| s.path == skill_path }
+
+      skill = Skill.load(skill_path)
+      @skills << skill
+      skill
+    end
   end
 end

data/lib/antigravity/config.rb

@@ -4,11 +4,22 @@ module Antigravity
   class Config
     attr_accessor :api_key, :default_model, :harness_path, :log_level
 
+    # Timeouts (seconds) — aggressive for dev, override via ENV for prod.
+    #   ANTIGRAVITY_TIMEOUT_LLM       — per-message wait for LLM response (default: 20s)
+    #   ANTIGRAVITY_TIMEOUT_WS        — WebSocket connect/handshake (default: 3s)
+    #   ANTIGRAVITY_TIMEOUT_HANDSHAKE — stdio binary handshake (default: 5s)
+    attr_accessor :timeout_llm, :timeout_ws, :timeout_handshake
+
     def initialize
       @api_key = ENV["GEMINI_API_KEY"]
-      @default_model = "gemini-flash-latest"
+      @default_model = ENV["GEMINI_MODEL"] || ENV["ANTIGRAVITY_MODEL"] || "gemini-3.6-flash"
       @harness_path = ENV["ANTIGRAVITY_HARNESS_PATH"] || File.expand_path("~/.antigravity/bin/localharness")
       @log_level = :info
+
+      # Timeouts: aggressive for dev, relax via ENV for production
+      @timeout_llm       = (ENV["ANTIGRAVITY_TIMEOUT_LLM"]       || 40).to_i
+      @timeout_ws        = (ENV["ANTIGRAVITY_TIMEOUT_WS"]        || 3).to_i
+      @timeout_handshake = (ENV["ANTIGRAVITY_TIMEOUT_HANDSHAKE"] || 5).to_i
     end
 
     def api_key?

data/lib/antigravity/connection/binary_fetcher.rb

@@ -0,0 +1,163 @@
+# frozen_string_literal: true
+
+require 'open3'
+require 'fileutils'
+require 'tmpdir'
+require 'net/http'
+require 'json'
+require 'uri'
+
+module Antigravity
+  module Connection
+    # Downloads and extracts the localharness binary from the official
+    # google-antigravity PyPI wheel. Used when the binary isn't already
+    # installed via Antigravity.app or a manual setup.
+    module BinaryFetcher
+      PYPI_PACKAGE = 'google-antigravity'
+      INSTALL_DIR  = File.expand_path('~/.antigravity/bin')
+      BINARY_NAME  = 'localharness'
+
+      # Platform mapping: Ruby's RUBY_PLATFORM -> PyPI wheel platform tag
+      PLATFORM_MAP = {
+        /darwin.*arm/i   => 'macosx_11_0_arm64',
+        /darwin.*x86/i   => 'macosx_10_15_x86_64',
+        /darwin/i        => 'macosx_11_0_arm64',    # default macOS = ARM
+        /linux.*x86_64/i => 'manylinux2014_x86_64',
+        /linux.*aarch/i  => 'manylinux2014_aarch64',
+      }.freeze
+
+      class << self
+        # Fetch and install the binary. Returns the installed path.
+        # @param quiet [Boolean] suppress progress output
+        # @return [String] path to the installed binary
+        def fetch!(quiet: false)
+          platform = detect_platform
+          say("🔍 Detecting platform: #{platform}", quiet: quiet)
+
+          # Step 1: Find the wheel URL from PyPI
+          say("📡 Querying PyPI for #{PYPI_PACKAGE}...", quiet: quiet)
+          wheel_url = find_wheel_url(platform)
+
+          # Step 2: Download the wheel
+          say("⏳ Downloading binary... this will take some time (~50MB wheel)", quiet: quiet)
+          wheel_path = download_wheel(wheel_url, quiet: quiet)
+
+          # Step 3: Extract the binary
+          say("📦 Extracting localharness binary...", quiet: quiet)
+          binary_path = extract_binary(wheel_path)
+
+          # Step 4: Make executable
+          FileUtils.chmod(0o755, binary_path)
+          say("✅ Installed localharness at #{binary_path}", quiet: quiet)
+
+          binary_path
+        ensure
+          # Clean up temp wheel
+          FileUtils.rm_f(wheel_path) if wheel_path && File.exist?(wheel_path.to_s)
+        end
+
+        # Check if the binary is already installed via fetch
+        def installed?
+          path = File.join(INSTALL_DIR, BINARY_NAME)
+          File.executable?(path)
+        end
+
+        def installed_path
+          File.join(INSTALL_DIR, BINARY_NAME)
+        end
+
+        private
+
+        def detect_platform
+          PLATFORM_MAP.each do |pattern, tag|
+            return tag if RUBY_PLATFORM.match?(pattern)
+          end
+          raise HarnessNotFoundError,
+            "Unsupported platform: #{RUBY_PLATFORM}. Cannot auto-download localharness."
+        end
+
+        def find_wheel_url(platform)
+          uri = URI("https://pypi.org/pypi/#{PYPI_PACKAGE}/json")
+          response = Net::HTTP.get(uri)
+          data = JSON.parse(response)
+
+          # Find latest version's wheel for our platform
+          urls = data['urls'] || []
+          wheel = urls.find { |u| u['filename']&.include?(platform) && u['filename']&.end_with?('.whl') }
+
+          unless wheel
+            # Try all versions for the platform
+            versions = data['releases']&.keys&.sort_by { |v| Gem::Version.new(v) rescue v }&.reverse
+            versions&.each do |ver|
+              files = data.dig('releases', ver) || []
+              wheel = files.find { |u| u['filename']&.include?(platform) && u['filename']&.end_with?('.whl') }
+              break if wheel
+            end
+          end
+
+          raise HarnessNotFoundError,
+            "No wheel found for platform #{platform} on PyPI. Install Antigravity.app manually." unless wheel
+
+          wheel['url']
+        end
+
+        def download_wheel(url, quiet: false)
+          dest = File.join(Dir.tmpdir, "antigravity-wheel-#{$$}.whl")
+          uri = URI(url)
+
+          Net::HTTP.start(uri.host, uri.port, use_ssl: uri.scheme == 'https') do |http|
+            request = Net::HTTP::Get.new(uri)
+            http.request(request) do |response|
+              total = response['content-length']&.to_i
+              downloaded = 0
+
+              File.open(dest, 'wb') do |f|
+                response.read_body do |chunk|
+                  f.write(chunk)
+                  downloaded += chunk.bytesize
+                  if total && total > 0 && !quiet
+                    pct = (downloaded * 100.0 / total).round(1)
+                    print "\r  ⏳ #{pct}% (#{(downloaded / 1024.0 / 1024).round(1)} MB / #{(total / 1024.0 / 1024).round(1)} MB)"
+                  end
+                end
+              end
+              puts unless quiet
+            end
+          end
+
+          dest
+        end
+
+        def extract_binary(wheel_path)
+          FileUtils.mkdir_p(INSTALL_DIR)
+
+          # Wheels are just ZIP files. Extract the binary.
+          # Look for: google/antigravity/bin/localharness (or language_server)
+          extract_dir = Dir.mktmpdir('agy-extract-')
+          system('unzip', '-q', '-o', wheel_path, '-d', extract_dir)
+
+          # Search for the binary inside (prefer localharness over language_server)
+          candidates = Dir.glob("#{extract_dir}/**/localharness") +
+                       Dir.glob("#{extract_dir}/**/language_server")
+
+          binary = candidates.find { |f| File.file?(f) && !File.directory?(f) }
+
+          unless binary
+            FileUtils.rm_rf(extract_dir)
+            raise HarnessNotFoundError,
+              "Could not find localharness binary inside the wheel. Contents: #{Dir.glob("#{extract_dir}/**/*").join(', ')}"
+          end
+
+          dest_path = File.join(INSTALL_DIR, BINARY_NAME)
+          FileUtils.cp(binary, dest_path)
+          FileUtils.rm_rf(extract_dir)
+          dest_path
+        end
+
+        def say(msg, quiet: false)
+          $stderr.puts msg unless quiet
+        end
+      end
+    end
+  end
+end

data/lib/antigravity/connection/local_connection.rb

@@ -0,0 +1,184 @@
+# frozen_string_literal: true
+
+require 'open3'
+require 'fileutils'
+require 'tmpdir'
+
+module Antigravity
+  module Connection
+    # Manages the full lifecycle of communicating with the localharness Go binary.
+    #
+    # Lifecycle: discover → spawn → stdio handshake → WebSocket → events → shutdown
+    #
+    # The localharness binary is the core Go engine that powers Antigravity.
+    # It handles model calls, tool orchestration, and agent logic.
+    # This class is the Ruby adapter that speaks its protocol.
+    class LocalConnection < Base
+      # Known locations for the localharness binary (checked in order).
+      # IMPORTANT: The standalone `localharness` binary (from PyPI wheel) uses
+      # the stdin/stdout protobuf handshake. The `language_server` binary
+      # (Antigravity.app) does NOT — it's the IDE server mode.
+      BINARY_SEARCH_PATHS = [
+        File.expand_path('~/.antigravity/bin/localharness'),
+        File.expand_path('~/.local/bin/localharness'),
+      ].freeze
+
+      HARNESS_SUBCOMMAND = 'localharness'
+
+      attr_reader :port, :api_key, :pid
+
+      def initialize(binary_path: nil)
+        @binary_path = binary_path || self.class.find_binary!
+        @port = nil
+        @api_key = nil
+        @pid = nil
+        @stdin = nil
+        @stdout = nil
+        @stderr = nil
+        @connected = false
+      end
+
+      # --- Binary Discovery ---
+
+      # Finds the localharness binary. Search order:
+      #   1. ANTIGRAVITY_HARNESS_PATH env var
+      #   2. Antigravity.app (macOS)
+      #   3. ~/.antigravity/bin/
+      #   4. PATH lookup via `which`
+      #   5. Auto-download from PyPI (unless ANTIGRAVITY_AUTO_DOWNLOAD=false)
+      #
+      # @return [String] absolute path to the binary
+      # @raise [HarnessNotFoundError] if not found anywhere
+      def self.find_binary!
+        # 1. Env var override
+        if (env_path = ENV['ANTIGRAVITY_HARNESS_PATH'])
+          return env_path if File.executable?(env_path)
+
+          raise HarnessNotFoundError,
+            "ANTIGRAVITY_HARNESS_PATH=#{env_path} is not executable"
+        end
+
+        # 2. Known paths
+        BINARY_SEARCH_PATHS.each do |path|
+          return path if File.executable?(path)
+        end
+
+        # 3. PATH lookup
+        which_result = `which localharness 2>/dev/null`.strip
+        return which_result unless which_result.empty?
+
+        # 4. Auto-download from PyPI wheel (opt-out with ANTIGRAVITY_AUTO_DOWNLOAD=false)
+        auto_dl = ENV.fetch('ANTIGRAVITY_AUTO_DOWNLOAD', 'true').downcase
+        unless %w[false 0 no].include?(auto_dl)
+          $stderr.puts "⏳ localharness not found locally. Downloading from PyPI... this may take a minute."
+          return BinaryFetcher.fetch!
+        end
+
+        raise HarnessNotFoundError,
+          "Could not find localharness binary. Install Antigravity.app, set ANTIGRAVITY_HARNESS_PATH, or allow auto-download."
+      end
+
+      # --- Connection Lifecycle ---
+
+      def connect!
+        spawn_process!
+        perform_handshake!
+        connect_websocket!
+        @connected = true
+        self
+      end
+
+      def connected?
+        @connected && process_alive? && @ws_client&.connected?
+      end
+
+      def disconnect!
+        @connected = false
+        @ws_client&.close
+        kill_process!
+      end
+
+      # Expose the WebSocket client for Conversation to use
+      def ws_client
+        @ws_client
+      end
+
+      private
+
+      def spawn_process!
+        # Only pass 'localharness' subcommand when binary is language_server
+        # (Antigravity.app). The standalone localharness binary needs no subcommand.
+        cmd = if @binary_path.end_with?('language_server')
+                [@binary_path, HARNESS_SUBCOMMAND]
+              else
+                [@binary_path]
+              end
+
+        @stdin, @stdout, @stderr, @wait_thread = Open3.popen3(*cmd)
+        @pid = @wait_thread.pid
+      rescue Errno::ENOENT => e
+        raise HarnessNotFoundError, "Failed to spawn localharness: #{e.message}"
+      end
+
+      def perform_handshake!
+        storage_dir = File.join(Dir.tmpdir, "antigravity-ruby-#{$$}")
+        FileUtils.mkdir_p(storage_dir)
+
+        # Send InputConfig via stdin
+        input_config = Protocol.encode_input_config(
+          storage_directory: storage_dir,
+          bind_address: 'localhost'
+        )
+        @stdin.write(input_config)
+        @stdin.flush
+
+        # Read OutputConfig from stdout
+        frame = Protocol.read_length_prefixed(@stdout, timeout: Antigravity.config.timeout_handshake)
+        output = Protocol.decode_output_config(frame)
+        @port = output[:port]
+        @api_key = output[:api_key]
+
+        raise HarnessHandshakeError, "Harness returned port=0" if @port == 0
+      rescue IOError, Errno::EPIPE, ProtocolError => e
+        # Capture stderr for diagnostics
+        stderr_output = begin
+          @stderr&.read_nonblock(4096)
+        rescue EOFError, IOError, Errno::EAGAIN
+          nil
+        end
+        kill_process!
+        detail = stderr_output ? " Stderr: #{stderr_output.strip}" : ''
+        raise HarnessHandshakeError, "Stdio handshake failed: #{e.message}#{detail}"
+      end
+
+      def connect_websocket!
+        @ws_client = WebSocketClient.new(port: @port, api_key: @api_key)
+        @ws_client.connect!
+      rescue => e
+        kill_process!
+        raise ProtocolError, "WebSocket connection failed: #{e.message}"
+      end
+
+      def process_alive?
+        return false unless @pid
+
+        Process.kill(0, @pid)
+        true
+      rescue Errno::ESRCH, Errno::EPERM
+        false
+      end
+
+      def kill_process!
+        [@stdin, @stdout, @stderr].each { |io| io&.close rescue nil }
+        Process.kill('TERM', @pid) if @pid && process_alive?
+        @wait_thread&.join(5)
+        Process.kill('KILL', @pid) if @pid && process_alive?
+      rescue Errno::ESRCH, Errno::EPERM
+        # Already dead, fine
+      ensure
+        @pid = nil
+        @stdin = @stdout = @stderr = nil
+      end
+    end
+  end
+end