mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-04 07:27:51 +00:00
fix(cli): renew an expired login once, show the proxy's 401 reason, read a models file once, try each BROWSER entry
This commit is contained in:
parent
07cafbc198
commit
c5ed701362
3 changed files with 253 additions and 35 deletions
|
|
@ -1,8 +1,9 @@
|
|||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
import webbrowser
|
||||
from collections.abc import Callable, Generator, Mapping
|
||||
from collections.abc import Callable, Generator, Mapping, Sequence
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass, field
|
||||
from http import HTTPStatus
|
||||
|
|
@ -15,9 +16,10 @@ from filelock import FileLock, Timeout
|
|||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
from litellm.litellm_core_utils.cli_token_utils import get_cli_token_file_path, load_cli_token
|
||||
from litellm.litellm_core_utils.secret_redaction import redact_string
|
||||
|
||||
from ...exceptions import UnauthorizedError
|
||||
from .auth import context_secret_vault, get_stored_api_key, replace_stored_token, storage_notice
|
||||
from .auth import context_secret_vault, replace_stored_token, storage_notice
|
||||
from .pkce_login import LOGIN_TIMEOUT_SECONDS, PkceCredential, PkceFailure, pkce_token_record, run_pkce_login
|
||||
|
||||
|
||||
|
|
@ -49,12 +51,15 @@ def _ssh_session_without_a_display(environ: Mapping[str, str]) -> bool:
|
|||
return over_ssh and not any(environ.get(name) for name in _DISPLAY_VARIABLES)
|
||||
|
||||
|
||||
_OPEN_URL_IN_BROWSER: Final = "import sys, webbrowser; sys.exit(0 if webbrowser.get().open(sys.argv[1]) else 1)"
|
||||
_OPEN_URL_IN_BROWSER: Final = (
|
||||
"import sys, webbrowser; "
|
||||
"sys.exit(0 if any(webbrowser.get(name).open(sys.argv[1]) for name in sys.argv[2:] or [None]) else 1)"
|
||||
)
|
||||
|
||||
|
||||
def open_browser_detached(url: str) -> PkceFailure | None:
|
||||
def open_browser_detached(url: str, browsers: Sequence[str]) -> PkceFailure | None:
|
||||
launcher: Final = subprocess.run(
|
||||
(sys.executable, "-I", "-c", _OPEN_URL_IN_BROWSER, url),
|
||||
(sys.executable, "-I", "-c", _OPEN_URL_IN_BROWSER, url, *browsers),
|
||||
stdin=subprocess.DEVNULL,
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.DEVNULL,
|
||||
|
|
@ -67,11 +72,40 @@ def _live_environ() -> Mapping[str, str]:
|
|||
return os.environ
|
||||
|
||||
|
||||
def _listed_browsers(environ: Mapping[str, str]) -> tuple[str, ...]:
|
||||
return tuple(name for name in environ.get("BROWSER", "").split(os.pathsep) if name)
|
||||
|
||||
|
||||
def _locate_browser(
|
||||
find_browser: Callable[[str | None], webbrowser.BaseBrowser], name: str | None
|
||||
) -> webbrowser.BaseBrowser | None:
|
||||
try:
|
||||
return find_browser(name)
|
||||
except webbrowser.Error:
|
||||
return None
|
||||
|
||||
|
||||
def _shows_a_window(browser: webbrowser.BaseBrowser) -> bool:
|
||||
return Path(browser.name).name not in _TEXT_MODE_BROWSERS
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PkceBrowserSignIn:
|
||||
find_browser: Callable[[], webbrowser.BaseBrowser] = webbrowser.get
|
||||
find_browser: Callable[[str | None], webbrowser.BaseBrowser] = webbrowser.get
|
||||
environ: Mapping[str, str] = field(default_factory=_live_environ)
|
||||
open_browser: Callable[[str], object] = open_browser_detached
|
||||
open_browser: Callable[[str, Sequence[str]], object] = open_browser_detached
|
||||
|
||||
def _browsers_to_open(self) -> tuple[str, ...] | PkceFailure:
|
||||
names: Final = _listed_browsers(self.environ) or (None,)
|
||||
located: Final = tuple(
|
||||
(name, browser) for name in names if (browser := _locate_browser(self.find_browser, name)) is not None
|
||||
)
|
||||
if not located:
|
||||
return PkceFailure("no web browser is available on this machine")
|
||||
windowed: Final = tuple(name for name, browser in located if _shows_a_window(browser))
|
||||
if not windowed:
|
||||
return PkceFailure("the web browser is a text-mode one, which lite cannot show while a command waits")
|
||||
return tuple(name for name in windowed if name is not None)
|
||||
|
||||
def unavailable_because(self) -> PkceFailure | None:
|
||||
if _runs_in_ci(self.environ):
|
||||
|
|
@ -80,21 +114,21 @@ class PkceBrowserSignIn:
|
|||
return PkceFailure(
|
||||
"this is an SSH session without a display, so a browser would open where you cannot see it"
|
||||
)
|
||||
try:
|
||||
browser: Final = self.find_browser()
|
||||
except webbrowser.Error:
|
||||
return PkceFailure("no web browser is available on this machine")
|
||||
if Path(browser.name).name in _TEXT_MODE_BROWSERS:
|
||||
return PkceFailure("the web browser is a text-mode one, which lite cannot show while a command waits")
|
||||
return None
|
||||
browsers: Final = self._browsers_to_open()
|
||||
return browsers if isinstance(browsers, PkceFailure) else None
|
||||
|
||||
def __call__(self, base_url: str, echo: Callable[[str], None]) -> PkceCredential | PkceFailure:
|
||||
browsers: Final = self._browsers_to_open()
|
||||
if isinstance(browsers, PkceFailure):
|
||||
return browsers
|
||||
echo(
|
||||
f"lite is opening your browser to sign in and waits up to {LOGIN_TIMEOUT_SECONDS // 60} minutes "
|
||||
"for you to approve it. "
|
||||
"A shell or agent that stops commands sooner should run this one with a longer timeout."
|
||||
)
|
||||
return run_pkce_login(base_url, requests.Session(), open_browser=self.open_browser, echo=echo)
|
||||
return run_pkce_login(
|
||||
base_url, requests.Session(), open_browser=lambda url: self.open_browser(url, browsers), echo=echo
|
||||
)
|
||||
|
||||
|
||||
class _SignInContext(TypedDict):
|
||||
|
|
@ -156,15 +190,30 @@ def _echo_to_stderr(message: str) -> None:
|
|||
click.echo(message, err=True)
|
||||
|
||||
|
||||
def _newer_stored_key(ctx: click.Context, base_url: str, rejected_key: str | None) -> str | None:
|
||||
login: Final = load_cli_token(vault=context_secret_vault(ctx))
|
||||
if login is None or login.base_url != base_url.rstrip("/") or not login.key or login.key == rejected_key:
|
||||
return None
|
||||
expired: Final = login.expires_at is not None and time.time() >= login.expires_at
|
||||
return None if expired else login.key
|
||||
|
||||
|
||||
def _attempt_with_newer_stored_key(
|
||||
run: Callable[[click.Context], object], ctx: click.Context, base_url: str, rejected_key: str | None
|
||||
) -> _Completed | _Unauthorized | None:
|
||||
stored_key: Final = get_stored_api_key(expected_base_url=base_url, vault=context_secret_vault(ctx))
|
||||
if stored_key is None or stored_key == rejected_key:
|
||||
stored_key: Final = _newer_stored_key(ctx, base_url, rejected_key)
|
||||
if stored_key is None:
|
||||
return None
|
||||
return _attempt_with_key(run, ctx, stored_key)
|
||||
|
||||
|
||||
def _proxy_reason(error: UnauthorizedError) -> str:
|
||||
original: Final = error.orig_exception
|
||||
response: Final = original.response if isinstance(original, requests.HTTPError) else None
|
||||
body: Final = response.text.strip() if response is not None else ""
|
||||
return redact_string(body) if body else str(error)
|
||||
|
||||
|
||||
def _sign_in_failure(base_url: str, failure: PkceFailure) -> str:
|
||||
if failure.proxy_predates_pkce:
|
||||
return (
|
||||
|
|
@ -254,13 +303,15 @@ def _login_for_another_proxy(ctx: click.Context, base_url: str) -> str | None:
|
|||
return login.base_url
|
||||
|
||||
|
||||
def _rerun_after_sign_in(run: Callable[[click.Context], object], ctx: click.Context, context: _SignInContext) -> object:
|
||||
def _rerun_after_sign_in(
|
||||
run: Callable[[click.Context], object], ctx: click.Context, context: _SignInContext, error: UnauthorizedError
|
||||
) -> object:
|
||||
base_url: Final = context["base_url"]
|
||||
if not context.get("api_key_from_token_file", False):
|
||||
raise click.ClickException(
|
||||
f"The proxy at {base_url} refused `{ctx.command_path}` with HTTP 401 for the API key from --api-key "
|
||||
"or LITELLM_PROXY_API_KEY. Check that the key is valid and has the role this command needs, "
|
||||
"or unset both to sign in through your browser instead."
|
||||
f"or unset both to sign in through your browser instead. The proxy said: {_proxy_reason(error)}"
|
||||
)
|
||||
sign_in: Final = context.get("browser_sign_in") or PkceBrowserSignIn()
|
||||
unavailable: Final = sign_in.unavailable_because()
|
||||
|
|
@ -271,7 +322,8 @@ def _rerun_after_sign_in(run: Callable[[click.Context], object], ctx: click.Cont
|
|||
if rejected_key is not None and _proxy_accepts(base_url, rejected_key):
|
||||
raise click.ClickException(
|
||||
f"The proxy at {base_url} accepts your stored login but refused `{ctx.command_path}` with HTTP 401, "
|
||||
"so it likely lacks the role this command needs. Set LITELLM_PROXY_API_KEY to a key that has it."
|
||||
"so it likely lacks the role this command needs. Set LITELLM_PROXY_API_KEY to a key that has it. "
|
||||
f"The proxy said: {_proxy_reason(error)}"
|
||||
)
|
||||
with _sign_in_lock_held(context.get("sign_in_lock", _sign_in_lock)()) as waited_for_another_sign_in:
|
||||
return _rerun_once_signed_in(run, ctx, base_url, sign_in, rejected_key, waited_for_another_sign_in)
|
||||
|
|
@ -304,11 +356,11 @@ def _rerun_once_signed_in(
|
|||
match _attempt_with_key(run, ctx, _sign_in(ctx, base_url, sign_in)):
|
||||
case _Completed(value=value):
|
||||
return value
|
||||
case _Unauthorized():
|
||||
case _Unauthorized(error=error):
|
||||
raise click.ClickException(
|
||||
f"The proxy at {base_url} refused `{ctx.command_path}` with HTTP 401 again after the browser sign-in, "
|
||||
"so the signed-in user likely lacks the role this command needs. "
|
||||
"Set LITELLM_PROXY_API_KEY to a key that has it."
|
||||
f"Set LITELLM_PROXY_API_KEY to a key that has it. The proxy said: {_proxy_reason(error)}"
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -322,7 +374,7 @@ class SignInOnUnauthorizedCommand(click.Command):
|
|||
context: Final = _sign_in_context(ctx)
|
||||
if context is None:
|
||||
raise error
|
||||
return _rerun_after_sign_in(run, ctx, context)
|
||||
return _rerun_after_sign_in(run, ctx, context, error)
|
||||
|
||||
|
||||
class SignInOnUnauthorizedGroup(click.Group):
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
# stdlib imports
|
||||
import re
|
||||
from collections import defaultdict
|
||||
from collections.abc import Callable
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal
|
||||
|
|
@ -45,7 +45,7 @@ class ModelYamlInfo:
|
|||
return ", ".join(self.access_groups) if self.access_groups else ""
|
||||
|
||||
|
||||
def _get_model_info_obj_from_yaml(model: dict[str, Any]) -> ModelYamlInfo:
|
||||
def _get_model_info_obj_from_yaml(model: Mapping[str, Any]) -> ModelYamlInfo:
|
||||
"""Extract model info from a model dict and return as ModelYamlInfo dataclass."""
|
||||
model_name: Final[str] = model["model_name"]
|
||||
model_params: Final[dict[str, Any]] = model["litellm_params"]
|
||||
|
|
@ -392,6 +392,12 @@ def get_model_list_from_yaml_file(yaml_file: str) -> list[dict[str, object]]:
|
|||
return model_list
|
||||
|
||||
|
||||
def _model_list_from_yaml_argument(
|
||||
_ctx: click.Context, _param: click.Parameter, yaml_file: str
|
||||
) -> Sequence[Mapping[str, object]]:
|
||||
return get_model_list_from_yaml_file(yaml_file)
|
||||
|
||||
|
||||
def _get_filtered_model_list(model_list, only_models_matching_regex, only_access_groups_matching_regex):
|
||||
"""Return a list of models that pass the filter criteria."""
|
||||
model_regex: Final = re.compile(only_models_matching_regex) if only_models_matching_regex else None
|
||||
|
|
@ -407,7 +413,12 @@ def _import_models_get_table_title(dry_run: bool) -> str:
|
|||
|
||||
|
||||
@models.command("import")
|
||||
@click.argument("yaml_file", type=click.Path(exists=True, dir_okay=False, readable=True))
|
||||
@click.argument(
|
||||
"model_list",
|
||||
metavar="YAML_FILE",
|
||||
type=click.Path(exists=True, dir_okay=False, readable=True),
|
||||
callback=_model_list_from_yaml_argument,
|
||||
)
|
||||
@click.option(
|
||||
"--dry-run",
|
||||
is_flag=True,
|
||||
|
|
@ -426,7 +437,7 @@ def _import_models_get_table_title(dry_run: bool) -> str:
|
|||
@click.pass_context
|
||||
def import_models(
|
||||
ctx: click.Context,
|
||||
yaml_file: str,
|
||||
model_list: Sequence[Mapping[str, object]],
|
||||
dry_run: bool,
|
||||
only_models_matching_regex: str | None,
|
||||
only_access_groups_matching_regex: str | None,
|
||||
|
|
@ -434,7 +445,6 @@ def import_models(
|
|||
"""Import models from a YAML file and add them to the proxy."""
|
||||
provider_counts: Final[dict[str, int]] = defaultdict(int)
|
||||
added_models: Final[list[ModelYamlInfo]] = []
|
||||
model_list: Final = get_model_list_from_yaml_file(yaml_file)
|
||||
filtered_model_list: Final = _get_filtered_model_list(
|
||||
model_list, only_models_matching_regex, only_access_groups_matching_regex
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
import json
|
||||
import os
|
||||
import webbrowser
|
||||
from collections.abc import Callable, Sequence
|
||||
from dataclasses import replace
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Final, Protocol
|
||||
|
||||
|
|
@ -23,6 +25,9 @@ BASE_URL: Final = "https://proxy.example.com"
|
|||
SIGNED_IN_KEY: Final = "sk-signed-in"
|
||||
GENERATED: Final = {"key": "sk-generated", "key_alias": "experiment-1"}
|
||||
REJECTED: Final = {"error": {"message": "Authentication Error, No api key passed in.", "type": "auth_error"}}
|
||||
INVALID_KEY: Final = {
|
||||
"error": {"message": "Authentication Error, Invalid API key provided", "type": "token_not_found_in_db"}
|
||||
}
|
||||
SIGNED_IN_AT: Final = 1_790_000_000.0
|
||||
EXPIRES_AT: Final = 4_102_444_800.0
|
||||
|
||||
|
|
@ -151,6 +156,28 @@ def test_a_stored_login_the_proxy_rejects_is_replaced_by_one_sign_in_with_only_j
|
|||
assert get_stored_api_key(BASE_URL, vault=vault) == SIGNED_IN_KEY
|
||||
|
||||
|
||||
def test_an_expired_login_whose_renewal_the_proxy_refuses_is_renewed_once_and_then_signs_in(
|
||||
requests_mock: Mocker, secret_vault_factory: type[FakeSecretVault]
|
||||
) -> None:
|
||||
vault: Final = secret_vault_factory()
|
||||
save_token(pkce_token_record(BASE_URL, replace(_credential("sk-expired"), expires_at=SIGNED_IN_AT)), vault=vault)
|
||||
renewal: Final = requests_mock.post(f"{BASE_URL}/v1/cli/token", status_code=400, json={"error": "invalid_grant"})
|
||||
requests_mock.post(f"{BASE_URL}/v1/cli/revoke", status_code=200)
|
||||
route: Final = requests_mock.post(
|
||||
f"{BASE_URL}/key/generate", [{"status_code": 401, "json": REJECTED}, {"status_code": 200, "json": GENERATED}]
|
||||
)
|
||||
sign_in: Final = FakeBrowserSignIn(_credential())
|
||||
|
||||
result: Final = _lite(["keys", "generate"], vault, sign_in)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert json.loads(result.stdout) == GENERATED
|
||||
assert renewal.call_count == 1
|
||||
assert result.stderr.count("Could not renew the key") == 1
|
||||
assert sign_in.calls == [BASE_URL]
|
||||
assert _bearers(route) == [None, f"Bearer {SIGNED_IN_KEY}"]
|
||||
|
||||
|
||||
def test_a_login_for_another_proxy_is_never_replaced_by_an_automatic_sign_in(
|
||||
requests_mock: Mocker, secret_vault_factory: type[FakeSecretVault]
|
||||
) -> None:
|
||||
|
|
@ -172,7 +199,7 @@ def test_a_login_for_another_proxy_is_never_replaced_by_an_automatic_sign_in(
|
|||
|
||||
|
||||
def test_the_browser_sign_in_fails_at_once_on_a_machine_without_a_browser(requests_mock: Mocker) -> None:
|
||||
def no_browser() -> webbrowser.BaseBrowser:
|
||||
def no_browser(_name: str | None) -> webbrowser.BaseBrowser:
|
||||
raise webbrowser.Error("could not locate runnable browser")
|
||||
|
||||
outcome: Final = PkceBrowserSignIn(find_browser=no_browser, environ={}).unavailable_because()
|
||||
|
|
@ -183,7 +210,7 @@ def test_the_browser_sign_in_fails_at_once_on_a_machine_without_a_browser(reques
|
|||
|
||||
@pytest.mark.parametrize("ci", ["true", "1", "True"])
|
||||
def test_the_browser_sign_in_never_looks_for_a_browser_in_ci(requests_mock: Mocker, ci: str) -> None:
|
||||
def browser_lookup() -> webbrowser.BaseBrowser:
|
||||
def browser_lookup(_name: str | None) -> webbrowser.BaseBrowser:
|
||||
raise AssertionError("looked for a browser while CI is set")
|
||||
|
||||
outcome: Final = PkceBrowserSignIn(find_browser=browser_lookup, environ={"CI": ci}).unavailable_because()
|
||||
|
|
@ -202,13 +229,14 @@ def test_the_browser_sign_in_never_looks_for_a_browser_in_ci(requests_mock: Mock
|
|||
("w3m", {"TERM": "xterm"}),
|
||||
("lynx", {"BROWSER": "lynx"}),
|
||||
("/usr/bin/w3m", {"BROWSER": "/usr/bin/w3m"}),
|
||||
("lynx", {"BROWSER": f"lynx{os.pathsep}/usr/bin/w3m"}),
|
||||
],
|
||||
)
|
||||
def test_a_text_mode_browser_fails_at_once_even_when_browser_picks_it(
|
||||
requests_mock: Mocker, name: str, environ: dict[str, str]
|
||||
) -> None:
|
||||
outcome: Final = PkceBrowserSignIn(
|
||||
find_browser=lambda: webbrowser.GenericBrowser(name), environ=environ
|
||||
find_browser=lambda listed: webbrowser.GenericBrowser(listed or name), environ=environ
|
||||
).unavailable_because()
|
||||
|
||||
assert outcome == PkceFailure("the web browser is a text-mode one, which lite cannot show while a command waits")
|
||||
|
|
@ -219,7 +247,7 @@ def test_a_text_mode_browser_fails_at_once_even_when_browser_picks_it(
|
|||
def test_an_ssh_session_without_a_display_never_opens_a_browser_on_the_remote_screen(
|
||||
requests_mock: Mocker, ssh_variable: str
|
||||
) -> None:
|
||||
def browser_lookup() -> webbrowser.BaseBrowser:
|
||||
def browser_lookup(_name: str | None) -> webbrowser.BaseBrowser:
|
||||
raise AssertionError("looked for a browser in an SSH session without a display")
|
||||
|
||||
outcome: Final = PkceBrowserSignIn(
|
||||
|
|
@ -248,7 +276,9 @@ def test_a_browser_the_user_picked_starts_the_sign_in_and_names_how_long_it_wait
|
|||
) -> None:
|
||||
discovery: Final = requests_mock.get(f"{BASE_URL}/.well-known/litellm-cli-auth", status_code=404)
|
||||
echoed: list[str] = []
|
||||
sign_in: Final = PkceBrowserSignIn(find_browser=lambda: webbrowser.GenericBrowser("firefox"), environ=environ)
|
||||
sign_in: Final = PkceBrowserSignIn(
|
||||
find_browser=lambda _name: webbrowser.GenericBrowser("firefox"), environ=environ
|
||||
)
|
||||
|
||||
assert sign_in.unavailable_because() is None
|
||||
outcome: Final = sign_in(BASE_URL, echoed.append)
|
||||
|
|
@ -262,6 +292,46 @@ def test_a_browser_the_user_picked_starts_the_sign_in_and_names_how_long_it_wait
|
|||
|
||||
|
||||
|
||||
def test_a_browser_list_hands_the_launcher_every_listed_browser_with_a_window_in_order(
|
||||
requests_mock: Mocker,
|
||||
) -> None:
|
||||
requests_mock.get(f"{BASE_URL}/.well-known/litellm-cli-auth", json=_cli_auth_contract())
|
||||
requests_mock.post(f"{BASE_URL}/register", status_code=201, json={"client_id": "cli-client"})
|
||||
launched: list[tuple[str, Sequence[str]]] = []
|
||||
|
||||
def launch(url: str, browsers: Sequence[str]) -> PkceFailure:
|
||||
launched.append((url, browsers))
|
||||
return PkceFailure("the web browser did not open")
|
||||
|
||||
sign_in: Final = PkceBrowserSignIn(
|
||||
find_browser=lambda name: webbrowser.GenericBrowser(name or "xdg-open"),
|
||||
environ={"BROWSER": os.pathsep.join(("lynx", "firefox", "/usr/bin/w3m", "chromium"))},
|
||||
open_browser=launch,
|
||||
)
|
||||
|
||||
assert sign_in.unavailable_because() is None
|
||||
assert sign_in(BASE_URL, lambda _line: None) == PkceFailure("the web browser did not open")
|
||||
assert [browsers for _url, browsers in launched] == [("firefox", "chromium")]
|
||||
assert launched[0][0].startswith(f"{BASE_URL}/authorize?")
|
||||
|
||||
|
||||
def _cli_auth_contract() -> dict[str, object]:
|
||||
return {
|
||||
"contract_version": 1,
|
||||
"issuer": BASE_URL,
|
||||
"authorization_endpoint": f"{BASE_URL}/authorize",
|
||||
"token_endpoint": f"{BASE_URL}/token",
|
||||
"registration_endpoint": f"{BASE_URL}/register",
|
||||
"revocation_endpoint": f"{BASE_URL}/revoke",
|
||||
"resource": BASE_URL,
|
||||
"response_types_supported": ["code"],
|
||||
"grant_types_supported": ["authorization_code", "refresh_token"],
|
||||
"code_challenge_methods_supported": ["S256"],
|
||||
"token_endpoint_auth_methods_supported": ["none"],
|
||||
"revocation_endpoint_auth_methods_supported": ["none"],
|
||||
}
|
||||
|
||||
|
||||
def _browser_script(path: Path, ran: Path, exit_code: int) -> Path:
|
||||
path.write_text(f'#!/bin/sh\n: > "{ran}"\nexit {exit_code}\n')
|
||||
path.chmod(0o755)
|
||||
|
|
@ -286,10 +356,36 @@ def test_the_browser_launcher_opens_only_the_browser_lite_picked_and_never_a_hid
|
|||
monkeypatch.delenv("DISPLAY", raising=False)
|
||||
monkeypatch.delenv("WAYLAND_DISPLAY", raising=False)
|
||||
|
||||
assert open_browser_detached(f"{BASE_URL}/v1/cli/authorize?state=abc") == expected
|
||||
assert open_browser_detached(f"{BASE_URL}/v1/cli/authorize?state=abc", ()) == expected
|
||||
assert picked_ran.exists()
|
||||
assert not console_ran.exists()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("second_exit_code", "expected"),
|
||||
[(0, None), (1, PkceFailure("the web browser did not open"))],
|
||||
ids=["second-browser-opens", "every-browser-fails"],
|
||||
)
|
||||
def test_the_browser_launcher_tries_each_listed_browser_in_order_and_never_a_console_one(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, second_exit_code: int, expected: PkceFailure | None
|
||||
) -> None:
|
||||
first_ran: Final = tmp_path / "first-browser-ran"
|
||||
second_ran: Final = tmp_path / "second-browser-ran"
|
||||
console_ran: Final = tmp_path / "console-browser-ran"
|
||||
first: Final = str(_browser_script(tmp_path / "first-browser", first_ran, 1))
|
||||
second: Final = str(_browser_script(tmp_path / "second-browser", second_ran, second_exit_code))
|
||||
_browser_script(tmp_path / "www-browser", console_ran, 0)
|
||||
monkeypatch.setenv("PATH", str(tmp_path))
|
||||
monkeypatch.setenv("TERM", "xterm")
|
||||
monkeypatch.setenv("BROWSER", os.pathsep.join((first, second)))
|
||||
monkeypatch.delenv("DISPLAY", raising=False)
|
||||
monkeypatch.delenv("WAYLAND_DISPLAY", raising=False)
|
||||
|
||||
assert open_browser_detached(f"{BASE_URL}/v1/cli/authorize?state=abc", (first, second)) == expected
|
||||
assert first_ran.exists()
|
||||
assert second_ran.exists()
|
||||
assert not console_ran.exists()
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("command", "method", "path", "body"),
|
||||
[
|
||||
|
|
@ -341,6 +437,37 @@ def test_models_import_signs_in_instead_of_reporting_models_it_never_added(
|
|||
assert _bearers(route) == [None, f"Bearer {SIGNED_IN_KEY}"]
|
||||
|
||||
|
||||
def _rejected_once_it_is_read(models_file: Path):
|
||||
def respond(request, context):
|
||||
models_file.unlink()
|
||||
context.status_code = 401
|
||||
return REJECTED
|
||||
|
||||
return respond
|
||||
|
||||
|
||||
def test_models_import_reads_its_file_once_so_a_piped_file_still_imports_after_the_sign_in(
|
||||
requests_mock: Mocker, secret_vault_factory: type[FakeSecretVault], tmp_path: Path
|
||||
) -> None:
|
||||
models_file: Final = tmp_path / "models.yaml"
|
||||
models_file.write_text(
|
||||
"model_list:\n - model_name: gpt-5.4-mini\n litellm_params:\n model: openai/gpt-5.4-mini\n"
|
||||
)
|
||||
route: Final = requests_mock.post(
|
||||
f"{BASE_URL}/model/new",
|
||||
[{"json": _rejected_once_it_is_read(models_file)}, {"status_code": 200, "json": {}}],
|
||||
)
|
||||
sign_in: Final = FakeBrowserSignIn(_credential())
|
||||
|
||||
result: Final = _lite(["models", "import", str(models_file)], secret_vault_factory(), sign_in)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert sign_in.calls == [BASE_URL]
|
||||
assert [request.json()["model_name"] for request in route.request_history] == ["gpt-5.4-mini", "gpt-5.4-mini"]
|
||||
assert _bearers(route) == [None, f"Bearer {SIGNED_IN_KEY}"]
|
||||
assert "gpt-5.4-mini" in result.stdout
|
||||
|
||||
|
||||
def test_a_401_partway_through_models_import_stops_without_signing_in_or_adding_models_twice(
|
||||
requests_mock: Mocker, secret_vault_factory: type[FakeSecretVault], tmp_path: Path
|
||||
) -> None:
|
||||
|
|
@ -412,6 +539,33 @@ def test_a_rejected_explicit_key_fails_cleanly_without_opening_a_browser(
|
|||
assert vault.writes == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("command", "method", "path"),
|
||||
[(["users", "list"], "GET", "/user/list"), (["keys", "generate"], "POST", "/key/generate")],
|
||||
ids=["users", "keys"],
|
||||
)
|
||||
def test_a_rejected_explicit_key_names_the_reason_the_proxy_gave_without_echoing_the_key(
|
||||
requests_mock: Mocker,
|
||||
secret_vault_factory: type[FakeSecretVault],
|
||||
command: list[str],
|
||||
method: str,
|
||||
path: str,
|
||||
) -> None:
|
||||
passed_key: Final = "sk-" + "x" * 48
|
||||
echoed_key: Final = {"error": {**INVALID_KEY["error"], "param": f"Received API Key = {passed_key}"}}
|
||||
requests_mock.request(method, f"{BASE_URL}{path}", status_code=401, json=echoed_key)
|
||||
sign_in: Final = FakeBrowserSignIn(_credential())
|
||||
|
||||
result: Final = _lite(command, secret_vault_factory(), sign_in, ["--api-key", passed_key])
|
||||
|
||||
assert result.exit_code == 1
|
||||
assert "for the API key from --api-key or LITELLM_PROXY_API_KEY" in result.stderr
|
||||
assert "Authentication Error, Invalid API key provided" in result.stderr
|
||||
assert "token_not_found_in_db" in result.stderr
|
||||
assert passed_key not in result.output
|
||||
assert sign_in.calls == []
|
||||
|
||||
|
||||
def test_a_failed_sign_in_names_the_reason_and_points_at_lite_login(
|
||||
requests_mock: Mocker, secret_vault_factory: type[FakeSecretVault]
|
||||
) -> None:
|
||||
|
|
@ -445,7 +599,7 @@ def test_a_proxy_without_browser_sign_in_is_named_as_too_old_instead_of_pointing
|
|||
requests_mock.get(f"{BASE_URL}/.well-known/litellm-cli-auth", status_code=discovery_status)
|
||||
|
||||
sign_in_with_a_picked_browser: Final = PkceBrowserSignIn(
|
||||
find_browser=lambda: webbrowser.GenericBrowser("firefox"), environ={"BROWSER": "firefox"}
|
||||
find_browser=lambda _name: webbrowser.GenericBrowser("firefox"), environ={"BROWSER": "firefox"}
|
||||
)
|
||||
|
||||
result: Final = CliRunner().invoke(
|
||||
|
|
@ -469,6 +623,7 @@ def test_a_signed_in_key_the_proxy_also_rejects_is_not_followed_by_another_sign_
|
|||
assert result.exit_code == 1
|
||||
assert isinstance(result.exception, SystemExit)
|
||||
assert "refused `cli keys generate` with HTTP 401 again after the browser sign-in" in result.stderr
|
||||
assert REJECTED["error"]["message"] in result.stderr
|
||||
assert sign_in.calls == [BASE_URL]
|
||||
assert _bearers(route) == [None, f"Bearer {SIGNED_IN_KEY}"]
|
||||
|
||||
|
|
@ -544,6 +699,7 @@ def test_a_stored_login_the_proxy_accepts_but_a_command_refuses_reports_a_missin
|
|||
assert isinstance(result.exception, SystemExit)
|
||||
assert "accepts your stored login but refused `cli credentials list` with HTTP 401" in result.stderr
|
||||
assert "Set LITELLM_PROXY_API_KEY to a key that has it" in result.stderr
|
||||
assert REJECTED["error"]["message"] in result.stderr
|
||||
assert sign_in.calls == []
|
||||
assert route.call_count == 1
|
||||
assert _bearers(key_check) == ["Bearer sk-member"]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue