lex-llm 0.3.1 → 0.4.3

data/lib/legion/extensions/llm/fleet/publish_safety.rb

@@ -0,0 +1,123 @@
+# frozen_string_literal: true
+
+module Legion
+  module Extensions
+    module Llm
+      module Fleet
+        # Publish-result helpers kept local to fleet messages so they work with older legion-transport releases.
+        module PublishSafety
+          private
+
+          def return_publish_result?(options)
+            options[:return_result] == true || options[:mandatory] == true || options[:publisher_confirm] == true ||
+              options[:spool] == false
+          end
+
+          def install_return_listener(exchange_dest, options, return_state)
+            return unless options[:mandatory] == true
+
+            return_channel = publish_channel(exchange_dest)
+            return unless return_channel.respond_to?(:on_return)
+
+            expected_correlation_id = correlation_id
+            expected_message_id = message_id
+            return_channel.on_return do |return_info, properties, _content|
+              next unless returned_message_matches?(
+                properties,
+                correlation_id: expected_correlation_id,
+                message_id: expected_message_id
+              )
+
+              record_return!(return_state, return_info)
+            end
+          end
+
+          def returned_message_matches?(properties, correlation_id:, message_id:)
+            return false if property_mismatch?(properties, :correlation_id, correlation_id)
+            return false if property_mismatch?(properties, :message_id, message_id)
+
+            true
+          end
+
+          def property_mismatch?(properties, key, expected)
+            return false unless expected
+            return false unless properties.respond_to?(key)
+
+            value = properties.public_send(key)
+            value && value != expected
+          end
+
+          def record_return!(return_state, return_info)
+            return_state[:returned] = true
+            return_state[:reply_code] = return_info.reply_code if return_info.respond_to?(:reply_code)
+            return_state[:reply_text] = return_info.reply_text if return_info.respond_to?(:reply_text)
+          end
+
+          def prepare_publisher_confirms(exchange_dest, options)
+            return unless options[:publisher_confirm] == true
+
+            confirm_channel = publish_channel(exchange_dest)
+            confirm_channel.confirm_select if confirm_channel.respond_to?(:confirm_select)
+          end
+
+          def publish_result(exchange_dest, options, return_state)
+            status = confirm_publish(exchange_dest, options)
+            status = :unroutable if return_state[:returned]
+            {
+              status: status,
+              accepted: status == :accepted,
+              exchange: exchange_name(exchange_dest),
+              routing_key: options[:routing_key] || routing_key || '',
+              message_id: message_id,
+              return_reply_code: return_state[:reply_code],
+              return_reply_text: return_state[:reply_text],
+              correlation_id: correlation_id
+            }.compact
+          end
+
+          def publish_failure_result(status, error, options)
+            {
+              status: status,
+              accepted: false,
+              error_class: error.class.name,
+              error: error.message,
+              routing_key: options[:routing_key] || routing_key || '',
+              message_id: message_id,
+              correlation_id: correlation_id
+            }.compact
+          end
+
+          def confirm_publish(exchange_dest, options)
+            return :accepted unless options[:publisher_confirm] == true
+
+            confirm_channel = publish_channel(exchange_dest)
+            return :accepted unless confirm_channel.respond_to?(:wait_for_confirms)
+
+            timeout = options[:publish_confirm_timeout_ms]
+            confirmed = if timeout
+                          confirm_channel.wait_for_confirms(timeout.to_f / 1000.0)
+                        else
+                          confirm_channel.wait_for_confirms
+                        end
+            confirmed == false ? :nacked : :accepted
+          rescue Timeout::Error => e
+            handle_exception(e, level: :warn, handled: true, operation: 'llm.fleet.publish.confirm')
+            :confirm_timeout
+          end
+
+          def publish_channel(exchange_dest)
+            return exchange_dest.channel if exchange_dest.respond_to?(:channel)
+
+            channel
+          end
+
+          def exchange_name(exchange_dest)
+            return exchange_dest.name if exchange_dest.respond_to?(:name)
+
+            exchange_dest.to_s
+          end
+        end
+      end
+    end
+  end
+end

data/lib/legion/extensions/llm/fleet/settings.rb

@@ -0,0 +1,66 @@
+# frozen_string_literal: true
+
+module Legion
+  module Extensions
+    module Llm
+      module Fleet
+        # Reads fleet settings from Legion::Settings when available, falling back to lex-llm defaults.
+        module Settings
+          module_function
+
+          def value(*path, default:)
+            configured_llm_settings.each do |configured|
+              found = dig(configured, *path)
+              return found unless found.nil?
+            end
+
+            fallback = dig(default_settings, *path)
+            fallback.nil? ? default : fallback
+          end
+
+          def configured_llm_settings
+            return [] unless defined?(::Legion::Settings) && ::Legion::Settings.respond_to?(:[])
+
+            configured = []
+            extensions = safe_fetch(::Legion::Settings, :extensions)
+            extension_llm = dig(extensions, :llm)
+            configured << extension_llm if extension_llm.respond_to?(:key?)
+
+            llm = safe_fetch(::Legion::Settings, :llm)
+            configured << llm if llm.respond_to?(:key?)
+            configured
+          rescue StandardError
+            []
+          end
+
+          def dig(hash, *keys)
+            keys.reduce(hash) do |current, key|
+              break nil unless current.respond_to?(:key?)
+
+              symbol_key = key.respond_to?(:to_sym) ? key.to_sym : key
+              string_key = key.to_s
+              if current.key?(symbol_key)
+                current[symbol_key]
+              elsif current.key?(string_key)
+                current[string_key]
+              end
+            end
+          end
+
+          def safe_fetch(source, key)
+            source[key] || source[key.to_s]
+          rescue StandardError
+            nil
+          end
+
+          def default_settings
+            return ::Legion::Extensions::Llm.default_settings if
+              ::Legion::Extensions::Llm.respond_to?(:default_settings)
+
+            {}
+          end
+        end
+      end
+    end
+  end
+end

data/lib/legion/extensions/llm/fleet/token_error.rb

@@ -0,0 +1,11 @@
+# frozen_string_literal: true
+
+module Legion
+  module Extensions
+    module Llm
+      module Fleet
+        class TokenError < StandardError; end
+      end
+    end
+  end
+end

data/lib/legion/extensions/llm/fleet/token_validator.rb

@@ -0,0 +1,205 @@
+# frozen_string_literal: true
+
+require 'concurrent'
+require 'time'
+
+require_relative 'settings'
+require_relative 'token_error'
+
+module Legion
+  module Extensions
+    module Llm
+      module Fleet
+        # Verifies responder-side fleet JWTs and prevents replay on provider nodes.
+        module TokenValidator
+          @seen_jtis = Concurrent::Map.new
+          @replay_mutex = Mutex.new
+
+          module_function
+
+          def validate!(token:, envelope:, record_replay: true)
+            raise TokenError, 'fleet token is required' if token.to_s.empty?
+
+            claims = symbolize_keys(jwt_module.verify(
+                                      token,
+                                      verification_key: signing_key,
+                                      issuer: issuer,
+                                      algorithm: algorithm,
+                                      verify_issuer: false
+                                    ))
+            validate_registered_claims!(claims)
+            validate_request_expiry!(claims)
+            validate_envelope_claims!(claims, symbolize_keys(envelope || {}))
+            record_replay ? reserve_replay!(claims[:jti]) : ensure_not_replayed!(claims[:jti])
+            claims
+          rescue TokenError
+            raise
+          rescue StandardError => e
+            raise TokenError, "fleet token verification failed: #{e.message}"
+          end
+
+          def reset_replay_cache!
+            @seen_jtis = Concurrent::Map.new
+            @replay_mutex = Mutex.new
+          end
+
+          def validate_registered_claims!(claims)
+            now = Time.now.to_i
+            raise TokenError, 'fleet token issuer mismatch' unless accepted_issuer?(claims[:iss])
+            raise TokenError, 'fleet token audience mismatch' unless claims[:aud].to_s == audience
+            if claims[:exp].nil? || claims[:exp].to_i + clock_skew_seconds <= now
+              raise TokenError,
+                    'fleet token expired'
+            end
+            if claims[:nbf].nil? || claims[:nbf].to_i - clock_skew_seconds > now
+              raise TokenError,
+                    'fleet token not yet valid'
+            end
+            raise TokenError, 'fleet token missing jti' if claims[:jti].to_s.empty?
+          end
+
+          def validate_request_expiry!(claims)
+            expires_at = claims[:expires_at]
+            raise TokenError, 'fleet request expires_at is required' if expires_at.to_s.empty?
+
+            expires = Time.iso8601(expires_at.to_s)
+            raise TokenError, 'fleet request expired' if expires + clock_skew_seconds <= Time.now.utc
+          rescue ArgumentError
+            raise TokenError, 'fleet request expires_at is invalid'
+          end
+
+          def validate_envelope_claims!(claims, envelope)
+            %i[
+              request_id correlation_id idempotency_key operation provider provider_instance
+              model reply_to message_context params caller trace_context timeout_seconds expires_at
+            ].each do |key|
+              expected = canonical_value(envelope[key])
+              actual = canonical_value(claims[key])
+              raise TokenError, "fleet token #{key} claim mismatch" unless actual == expected
+            end
+          end
+
+          def reserve_replay!(jti)
+            @replay_mutex.synchronize do
+              now = Time.now.to_i
+              purge_replay_cache_locked!(now)
+              existing = @seen_jtis[jti.to_s]
+              raise TokenError, 'fleet token replay detected' if active_replay?(existing, now)
+
+              @seen_jtis[jti.to_s] = replay_entry(:inflight, now)
+            end
+          end
+
+          def mark_replay!(jti)
+            @replay_mutex.synchronize do
+              @seen_jtis[jti.to_s] = replay_entry(:complete)
+            end
+          end
+
+          def release_replay!(jti)
+            @replay_mutex.synchronize do
+              entry = @seen_jtis[jti.to_s]
+              @seen_jtis.delete(jti.to_s) if entry.nil? || entry[:state] == :inflight
+            end
+          end
+
+          def ensure_not_replayed!(jti)
+            @replay_mutex.synchronize do
+              now = Time.now.to_i
+              purge_replay_cache_locked!(now)
+              raise TokenError, 'fleet token replay detected' if active_replay?(@seen_jtis[jti.to_s], now)
+            end
+          end
+
+          def purge_replay_cache!
+            @replay_mutex.synchronize { purge_replay_cache_locked!(Time.now.to_i) }
+          end
+
+          def purge_replay_cache_locked!(now)
+            @seen_jtis.each_pair { |jti, entry| @seen_jtis.delete(jti) unless active_replay?(entry, now) }
+          end
+
+          def active_replay?(entry, now)
+            entry && entry[:expires_at] > now
+          end
+
+          def replay_entry(state, now = Time.now.to_i)
+            { state: state, expires_at: now + replay_ttl_seconds }
+          end
+
+          def replay_ttl_seconds
+            ttl = Settings.value(:fleet, :auth, :replay_ttl_seconds, default: 600).to_i
+            ttl.positive? ? ttl : 600
+          end
+
+          def accepted_issuer?(value)
+            accepted_issuers.map(&:to_s).include?(value.to_s)
+          end
+
+          def accepted_issuers
+            issuers = Settings.value(:fleet, :auth, :accepted_issuers, default: [issuer])
+            issuers = [issuer] if Array(issuers).empty?
+            Array(issuers)
+          end
+
+          def clock_skew_seconds
+            Settings.value(:fleet, :auth, :max_clock_skew_seconds, default: 30).to_i
+          end
+
+          def issuer
+            Settings.value(:fleet, :auth, :issuer, default: 'legion-llm')
+          end
+
+          def audience
+            Settings.value(:fleet, :auth, :audience, default: 'lex-llm-fleet-worker')
+          end
+
+          def algorithm
+            Settings.value(:fleet, :auth, :algorithm, default: 'HS256')
+          end
+
+          def signing_key
+            if defined?(::Legion::Crypt) && ::Legion::Crypt.respond_to?(:cluster_secret)
+              return ::Legion::Crypt.cluster_secret
+            end
+
+            raise TokenError, 'no signing key available - Legion::Crypt not initialized'
+          rescue TokenError
+            raise
+          rescue StandardError => e
+            raise TokenError, "no signing key available: #{e.message}"
+          end
+
+          def jwt_module
+            return ::Legion::Crypt::JWT if defined?(::Legion::Crypt::JWT) && ::Legion::Crypt::JWT.respond_to?(:verify)
+
+            raise TokenError, 'Legion::Crypt::JWT.verify unavailable'
+          end
+
+          def symbolize_keys(hash)
+            return {} unless hash.respond_to?(:each)
+
+            hash.each_with_object({}) do |(key, value), result|
+              result[key.respond_to?(:to_sym) ? key.to_sym : key] = value
+            end
+          end
+
+          def canonical_value(value)
+            case value
+            when Hash
+              value.each_with_object({}) do |(key, child), result|
+                result[key.to_s] = canonical_value(child)
+              end.sort.to_h
+            when Array
+              value.map { |child| canonical_value(child) }
+            when Symbol
+              value.to_s
+            else
+              value
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/legion/extensions/llm/fleet/worker_execution.rb

@@ -0,0 +1,165 @@
+# frozen_string_literal: true
+
+require 'concurrent'
+
+require_relative 'settings'
+require_relative 'token_validator'
+
+module Legion
+  module Extensions
+    module Llm
+      module Fleet
+        # Applies responder-side policy and dispatches a fleet request to a local lex-llm provider.
+        module WorkerExecution
+          class PolicyError < StandardError; end
+
+          @idempotency_keys = Concurrent::Map.new
+          @idempotency_mutex = Mutex.new
+
+          module_function
+
+          def call(envelope:, provider:)
+            claims = nil
+            idempotency_key = nil
+            claims = validate_identity!(envelope)
+            validate_policy!(envelope)
+            idempotency_key = validate_idempotency!(envelope)
+            response = dispatch_local_provider!(envelope: envelope, provider: provider)
+            mark_idempotency_success!(idempotency_key) if idempotency_key
+            TokenValidator.mark_replay!(claims[:jti]) if claims.is_a?(Hash)
+            response
+          rescue TokenError => e
+            release_idempotency!(idempotency_key) if idempotency_key
+            release_replay!(claims)
+            raise PolicyError, e.message
+          rescue StandardError
+            release_idempotency!(idempotency_key) if idempotency_key
+            release_replay!(claims)
+            raise
+          end
+
+          def validate_identity!(envelope)
+            return true unless responder_setting(:require_auth, default: true)
+
+            TokenValidator.validate!(token: envelope_value(envelope, :signed_token), envelope: envelope)
+          end
+
+          def validate_policy!(_envelope)
+            return true unless responder_setting(:require_policy, default: false)
+
+            raise PolicyError, 'fleet responder policy enforcement unavailable'
+          end
+
+          def validate_idempotency!(envelope)
+            return nil unless responder_setting(:require_idempotency, default: true)
+
+            key = envelope_value(envelope, :idempotency_key)
+            raise PolicyError, 'fleet idempotency_key is required' if key.to_s.empty?
+
+            reserve_idempotency_key!(key.to_s)
+            key.to_s
+          end
+
+          def dispatch_local_provider!(envelope:, provider:)
+            provider = provider.call(envelope) if provider.respond_to?(:call) && !provider.respond_to?(:chat)
+            operation = envelope_value(envelope, :operation).to_sym
+            params = normalize_hash(envelope_value(envelope, :params) || {})
+            model = envelope_value(envelope, :model)
+
+            case operation
+            when :chat
+              provider.chat(messages: params.fetch(:messages, []), model: model, **except(params, :messages))
+            when :stream
+              provider.stream_chat(messages: params.fetch(:messages, []), model: model, **except(params, :messages))
+            when :embed
+              provider.embed(text: params[:text], model: model, **except(params, :text))
+            when :count_tokens
+              provider.count_tokens(messages: params.fetch(:messages, []), model: model, **except(params, :messages))
+            else
+              raise PolicyError, "unsupported fleet operation: #{operation}"
+            end
+          end
+
+          def reset_idempotency_cache!
+            @idempotency_keys = Concurrent::Map.new
+            @idempotency_mutex = Mutex.new
+          end
+
+          def mark_idempotency_success!(key)
+            @idempotency_mutex.synchronize do
+              @idempotency_keys[key.to_s] = { state: :complete, expires_at: Time.now.to_i + idempotency_ttl_seconds }
+            end
+          end
+
+          def release_idempotency!(key)
+            @idempotency_mutex.synchronize { @idempotency_keys.delete(key.to_s) }
+          end
+
+          def release_replay!(claims)
+            return unless claims.is_a?(Hash) && claims[:jti]
+
+            TokenValidator.release_replay!(claims[:jti])
+          end
+
+          def purge_idempotency_cache!
+            @idempotency_mutex.synchronize do
+              now = Time.now.to_i
+              @idempotency_keys.each_pair do |key, entry|
+                @idempotency_keys.delete(key) if entry[:expires_at] <= now
+              end
+            end
+          end
+
+          def reserve_idempotency_key!(key)
+            @idempotency_mutex.synchronize do
+              now = Time.now.to_i
+              existing = @idempotency_keys[key]
+              raise PolicyError, 'duplicate fleet idempotency key' if existing && existing[:expires_at] > now
+
+              @idempotency_keys[key] = { state: :inflight, expires_at: now + idempotency_ttl_seconds }
+            end
+          end
+
+          def idempotency_ttl_seconds
+            ttl = responder_setting(:idempotency_ttl_seconds, default: 600).to_i
+            ttl.positive? ? ttl : 600
+          end
+
+          def responder_setting(key, default:)
+            value = Settings.value(:fleet, :responder, key, default: nil)
+            return auth_required? if key == :require_auth && value.nil?
+            return default if value.nil?
+
+            value
+          end
+
+          def auth_required?
+            Settings.value(:fleet, :auth, :require_signed_token, default: true) != false
+          end
+
+          def envelope_value(envelope, key)
+            return nil unless envelope.respond_to?(:key?)
+
+            envelope[key] || envelope[key.to_s]
+          end
+
+          def normalize_hash(hash)
+            return {} unless hash.respond_to?(:each)
+
+            hash.each_with_object({}) do |(key, value), result|
+              result[key.respond_to?(:to_sym) ? key.to_sym : key] = value
+            end
+          end
+
+          def except(hash, *keys)
+            exclusions = keys.map(&:to_sym)
+            hash.each_with_object({}) do |(key, value), result|
+              normalized_key = key.respond_to?(:to_sym) ? key.to_sym : key
+              result[normalized_key] = value unless exclusions.include?(normalized_key)
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/legion/extensions/llm/message.rb

@@ -80,12 +80,18 @@ module Legion
             content: content,
             model_id: model_id,
             tool_calls: tool_calls,
-            tool_call_id: tool_call_id,
-            thinking: thinking&.text,
-            thinking_signature: thinking&.signature
+            tool_call_id: tool_call_id
           }.merge(tokens ? tokens.to_h : {}).compact
         end
 
+        def to_internal_h
+          to_h.merge(
+            thinking: thinking&.text,
+            thinking_signature: thinking&.signature,
+            raw: raw
+          ).compact
+        end
+
         def instance_variables
           super - [:@raw]
         end