defmodule Gemini.RateLimiter.Manager do @moduledoc """ Central rate limiter manager that coordinates request submission. Wraps outbound requests with: - Rate limit checking and enforcement - Concurrency gating - Token budgeting - Retry handling with backoff Enabled by default. Use `disable_rate_limiter: true` to opt out. ## Features - ETS-based state for cross-process visibility - Per-model/location/metric tracking - Configurable concurrency limits with adaptive mode - Token budget estimation and tracking - Telemetry event emission ## Usage # Execute a request through the rate limiter {:ok, response} = Manager.execute( fn -> HTTP.post(path, body, opts) end, "gemini-flash-lite-latest", opts ) # Non-blocking mode returns immediately if rate limited case Manager.execute(fn -> ... end, model, non_blocking: true) do {:ok, response} -> handle_response(response) {:error, {:rate_limited, retry_at, details}} -> schedule_retry(retry_at) end ## Configuration Configure via application environment or per-request options: config :gemini_ex, :rate_limiter, max_concurrency_per_model: 4, max_attempts: 3, base_backoff_ms: 1000, profile: :prod Per-request overrides: Gemini.generate("Hello", [ disable_rate_limiter: true, # Bypass rate limiter non_blocking: true, # Return immediately if rate limited max_concurrency_per_model: 8 # Override concurrency ]) """ use GenServer alias Gemini.RateLimiter.{Config, State, ConcurrencyGate, RetryManager} alias Gemini.Telemetry @type execute_opts :: keyword() @type request_fn :: (-> {:ok, term()} | {:error, term()}) # Client API @doc """ Start the rate limiter manager. """ @spec start_link(keyword()) :: GenServer.on_start() def start_link(opts \\ []) do GenServer.start_link(__MODULE__, opts, name: __MODULE__) end @doc """ Execute a request through the rate limiter. ## Parameters - `request_fn` - Zero-arity function that makes the actual HTTP request - `model` - Model name for rate limit tracking - `opts` - Options for rate limiting and the underlying request ## Options - `:location` - Location for rate limit tracking (default: "us-central1") - `:disable_rate_limiter` - Bypass all rate limiting (default: false) - `:non_blocking` - Return immediately if rate limited (default: false) - `:max_concurrency_per_model` - Override concurrency limit - `:estimated_input_tokens` - Estimated tokens for budget checking - `:estimated_cached_tokens` - Estimated cached-context tokens for budget checking - `:token_budget_per_window` - Maximum tokens per window (nil = no limit) ## Returns - `{:ok, response}` - Request succeeded - `{:error, {:rate_limited, retry_at, details}}` - Rate limited - `{:error, {:transient_failure, attempts, last_error}}` - Transient failure - `{:error, term()}` - Other error """ @spec execute(request_fn(), String.t(), execute_opts()) :: {:ok, term()} | {:error, term()} def execute(request_fn, model, opts \\ []) do config = Config.build(opts) if Config.enabled?(config) do do_execute(request_fn, model, config, opts) else # Rate limiter disabled, execute directly request_fn.() end end @doc """ Execute a request, extracting and recording usage from the response. Similar to `execute/3` but also records token usage from successful responses. """ @spec execute_with_usage_tracking(request_fn(), String.t(), execute_opts()) :: {:ok, term()} | {:error, term()} def execute_with_usage_tracking(request_fn, model, opts \\ []) do result = execute(request_fn, model, opts) # Record usage from successful responses case result do {:ok, response} when is_map(response) -> record_usage_from_response(model, opts, response) _ -> :ok end result end @doc """ Check if a request would be rate limited without executing it. ## Returns - `:ok` - Request can proceed - `{:rate_limited, retry_at, details}` - Currently rate limited - `{:over_budget, usage}` - Would exceed token budget - `{:no_permits, available}` - No concurrency permits available """ @spec check_status(String.t(), execute_opts()) :: :ok | {:rate_limited, DateTime.t(), map()} | {:over_budget, map()} | {:no_permits, non_neg_integer()} def check_status(model, opts \\ []) do config = Config.build(opts) location = Keyword.get(opts, :location) state_key = State.build_key(model, location, :token_count) budget_status = check_token_budget(state_key, opts, config) cond do not Config.enabled?(config) -> :ok retry_until = State.get_retry_until(state_key) -> retry_state = State.get_retry_state(state_key) {:rate_limited, retry_until, %{ quota_metric: retry_state && retry_state.quota_metric, quota_id: retry_state && retry_state.quota_id }} Config.concurrency_enabled?(config) and ConcurrencyGate.available_permits(model, config) == 0 -> {:no_permits, 0} match?({:over_budget, _}, budget_status) -> {:over_budget, build_over_budget_status(budget_status, state_key)} true -> :ok end end @doc """ Get the current retry state for a model. """ @spec get_retry_state(String.t(), keyword()) :: State.retry_state() | nil def get_retry_state(model, opts \\ []) do location = Keyword.get(opts, :location) state_key = State.build_key(model, location, :token_count) State.get_retry_state(state_key) end defp build_over_budget_status({:over_budget, budget_ctx}, state_key) do budget_ctx |> Map.put_new(:usage, State.get_current_usage(state_key) || %{}) end defp build_over_budget_status(_status, state_key), do: State.get_current_usage(state_key) || %{} @doc """ Get current token usage for a model. """ @spec get_usage(String.t(), keyword()) :: State.usage_window() | nil def get_usage(model, opts \\ []) do location = Keyword.get(opts, :location) state_key = State.build_key(model, location, :token_count) State.get_current_usage(state_key) end @doc """ Reset all rate limiter state (useful for testing). """ @spec reset_all() :: :ok def reset_all do State.reset_all() ConcurrencyGate.reset_all() :ok end # GenServer callbacks @impl true def init(_opts) do # Initialize ETS tables State.init() ConcurrencyGate.init() {:ok, %{}} end # Private implementation defp do_execute(request_fn, model, config, opts) do location = Keyword.get(opts, :location) state_key = State.build_key(model, location, :token_count) concurrency_key = concurrency_key(model, opts) # Emit telemetry for request start start_time = System.monotonic_time() emit_request_start(model, opts) # Check token budget case check_token_budget(state_key, opts, config) do {:ok, _budget_ctx} -> execute_with_concurrency( request_fn, model, concurrency_key, state_key, config, start_time, opts ) {:over_budget, budget_ctx} -> handle_over_budget( state_key, config, start_time, opts, request_fn, model, concurrency_key, budget_ctx ) end end defp execute_with_concurrency( request_fn, model, concurrency_key, state_key, config, start_time, opts ) do # Acquire concurrency permit if enabled permit_result = if Config.concurrency_enabled?(config) do ConcurrencyGate.acquire(concurrency_key, config) else :ok end case permit_result do :ok -> try do result = execute_with_retry(request_fn, model, concurrency_key, state_key, config, opts) emit_request_complete(model, start_time, result, opts) result after if Config.concurrency_enabled?(config) do ConcurrencyGate.release(concurrency_key) end end {:error, :no_permit_available} -> emit_request_error(model, start_time, :no_permit_available, opts) {:error, {:rate_limited, nil, %{reason: :no_permit_available}}} {:error, :concurrency_disabled} -> # Concurrency disabled, proceed without permit result = execute_with_retry(request_fn, model, concurrency_key, state_key, config, opts) emit_request_complete(model, start_time, result, opts) result {:error, reason} -> emit_request_error(model, start_time, reason, opts) {:error, reason} end end defp execute_with_retry(request_fn, _model, concurrency_key, state_key, config, opts) do wrapped_fn = fn -> result = request_fn.() # Signal success/429 to adaptive concurrency case RetryManager.classify_response(result) do :success -> ConcurrencyGate.signal_success(concurrency_key, config) :rate_limited -> ConcurrencyGate.signal_429(concurrency_key, config) _ -> :ok end result end RetryManager.execute_with_retry(wrapped_fn, state_key, config, opts) end defp check_token_budget(state_key, opts, config) do # ADR-0001/0002: Use estimated tokens from opts, fall back to 0 estimated_input_tokens = Keyword.get(opts, :estimated_input_tokens, 0) estimated_cached_tokens = Keyword.get(opts, :estimated_cached_tokens, 0) estimated_total = estimated_input_tokens + estimated_cached_tokens # ADR-0002: Fall back to config.token_budget_per_window when not in opts budget = Keyword.get(opts, :token_budget_per_window, config.token_budget_per_window) usage = State.get_current_usage(state_key) cond do is_nil(budget) -> {:ok, %{ estimated_input_tokens: estimated_input_tokens, estimated_cached_tokens: estimated_cached_tokens, estimated_total_tokens: estimated_total, budget: budget, usage: usage }} estimated_total > budget -> {:over_budget, %{ reason: :over_budget, request_too_large: true, estimated_input_tokens: estimated_input_tokens, estimated_cached_tokens: estimated_cached_tokens, estimated_total_tokens: estimated_total, token_budget: budget, usage: usage }} usage && usage.input_tokens + usage.output_tokens + estimated_total > budget -> window_end = DateTime.add(usage.window_start, usage.window_duration_ms, :millisecond) {:over_budget, %{ reason: :over_budget, request_too_large: false, estimated_input_tokens: estimated_input_tokens, estimated_cached_tokens: estimated_cached_tokens, estimated_total_tokens: estimated_total, token_budget: budget, usage: usage, window_end: window_end }} true -> {:ok, %{ estimated_input_tokens: estimated_input_tokens, estimated_cached_tokens: estimated_cached_tokens, estimated_total_tokens: estimated_total, budget: budget, usage: usage }} end end defp handle_over_budget( state_key, config, start_time, opts, request_fn, model, concurrency_key, budget_ctx ) do retry_at = Map.get(budget_ctx, :window_end) base_details = rate_limit_details(budget_ctx) max_wait_ms = config.max_budget_wait_ms cond do Map.get(budget_ctx, :request_too_large) -> emit_rate_limit_error(state_key, :over_budget, start_time, base_details) {:error, {:rate_limited, nil, rate_limit_details(budget_ctx)}} config.non_blocking -> emit_rate_limit_error(state_key, :over_budget, start_time, base_details) {:error, {:rate_limited, retry_at, rate_limit_details(budget_ctx)}} true -> # Blocking mode: wait for window to clear once, then re-check if retry_at do wait_ms = max(0, DateTime.diff(retry_at, DateTime.utc_now(), :millisecond)) capped_wait = case max_wait_ms do nil -> {wait_ms, false} cap when wait_ms > cap -> {cap, true} _ -> {wait_ms, false} end {actual_wait_ms, capped?} = capped_wait wait_metadata = Map.merge(base_details, %{wait_ms: actual_wait_ms, wait_capped: capped?}) emit_rate_limit_wait(state_key, retry_at, :over_budget, wait_metadata) if actual_wait_ms > 0 do Process.sleep(actual_wait_ms) end end case check_token_budget(state_key, opts, config) do {:ok, _} -> execute_with_concurrency( request_fn, model, concurrency_key, state_key, config, start_time, opts ) {:over_budget, %{request_too_large: true} = over_again} -> emit_rate_limit_error( state_key, :over_budget, start_time, rate_limit_details(over_again) ) {:error, {:rate_limited, nil, rate_limit_details(over_again)}} {:over_budget, over_again} -> next_retry = Map.get(over_again, :window_end) emit_rate_limit_error( state_key, :over_budget, start_time, rate_limit_details(over_again) ) {:error, {:rate_limited, next_retry, rate_limit_details(over_again)}} end end end defp rate_limit_details(budget_ctx) do Map.merge( %{reason: :over_budget}, Map.take(budget_ctx, [ :estimated_input_tokens, :estimated_cached_tokens, :estimated_total_tokens, :token_budget, :request_too_large ]) ) end defp record_usage_from_response(model, opts, response) do location = Keyword.get(opts, :location) state_key = State.build_key(model, location, :token_count) # ADR-0002: Get window duration from config config = Config.build(opts) # Extract usage from response usage = extract_usage(response) if usage do State.record_usage( state_key, Map.get(usage, :input_tokens, 0), Map.get(usage, :output_tokens, 0), window_duration_ms: config.window_duration_ms ) end end defp extract_usage(response) do cond do Map.has_key?(response, :usage_metadata) -> cached_tokens = Map.get(response.usage_metadata, :cached_content_token_count, 0) %{ input_tokens: Map.get(response.usage_metadata, :prompt_token_count, 0) + cached_tokens, output_tokens: Map.get(response.usage_metadata, :candidates_token_count, 0) } Map.has_key?(response, "usageMetadata") -> cached_tokens = Map.get(response["usageMetadata"], "cachedContentTokenCount", 0) %{ input_tokens: Map.get(response["usageMetadata"], "promptTokenCount", 0) + cached_tokens, output_tokens: Map.get(response["usageMetadata"], "candidatesTokenCount", 0) } true -> nil end end # Telemetry helpers defp emit_request_start(model, opts) do metadata = %{ model: model, location: Keyword.get(opts, :location), system_time: System.system_time() } Telemetry.execute([:gemini, :rate_limit, :request, :start], %{}, metadata) end defp emit_request_complete(model, start_time, result, opts) do duration = Telemetry.calculate_duration(start_time) status = case result do {:ok, _} -> :success {:error, {:rate_limited, _, _}} -> :rate_limited {:error, {:transient_failure, _, _}} -> :transient_failure {:error, _} -> :error end metadata = %{ model: model, location: Keyword.get(opts, :location), status: status } Telemetry.execute([:gemini, :rate_limit, :request, :stop], %{duration: duration}, metadata) end defp emit_request_error(model, start_time, reason, opts) do duration = Telemetry.calculate_duration(start_time) metadata = %{ model: model, location: Keyword.get(opts, :location), reason: reason } Telemetry.execute( [:gemini, :rate_limit, :request, :error], %{duration: duration}, metadata ) end defp emit_rate_limit_wait(state_key, retry_at, reason, metadata) do {model, location, _metric} = state_key metadata = %{ model: model, location: location, retry_at: retry_at, reason: reason } |> Map.merge(metadata) Telemetry.execute([:gemini, :rate_limit, :wait], %{}, metadata) end defp emit_rate_limit_error(state_key, reason, start_time, metadata) do {model, location, _metric} = state_key duration = Telemetry.calculate_duration(start_time) metadata = %{ model: model, location: location, reason: reason } |> Map.merge(metadata) Telemetry.execute([:gemini, :rate_limit, :error], %{duration: duration}, metadata) end defp concurrency_key(model, opts) do case Keyword.get(opts, :concurrency_key) do nil -> model key -> "#{model}:#{to_string(key)}" end end end