antigravity-sdk 0.2.0 → 0.4.2

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

data/lib/antigravity/connection/websocket_client.rb

@@ -0,0 +1,159 @@
+# frozen_string_literal: true
+
+require 'socket'
+require 'websocket'
+require 'json'
+require 'securerandom'
+
+module Antigravity
+  module Connection
+    # Lightweight WebSocket client for the localharness.
+    # Uses the `websocket` gem for frame encoding/decoding over raw TCPSocket.
+    # No EventMachine, no threads-by-default — just blocking IO.
+    class WebSocketClient
+      attr_reader :port, :api_key, :connected
+
+      def initialize(port:, api_key:)
+        @port = port
+        @api_key = api_key
+        @socket = nil
+        @handshake = nil
+        @connected = false
+        @frame_buffer = WebSocket::Frame::Incoming::Client.new
+      end
+
+      # Open the WebSocket connection to ws://localhost:<port>/
+      def connect!
+        @socket = TCPSocket.new('127.0.0.1', @port)
+
+        # Build and send the HTTP upgrade handshake
+        @handshake = WebSocket::Handshake::Client.new(
+          url: "ws://127.0.0.1:#{@port}/",
+          headers: { 'x-goog-api-key' => @api_key }
+        )
+        @socket.write(@handshake.to_s)
+        @socket.flush
+
+        # Read the server's handshake response
+        loop do
+          line = @socket.gets
+          raise ProtocolError, 'EOF during WebSocket handshake' unless line
+
+          @handshake << line
+          break if @handshake.finished?
+        end
+
+        unless @handshake.valid?
+          raise ProtocolError, "WebSocket handshake rejected: #{@handshake.error}"
+        end
+
+        @connected = true
+        self
+      end
+
+      def connected?
+        @connected && @socket && !@socket.closed?
+      end
+
+      # Send a JSON message as a WebSocket text frame
+      def send_json(data)
+        json = data.is_a?(String) ? data : JSON.generate(data)
+        frame = WebSocket::Frame::Outgoing::Client.new(
+          data: json,
+          type: :text,
+          version: @handshake.version
+        )
+        @socket.write(frame.to_s)
+        @socket.flush
+      end
+
+      # Read the next JSON message. Blocks until a text frame arrives.
+      # Yields each message if a block is given (for streaming).
+      # Returns nil on connection close or idle_timeout.
+      def receive_json(timeout: Antigravity.config.timeout_llm, idle_timeout: nil)
+        deadline = Process.clock_gettime(Process::CLOCK_MONOTONIC) + timeout
+        last_activity = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+
+        loop do
+          now = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+          remaining = deadline - now
+          raise ProtocolError, 'WebSocket read timeout' if remaining <= 0
+
+          if idle_timeout && (now - last_activity) >= idle_timeout
+            return nil
+          end
+
+          select_time = [remaining, 0.5].min
+          if idle_timeout
+            idle_rem = idle_timeout - (now - last_activity)
+            select_time = [select_time, idle_rem].min if idle_rem > 0
+          end
+
+          ready = IO.select([@socket], nil, nil, [select_time, 0.05].max)
+          next unless ready
+
+          data = @socket.read_nonblock(16384, exception: false)
+          case data
+          when :wait_readable then next
+          when nil
+            @connected = false
+            return nil
+          end
+
+          last_activity = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+          @frame_buffer << data
+
+          while (frame = @frame_buffer.next)
+            case frame.type
+            when :text
+              return JSON.parse(frame.data, symbolize_names: true)
+            when :ping
+              pong = WebSocket::Frame::Outgoing::Client.new(
+                data: frame.data, type: :pong, version: @handshake.version
+              )
+              @socket.write(pong.to_s)
+            when :close
+              @connected = false
+              return nil
+            end
+          end
+        end
+      end
+
+      # Read messages in a loop, yielding each parsed JSON.
+      # Stops when block returns :stop, connection closes, or timeout.
+      def each_message(timeout: Antigravity.config.timeout_llm, &block)
+        loop do
+          # Block can return [:stop] or [:idle_timeout, seconds]
+          msg = receive_json(timeout: timeout, idle_timeout: @current_idle_timeout)
+          break unless msg
+
+          result = block.call(msg)
+          if result.is_a?(Array) && result.first == :idle_timeout
+            @current_idle_timeout = result.last
+          elsif result == :stop
+            @current_idle_timeout = nil
+            break
+          end
+        end
+      ensure
+        @current_idle_timeout = nil
+      end
+
+      def close
+        return unless @socket && !@socket.closed?
+
+        begin
+          close_frame = WebSocket::Frame::Outgoing::Client.new(
+            data: '', type: :close, version: @handshake&.version || 13
+          )
+          @socket.write(close_frame.to_s)
+        rescue IOError, Errno::EPIPE
+          # Already closed
+        end
+        @socket.close rescue nil
+        @connected = false
+      end
+    end
+  end
+end

data/lib/antigravity/conversation.rb

@@ -0,0 +1,302 @@
+# frozen_string_literal: true
+
+require 'json'
+require 'securerandom'
+
+module Antigravity
+  # Manages a multi-turn conversation with the harness over WebSocket.
+  # Handles the event loop: send InputEvent, receive OutputEvents, dispatch tools.
+  #
+  # Mirrors Python SDK's Conversation class:
+  #   - history, turn_count, conversation_id
+  #   - last_turn_usage, total_usage
+  class Conversation
+    attr_reader :conversation_id, :history, :turn_count, :total_usage, :last_turn_usage
+
+    def initialize(ws_client:, tool_runner: nil, hooks: nil)
+      @ws = ws_client
+      @tool_runner = tool_runner || ToolRunner.new
+      @hooks = hooks
+      @conversation_id = nil
+      @history = []
+      @turn_count = 0
+      @total_usage = empty_usage
+      @last_turn_usage = nil
+      @initialized = false
+    end
+
+    # Initialize the conversation session with the harness.
+    # Sends InitializeConversationEvent (protobuf JSON format), receives response.
+    def initialize_session!(harness_config:)
+      # harness_config IS the InitializeConversationEvent JSON:
+      # { config: { models: [...], workspaces: [...], ... } }
+      @ws.send_json(harness_config)
+
+      # Wait for InitializeConversationResponse
+      msg = @ws.receive_json(timeout: 30)
+      raise ProtocolError, 'No response to InitializeConversationEvent' unless msg
+
+      if (init_resp = msg[:initializeConversationResponse])
+        @conversation_id = init_resp[:cascadeId]
+        @initialized = true
+      else
+        # Some harness versions may wrap differently
+        @conversation_id = msg[:cascadeId] || msg.dig(:config, :cascadeId) || SecureRandom.uuid
+        @initialized = true
+      end
+
+      self
+    end
+
+    # Send a user message and collect the full response.
+    # Yields streaming chunks if a block is given.
+    #
+    # @param prompt [String] user message
+    # @return [Message] the complete response
+    def chat(prompt, timeout: Antigravity.config.timeout_llm, &block)
+      raise ProtocolError, 'Session not initialized' unless @initialized
+
+      @turn_count += 1
+      @last_turn_usage = empty_usage
+
+      # Send user input (protobuf InputEvent with user_input string field)
+      input_event = {
+        userInput: prompt
+      }
+      @ws.send_json(input_event)
+
+      # Collect response
+      collect_response(timeout: timeout, &block)
+    end
+
+    def initialized?
+      @initialized
+    end
+
+    # Generate a session summary hash (mirrors Python's metadata)
+    def session_summary(model: nil)
+      {
+        conversation_id: @conversation_id,
+        turn_count: @turn_count,
+        total_tokens: @total_usage[:total_token_count],
+        prompt_tokens: @total_usage[:prompt_token_count],
+        candidates_tokens: @total_usage[:candidates_token_count],
+        model: model
+      }
+    end
+
+    private
+
+    def collect_response(timeout: Antigravity.config.timeout_llm, &block)
+      text_parts = []
+      thinking_parts = []
+      steps = []
+      tool_calls_count = 0
+      finished = false
+      finished_at = nil
+
+      @ws.each_message(timeout: timeout) do |msg|
+        @hooks&.emit(:ws_message, msg)
+
+        if (step = msg[:stepUpdate])
+          step_record = parse_step(step)
+          steps << step_record
+
+          # Text delta — stream it (only from model aimed at user, not tool descriptions or error steps)
+          is_model_step = step[:source].to_s =~ /MODEL|model|3/
+          is_target_user = step[:target].to_s =~ /USER|user|1/
+          is_error_step = step[:state].to_s =~ /ERROR|error|4/ || step[:errorMessage]
+
+          if step[:textDelta] && !step[:textDelta].empty? && is_model_step && is_target_user && !is_error_step
+            text_parts << step[:textDelta]
+            chunk = Message.new(
+              content: step[:textDelta],
+              role: :assistant,
+              delta: true
+            )
+            block&.call(chunk)
+          end
+
+          # Thinking delta
+          if step[:thinkingDelta] && !step[:thinkingDelta].empty?
+            thinking_parts << step[:thinkingDelta]
+          end
+
+          # Custom tool action
+          if step[:customTool]
+            tool_calls_count += 1
+            handle_custom_tool(step)
+          end
+
+          # Harness built-in tool actions count
+          if step_record[:target] == :environment && step_record[:source] == :model
+            tool_calls_count += 1 unless step[:customTool]
+          end
+
+          # Finished? Model response to user with DONE state AND text collected.
+          # NOTE: DONE can arrive before text deltas, so we require text here.
+          # For tool-only turns (no text), FULLY_IDLE below is the authoritative stop.
+          if is_model_step && is_target_user && !is_error_step && step[:state] && step[:state].to_s =~ /DONE|done|2/ && !text_parts.empty?
+            finished = true
+          end
+        end
+
+        # Top-level tool call (custom tools are sent as separate messages, not in stepUpdate)
+        if (tool_call = msg[:toolCall])
+          tool_calls_count += 1
+          handle_tool_call(tool_call)
+        end
+
+        # Usage update
+        if (usage = msg[:usageUpdate])
+          update_usage(usage)
+        end
+
+        # Trajectory state: STATE_FULLY_IDLE / STATE_CANCELLED = turn complete (authoritative signal from harness)
+        # This is the MOST authoritative signal — always stop, even with empty text.
+        if (traj = msg[:trajectoryStateUpdate])
+          if traj[:state].to_s =~ /FULLY_IDLE|CANCELLED/
+            finished = true
+            finished_at = Time.now
+          end
+        end
+
+        # Stop conditions (in priority order):
+        # 1. Session end — always stop
+        # 2. Model response DONE or trajectory FULLY_IDLE / CANCELLED — turn complete
+        if msg.key?(:sessionEndResponse) || finished
+          :stop
+        elsif !text_parts.empty?
+          # If assistant has sent text response, allow 3s idle timeout for trailing metadata
+          [:idle_timeout, 3.0]
+        end
+      end
+
+      # Build final message
+      Message.new(
+        content: text_parts.join,
+        role: :assistant,
+        thinking: thinking_parts.join,
+        steps: steps,
+        tool_calls_count: tool_calls_count,
+        usage: @last_turn_usage.dup
+      )
+    end
+
+    def handle_custom_tool(step)
+      tool_data = step[:customTool]
+      tool_call = tool_data[:toolCall] || tool_data
+      handle_tool_call(tool_call)
+    end
+
+    # Handle a top-level toolCall message from the harness
+    # Format: {id: "...", name: "tool_name", argumentsJson: "{...}"}
+    def handle_tool_call(tool_call)
+      tool_id = tool_call[:id]
+      tool_name = tool_call[:name]
+
+      # Parse arguments from JSON string
+      args_json = tool_call[:argumentsJson] || tool_call[:arguments_json]
+      args = if args_json.is_a?(String) && !args_json.empty?
+               JSON.parse(args_json, symbolize_names: true)
+             elsif tool_call[:arguments].is_a?(Hash)
+               tool_call[:arguments]
+             else
+               {}
+             end
+
+      # Symbolize keys for Ruby kwargs
+      kwargs = args.transform_keys(&:to_sym)
+
+      begin
+        result = @tool_runner.execute(tool_name, **kwargs)
+      rescue ToolNotFoundError => e
+        result = { error: e.message }
+      end
+
+      # Send tool response back (protobuf InputEvent.tool_response format)
+      # The harness expects responseJson to be a JSON object (Python SDK wraps in {"result": ...})
+      result_dict = result.is_a?(Hash) ? result : { result: result.to_s }
+      tool_response = {
+        toolResponse: {
+          id: tool_id,
+          responseJson: JSON.generate(result_dict)
+        }
+      }
+      @ws.send_json(tool_response)
+    end
+
+    def parse_step(step)
+      {
+        step_index: step[:stepIndex],
+        state: parse_state(step[:state]),
+        source: parse_source(step[:source]),
+        target: parse_target(step[:target]),
+        text_delta: step[:textDelta],
+        text: step[:text],
+        thinking_delta: step[:thinkingDelta],
+        error: step[:errorMessage],
+        cascade_id: step[:cascadeId],
+        trajectory_id: step[:trajectoryId]
+      }
+    end
+
+    def parse_state(val)
+      case val.to_s
+      when /ACTIVE|1/ then :active
+      when /DONE|2/   then :done
+      when /WAITING|3/ then :waiting
+      when /ERROR|4/  then :error
+      else :unknown
+      end
+    end
+
+    def parse_source(val)
+      case val.to_s
+      when /SYSTEM|1/ then :system
+      when /USER|2/   then :user
+      when /MODEL|3/  then :model
+      else :unknown
+      end
+    end
+
+    def parse_target(val)
+      case val.to_s
+      when /USER|1/        then :user
+      when /MODEL|2/       then :model
+      when /ENVIRONMENT|3/ then :environment
+      else :unknown
+      end
+    end
+
+    def update_usage(usage)
+      # The harness sends: { total: { promptTokenCount: "3855", ... }, agents: [...] }
+      # Fall back to cumulativeUsage (legacy) or flat usage hash
+      meta = usage[:total] || usage[:cumulativeUsage] || usage
+      @last_turn_usage = {
+        prompt_token_count: meta[:promptTokenCount].to_i,
+        candidates_token_count: meta[:candidatesTokenCount].to_i,
+        thoughts_token_count: meta[:thoughtsTokenCount].to_i,
+        total_token_count: meta[:totalTokenCount].to_i,
+        cached_content_token_count: meta[:cachedContentTokenCount].to_i
+      }
+      # Accumulate into total
+      @last_turn_usage.each { |k, v| @total_usage[k] += v }
+    end
+
+    def empty_usage
+      {
+        prompt_token_count: 0,
+        candidates_token_count: 0,
+        thoughts_token_count: 0,
+        total_token_count: 0,
+        cached_content_token_count: 0
+      }
+    end
+
+    def next_seq
+      @seq_counter ||= 0
+      @seq_counter += 1
+    end
+  end
+end