lex-llm 0.4.16 → 0.5.0

data/spec/legion/extensions/llm/fleet/provider_responder_spec.rb

@@ -0,0 +1,120 @@
+# frozen_string_literal: true
+
+require 'legion/extensions/llm/fleet/provider_responder'
+
+RSpec.describe Legion::Extensions::Llm::Fleet::ProviderResponder do
+  let(:protocol) { Legion::Extensions::Llm::Fleet::Protocol }
+  let(:payload) do
+    {
+      request_id: 'req-1',
+      correlation_id: 'corr-1',
+      idempotency_key: 'idem-1',
+      operation: :chat,
+      provider: :ollama,
+      provider_instance: :default,
+      model: 'llama3',
+      params: { messages: [{ role: 'user', content: 'hello' }], temperature: 0.1 },
+      reply_to: 'reply.queue',
+      message_context: { conversation_id: 'conv-1' },
+      caller: { identity: 'user:matt' },
+      trace_context: { trace_id: 'trace-1' },
+      signed_token: 'unsigned',
+      timeout_seconds: 30,
+      expires_at: (Time.now.utc + 30).iso8601,
+      protocol_version: protocol::VERSION
+    }
+  end
+  let(:provider_instances) do
+    {
+      default: {
+        base_url: 'http://localhost:11434',
+        fleet: { respond_to_requests: true }
+      }
+    }
+  end
+  let(:provider_class) do
+    Class.new do
+      attr_reader :settings
+
+      def initialize(settings)
+        @settings = settings
+      end
+
+      def chat(messages:, model:, **params)
+        {
+          content: "chat #{model} #{messages.first[:content]}",
+          usage: { input_tokens: 1, output_tokens: 2 },
+          metadata: { temperature: params[:temperature], base_url: settings[:base_url] }
+        }
+      end
+    end
+  end
+
+  before do
+    Legion::Extensions::Llm::Fleet::WorkerExecution.reset_idempotency_cache!
+    allow(Legion::Extensions::Llm::Fleet::Settings).to receive(:value).and_call_original
+    allow(Legion::Extensions::Llm::Fleet::Settings).to receive(:value)
+      .with(:fleet, :auth, :require_signed_token, default: true).and_return(false)
+    allow(Legion::Extensions::Llm::Fleet::Settings).to receive(:value)
+      .with(:fleet, :responder, :require_idempotency, default: nil).and_return(false)
+  end
+
+  it 'does not require the legion-llm namespace on responder nodes' do
+    hide_const('Legion::LLM') if defined?(Legion::LLM)
+
+    expect(defined?(Legion::LLM)).to be_nil
+    expect(described_class).to respond_to(:call)
+  end
+
+  it 'builds a provider from the requested instance, dispatches the operation, and publishes a fleet response' do
+    response_message = instance_double(Legion::Extensions::Llm::Transport::Messages::FleetResponse, publish: true)
+    allow(Legion::Extensions::Llm::Transport::Messages::FleetResponse).to receive(:new).and_return(response_message)
+
+    response = described_class.call(
+      payload: payload,
+      provider_family: :ollama,
+      provider_class: provider_class,
+      provider_instances: provider_instances
+    )
+
+    expect(response).to include(content: 'chat llama3 hello')
+    expect(Legion::Extensions::Llm::Transport::Messages::FleetResponse).to have_received(:new).with(
+      hash_including(
+        request_id: 'req-1',
+        correlation_id: 'corr-1',
+        provider: :ollama,
+        provider_instance: :default,
+        model: 'llama3',
+        content: 'chat llama3 hello',
+        usage: { input_tokens: 1, output_tokens: 2 },
+        metadata: { temperature: 0.1, base_url: 'http://localhost:11434' }
+      )
+    )
+    expect(response_message).to have_received(:publish)
+  end
+
+  it 'rejects legacy fleet protocol fields before provider execution' do
+    expect do
+      described_class.call(
+        payload: payload.merge(request_type: 'chat'),
+        provider_family: :ollama,
+        provider_class: provider_class,
+        provider_instances: provider_instances
+      )
+    end.to raise_error(ArgumentError, /request_type/)
+  end
+
+  it 'reports whether a provider instance is enabled for fleet responses' do
+    expect(described_class.enabled_for?(provider_instances)).to be(true)
+    expect(described_class.enabled_for?(default: { fleet: { respond_to_requests: false } })).to be(false)
+  end
+
+  it 'raises an actionable configuration error when fleet transport messages cannot load' do
+    allow(Legion::Extensions::Llm::Transport::Messages).to receive(:const_get)
+      .with(:FleetResponse).and_raise(NameError, 'missing transport')
+
+    expect do
+      described_class.transport_message_class(:FleetResponse)
+    end.to raise_error(described_class::ConfigurationError, /fleet responder transport unavailable/)
+  end
+end

data/spec/legion/extensions/llm/fleet/token_validator_spec.rb

@@ -0,0 +1,163 @@
+# frozen_string_literal: true
+
+require 'legion/crypt'
+require 'legion/extensions/llm/fleet/token_validator'
+
+RSpec.describe Legion::Extensions::Llm::Fleet::TokenValidator do
+  let(:expires_at) { (Time.now.utc + 60).iso8601 }
+  let(:envelope) do
+    {
+      request_id: 'req-1',
+      correlation_id: 'corr-1',
+      idempotency_key: 'idem-1',
+      operation: :chat,
+      provider: :ollama,
+      provider_instance: :default,
+      model: 'llama3',
+      reply_to: 'reply.queue',
+      message_context: { conversation_id: 'conv-1' },
+      params: { messages: [{ role: 'user', content: 'hello' }] },
+      caller: { identity: 'user:matt' },
+      trace_context: { trace_id: 'trace-1' },
+      timeout_seconds: 30,
+      expires_at: expires_at
+    }
+  end
+  let(:claims) do
+    envelope.merge(
+      iss: 'legion-llm',
+      aud: 'lex-llm-fleet-worker',
+      exp: Time.now.to_i + 60,
+      nbf: Time.now.to_i - 1,
+      jti: 'jti-1'
+    )
+  end
+
+  before do
+    described_class.reset_replay_cache!
+    allow(Legion::Extensions::Llm::Fleet::Settings).to receive(:value).and_call_original
+    allow(Legion::Crypt).to receive(:cluster_secret).and_return('secret')
+    allow(Legion::Crypt::JWT).to receive(:verify).and_return(claims)
+  end
+
+  it 'verifies the signed token and validates matching envelope claims' do
+    expect(validate_token).to include(jti: 'jti-1')
+    expect(Legion::Crypt::JWT).to have_received(:verify).with(
+      'signed.jwt',
+      verification_key: 'secret',
+      issuer: 'legion-llm',
+      algorithm: 'HS256',
+      verify_issuer: true
+    )
+  end
+
+  it 'rejects missing tokens' do
+    expect do
+      described_class.validate!(token: nil, envelope: envelope)
+    end.to raise_error(Legion::Extensions::Llm::Fleet::TokenError, /token is required/)
+  end
+
+  {
+    iss: ['other-issuer', /issuer mismatch/],
+    aud: ['wrong-audience', /audience mismatch/],
+    exp: [Time.now.to_i - 300, /token expired/],
+    nbf: [Time.now.to_i + 300, /not yet valid/],
+    jti: ['', /missing jti/]
+  }.each do |claim, (value, message)|
+    it "rejects invalid registered claim #{claim}" do
+      allow(Legion::Crypt::JWT).to receive(:verify).and_return(claims.merge(claim => value))
+
+      expect { validate_token }.to raise_error(Legion::Extensions::Llm::Fleet::TokenError, message)
+    end
+  end
+
+  it 'rejects expired request envelopes' do
+    allow(Legion::Crypt::JWT).to receive(:verify)
+      .and_return(claims.merge(expires_at: (Time.now.utc - 31).iso8601))
+
+    expect { validate_token }.to raise_error(Legion::Extensions::Llm::Fleet::TokenError, /request expired/)
+  end
+
+  it 'rejects invalid request expiry timestamps' do
+    allow(Legion::Crypt::JWT).to receive(:verify).and_return(claims.merge(expires_at: 'not-time'))
+
+    expect { validate_token }.to raise_error(Legion::Extensions::Llm::Fleet::TokenError, /expires_at is invalid/)
+  end
+
+  %i[
+    request_id correlation_id idempotency_key operation provider provider_instance model reply_to
+    timeout_seconds expires_at
+  ].each do |claim|
+    it "rejects envelope claim mismatch for #{claim}" do
+      value = claim == :expires_at ? (Time.now.utc + 120).iso8601 : 'different'
+      allow(Legion::Crypt::JWT).to receive(:verify).and_return(claims.merge(claim => value))
+
+      expect do
+        validate_token
+      end.to raise_error(Legion::Extensions::Llm::Fleet::TokenError, /#{claim} claim mismatch/)
+    end
+  end
+
+  %i[message_context params caller trace_context].each do |claim|
+    it "rejects envelope hash mismatch for #{claim}" do
+      allow(Legion::Crypt::JWT).to receive(:verify).and_return(claims.merge(claim => 'different'))
+
+      expect do
+        validate_token
+      end.to raise_error(Legion::Extensions::Llm::Fleet::TokenError, /#{claim} hash mismatch/)
+    end
+  end
+
+  it 'wraps JWT verification failures with a token validation error' do
+    allow(Legion::Crypt::JWT).to receive(:verify).and_raise(StandardError, 'bad signature')
+
+    expect do
+      validate_token
+    end.to raise_error(Legion::Extensions::Llm::Fleet::TokenError, /verification failed: bad signature/)
+  end
+
+  it 'reserves replay tokens immediately and rejects duplicate validation' do
+    validate_token
+
+    expect { validate_token }.to raise_error(Legion::Extensions::Llm::Fleet::TokenError, /replay detected/)
+  end
+
+  it 'supports non-recording validation while still rejecting already recorded tokens' do
+    expect(validate_token(record_replay: false)).to include(jti: 'jti-1')
+    expect(validate_token(record_replay: false)).to include(jti: 'jti-1')
+
+    validate_token
+
+    expect do
+      validate_token(record_replay: false)
+    end.to raise_error(Legion::Extensions::Llm::Fleet::TokenError, /replay detected/)
+  end
+
+  it 'can release inflight replay reservations after failed dispatch' do
+    validate_token
+    described_class.release_replay!('jti-1')
+
+    expect(validate_token).to include(jti: 'jti-1')
+  end
+
+  it 'keeps completed replay reservations after release is requested' do
+    validate_token
+    described_class.mark_replay!('jti-1')
+    described_class.release_replay!('jti-1')
+
+    expect do
+      validate_token(record_replay: false)
+    end.to raise_error(Legion::Extensions::Llm::Fleet::TokenError, /replay detected/)
+  end
+
+  it 'uses a dedicated auth replay TTL setting' do
+    allow(Legion::Extensions::Llm::Fleet::Settings).to receive(:value)
+      .with(:fleet, :auth, :replay_ttl_seconds, default: 600).and_return(42)
+
+    expect(described_class.replay_ttl_seconds).to eq(42)
+  end
+
+  def validate_token(record_replay: true)
+    described_class.validate!(token: 'signed.jwt', envelope: envelope, record_replay: record_replay)
+  end
+end

data/spec/legion/extensions/llm/fleet/worker_execution_spec.rb

@@ -0,0 +1,128 @@
+# frozen_string_literal: true
+
+require 'legion/extensions/llm/fleet/worker_execution'
+
+RSpec.describe Legion::Extensions::Llm::Fleet::WorkerExecution do
+  let(:envelope) do
+    {
+      operation: :chat,
+      model: 'llama3',
+      params: { messages: [{ role: 'user', content: 'hello' }] },
+      signed_token: 'unsigned',
+      idempotency_key: 'idem-1'
+    }
+  end
+  let(:provider) do
+    Class.new do
+      def chat(messages:, model:, **)
+        { content: "#{model}:#{messages.first[:content]}" }
+      end
+    end.new
+  end
+
+  before do
+    described_class.reset_idempotency_cache!
+    allow(Legion::Extensions::Llm::Fleet::Settings).to receive(:value).and_call_original
+    allow(Legion::Extensions::Llm::Fleet::Settings).to receive(:value)
+      .with(:fleet, :auth, :require_signed_token, default: true).and_return(false)
+  end
+
+  it 'dispatches canonical chat operations directly to the local provider' do
+    allow(Legion::Extensions::Llm::Fleet::Settings).to receive(:value)
+      .with(:fleet, :responder, :require_idempotency, default: nil).and_return(false)
+
+    expect(described_class.call(envelope: envelope, provider: provider)).to eq(content: 'llama3:hello')
+  end
+
+  it 'unpacks legacy nested options before dispatching to the local provider' do
+    allow(Legion::Extensions::Llm::Fleet::Settings).to receive(:value)
+      .with(:fleet, :responder, :require_idempotency, default: nil).and_return(false)
+    captured_options = nil
+    capture_provider = Class.new do
+      define_method(:chat) do |messages:, model:, **options|
+        captured_options = options
+        { content: "#{model}:#{messages.first[:content]}" }
+      end
+    end.new
+    legacy_envelope = envelope.merge(
+      params: {
+        messages: [{ role: 'user', content: 'hello' }],
+        stream: false,
+        tools: { current: { name: 'current' } },
+        options: {
+          'stream' => true,
+          'system' => 'Use available tools.',
+          'tools' => { legacy: { name: 'legacy' } }
+        }
+      }
+    )
+
+    described_class.call(envelope: legacy_envelope, provider: capture_provider)
+
+    expect(captured_options).to include(
+      stream: false,
+      system: 'Use available tools.',
+      tools: { current: { name: 'current' } }
+    )
+    expect(captured_options).not_to have_key(:options)
+  end
+
+  it 'rejects duplicate idempotency keys before executing the provider again' do
+    allow(Legion::Extensions::Llm::Fleet::Settings).to receive(:value)
+      .with(:fleet, :responder, :require_idempotency, default: nil).and_return(true)
+
+    described_class.call(envelope: envelope, provider: provider)
+
+    expect do
+      described_class.call(envelope: envelope, provider: provider)
+    end.to raise_error(described_class::PolicyError, /duplicate fleet idempotency key/)
+  end
+
+  it 'replaces expired idempotency entries while reserving the new attempt' do
+    described_class.instance_variable_get(:@idempotency_keys)['idem-expired'] = {
+      state: :complete,
+      expires_at: Time.now.to_i - 1
+    }
+
+    expect { described_class.reserve_idempotency_key!('idem-expired') }.not_to raise_error
+
+    expect do
+      described_class.reserve_idempotency_key!('idem-expired')
+    end.to raise_error(described_class::PolicyError, /duplicate fleet idempotency key/)
+  end
+
+  it 'reserves token replay protection before provider dispatch and marks success after completion' do
+    allow(Legion::Extensions::Llm::Fleet::Settings).to receive(:value)
+      .with(:fleet, :auth, :require_signed_token, default: true).and_return(true)
+    allow(Legion::Extensions::Llm::Fleet::Settings).to receive(:value)
+      .with(:fleet, :responder, :require_idempotency, default: nil).and_return(false)
+    allow(Legion::Extensions::Llm::Fleet::TokenValidator).to receive(:validate!).and_return({ jti: 'jti-1' })
+    allow(Legion::Extensions::Llm::Fleet::TokenValidator).to receive(:mark_replay!)
+
+    described_class.call(envelope: envelope, provider: provider)
+
+    expect(Legion::Extensions::Llm::Fleet::TokenValidator).to have_received(:validate!)
+      .with(token: 'unsigned', envelope: envelope)
+    expect(Legion::Extensions::Llm::Fleet::TokenValidator).to have_received(:mark_replay!).with('jti-1')
+  end
+
+  it 'releases reserved token replay state when provider dispatch fails' do
+    failing_provider = Class.new do
+      def chat(**)
+        raise 'provider unavailable'
+      end
+    end.new
+    allow(Legion::Extensions::Llm::Fleet::Settings).to receive(:value)
+      .with(:fleet, :auth, :require_signed_token, default: true).and_return(true)
+    allow(Legion::Extensions::Llm::Fleet::Settings).to receive(:value)
+      .with(:fleet, :responder, :require_idempotency, default: nil).and_return(false)
+    allow(Legion::Extensions::Llm::Fleet::TokenValidator).to receive(:validate!).and_return({ jti: 'jti-2' })
+    allow(Legion::Extensions::Llm::Fleet::TokenValidator).to receive(:release_replay!)
+
+    expect do
+      described_class.call(envelope: envelope, provider: failing_provider)
+    end.to raise_error(RuntimeError, /provider unavailable/)
+
+    expect(Legion::Extensions::Llm::Fleet::TokenValidator).to have_received(:release_replay!).with('jti-2')
+  end
+end