path: root/spec/dispatch/adapter/claude/rate_limiter_spec.rb

# frozen_string_literal: true

require "webmock/rspec"
require "tmpdir"
require "fileutils"

RSpec.describe Dispatch::Adapter::Claude, "rate limiting" do
  let(:model_id)  { "claude-sonnet-4-6" }
  let(:api_key)   { "sk-ant-api03-test" }
  let(:base_url)  { "https://api.anthropic.com" }
  let(:tmpdir)    { Dir.mktmpdir("claude_rate_limit_test") }
  let(:store_path) { File.join(tmpdir, "claude_oauth.json") }
  let(:store) { Dispatch::Adapter::Claude::TokenStore.new(path: store_path) }

  after { FileUtils.rm_rf(tmpdir) }
  before { WebMock.disable_net_connect! }
  after  { WebMock.reset! }

  def make_adapter(min_request_interval: 0, rate_limit: nil, **opts)
    described_class.new(
      model: model_id,
      api_key: api_key,
      base_url: base_url,
      token_store: store,
      min_request_interval: min_request_interval,
      rate_limit: rate_limit,
      **opts
    )
  end

  let(:text_response) do
    {
      "id" => "msg_01",
      "type" => "message",
      "role" => "assistant",
      "model" => model_id,
      "stop_reason" => "end_turn",
      "content" => [{ "type" => "text", "text" => "Hello" }],
      "usage" => { "input_tokens" => 10, "output_tokens" => 5 }
    }
  end

  let(:messages) do
    [Dispatch::Adapter::Message.new(
      role: "user",
      content: [Dispatch::Adapter::TextBlock.new(text: "Hi")]
    )]
  end

  # ── Constructor accepts rate limit parameters ─────────────────────────────

  describe "constructor rate limit parameters" do
    it "accepts default min_request_interval of 1.0" do
      adapter = described_class.new(
        model: model_id,
        api_key: api_key,
        base_url: base_url,
        token_store: store
      )
      expect(adapter).to be_a(described_class)
    end

    it "accepts custom min_request_interval" do
      adapter = make_adapter(min_request_interval: 2.0)
      expect(adapter).to be_a(described_class)
    end

    it "accepts min_request_interval: 0 to disable cooldown" do
      adapter = make_adapter(min_request_interval: 0)
      expect(adapter).to be_a(described_class)
    end

    it "accepts rate_limit hash for sliding window" do
      adapter = make_adapter(rate_limit: { requests: 10, period: 60 })
      expect(adapter).to be_a(described_class)
    end

    it "raises ArgumentError for negative min_request_interval" do
      expect do
        make_adapter(min_request_interval: -1)
      end.to raise_error(ArgumentError)
    end

    it "raises ArgumentError for invalid rate_limit hash" do
      expect do
        make_adapter(rate_limit: { requests: 0, period: 60 })
      end.to raise_error(ArgumentError)
    end
  end

  # ── wait! is called before each request type ──────────────────────────────

  describe "#chat calls @rate_limiter.wait!" do
    before do
      stub_request(:post, "#{base_url}/v1/messages")
        .to_return(
          status: 200,
          body: JSON.generate(text_response),
          headers: { "Content-Type" => "application/json" }
        )
    end

    it "calls wait! once per chat request" do
      rate_limiter = instance_double(Dispatch::Adapter::RateLimiter)
      allow(Dispatch::Adapter::RateLimiter).to receive(:new).and_return(rate_limiter)
      allow(rate_limiter).to receive(:wait!)

      adapter = make_adapter
      adapter.chat(messages)

      expect(rate_limiter).to have_received(:wait!).once
    end

    it "calls wait! on each of multiple chat requests" do
      stub_request(:post, "#{base_url}/v1/messages")
        .to_return(
          status: 200,
          body: JSON.generate(text_response),
          headers: { "Content-Type" => "application/json" }
        )

      rate_limiter = instance_double(Dispatch::Adapter::RateLimiter)
      allow(Dispatch::Adapter::RateLimiter).to receive(:new).and_return(rate_limiter)
      allow(rate_limiter).to receive(:wait!)

      adapter = make_adapter
      adapter.chat(messages)
      adapter.chat(messages)

      expect(rate_limiter).to have_received(:wait!).exactly(2).times
    end
  end

  describe "#count_tokens calls @rate_limiter.wait!" do
    before do
      stub_request(:post, "#{base_url}/v1/messages/count_tokens")
        .to_return(
          status: 200,
          body: JSON.generate({ "input_tokens" => 42 }),
          headers: { "Content-Type" => "application/json" }
        )
    end

    it "calls wait! before counting tokens" do
      rate_limiter = instance_double(Dispatch::Adapter::RateLimiter)
      allow(Dispatch::Adapter::RateLimiter).to receive(:new).and_return(rate_limiter)
      allow(rate_limiter).to receive(:wait!)

      adapter = make_adapter
      adapter.count_tokens(messages)

      expect(rate_limiter).to have_received(:wait!).once
    end
  end

  describe "#usage_report calls @rate_limiter.wait!" do
    before do
      stub_request(:get, "#{base_url}/api/oauth/usage")
        .to_return(
          status: 200,
          body: JSON.generate({ "five_hour" => { "utilization" => 50.0 } }),
          headers: { "Content-Type" => "application/json" }
        )
      stub_request(:get, "#{base_url}/api/oauth/profile")
        .to_return(
          status: 200,
          body: JSON.generate({ "email" => "[email protected]", "account_id" => "acct" }),
          headers: { "Content-Type" => "application/json" }
        )
    end

    it "calls wait! before fetching usage (OAuth)" do
      rate_limiter = instance_double(Dispatch::Adapter::RateLimiter)
      allow(Dispatch::Adapter::RateLimiter).to receive(:new).and_return(rate_limiter)
      allow(rate_limiter).to receive(:wait!)

      oauth_adapter = described_class.new(
        model: model_id,
        api_key: "sk-ant-oat-test",
        base_url: base_url,
        token_store: store,
        is_oauth: true,
        min_request_interval: 0
      )
      oauth_adapter.usage_report

      expect(rate_limiter).to have_received(:wait!).once
    end
  end

  describe "#list_models calls @rate_limiter.wait!" do
    before do
      stub_request(:get, "#{base_url}/v1/models?limit=200")
        .to_return(
          status: 200,
          body: JSON.generate({ "data" => [{ "id" => model_id, "display_name" => "Sonnet" }] }),
          headers: { "Content-Type" => "application/json" }
        )
    end

    it "calls wait! on list_models (API fetch)" do
      rate_limiter = instance_double(Dispatch::Adapter::RateLimiter)
      allow(Dispatch::Adapter::RateLimiter).to receive(:new).and_return(rate_limiter)
      allow(rate_limiter).to receive(:wait!)

      adapter = make_adapter
      adapter.list_models

      expect(rate_limiter).to have_received(:wait!).once
    end

    it "does not call wait! on a cached list_models call" do
      rate_limiter = instance_double(Dispatch::Adapter::RateLimiter)
      allow(Dispatch::Adapter::RateLimiter).to receive(:new).and_return(rate_limiter)
      allow(rate_limiter).to receive(:wait!)

      adapter = make_adapter
      adapter.list_models  # first call: fetches from API
      adapter.list_models  # second call: from cache

      # Only 1 wait! for the first call; the cached call is free
      expect(rate_limiter).to have_received(:wait!).once
    end
  end

  # ── Two consecutive chat calls with real rate limiter ─────────────────────

  describe "two consecutive chat calls with min_request_interval: 1.0" do
    subject(:adapter) do
      described_class.new(
        model: model_id,
        api_key: api_key,
        base_url: base_url,
        token_store: store,
        min_request_interval: 1.0
      )
    end

    before do
      stub_request(:post, "#{base_url}/v1/messages")
        .to_return(
          status: 200,
          body: JSON.generate(text_response),
          headers: { "Content-Type" => "application/json" }
        ).times(2)
    end

    it "total elapsed for two calls is >= 1.0 s" do
      t0 = Time.now.to_f
      adapter.chat(messages)
      adapter.chat(messages)
      elapsed = Time.now.to_f - t0
      expect(elapsed).to be >= 1.0
    end
  end

  # ── min_request_interval: 0 disables throttling ───────────────────────────

  describe "min_request_interval: 0 disables throttling" do
    subject(:adapter) { make_adapter(min_request_interval: 0) }

    before do
      stub_request(:post, "#{base_url}/v1/messages")
        .to_return(
          status: 200,
          body: JSON.generate(text_response),
          headers: { "Content-Type" => "application/json" }
        ).times(2)
    end

    it "two consecutive calls complete in under 0.5 s" do
      t0 = Time.now.to_f
      adapter.chat(messages)
      adapter.chat(messages)
      elapsed = Time.now.to_f - t0
      expect(elapsed).to be < 0.5
    end
  end

  # ── Rate limit file location ───────────────────────────────────────────────

  describe "rate limit file location" do
    it "stores the rate limit file in the same directory as the token store" do
      stub_request(:post, "#{base_url}/v1/messages")
        .to_return(
          status: 200,
          body: JSON.generate(text_response),
          headers: { "Content-Type" => "application/json" }
        )

      adapter = make_adapter(min_request_interval: 1.0)
      adapter.chat(messages)

      rate_limit_path = File.join(tmpdir, "claude_rate_limit")
      expect(File.exist?(rate_limit_path)).to be true
    end
  end

  # ── DEFAULT_MIN_REQUEST_INTERVAL constant ─────────────────────────────────

  describe "DEFAULT_MIN_REQUEST_INTERVAL" do
    it "is 1.0" do
      expect(described_class::DEFAULT_MIN_REQUEST_INTERVAL).to eq(1.0)
    end
  end
end