Class: Antigravity::Connection::WebSocketClient

Inherits:
Object
  • Object
show all
Defined in:
lib/antigravity/connection/websocket_client.rb

Overview

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.

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(port:, api_key:) ⇒ WebSocketClient

Returns a new instance of WebSocketClient.



16
17
18
19
20
21
22
23
 
# File 'lib/antigravity/connection/websocket_client.rb', line 16

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

Instance Attribute Details

#api_key ⇒ Object (readonly)

Returns the value of attribute api_key.



14
15
16
 
# File 'lib/antigravity/connection/websocket_client.rb', line 14

def api_key
  @api_key
end

#connected ⇒ Object (readonly)

Returns the value of attribute connected.



14
15
16
 
# File 'lib/antigravity/connection/websocket_client.rb', line 14

def connected
  @connected
end

#port ⇒ Object (readonly)

Returns the value of attribute port.



14
15
16
 
# File 'lib/antigravity/connection/websocket_client.rb', line 14

def port
  @port
end

Instance Method Details

#close ⇒ Object 



143
144
145
146
147
148
149
150
151
152
153
154
155
156
 
# File 'lib/antigravity/connection/websocket_client.rb', line 143

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

#connect! ⇒ Object

Open the WebSocket connection to ws://localhost:/



26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
 
# File 'lib/antigravity/connection/websocket_client.rb', line 26

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

#connected? ⇒ Boolean

Returns:

  • (Boolean) 


54
55
56
 
# File 'lib/antigravity/connection/websocket_client.rb', line 54

def connected?
  @connected && @socket && !@socket.closed?
end

#each_message(timeout: Antigravity.config.timeout_llm, &block) ⇒ Object

Read messages in a loop, yielding each parsed JSON. Stops when block returns :stop, connection closes, or timeout.



125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
 
# File 'lib/antigravity/connection/websocket_client.rb', line 125

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

#receive_json(timeout: Antigravity.config.timeout_llm, idle_timeout: nil) ⇒ Object

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.



73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
 
# File 'lib/antigravity/connection/websocket_client.rb', line 73

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

#send_json(data) ⇒ Object

Send a JSON message as a WebSocket text frame



59
60
61
62
63
64
65
66
67
68
 
# File 'lib/antigravity/connection/websocket_client.rb', line 59

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