diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index 82708d412c9..63d6f2c95e8 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -413,6 +413,29 @@ def _map_openai_exception( response=getattr(original_exception, "response", None), litellm_debug_info=extra_information, ) + elif original_exception.status_code == 403: + from litellm.llms.xai.oauth import is_xai_spending_limit_error + + if is_xai_spending_limit_error( + custom_llm_provider=custom_llm_provider, + status_code=original_exception.status_code, + error_str=error_str, + ): + raise RateLimitError( + message=f"RateLimitError: {exception_provider} - {message}", + model=model, + llm_provider=custom_llm_provider, + response=getattr(original_exception, "response", None), + litellm_debug_info=extra_information, + ) + raise APIError( + status_code=original_exception.status_code, + message=f"APIError: {exception_provider} - {message}", + llm_provider=custom_llm_provider, + model=model, + request=getattr(original_exception, "request", None), + litellm_debug_info=extra_information, + ) elif original_exception.status_code == 404: raise NotFoundError( message=f"NotFoundError: {exception_provider} - {message}", diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index edd2e88f95c..b231eadd7f0 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -57,6 +57,7 @@ OPTIONAL_KWARGS_KEYS: Final = ( "itpm", "otpm", "use_xai_oauth", + "xai_oauth_token_file", } ) | AWS_CREDENTIAL_KWARGS_KEYS diff --git a/litellm/llms/xai/chat/transformation.py b/litellm/llms/xai/chat/transformation.py index 747ee0b6c49..6e6353d015e 100644 --- a/litellm/llms/xai/chat/transformation.py +++ b/litellm/llms/xai/chat/transformation.py @@ -68,8 +68,10 @@ class XAIChatConfig(OpenAIGPTConfig): dynamic_api_key: Final = XAIModelInfo.get_api_key(api_key) if should_use_xai_oauth(litellm_params) and not dynamic_api_key: + raw_token_file: Final = (litellm_params or {}).get("xai_oauth_token_file") + token_file: Final = raw_token_file if isinstance(raw_token_file, str) else None try: - headers["Authorization"] = f"Bearer {XAIOAuthAuthenticator().get_access_token()}" + headers["Authorization"] = f"Bearer {XAIOAuthAuthenticator(auth_file=token_file).get_access_token()}" except XAIOAuthError as exc: raise AuthenticationError( model=model, @@ -103,7 +105,9 @@ class XAIChatConfig(OpenAIGPTConfig): dynamic_api_key: Final = XAIModelInfo.get_api_key(api_key) if should_use_xai_oauth(litellm_params) and not dynamic_api_key: - api_base = XAIOAuthAuthenticator().get_api_base() + raw_token_file: Final = (litellm_params or {}).get("xai_oauth_token_file") + token_file: Final = raw_token_file if isinstance(raw_token_file, str) else None + api_base = XAIOAuthAuthenticator(auth_file=token_file).get_api_base() return super().get_complete_url( api_base=api_base, diff --git a/litellm/llms/xai/oauth.py b/litellm/llms/xai/oauth.py index e8196ec6cb9..3476228fd85 100644 --- a/litellm/llms/xai/oauth.py +++ b/litellm/llms/xai/oauth.py @@ -2,6 +2,7 @@ import base64 import hashlib import json import os +import re import secrets import sys import threading @@ -30,6 +31,7 @@ XAI_OAUTH_REDIRECT_PORT: Final = 56121 XAI_OAUTH_REDIRECT_PATH: Final = "/callback" XAI_OAUTH_EXPIRY_SKEW_SECONDS: Final = 120 XAI_OAUTH_CALLBACK_TIMEOUT_SECONDS: Final = 180 +_XAI_OAUTH_ACCOUNT_NAME_RE: Final = re.compile(r"^[A-Za-z0-9_-]+$") _XAI_OAUTH_REFRESH_LOCK: Final = threading.Lock() @@ -75,6 +77,30 @@ class XAIOAuthLoginRequiredError(XAIOAuthError): pass +def _default_xai_oauth_token_dir() -> str: + return get_secret_str("XAI_OAUTH_TOKEN_DIR") or os.path.expanduser("~/.config/litellm/xai_oauth") + + +def resolve_xai_oauth_auth_file(auth_file: str | None, token_dir: str) -> str: + requested: Final = auth_file or get_secret_str("XAI_OAUTH_AUTH_FILE") or "auth.json" + candidate: Final = requested if os.path.isabs(requested) else os.path.join(token_dir, requested) + resolved: Final = os.path.realpath(candidate) + allowed: Final = os.path.realpath(token_dir) + if resolved == allowed or resolved.startswith(allowed + os.sep): + return resolved + raise XAIOAuthError("xAI OAuth auth file must stay inside the token directory") + + +def oauth_auth_file_for_account(account: str, token_dir: str) -> str: + if not _XAI_OAUTH_ACCOUNT_NAME_RE.fullmatch(account): + raise ValueError("xAI OAuth account must match ^[A-Za-z0-9_-]+$") + return resolve_xai_oauth_auth_file(f"auth-{account}.json", token_dir) + + +def is_xai_spending_limit_error(*, custom_llm_provider: str, status_code: int, error_str: str) -> bool: + return custom_llm_provider == "xai" and status_code == 403 and "spending-limit" in error_str + + class _CallbackHandler(BaseHTTPRequestHandler): server: "_CallbackServer" # pyright: ignore[reportIncompatibleVariableOverride] # stdlib stubs type server as BaseServer; _CallbackServer is the only server this handler is registered on @@ -121,9 +147,14 @@ class _CallbackServer(HTTPServer): class XAIOAuthAuthenticator: - def __init__(self, http_client: httpx.Client | HTTPHandler | None = None) -> None: - self.token_dir = get_secret_str("XAI_OAUTH_TOKEN_DIR") or os.path.expanduser("~/.config/litellm/xai_oauth") - self.auth_file = os.path.join(self.token_dir, get_secret_str("XAI_OAUTH_AUTH_FILE") or "auth.json") + def __init__( + self, + http_client: httpx.Client | HTTPHandler | None = None, + auth_file: str | None = None, + token_dir: str | None = None, + ) -> None: + self.token_dir = token_dir or _default_xai_oauth_token_dir() + self.auth_file = resolve_xai_oauth_auth_file(auth_file, self.token_dir) self.http_client = http_client def get_api_base(self) -> str: diff --git a/litellm/llms/xai/responses/transformation.py b/litellm/llms/xai/responses/transformation.py index 646d6798783..a1e53da1ab3 100644 --- a/litellm/llms/xai/responses/transformation.py +++ b/litellm/llms/xai/responses/transformation.py @@ -231,8 +231,9 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): ) if should_use_xai_oauth(litellm_params.model_dump()): + token_file: Final = litellm_params.xai_oauth_token_file try: - api_key = XAIOAuthAuthenticator().get_access_token() + api_key = XAIOAuthAuthenticator(auth_file=token_file).get_access_token() except XAIOAuthError as exc: raise AuthenticationError( model=model, @@ -268,7 +269,9 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): api_key: Final = XAIModelInfo.get_api_key(litellm_params.get("api_key"), legacy_generic_before_env=True) if should_use_xai_oauth(litellm_params) and not api_key: - api_base = XAIOAuthAuthenticator().get_api_base() + raw_token_file: Final = litellm_params.get("xai_oauth_token_file") + token_file: Final = raw_token_file if isinstance(raw_token_file, str) else None + api_base = XAIOAuthAuthenticator(auth_file=token_file).get_api_base() else: api_base = api_base or litellm.api_base or get_secret_str("XAI_API_BASE") or XAI_API_BASE diff --git a/litellm/main.py b/litellm/main.py index 17edafcdfca..f1ce8f42b8b 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -5559,6 +5559,7 @@ def completion( tpm=kwargs.get("tpm"), rpm=kwargs.get("rpm"), use_xai_oauth=kwargs.get("use_xai_oauth", False), + xai_oauth_token_file=kwargs.get("xai_oauth_token_file"), gigachat_scope=kwargs.get("gigachat_scope"), gigachat_auth_url=kwargs.get("gigachat_auth_url"), gigachat_access_token=kwargs.get("gigachat_access_token"), diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index be65c3b39ec..62ab58940ce 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -326,6 +326,11 @@ _BANNED_REQUEST_BODY_PARAMS: Final[tuple[str, ...]] = ( # caller-supplied value is the same exfil shape as # ``aws_web_identity_token`` on the Bedrock path. "azure_ad_token", + # xAI SuperGrok OAuth token file. The xAI transformer reads + # ``xai_oauth_token_file`` and authenticates as that local account, + # so a caller-supplied path is the same credential-selector shape as + # ``aws_profile_name`` on the Bedrock path. + "xai_oauth_token_file", # Endpoint-targeting fields that retarget the outbound request or # an observability callback. An attacker-controlled value either # exfiltrates the request payload (incl. messages + admin-set diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index c2d60cd5488..a92468a1897 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -1003,10 +1003,22 @@ def run_server( prometheus_metrics_port: int | None, ): if cli_args: - if cli_args == ("xai-oauth", "login"): - from litellm.llms.xai.oauth import XAIOAuthAuthenticator + if len(cli_args) >= 2 and cli_args[0] == "xai-oauth" and cli_args[1] == "login": + from litellm.llms.xai.oauth import ( + XAIOAuthAuthenticator, + oauth_auth_file_for_account, + ) + from litellm.secret_managers.main import get_secret_str - authenticator: Final = XAIOAuthAuthenticator() + account: Final = cli_args[2] if len(cli_args) >= 3 else None + token_dir: Final = get_secret_str("XAI_OAUTH_TOKEN_DIR") or os.path.expanduser( + "~/.config/litellm/xai_oauth" + ) + authenticator: Final = ( + XAIOAuthAuthenticator(auth_file=oauth_auth_file_for_account(account, token_dir), token_dir=token_dir) + if account + else XAIOAuthAuthenticator() + ) auth_data: Final = authenticator.login() click.echo(f"xAI OAuth login successful. Credentials saved to {authenticator.auth_file}.") if auth_data.get("expires_at"): diff --git a/litellm/types/router.py b/litellm/types/router.py index fc09c40fe08..57248e7b87c 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -349,6 +349,14 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): default=False, description="Use stored xAI OAuth credentials when no xAI API key is configured.", ) + xai_oauth_token_file: str | None = Field( + default=None, + description=( + "Per-deployment xAI OAuth token file path. When set with use_xai_oauth=True, " + "the request reads that account's OAuth credentials instead of the global " + "token file, enabling multi-account SuperGrok routing." + ), + ) model_config = ConfigDict(extra="allow", arbitrary_types_allowed=True) merge_reasoning_content_in_choices: bool | None = False model_info: dict | None = None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index e3ea37dc0c8..0a64b9dc05c 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3805,6 +3805,7 @@ all_litellm_params = ( "enable_tag_filtering", "enable_json_schema_validation", "use_xai_oauth", + "xai_oauth_token_file", "auto_router_config_path", "auto_router_config", "auto_router_default_model", diff --git a/tests/test_litellm/llms/xai/test_xai_oauth_multi_account.py b/tests/test_litellm/llms/xai/test_xai_oauth_multi_account.py new file mode 100644 index 00000000000..254b6cd151d --- /dev/null +++ b/tests/test_litellm/llms/xai/test_xai_oauth_multi_account.py @@ -0,0 +1,287 @@ +import json +import os +import time + +import pytest +from click.testing import CliRunner + +import litellm +from litellm.litellm_core_utils.exception_mapping_utils import _map_openai_exception +from litellm.llms.xai.chat.transformation import XAIChatConfig +from litellm.llms.xai.oauth import XAIOAuthAuthenticator, XAIOAuthError +from litellm.llms.xai.responses.transformation import XAIResponsesAPIConfig +from litellm.types.router import GenericLiteLLMParams + + +def test_authenticator_accepts_explicit_absolute_auth_file(tmp_path): + custom = tmp_path / "custom-dir" / "alice.json" + custom.parent.mkdir(parents=True) + custom.write_text( + json.dumps( + { + "access_token": "alice-token", + "refresh_token": "refresh-token", + "expires_at": time.time() + 3600, + } + ) + ) + + auth = XAIOAuthAuthenticator(auth_file=str(custom), token_dir=str(custom.parent)) + + assert auth.auth_file == os.path.realpath(str(custom)) + assert auth.get_access_token() == "alice-token" + + +def test_authenticator_relative_auth_file_joins_token_dir(tmp_path, monkeypatch): + token_dir = tmp_path / "xai_oauth" + token_dir.mkdir() + (token_dir / "auth-bob.json").write_text( + json.dumps( + { + "access_token": "bob-token", + "refresh_token": "refresh-token", + "expires_at": time.time() + 3600, + } + ) + ) + monkeypatch.setenv("XAI_OAUTH_TOKEN_DIR", str(token_dir)) + + auth = XAIOAuthAuthenticator(auth_file="auth-bob.json") + + assert auth.auth_file == os.path.realpath(str(token_dir / "auth-bob.json")) + assert auth.get_access_token() == "bob-token" + + +def test_authenticator_rejects_relative_path_escaping_token_dir(tmp_path, monkeypatch): + token_dir = tmp_path / "xai_oauth" + token_dir.mkdir() + monkeypatch.setenv("XAI_OAUTH_TOKEN_DIR", str(token_dir)) + + with pytest.raises(XAIOAuthError, match="token directory"): + XAIOAuthAuthenticator(auth_file="../secret.json") + + +def test_authenticator_rejects_absolute_path_outside_token_dir(tmp_path, monkeypatch): + token_dir = tmp_path / "xai_oauth" + token_dir.mkdir() + outside = tmp_path / "outside.json" + outside.write_text("{}") + monkeypatch.setenv("XAI_OAUTH_TOKEN_DIR", str(token_dir)) + + with pytest.raises(XAIOAuthError, match="token directory"): + XAIOAuthAuthenticator(auth_file=str(outside)) + + +def test_authenticator_rejects_dotdot_absolute_path_outside_token_dir(tmp_path, monkeypatch): + token_dir = tmp_path / "xai_oauth" + token_dir.mkdir() + monkeypatch.setenv("XAI_OAUTH_TOKEN_DIR", str(token_dir)) + escaped = str((token_dir / ".." / "outside.json").resolve()) + + with pytest.raises(XAIOAuthError, match="token directory"): + XAIOAuthAuthenticator(auth_file=escaped) + + +def test_oauth_auth_file_for_account_builds_path_inside_token_dir(tmp_path): + from litellm.llms.xai.oauth import oauth_auth_file_for_account + + assert oauth_auth_file_for_account("alice", str(tmp_path)) == os.path.realpath( + os.path.join(str(tmp_path), "auth-alice.json") + ) + assert oauth_auth_file_for_account("Bob_1-2", str(tmp_path)) == os.path.realpath( + os.path.join(str(tmp_path), "auth-Bob_1-2.json") + ) + + +@pytest.mark.parametrize( + "account", + ["", ".", "..", "foo/bar", "../alice", "alice.json", "foo\\bar", "alice/../bob"], +) +def test_oauth_auth_file_for_account_rejects_unsafe_names(account, tmp_path): + from litellm.llms.xai.oauth import oauth_auth_file_for_account + + with pytest.raises(ValueError, match="account"): + oauth_auth_file_for_account(account, str(tmp_path)) + + +def test_authenticator_explicit_auth_file_overrides_env(tmp_path, monkeypatch): + env_file = tmp_path / "env.json" + env_file.write_text( + json.dumps( + { + "access_token": "env-token", + "refresh_token": "r", + "expires_at": time.time() + 3600, + } + ) + ) + explicit_file = tmp_path / "explicit.json" + explicit_file.write_text( + json.dumps( + { + "access_token": "explicit-token", + "refresh_token": "r", + "expires_at": time.time() + 3600, + } + ) + ) + monkeypatch.setenv("XAI_OAUTH_AUTH_FILE", str(env_file)) + + auth = XAIOAuthAuthenticator(auth_file=str(explicit_file), token_dir=str(tmp_path)) + + assert auth.get_access_token() == "explicit-token" + + +def test_chat_config_uses_per_deployment_token_file(tmp_path, monkeypatch): + monkeypatch.setenv("XAI_OAUTH_TOKEN_DIR", str(tmp_path)) + alice_file = tmp_path / "auth-alice.json" + alice_file.write_text( + json.dumps( + { + "access_token": "alice-chat-token", + "refresh_token": "refresh-token", + "expires_at": time.time() + 3600, + } + ) + ) + + headers = XAIChatConfig().validate_environment( + headers={}, + model="grok-4", + messages=[], + optional_params={}, + litellm_params={ + "use_xai_oauth": True, + "xai_oauth_token_file": str(alice_file), + }, + api_key=None, + ) + + assert headers["Authorization"] == "Bearer alice-chat-token" + + +def test_chat_config_multi_deployment_token_isolation(tmp_path, monkeypatch): + monkeypatch.setenv("XAI_OAUTH_TOKEN_DIR", str(tmp_path)) + alice = tmp_path / "auth-alice.json" + alice.write_text( + json.dumps( + { + "access_token": "alice-token", + "refresh_token": "r", + "expires_at": time.time() + 3600, + } + ) + ) + bob = tmp_path / "auth-bob.json" + bob.write_text( + json.dumps( + { + "access_token": "bob-token", + "refresh_token": "r", + "expires_at": time.time() + 3600, + } + ) + ) + + headers_alice = XAIChatConfig().validate_environment( + headers={}, + model="grok-4", + messages=[], + optional_params={}, + litellm_params={ + "use_xai_oauth": True, + "xai_oauth_token_file": str(alice), + }, + api_key=None, + ) + headers_bob = XAIChatConfig().validate_environment( + headers={}, + model="grok-4", + messages=[], + optional_params={}, + litellm_params={ + "use_xai_oauth": True, + "xai_oauth_token_file": str(bob), + }, + api_key=None, + ) + + assert headers_alice["Authorization"] == "Bearer alice-token" + assert headers_bob["Authorization"] == "Bearer bob-token" + + +def test_responses_config_uses_per_deployment_token_file(tmp_path, monkeypatch): + monkeypatch.setenv("XAI_OAUTH_TOKEN_DIR", str(tmp_path)) + bob = tmp_path / "auth-bob.json" + bob.write_text( + json.dumps( + { + "access_token": "bob-responses-token", + "refresh_token": "r", + "expires_at": time.time() + 3600, + } + ) + ) + + headers = XAIResponsesAPIConfig().validate_environment( + headers={}, + model="grok-4", + litellm_params=GenericLiteLLMParams( + use_xai_oauth=True, + xai_oauth_token_file=str(bob), + ), + ) + + assert headers["Authorization"] == "Bearer bob-responses-token" + + +def test_proxy_cli_xai_oauth_login_rejects_unsafe_account(monkeypatch, tmp_path): + from litellm.proxy.proxy_cli import run_server + + monkeypatch.setenv("XAI_OAUTH_TOKEN_DIR", str(tmp_path)) + result = CliRunner().invoke(run_server, ["xai-oauth", "login", "../alice"]) + + assert result.exit_code != 0 + assert result.exception is not None + + +def test_spending_limit_403_maps_to_rate_limit_error(): + class FakeHTTPError(Exception): + def __init__(self): + self.status_code = 403 + self.message = "personal-team-blocked:spending-limit" + self.response = None + super().__init__(self.message) + + with pytest.raises(litellm.RateLimitError): + _map_openai_exception( + model="grok-4", + original_exception=FakeHTTPError(), + custom_llm_provider="xai", + error_str="personal-team-blocked:spending-limit", + exception_type="APIStatusError", + exception_provider="XaiException", + extra_information="", + ) + + +def test_spending_limit_403_on_non_xai_does_not_map_to_rate_limit_error(): + class FakeHTTPError(Exception): + def __init__(self): + self.status_code = 403 + self.message = "personal-team-blocked:spending-limit" + self.response = None + super().__init__(self.message) + + with pytest.raises(litellm.APIError) as exc_info: + _map_openai_exception( + model="gpt-4", + original_exception=FakeHTTPError(), + custom_llm_provider="openai", + error_str="personal-team-blocked:spending-limit", + exception_type="APIStatusError", + exception_provider="OpenAIException", + extra_information="", + ) + assert not isinstance(exc_info.value, litellm.RateLimitError) + assert exc_info.value.status_code == 403 diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index cdf1f897707..d5328822e32 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -3586,3 +3586,53 @@ class TestIsRequestBodySafeBlocksAwsIdentitySelectors: ) is True ) + + +class TestIsRequestBodySafeBlocksXaiOauthTokenFile: + """A caller must not select another local xAI OAuth account. Router kwargs + override deployment params, so ``xai_oauth_token_file`` in the request body + would authenticate as any token file readable on the proxy host. + """ + + def test_xai_oauth_token_file_is_in_banned_set(self): + from litellm.proxy.auth.auth_utils import _BANNED_REQUEST_BODY_PARAMS + + assert "xai_oauth_token_file" in set(_BANNED_REQUEST_BODY_PARAMS) + + def test_xai_oauth_token_file_in_request_body_is_rejected(self): + with pytest.raises(ValueError, match="xai_oauth_token_file"): + is_request_body_safe( + request_body={ + "model": "grok-4", + "xai_oauth_token_file": "/etc/passwd", + }, + general_settings={}, + llm_router=None, + model="grok-4", + ) + + def test_xai_oauth_token_file_under_extra_body_is_rejected(self): + with pytest.raises(ValueError, match="xai_oauth_token_file"): + is_request_body_safe( + request_body={ + "model": "grok-4", + "extra_body": {"xai_oauth_token_file": "/etc/passwd"}, + }, + general_settings={}, + llm_router=None, + model="grok-4", + ) + + def test_xai_oauth_token_file_allowed_under_proxy_wide_opt_in(self): + assert ( + is_request_body_safe( + request_body={ + "model": "grok-4", + "xai_oauth_token_file": "auth-alice.json", + }, + general_settings={"allow_client_side_credentials": True}, + llm_router=None, + model="grok-4", + ) + is True + ) diff --git a/tests/test_litellm/proxy/auth/test_banned_params_extra_body.py b/tests/test_litellm/proxy/auth/test_banned_params_extra_body.py index e87b206a40a..6c621b03cba 100644 --- a/tests/test_litellm/proxy/auth/test_banned_params_extra_body.py +++ b/tests/test_litellm/proxy/auth/test_banned_params_extra_body.py @@ -28,6 +28,7 @@ from litellm.proxy.auth.auth_utils import is_request_body_safe # noqa: E402 "base_url", "vertex_credentials", "azure_ad_token", + "xai_oauth_token_file", ], ) def test_banned_param_under_extra_body_is_rejected(banned_param):