bynn-decisionnode 0.4.1

data/lib/decisionnode/session.rb

@@ -0,0 +1,286 @@
+# frozen_string_literal: true
+
+require "json"
+require "uri"
+
+class DecisionNode
+  # A frame message's usage: its token counts and its cost, or nil when the
+  # message has none.
+  module FrameUsage
+    private
+
+    def frame_usage(usage)
+      return nil unless usage.is_a?(Hash)
+
+      Usage.new(usage["input_tokens"], usage["output_tokens"], nil, nil, nil, nil, Cost.parse(usage["cost"]))
+    end
+  end
+  private_constant :FrameUsage
+
+  # The reply to one session frame.
+  class Reply
+    include FrameUsage
+
+    # The `seq` of the frame answered.
+    attr_reader :seq
+    # Question id to typed Answer, as in a Decision.
+    attr_reader :answers
+    # The `seq` of each frame this frame superseded, in arrival order, or nil.
+    attr_reader :dropped
+    # Token counts of this frame and what it cost, a Usage.
+    attr_reader :usage
+    # The worker's time for this frame, in milliseconds.
+    attr_reader :latency_ms
+    # The safety check on this frame, as the server sent it (a Hash), or nil.
+    attr_reader :safety
+    # The message as received.
+    attr_reader :raw
+
+    def initialize(raw)
+      @raw = raw
+      @seq = raw["seq"]
+      @answers = Answer.map(raw["answers"])
+      @dropped = raw["dropped"]
+      @usage = frame_usage(raw["usage"])
+      @latency_ms = raw["latency_ms"]
+      @safety = raw["safety"]
+    end
+
+    # The answer to the question with this id.
+    def [](id)
+      answers.fetch(id.to_s) { raise KeyError, "no answer with the id #{id.to_s.inspect}" }
+    end
+  end
+
+  # A frame the session could not answer. A value, never raised: the session
+  # stays open and the next frame can be sent.
+  class FrameError
+    include FrameUsage
+
+    # The frame's `seq`, or nil when the message could not be read as a frame.
+    attr_reader :seq
+    attr_reader :error_type, :message, :category
+    # The wait the server asks for before the next frame, for `rate_limit_error`
+    # and `overloaded_error`.
+    attr_reader :retry_after_ms
+    # What a refused frame billed (a Usage), and its safety result (a Hash).
+    attr_reader :usage, :safety
+    attr_reader :raw
+
+    def initialize(raw)
+      @raw = raw
+      @seq = raw["seq"]
+      @error_type = raw.dig("error", "error_type")
+      @message = raw.dig("error", "message")
+      @category = raw.dig("error", "category")
+      @retry_after_ms = raw["retry_after_ms"]
+      @usage = frame_usage(raw["usage"])
+      @safety = raw["safety"]
+    end
+  end
+
+  # The end of a session's stream, with the close code. It is the last event
+  # of Session#replies, and is raised by Session#connect when the server
+  # refuses the socket before it opens.
+  class SessionClosed < StandardError
+    # Codes after which a new session with the same body is the remedy:
+    # 1001, 1006 (the connection was lost without a close frame), 1011, 4404,
+    # 4408, 4409 and 4410.
+    REOPEN_CODES = [1001, 1006, 1011, 4404, 4408, 4409, 4410].freeze
+    # The server is at its socket limit; the session is kept, connect again.
+    TRY_AGAIN = 1013
+
+    attr_reader :code, :reason
+
+    def initialize(code, reason = "")
+      @code = code
+      @reason = reason.to_s
+      super("session closed with #{code}#{@reason.empty? ? '' : " (#{@reason})"}")
+    end
+
+    # True for the codes after which a new session with the same body is the
+    # remedy; false otherwise, 1013 included (the session is kept).
+    def reopen
+      REOPEN_CODES.include?(code)
+    end
+    alias reopen? reopen
+  end
+
+  # An open session: one state, a fixed set of questions, and a stream of
+  # frames answered as they arrive. Made by DecisionNode#open_session.
+  #
+  #   s = dn.open_session(state: { mission: "Survey" }, questions: { abort: DecisionNode.truth("Abort?") })
+  #   s.connect
+  #   s.send(frame: { battery: 0.31 })
+  #   s.replies.each { |event| ... } # Reply, FrameError, then SessionClosed
+  #   s.end
+  #
+  # Every field of the open response is available by name (`session_id`,
+  # `model`, `stream_url`, `window`, `ttl_seconds`, `expires_at`, `limits`,
+  # `questions`, `prefix_tokens`, `usage`, `safety`) or as `session["field"]`.
+  # The socket is safe to send on from several threads; read the replies from
+  # one.
+  class Session
+    # The open response as received.
+    attr_reader :raw
+    # The answers to the session's questions on the state alone, typed.
+    attr_reader :answers
+
+    FIELDS = %w[session_id model stream_url window ttl_seconds expires_at limits questions prefix_tokens usage safety].freeze
+
+    FIELDS.each do |name|
+      define_method(name) { @raw[name] }
+    end
+
+    def initialize(client, raw)
+      @client = client
+      @raw = raw.is_a?(Hash) ? raw : {}
+      @answers = Answer.map(@raw["answers"])
+      @seq = 0
+      @seq_lock = Mutex.new
+      @stream = nil
+      @closed = nil
+    end
+
+    def id
+      session_id
+    end
+
+    def [](key)
+      @raw[key.to_s]
+    end
+
+    def inspect
+      "#<#{self.class.name} #{session_id.inspect} model=#{model.inspect}>"
+    end
+
+    # Opens the socket: the key in the Authorization header (never in the URL)
+    # and the client's custom headers on the handshake. Returns self. A server
+    # that closes the socket before it opens raises SessionClosed; a refusal
+    # sent right after it opens is the first event of #replies. Connecting
+    # again after a close (for example after 1013) opens a new socket to the
+    # same session.
+    def connect
+      @stream&.close_socket
+      @closed = nil
+      @stream = @client.__send__(:session_stream, stream_url)
+      self
+    end
+
+    # Sends one frame and returns its `seq`; never waits for a reply.
+    #
+    # frame::     text or any JSON value (the reading).
+    # image::     one Image: sent inline under the client's inline threshold,
+    #             else uploaded first and sent by `upload_id`. A URL image is
+    #             refused: frames take bytes or an upload.
+    # questions:: the ids of the session's questions to ask (default all).
+    # reset::     true empties the window before this frame is answered.
+    # seq::       the frame's number; by default a counter from 1.
+    #
+    # A frame that breaks a fixed rule (no frame and no image, an empty frame,
+    # an unknown question id, a seq out of range, a URL image, an image the
+    # API would refuse) raises BadRequestError; nothing is written.
+    def send_frame(frame: nil, image: nil, questions: nil, reset: nil, seq: nil)
+      raise ArgumentError, "connect the session before sending" if @stream.nil?
+
+      # The frame's fixed rules first: a fault raises and nothing is written,
+      # and the refused frame takes no seq.
+      check_frame(frame: frame, image: image, questions: questions, reset: reset, seq: seq)
+      message = { "seq" => seq.nil? ? next_seq : seq }
+      message["frame"] = frame unless frame.nil?
+      message["image"] = @client.__send__(:frame_image, image) unless image.nil?
+      message["questions"] = questions.map(&:to_s) unless questions.nil?
+      message["reset"] = reset unless reset.nil?
+      @stream.text(@client.__send__(:encode, message))
+      message["seq"]
+    end
+
+    # Runs every check #send runs on a frame before writing it, alone: raises
+    # the BadRequestError #send would raise (held /v1/models limits included)
+    # and returns nil for a frame #send would send. Nothing is written, nothing
+    # is uploaded, no seq is used, and the session need not be connected.
+    def check_frame(frame: nil, image: nil, questions: nil, reset: nil, seq: nil)
+      Checks.frame!(frame: frame, image: image, questions: questions, reset: reset, seq: seq,
+                    known_questions: @raw["questions"], limits: @client.__send__(:limits_for, model))
+      nil
+    end
+
+    # Session#send with keyword arguments sends a frame (see #send_frame);
+    # with a method name it is Object#send.
+    def send(*args, **kwargs, &block)
+      return super unless args.empty? && block.nil?
+
+      send_frame(**kwargs)
+    end
+
+    # The next event: a Reply, a FrameError, or the SessionClosed that ends
+    # the stream (returned again on every later call). nil when nothing
+    # arrived within `timeout` seconds (nil waits without limit).
+    def next_event(timeout: nil)
+      return @closed if @closed
+      raise ArgumentError, "connect the session before reading" if @stream.nil?
+
+      event = @stream.events.pop(timeout: timeout)
+      return nil if event.nil?
+
+      kind, a, b = event
+      if kind == :close
+        @closed = SessionClosed.new(a, b)
+        return @closed
+      end
+      parse_message(a)
+    end
+
+    # Each event in turn: Reply and FrameError values, then the SessionClosed
+    # that ends the stream. Without a block, an Enumerator.
+    def replies(timeout: nil)
+      return enum_for(:replies, timeout: timeout) unless block_given?
+
+      loop do
+        event = next_event(timeout: timeout)
+        return nil if event.nil?
+
+        yield event
+        return nil if event.is_a?(SessionClosed)
+      end
+    end
+
+    # Closes the socket with 1000. The session ends on the server.
+    def close
+      @stream&.close(1000)
+      nil
+    end
+
+    # Ends the session (`DELETE /v1/sessions/{id}`); the server then closes the
+    # socket with 1000. A session that is already gone (404) counts as ended.
+    # Returns the response body, or nil after a 404.
+    def end_session
+      body, = @client.__send__(:request, "DELETE", "/v1/sessions/#{URI.encode_www_form_component(session_id.to_s)}", nil)
+      body
+    rescue NotFoundError
+      nil
+    end
+    alias end end_session
+
+    # Closes the socket without a close frame. For the client's own use.
+    def drop
+      @stream&.close_socket
+    end
+
+    private
+
+    def next_seq
+      @seq_lock.synchronize { @seq += 1 }
+    end
+
+    def parse_message(text)
+      raw = JSON.parse(text, max_nesting: false)
+      raise JSON::ParserError, "not an object" unless raw.is_a?(Hash)
+
+      raw.key?("error") ? FrameError.new(raw) : Reply.new(raw)
+    rescue JSON::ParserError
+      FrameError.new({ "seq" => nil, "error" => { "error_type" => "invalid_message",
+                                                  "message" => "the server sent a message that is not a JSON object" } })
+    end
+  end
+end

data/lib/decisionnode/session_stream.rb

@@ -0,0 +1,255 @@
+# frozen_string_literal: true
+
+require "openssl"
+require "socket"
+require "uri"
+require "websocket/driver"
+
+class DecisionNode
+  # The socket of a session: RFC 6455 through websocket-driver over a TCP or
+  # TLS socket from the standard library, with one reader thread. Pings are
+  # answered at once; a socket that receives nothing, pings included, for
+  # IDLE_LIMIT seconds is treated as lost (1006). Used by Session; not used directly.
+  #
+  # The socket follows the client's transport where a WebSocket can: a `wss`
+  # URL goes through the transport's proxy (else the proxy the environment
+  # names for Net::HTTP, `no_proxy` honoured) with an HTTP CONNECT tunnel, and
+  # TLS takes the transport's `net_http:` trust and client certificate
+  # settings (ca_file, ca_path, cert_store, cert, key, ...), with peer and
+  # hostname checks always on.
+  class SessionStream
+    # The server pings every 20 s and drops a socket after 45 s of silence.
+    IDLE_LIMIT = 65
+    MAX_MESSAGE = 16 * 1024 * 1024
+
+    attr_reader :url, :events
+
+    # The `net_http:` settings that apply to the session's TLS context.
+    TLS_SETTINGS = %w[ca_file ca_path cert_store cert key extra_chain_cert ssl_version min_version max_version
+                      ciphers verify_depth].freeze
+
+    # proxy:: a proxy URI, :env to take the one the environment names, or nil for none.
+    # tls::   settings by name (TLS_SETTINGS); others are ignored.
+    def initialize(url, headers, timeout, proxy: :env, tls: {})
+      @url = url
+      @proxy = proxy
+      @tls = tls
+      @events = Queue.new
+      @write_lock = Mutex.new
+      @driver_lock = Mutex.new
+      @finish_lock = Mutex.new
+      @opened = Queue.new
+      @done = false
+      @io = open_socket(URI.parse(url), timeout)
+      @driver = Driver.new(self, max_length: MAX_MESSAGE)
+      headers.each { |name, value| @driver.set_header(name, value) }
+      @driver.on(:open) { @opened << [:open] }
+      @driver.on(:message) { |e| @events << [:message, e.data] }
+      @driver.on(:close) { |e| finish(@driver.received_close_code || e.code || 1005, e.reason.to_s) }
+      @driver_lock.synchronize { @driver.start }
+      @reader = Thread.new { read_loop }
+      @reader.report_on_exception = false
+      await_open(timeout)
+    end
+
+    # Sends one text message.
+    def text(data)
+      @driver_lock.synchronize { @driver.text(data) }
+    end
+
+    def close(code)
+      @driver_lock.synchronize { @driver.close("", code) }
+    end
+
+    # Called by the driver.
+    def write(data)
+      @write_lock.synchronize { @io.write(data) }
+    rescue IOError, SystemCallError, OpenSSL::SSL::SSLError
+      finish(1006, "")
+    end
+
+    def close_socket
+      @done = true
+      @io.close unless @io.closed?
+    rescue IOError, SystemCallError, OpenSSL::SSL::SSLError
+      nil
+    end
+
+    private
+
+    LOOPBACK = %w[localhost 127.0.0.1 ::1].freeze
+
+    def open_socket(uri, timeout)
+      host = uri.host.to_s.delete_prefix("[").delete_suffix("]")
+      case uri.scheme&.downcase
+      when "wss"
+        port = uri.port || 443
+        proxy = proxy_for(host, port)
+        tcp = proxy ? self.class.tunnel(proxy, host, port, timeout) : Socket.tcp(host, port, connect_timeout: timeout)
+        tls(tcp, host, timeout)
+      when "ws"
+        raise ArgumentError, "the stream URL must use wss, got #{uri}" unless LOOPBACK.include?(host.downcase)
+
+        tcp = Socket.tcp(host, uri.port || 80, connect_timeout: timeout)
+        tcp.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_NODELAY, 1)
+        tcp
+      else
+        raise ArgumentError, "the stream URL must use wss, got #{uri}"
+      end
+    rescue Errno::ETIMEDOUT, IO::TimeoutError => e
+      raise DecisionNode::TimeoutError, "no connection to the session stream within #{timeout} s (#{e.class})"
+    rescue SocketError, SystemCallError => e
+      raise DecisionNode::ConnectionError, "cannot connect to the session stream: #{e.message} (#{e.class})"
+    end
+
+    # The proxy for a wss host: the transport's, else the environment's as
+    # Net::HTTP finds it (no_proxy honoured), else none.
+    def proxy_for(host, port)
+      return @proxy unless @proxy == :env
+
+      URI::HTTP.new("http", nil, host, port, nil, nil, nil, nil, nil).find_proxy
+    end
+
+    def tls(tcp, host, timeout)
+      tcp.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_NODELAY, 1)
+      ssl = OpenSSL::SSL::SSLSocket.new(tcp, self.class.tls_context(@tls))
+      ssl.hostname = host
+      ssl.sync_close = true
+      deadline = Process.clock_gettime(Process::CLOCK_MONOTONIC) + timeout
+      loop do
+        case ssl.connect_nonblock(exception: false)
+        when :wait_readable then wait(tcp, :r, deadline)
+        when :wait_writable then wait(tcp, :w, deadline)
+        else break
+        end
+      end
+      ssl.post_connection_check(host)
+      ssl
+    rescue OpenSSL::SSL::SSLError => e
+      tcp.close
+      raise DecisionNode::ConnectionError, "TLS to the session stream failed: #{e.message}"
+    end
+
+    # A TLS context with the given settings, certificate and hostname checks on.
+    def self.tls_context(settings)
+      params = settings.to_h { |k, v| [k.to_s.delete_suffix("="), v] }.select { |k, _| TLS_SETTINGS.include?(k) }
+      context = OpenSSL::SSL::SSLContext.new
+      context.set_params(params.transform_keys(&:to_sym).merge(verify_mode: OpenSSL::SSL::VERIFY_PEER, verify_hostname: true))
+      context
+    end
+
+    # A TCP socket to host:port through an HTTP proxy's CONNECT tunnel.
+    def self.tunnel(proxy, host, port, timeout)
+      unless proxy.scheme.to_s.casecmp?("http")
+        raise DecisionNode::ConnectionError, "a session socket goes through an http:// proxy only, got #{proxy.scheme}://"
+      end
+
+      tcp = Socket.tcp(proxy.host, proxy.port || 80, connect_timeout: timeout)
+      target = host.include?(":") ? "[#{host}]:#{port}" : "#{host}:#{port}"
+      head = +"CONNECT #{target} HTTP/1.1\r\nHost: #{target}\r\n"
+      if proxy.user
+        login = [proxy.user, proxy.password.to_s].map { |part| URI.decode_www_form_component(part) }.join(":")
+        head << "Proxy-Authorization: Basic #{[login].pack('m0')}\r\n"
+      end
+      tcp.write("#{head}\r\n")
+      answer = read_head(tcp, Process.clock_gettime(Process::CLOCK_MONOTONIC) + timeout)
+      status = answer[%r{\AHTTP/1\.[01] (\d{3})}, 1].to_i
+      return tcp if status == 200
+
+      tcp.close
+      raise DecisionNode::ConnectionError, "the proxy refused the tunnel to the session stream with HTTP #{status}"
+    rescue IOError, SystemCallError => e
+      tcp&.close
+      raise DecisionNode::ConnectionError, "the proxy tunnel to the session stream failed: #{e.message} (#{e.class})"
+    end
+
+    # The proxy's answer up to the blank line, read a byte at a time so nothing
+    # of the tunnel is taken.
+    def self.read_head(io, deadline)
+      head = +""
+      until head.end_with?("\r\n\r\n")
+        left = deadline - Process.clock_gettime(Process::CLOCK_MONOTONIC)
+        raise DecisionNode::TimeoutError, "the proxy did not answer in time" unless left.positive? && io.wait_readable(left)
+
+        byte = io.read_nonblock(1, exception: false)
+        next if byte == :wait_readable
+        raise DecisionNode::ConnectionError, "the proxy closed the connection" if byte.nil?
+
+        head << byte
+        raise DecisionNode::ConnectionError, "the proxy's answer is too long" if head.bytesize > 16_384
+      end
+      head
+    end
+
+    def wait(io, dir, deadline)
+      left = deadline - Process.clock_gettime(Process::CLOCK_MONOTONIC)
+      ready = left.positive? && (dir == :r ? io.wait_readable(left) : io.wait_writable(left))
+      raise DecisionNode::TimeoutError, "no TLS handshake with the session stream in time" unless ready
+    end
+
+    def await_open(timeout)
+      first = @opened.pop(timeout: timeout)
+      if first.nil?
+        close_socket
+        raise DecisionNode::TimeoutError, "the session stream did not open within #{timeout} s"
+      end
+      return if first.first == :open
+
+      _, code, reason = first
+      status = @driver.status
+      if status && status != 101
+        raise DecisionNodeError.new("the session stream refused the upgrade with HTTP #{status}", status: status)
+      end
+      raise SessionClosed.new(code, reason)
+    end
+
+    def read_loop
+      until @done
+        if !(@io.respond_to?(:pending) && @io.pending.positive?) && !@io.to_io.wait_readable(IDLE_LIMIT)
+          finish(1006, "no traffic for #{IDLE_LIMIT} s")
+          break
+        end
+        data = @io.read_nonblock(65_536, exception: false)
+        next if data == :wait_readable || data == :wait_writable
+
+        if data.nil?
+          finish(1006, "")
+          break
+        end
+        @driver_lock.synchronize { @driver.parse(data) }
+      end
+    rescue IOError, SystemCallError, OpenSSL::SSL::SSLError
+      finish(1006, "")
+    end
+
+    # Records the close once, closes the socket and wakes #connect.
+    def finish(code, reason)
+      first = @finish_lock.synchronize do
+        next false if @finished
+
+        @finished = true
+      end
+      return unless first
+
+      @events << [:close, code, reason]
+      @opened << [:close, code, reason]
+      close_socket
+    end
+
+    # The driver, keeping the close code the server sent: websocket-driver
+    # replaces codes it does not list (1012 to 1014, 1013 among them) with 1002.
+    class Driver < ::WebSocket::Driver::Client
+      attr_reader :received_close_code
+
+      private
+
+      def emit_frame(buffer)
+        if @frame.opcode == OPCODES[:close] && buffer.bytesize >= 2
+          payload = ::WebSocket::Mask.mask(buffer.dup, @frame.masking_key)
+          @received_close_code = payload.unpack1("S>")
+        end
+        super
+      end
+    end
+  end
+end

data/lib/decisionnode/supervisor.rb

@@ -0,0 +1,107 @@
+# frozen_string_literal: true
+
+class DecisionNode
+  # The loop behind DecisionNode#supervise: one feed, sessions reopened as the
+  # close codes say. Not used directly.
+  class Supervisor
+    # Codes after which no backoff is needed before the next session: the
+    # session went idle (4408) or reached its end of life (4410).
+    AT_ONCE = [4408, 4410].freeze
+
+    def initialize(client, body, handler, max_sessions, drain_timeout)
+      @client = client
+      @body = body
+      @handler = handler
+      @max_sessions = max_sessions
+      @drain_timeout = drain_timeout
+      @session = nil  # the session being fed, nil when none is connected
+      @kept = nil     # a session closed with 1013, to connect to again
+      @last = nil     # the last SessionClosed seen
+      @connects = 0
+      @failures = 0
+      @answered = 0
+      @pending = 0
+    end
+
+    def run(feed)
+      feed.each do |reading|
+        drain(0)
+        until @session
+          return @last if stopped? || @connects >= @max_sessions
+
+          pause_before_connect
+          connect
+        end
+        @session.send_frame(**reading.to_h { |k, v| [k.to_sym, v] })
+        @pending += 1
+      end
+      drain(@drain_timeout)
+      if @session
+        @session.end_session
+        @session.close
+      end
+      @last
+    end
+
+    private
+
+    def stopped?
+      @last && !@last.reopen && @last.code != SessionClosed::TRY_AGAIN
+    end
+
+    # A capped backoff with jitter, 1 s doubling to 30 s, reset after a
+    # session that answered frames; none after 4408 and 4410.
+    def pause_before_connect
+      return if @last.nil? || AT_ONCE.include?(@last.code)
+
+      @failures = 0 if @answered.positive?
+      base = [DecisionNode::RECONNECT_CAP, DecisionNode::RECONNECT_START * (2**@failures)].min
+      @client.__send__(:pause, base * (0.5 + (rand * 0.5)))
+      @failures += 1
+    end
+
+    def connect
+      session = if @last&.code == SessionClosed::TRY_AGAIN && @kept
+                  @kept
+                else
+                  @client.open_session(**@body)
+                end
+      @connects += 1
+      @answered = 0
+      @pending = 0
+      @last = nil
+      @kept = nil
+      session.connect
+      @session = session
+    rescue SessionClosed => e
+      closed(e, session)
+    end
+
+    def closed(event, session)
+      @last = event
+      @kept = event.code == SessionClosed::TRY_AGAIN ? session : nil
+      @session = nil
+    end
+
+    # Hands every event that arrived to the handler. With a wait, waits up to
+    # that many seconds for the outstanding replies.
+    def drain(wait)
+      deadline = Process.clock_gettime(Process::CLOCK_MONOTONIC) + wait
+      while @session
+        timeout = wait.zero? ? 0 : [deadline - Process.clock_gettime(Process::CLOCK_MONOTONIC), 0.01].max
+        event = @session.next_event(timeout: timeout)
+        return if event.nil?
+
+        if event.is_a?(SessionClosed)
+          closed(event, @session)
+          return
+        end
+        dropped = event.is_a?(Reply) ? (event.dropped || []).size : 0
+        @pending = [@pending - 1 - dropped, 0].max
+        @answered += 1 if event.is_a?(Reply)
+        @handler.call(event)
+        return if wait.positive? && @pending.zero?
+      end
+    end
+  end
+end