scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 101

Merge 53b12558fe into 559247fa84

Browse source
This commit is contained in:
HX 2026-09-13 04:10:32 +08:00 • committed by GitHub
commit 343e3a0e49
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
13 changed files with 437 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

23
litellm/litellm_core_utils/exception_mapping_utils.py
View file

@ -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}",

1
litellm/litellm_core_utils/get_litellm_params.py
View file

@ -57,6 +57,7 @@ OPTIONAL_KWARGS_KEYS: Final = (
"itpm",
"otpm",
"use_xai_oauth",
"xai_oauth_token_file",
}
)
| AWS_CREDENTIAL_KWARGS_KEYS

8
litellm/llms/xai/chat/transformation.py
View file

@ -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,

37
litellm/llms/xai/oauth.py
View file

@ -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:

7
litellm/llms/xai/responses/transformation.py
View file

@ -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

1
litellm/main.py
View file

@ -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"),

5
litellm/proxy/auth/auth_utils.py
View file

@ -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

18
litellm/proxy/proxy_cli.py
View file

@ -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"):

8
litellm/types/router.py
View file

@ -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

1
litellm/types/utils.py
View file

@ -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",

287
tests/test_litellm/llms/xai/test_xai_oauth_multi_account.py Normal file
View file

@ -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

50
tests/test_litellm/proxy/auth/test_auth_utils.py
View file

@ -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
)

1
tests/test_litellm/proxy/auth/test_banned_params_extra_body.py
View file

@ -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):