antigravity-sdk 0.3.0 → 0.4.2

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

data/lib/antigravity/emojis.rb

@@ -2,6 +2,7 @@
 
 module Antigravity
   EMOJIS = {
+    antigravity: "🛰️",
     gem: "💎",
     agent: "🕵️‍♂️",
     prompt: "💬",
@@ -13,6 +14,9 @@ module Antigravity
     sidecar: "🚗",
     logger: "🪵",
     skill: "📁",
+    check: "✅",
+    magnifying: "🔍",
+    workspace: "📂",
     message: "💬",
     guard: "🛡️",
     test: "🧪",

data/lib/antigravity/errors.rb

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+module Antigravity
+  # Base error for all Antigravity SDK errors
+  class Error < StandardError; end
+
+  # Raised when the localharness binary cannot be found
+  class HarnessNotFoundError < Error; end
+
+  # Raised when the stdio handshake with localharness fails
+  class HarnessHandshakeError < Error; end
+
+  # Raised when protobuf encoding/decoding fails
+  class ProtocolError < Error; end
+
+  # Raised when a tool callback fails
+  class ToolError < Error; end
+
+  # Raised when trying to execute an unregistered tool
+  class ToolNotFoundError < ToolError; end
+
+  # Raised when WebSocket connection fails
+  class ConnectionError < Error; end
+
+  # Raised when required configuration (e.g., GEMINI_API_KEY) is missing
+  class ConfigError < Error; end
+end

data/lib/antigravity/guards.rb

@@ -1,31 +1,48 @@
 # frozen_string_literal: true
 
-require "logger"
+require "json"
 require "fileutils"
 
 module Antigravity
   module Guards
-    # Configurable opt-in logger guard supporting file loggers and Rails.logger
+    # Dual-output logger guard:
+    #   1. JSONL (log/antigravity.jsonl) — structured, machine-parseable, full data
+    #   2. Compact log (log/antigravity.log) — human-readable one-liners with byte sizes
+    # Falls back to Rails.logger for both if available.
     class AgentLogger
-      attr_reader :logger, :target_description
+      attr_reader :target_description
 
-      def initialize(log_target = nil, level: Logger::INFO, silent_notice: false)
-        resolved_target = resolve_log_target(log_target)
+      def initialize(log_target = nil, level: :info, silent_notice: false)
+        resolved = resolve_log_target(log_target)
 
-        if resolved_target.is_a?(String)
-          dir = File.dirname(resolved_target)
+        if resolved.is_a?(String)
+          dir = File.dirname(resolved)
           FileUtils.mkdir_p(dir) unless Dir.exist?(dir)
-          @logger = ::Logger.new(resolved_target)
-          @target_description = resolved_target
-        elsif resolved_target.respond_to?(:info)
-          @logger = resolved_target
+
+          # Fat JSONL log
+          @jsonl = File.open(resolved, 'a')
+          @jsonl.sync = true
+
+          # Skinny compact log (same dir, .log extension)
+          compact_path = resolved.sub(/\.jsonl$/, '.log')
+          @compact = File.open(compact_path, 'a')
+          @compact.sync = true
+
+          @rails_logger = nil
+          @target_description = resolved
+        elsif resolved.respond_to?(:info)
+          @jsonl = nil
+          @compact = nil
+          @rails_logger = resolved
           @target_description = "Rails.logger"
         else
-          @logger = ::Logger.new($stdout)
+          @jsonl = $stdout
+          @compact = nil
+          @rails_logger = nil
           @target_description = "$stdout"
         end
 
-        @logger.level = level if @logger.respond_to?(:level=)
+        @level = level
 
         unless silent_notice
           puts "#{Antigravity.emoji(:logger)} Logging to #{@target_description}"
@@ -33,25 +50,49 @@ module Antigravity
       end
 
       def before_prompt(prompt_text)
-        @logger.info("#{Antigravity.emoji(:prompt)} [Antigravity::Prompt] User: '#{prompt_text}'")
+        size = prompt_text.to_s.bytesize
+        log_jsonl('prompt', { user_input: prompt_text })
+        log_compact("#{Antigravity.emoji(:prompt)} PROMPT #{size}B | #{prompt_text.to_s[0, 80]}")
       end
 
       def after_response(response)
-        @logger.info("#{Antigravity.emoji(:response)} [Antigravity::Response] Assistant (#{response.model_id}): #{response.content.strip}")
+        content = response.content&.strip || ''
+        log_jsonl('response', {
+          model: response.model_id,
+          content: content,
+          tokens: response.usage[:total_token_count],
+          tool_calls: response.tool_calls_count,
+          steps: response.steps&.length
+        })
+        log_compact("#{Antigravity.emoji(:response)} RESPONSE #{content.bytesize}B | " \
+                    "tokens=#{response.usage[:total_token_count]} " \
+                    "tools=#{response.tool_calls_count} " \
+                    "steps=#{response.steps&.length} " \
+                    "model=#{response.model_id}")
       end
 
       def before_tool_call(tool_name, params)
-        @logger.info("#{Antigravity.emoji(:tool)} [Antigravity::Tool] Executing '#{tool_name}' with params: #{params.inspect}")
+        params_size = params.to_s.bytesize
+        log_jsonl('tool_call', { tool: tool_name, params: params })
+        log_compact("#{Antigravity.emoji(:tool)} TOOL_CALL #{tool_name} params=#{params_size}B")
       end
 
       def after_tool_call(tool_name, params, result)
-        prefix = result.to_s.include?("TOOL BLOCKED") ? Antigravity.emoji(:tool_blocked) : Antigravity.emoji(:tool_result)
-        @logger.info("#{prefix} [Antigravity::Tool] Result for '#{tool_name}': #{result}")
+        blocked = result.to_s.include?("TOOL BLOCKED")
+        result_size = result.to_s.bytesize
+        log_jsonl('tool_result', {
+          tool: tool_name,
+          result: result.to_s[0, 500],
+          blocked: blocked
+        })
+        status = blocked ? 'BLOCKED' : 'OK'
+        log_compact("#{blocked ? Antigravity.emoji(:tool_blocked) : Antigravity.emoji(:tool_result)} TOOL_RESULT #{tool_name} #{status} result=#{result_size}B")
         result
       end
 
       def sidecar_event(event_type, payload)
-        @logger.info("#{Antigravity.emoji(:sidecar)} [Antigravity::Sidecar] Event :#{event_type} payload: #{payload.inspect}")
+        log_jsonl('sidecar', { type: event_type.to_s, payload: payload })
+        log_compact("#{Antigravity.emoji(:sidecar)} SIDECAR :#{event_type}")
       end
 
       def attach_to(agent)
@@ -64,17 +105,36 @@ module Antigravity
 
       private
 
+      def ts
+        Time.now.utc.strftime('%Y-%m-%dT%H:%M:%S.%3NZ')
+      end
+
+      def log_jsonl(event, data)
+        if @rails_logger
+          @rails_logger.info("[Antigravity] #{event}: #{data.inspect}")
+        elsif @jsonl
+          entry = { ts: ts, event: event, pid: Process.pid }.merge(data.compact)
+          @jsonl.puts(JSON.generate(entry))
+        end
+      end
+
+      def log_compact(line)
+        if @compact
+          @compact.puts("#{ts} #{line}")
+        end
+      end
+
       def resolve_log_target(target)
         return target if target
 
         if defined?(Rails) && Rails.respond_to?(:logger) && Rails.logger
           Rails.logger
         elsif ENV["RAILS_ENV"] && !ENV["RAILS_ENV"].empty?
-          "log/#{ENV['RAILS_ENV']}.log"
+          "log/#{ENV['RAILS_ENV']}.jsonl"
         elsif ENV["RACK_ENV"] && !ENV["RACK_ENV"].empty?
-          "log/#{ENV['RACK_ENV']}.log"
+          "log/#{ENV['RACK_ENV']}.jsonl"
         else
-          "log/antigravity.log"
+          "log/antigravity.jsonl"
         end
       end
     end

data/lib/antigravity/hooks.rb

@@ -9,6 +9,7 @@ module Antigravity
       @post_response_hooks = []
       @pre_tool_hooks = []
       @post_tool_hooks = []
+      @listeners = Hash.new { |h, k| h[k] = [] }
     end
 
     def before_prompt(&block)
@@ -29,6 +30,17 @@ module Antigravity
       @post_tool_hooks << block if block_given?
     end
 
+    # Generic event system — subscribe to any named event.
+    # Usage: hooks.on(:ws_message) { |msg| puts msg.keys }
+    def on(event, &block)
+      @listeners[event.to_sym] << block if block_given?
+    end
+
+    # Emit a named event to all subscribers.
+    def emit(event, *args)
+      @listeners[event.to_sym].each { |cb| cb.call(*args) }
+    end
+
     def run_pre_prompt(prompt_text)
       @pre_prompt_hooks.each { |hook| hook.call(prompt_text) }
     end