diff --git a/src/rotator_library/client.py b/src/rotator_library/client.py index 93246cba..34f64bc5 100644 --- a/src/rotator_library/client.py +++ b/src/rotator_library/client.py @@ -65,6 +65,43 @@ def __init__(self, message, data=None): self.data = data +OPENROUTER_ATTRIBUTION_HEADER_KEYS = ( + "HTTP-Referer", + "X-OpenRouter-Title", + "X-Title", + "X-OpenRouter-Categories", +) + + +def _get_request_header(request: Optional[Any], key: str) -> Optional[str]: + if request is None: + return None + headers = getattr(request, "headers", None) + if headers is None: + return None + if hasattr(headers, "get"): + value = headers.get(key) + if value is None: + value = headers.get(key.lower()) + return value + if isinstance(headers, dict): + return headers.get(key) or headers.get(key.lower()) + return None + + +def _merge_openrouter_extra_headers( + litellm_kwargs: Dict[str, Any], request: Optional[Any] +) -> Dict[str, Any]: + extra_headers = dict(litellm_kwargs.get("extra_headers") or {}) + for key in OPENROUTER_ATTRIBUTION_HEADER_KEYS: + value = _get_request_header(request, key) + if value and key not in extra_headers: + extra_headers[key] = value + if extra_headers: + litellm_kwargs["extra_headers"] = extra_headers + return litellm_kwargs + + class RotatingClient: """ A client that intelligently rotates and retries API keys using LiteLLM, @@ -1698,6 +1735,8 @@ async def _execute_with_retry( ] litellm_kwargs = sanitize_request_payload(litellm_kwargs, model) + if provider == "openrouter": + litellm_kwargs = _merge_openrouter_extra_headers(litellm_kwargs, request) for attempt in range(self.max_retries): try: @@ -2431,6 +2470,8 @@ async def _streaming_acompletion_with_retry( ] litellm_kwargs = sanitize_request_payload(litellm_kwargs, model) + if provider == "openrouter": + litellm_kwargs = _merge_openrouter_extra_headers(litellm_kwargs, request) # If the provider is 'qwen_code', set the custom provider to 'qwen' # and strip the prefix from the model name for LiteLLM. diff --git a/tests/test_client_routing_policy.py b/tests/test_client_routing_policy.py index 284734f8..68fd2a00 100644 --- a/tests/test_client_routing_policy.py +++ b/tests/test_client_routing_policy.py @@ -1,3 +1,4 @@ +import importlib import sys import asyncio import random @@ -9,8 +10,13 @@ ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(ROOT / "src")) -from rotator_library.client import RotatingClient -from rotator_library.routing_policy import RoutingPolicy, RoutingPolicyError +client_module = importlib.import_module("rotator_library.client") +routing_policy_module = importlib.import_module("rotator_library.routing_policy") + +RotatingClient = getattr(client_module, "RotatingClient") +_merge_openrouter_extra_headers = getattr(client_module, "_merge_openrouter_extra_headers") +RoutingPolicy = getattr(routing_policy_module, "RoutingPolicy") +RoutingPolicyError = getattr(routing_policy_module, "RoutingPolicyError") def test_client_helper_rewrites_weighted_router_model(): @@ -125,3 +131,42 @@ def test_client_helper_rewrites_weighted_qwen3_5_model(): assert decision is not None assert decision.strategy == "weighted" assert decision.excluded_providers == ["opencode_go"] + + +def test_merge_openrouter_extra_headers_copies_attribution_from_request(): + class Request: + headers = { + "HTTP-Referer": "https://opencode.ai", + "X-OpenRouter-Title": "OpenCode/opencode-router", + "X-Title": "OpenCode/opencode-router", + "X-OpenRouter-Categories": "cli-agent", + } + + kwargs = _merge_openrouter_extra_headers({"messages": []}, Request()) + + assert kwargs["extra_headers"]["HTTP-Referer"] == "https://opencode.ai" + assert kwargs["extra_headers"]["X-OpenRouter-Title"] == "OpenCode/opencode-router" + assert kwargs["extra_headers"]["X-Title"] == "OpenCode/opencode-router" + assert kwargs["extra_headers"]["X-OpenRouter-Categories"] == "cli-agent" + + +def test_merge_openrouter_extra_headers_preserves_existing_values(): + class Request: + headers = { + "HTTP-Referer": "https://opencode.ai", + "X-OpenRouter-Title": "OpenCode/opencode-router", + } + + kwargs = _merge_openrouter_extra_headers( + { + "extra_headers": { + "HTTP-Referer": "https://custom.example", + "Existing": "value", + } + }, + Request(), + ) + + assert kwargs["extra_headers"]["HTTP-Referer"] == "https://custom.example" + assert kwargs["extra_headers"]["X-OpenRouter-Title"] == "OpenCode/opencode-router" + assert kwargs["extra_headers"]["Existing"] == "value"