mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-04 07:27:51 +00:00
fix(cli): check for a browser before the key check and lock, cap the sign-in wait, and sign in without the lock when the login dir is read-only
This commit is contained in:
mateo-berri
parent
a4055fdd64
commit
e63339c701
3 changed files with 251 additions and 80 deletions
|
|
@ -57,11 +57,11 @@ When a command that calls the proxy (`lite keys`, `lite models`, `lite model-gro
|
|||
lite --base-url https://your-proxy.example.com keys generate --models gpt-5.4-mini --key-alias experiment-1
|
||||
```
|
||||
|
||||
The sign-in is the browser flow of `lite login --pkce`, so you pick your team on the consent page, the credential is stored where `lite login` stores it, and its refresh token renews the key on its own from then on. Every line about the sign-in goes to stderr, so stdout carries only the command's own output, such as the JSON from `lite keys generate`. If another `lite` process is already signing in, the command waits for it to finish and then uses the login it stored, so parallel commands open one browser tab. If another `lite` process stored a newer login for the same proxy in the meantime, that key is tried first and no browser opens. `lite` keeps one stored login, so when that login is for a different proxy, the command fails and tells you to run `lite --base-url <proxy> login --pkce` to switch, instead of replacing and revoking the other proxy's login on its own
|
||||
The sign-in is the browser flow of `lite login --pkce`, so you pick your team on the consent page, the credential is stored where `lite login` stores it, and its refresh token renews the key on its own from then on. Every line about the sign-in goes to stderr, so stdout carries only the command's own output, such as the JSON from `lite keys generate`. If another `lite` process is already signing in, the command waits for it to finish and then uses the login it stored, so parallel commands open one browser tab. When that sign-in ends without a login the proxy accepts, or is still waiting after 6 minutes, the command fails instead of opening a second tab. If another `lite` process stored a newer login for the same proxy in the meantime, that key is tried first and no browser opens. `lite` keeps one stored login, so when that login is for a different proxy, the command fails and tells you to run `lite --base-url <proxy> login --pkce` to switch, instead of replacing and revoking the other proxy's login on its own
|
||||
|
||||
Nothing signs in ahead of time. A proxy that accepts the request without a key never opens a browser, and neither does `--help`. The proxy also answers 401 when a valid key lacks the role a command needs, such as a non-admin key running `lite credentials list`, so before it signs in, `lite` checks a rejected stored login against `/key/info`. When the proxy accepts the login there, the command fails with an error that says the login likely lacks the role, and no browser opens. A key you passed yourself is never swapped out or checked again, and a 401 for it fails with an error that names `--api-key` and `LITELLM_PROXY_API_KEY`. When the browser sign-in cannot finish (the machine has no web browser or only a text-mode one such as `w3m` or `lynx`, the shell runs over SSH with no display and `BROWSER` is unset, the browser fails to open, the proxy does not serve `/.well-known/litellm-cli-auth`, you deny the request, or nobody approves it within 5 minutes), the command fails and tells you to run `lite login --pkce` from a terminal on a machine with a browser, `lite login` on one without, or set `LITELLM_PROXY_API_KEY`. `lite` prints that 5-minute wait on stderr when the sign-in starts, so an agent whose shell tool stops commands sooner knows to run the command again with a longer timeout
|
||||
|
||||
When the `CI` environment variable is set to anything but empty, `0`, or `false`, `lite` never opens a browser and the command fails at once with the same message. Set `CI=1` for a script or cron job that should fail instead of waiting for a sign-in
|
||||
When the `CI` environment variable is set to anything but empty, `0`, or `false`, `lite` never opens a browser and the command fails at once with the same message, without checking the stored login against the proxy first. Set `CI=1` for a script or cron job that should fail instead of waiting for a sign-in
|
||||
|
||||
`lite chat`, `lite encryption migrate`, `lite http request`, `lite keys import`, and `lite teams assign-key` do not sign in on their own. `lite chat` and `lite teams assign-key` are interactive, `lite encryption migrate` is a one-time admin migration you run with an admin key, `lite keys import` also talks to a second proxy, and `lite http request` cannot tell a 401 from the proxy apart from a provider's 401 that the proxy passes through
|
||||
|
||||
|
|
|
|||
|
|
@ -2,16 +2,16 @@ import os
|
|||
import subprocess
|
||||
import sys
|
||||
import webbrowser
|
||||
from collections.abc import Callable, Mapping
|
||||
from contextlib import AbstractContextManager
|
||||
from dataclasses import dataclass
|
||||
from collections.abc import Callable, Generator, Mapping
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass, field
|
||||
from http import HTTPStatus
|
||||
from pathlib import Path
|
||||
from typing import Final, Literal, Protocol
|
||||
|
||||
import click
|
||||
import requests
|
||||
from filelock import FileLock
|
||||
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
|
||||
|
|
@ -22,13 +22,22 @@ from .pkce_login import LOGIN_TIMEOUT_SECONDS, PkceCredential, PkceFailure, pkce
|
|||
|
||||
|
||||
class BrowserSignIn(Protocol):
|
||||
def unavailable_because(self) -> PkceFailure | None: ...
|
||||
|
||||
def __call__(self, base_url: str, echo: Callable[[str], None]) -> PkceCredential | PkceFailure: ...
|
||||
|
||||
|
||||
class SignInLock(Protocol):
|
||||
def acquire(self, timeout: float) -> object: ...
|
||||
|
||||
def release(self) -> object: ...
|
||||
|
||||
|
||||
_TEXT_MODE_BROWSERS: Final = frozenset({"www-browser", "links", "elinks", "lynx", "w3m"})
|
||||
_SSH_SESSION_VARIABLES: Final = ("SSH_CONNECTION", "SSH_TTY")
|
||||
_DISPLAY_VARIABLES: Final = ("DISPLAY", "WAYLAND_DISPLAY")
|
||||
_SIGNED_IN_LOGIN_WITHOUT_KEY_ROW: Final = HTTPStatus.NOT_FOUND
|
||||
_WAIT_FOR_ANOTHER_SIGN_IN_SECONDS: Final = LOGIN_TIMEOUT_SECONDS + 60
|
||||
|
||||
|
||||
def _runs_in_ci(environ: Mapping[str, str]) -> bool:
|
||||
|
|
@ -54,28 +63,38 @@ def open_browser_detached(url: str) -> PkceFailure | None:
|
|||
return None if launcher.returncode == 0 else PkceFailure("the web browser did not open")
|
||||
|
||||
|
||||
def pkce_browser_sign_in(
|
||||
base_url: str,
|
||||
echo: Callable[[str], None],
|
||||
find_browser: Callable[[], webbrowser.BaseBrowser] = webbrowser.get,
|
||||
environ: Mapping[str, str] = os.environ,
|
||||
open_browser: Callable[[str], object] = open_browser_detached,
|
||||
) -> PkceCredential | PkceFailure:
|
||||
if _runs_in_ci(environ):
|
||||
return PkceFailure("the CI environment variable is set, so lite does not open a browser")
|
||||
if "BROWSER" not in environ and _ssh_session_without_a_display(environ):
|
||||
return PkceFailure("this is an SSH session without a display, so a browser would open where you cannot see it")
|
||||
try:
|
||||
browser: Final = 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")
|
||||
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=open_browser, echo=echo)
|
||||
def _live_environ() -> Mapping[str, str]:
|
||||
return os.environ
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PkceBrowserSignIn:
|
||||
find_browser: Callable[[], webbrowser.BaseBrowser] = webbrowser.get
|
||||
environ: Mapping[str, str] = field(default_factory=_live_environ)
|
||||
open_browser: Callable[[str], object] = open_browser_detached
|
||||
|
||||
def unavailable_because(self) -> PkceFailure | None:
|
||||
if _runs_in_ci(self.environ):
|
||||
return PkceFailure("the CI environment variable is set, so lite does not open a browser")
|
||||
if "BROWSER" not in self.environ and _ssh_session_without_a_display(self.environ):
|
||||
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
|
||||
|
||||
def __call__(self, base_url: str, echo: Callable[[str], None]) -> PkceCredential | PkceFailure:
|
||||
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)
|
||||
|
||||
|
||||
class _SignInContext(TypedDict):
|
||||
|
|
@ -83,7 +102,7 @@ class _SignInContext(TypedDict):
|
|||
api_key: ReadOnly[NotRequired[str | None]]
|
||||
api_key_from_token_file: ReadOnly[NotRequired[bool]]
|
||||
browser_sign_in: ReadOnly[NotRequired[BrowserSignIn]]
|
||||
sign_in_lock: ReadOnly[NotRequired[Callable[[], AbstractContextManager[object]]]]
|
||||
sign_in_lock: ReadOnly[NotRequired[Callable[[], SignInLock]]]
|
||||
|
||||
|
||||
class _SignInContextView(Protocol):
|
||||
|
|
@ -159,10 +178,8 @@ def _sign_in_failure(base_url: str, failure: PkceFailure) -> str:
|
|||
)
|
||||
|
||||
|
||||
def _sign_in(ctx: click.Context, context: _SignInContext) -> str:
|
||||
base_url: Final = context["base_url"]
|
||||
def _sign_in(ctx: click.Context, base_url: str, sign_in: BrowserSignIn) -> str:
|
||||
_echo_to_stderr(f"The proxy at {base_url} needs a sign-in.")
|
||||
sign_in: Final = context.get("browser_sign_in") or pkce_browser_sign_in
|
||||
credential: Final = sign_in(base_url, _echo_to_stderr)
|
||||
if isinstance(credential, PkceFailure):
|
||||
raise click.ClickException(_sign_in_failure(base_url, credential))
|
||||
|
|
@ -184,10 +201,52 @@ def _proxy_accepts(base_url: str, api_key: str) -> bool:
|
|||
return response.ok or response.status_code == _SIGNED_IN_LOGIN_WITHOUT_KEY_ROW
|
||||
|
||||
|
||||
def _sign_in_lock() -> FileLock:
|
||||
def _sign_in_lock() -> SignInLock:
|
||||
return FileLock(Path(get_cli_token_file_path()).with_name("sign-in.lock"))
|
||||
|
||||
|
||||
def _try_to_lock(lock: SignInLock, timeout: float) -> Literal["held", "busy", "unlockable"]:
|
||||
try:
|
||||
lock.acquire(timeout=timeout)
|
||||
except Timeout:
|
||||
return "busy"
|
||||
except OSError:
|
||||
return "unlockable"
|
||||
return "held"
|
||||
|
||||
|
||||
def _lock_sign_in(lock: SignInLock) -> Literal["held", "waited", "unlockable"]:
|
||||
first_try: Final = _try_to_lock(lock, timeout=0)
|
||||
if first_try != "busy":
|
||||
return first_try
|
||||
_echo_to_stderr(
|
||||
"Another lite command is waiting for a browser sign-in, so this one waits for it instead of opening a second one."
|
||||
)
|
||||
match _try_to_lock(lock, timeout=_WAIT_FOR_ANOTHER_SIGN_IN_SECONDS):
|
||||
case "held":
|
||||
return "waited"
|
||||
case "unlockable":
|
||||
return "unlockable"
|
||||
case "busy":
|
||||
raise click.ClickException(
|
||||
"Another lite command has been waiting for a browser sign-in for over "
|
||||
f"{_WAIT_FOR_ANOTHER_SIGN_IN_SECONDS // 60} minutes, so this one stopped waiting. "
|
||||
"Finish or stop that sign-in, then run this command again."
|
||||
)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _sign_in_lock_held(lock: SignInLock) -> Generator[bool]:
|
||||
outcome: Final = _lock_sign_in(lock)
|
||||
if outcome == "unlockable":
|
||||
yield False
|
||||
return
|
||||
try:
|
||||
yield outcome == "waited"
|
||||
finally:
|
||||
lock.release()
|
||||
|
||||
|
||||
def _login_for_another_proxy(ctx: click.Context, base_url: str) -> str | None:
|
||||
login: Final = load_cli_token(vault=context_secret_vault(ctx))
|
||||
if login is None or not login.base_url or login.base_url == base_url.rstrip("/"):
|
||||
|
|
@ -203,20 +262,29 @@ def _rerun_after_sign_in(run: Callable[[click.Context], object], ctx: click.Cont
|
|||
"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."
|
||||
)
|
||||
sign_in: Final = context.get("browser_sign_in") or PkceBrowserSignIn()
|
||||
unavailable: Final = sign_in.unavailable_because()
|
||||
if unavailable is not None:
|
||||
_echo_to_stderr(f"The proxy at {base_url} needs a sign-in.")
|
||||
raise click.ClickException(_sign_in_failure(base_url, unavailable))
|
||||
rejected_key: Final = context.get("api_key")
|
||||
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."
|
||||
)
|
||||
with context.get("sign_in_lock", _sign_in_lock)():
|
||||
return _rerun_once_signed_in(run, ctx, context, rejected_key)
|
||||
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)
|
||||
|
||||
|
||||
def _rerun_once_signed_in(
|
||||
run: Callable[[click.Context], object], ctx: click.Context, context: _SignInContext, rejected_key: str | None
|
||||
run: Callable[[click.Context], object],
|
||||
ctx: click.Context,
|
||||
base_url: str,
|
||||
sign_in: BrowserSignIn,
|
||||
rejected_key: str | None,
|
||||
waited_for_another_sign_in: bool,
|
||||
) -> object:
|
||||
base_url: Final = context["base_url"]
|
||||
retried: Final = _attempt_with_newer_stored_key(run, ctx, base_url, rejected_key)
|
||||
if isinstance(retried, _Completed):
|
||||
return retried.value
|
||||
|
|
@ -227,7 +295,13 @@ def _rerun_once_signed_in(
|
|||
f"so lite did not replace it on its own. Run `lite --base-url {base_url} login --pkce` to switch proxies, "
|
||||
"or set LITELLM_PROXY_API_KEY."
|
||||
)
|
||||
match _attempt_with_key(run, ctx, _sign_in(ctx, context)):
|
||||
if waited_for_another_sign_in:
|
||||
raise click.ClickException(
|
||||
"The browser sign-in another lite command started ended while this one waited, "
|
||||
f"and the proxy at {base_url} still refuses `{ctx.command_path}`, so lite did not open a second one. "
|
||||
"Run the command again to sign in."
|
||||
)
|
||||
match _attempt_with_key(run, ctx, _sign_in(ctx, base_url, sign_in)):
|
||||
case _Completed(value=value):
|
||||
return value
|
||||
case _Unauthorized():
|
||||
|
|
|
|||
|
|
@ -7,11 +7,12 @@ from typing import Final
|
|||
import pytest
|
||||
import requests
|
||||
from click.testing import CliRunner, Result
|
||||
from filelock import Timeout
|
||||
from requests_mock import Mocker
|
||||
|
||||
from litellm.proxy.client.cli import cli
|
||||
from litellm.proxy.client.cli.commands.auth import CliTokenData, get_stored_api_key, save_token
|
||||
from litellm.proxy.client.cli.commands.auto_sign_in import pkce_browser_sign_in
|
||||
from litellm.proxy.client.cli.commands.auto_sign_in import PkceBrowserSignIn
|
||||
from litellm.proxy.client.cli.commands.pkce_login import PkceCredential, PkceFailure, pkce_token_record
|
||||
from tests.unit.litellm_core_utils.fake_secret_vault import FakeSecretVault
|
||||
|
||||
|
|
@ -28,6 +29,9 @@ class FakeBrowserSignIn:
|
|||
self.outcome = outcome
|
||||
self.calls: list[str] = []
|
||||
|
||||
def unavailable_because(self) -> PkceFailure | None:
|
||||
return None
|
||||
|
||||
def __call__(self, base_url: str, echo: Callable[[str], None]) -> PkceCredential | PkceFailure:
|
||||
self.calls.append(base_url)
|
||||
echo(f"Opening browser to: {base_url}/v1/cli/authorize?state=fake")
|
||||
|
|
@ -160,13 +164,10 @@ def test_the_browser_sign_in_fails_at_once_on_a_machine_without_a_browser(reques
|
|||
def no_browser() -> webbrowser.BaseBrowser:
|
||||
raise webbrowser.Error("could not locate runnable browser")
|
||||
|
||||
echoed: list[str] = []
|
||||
|
||||
outcome: Final = pkce_browser_sign_in(BASE_URL, echoed.append, find_browser=no_browser, environ={})
|
||||
outcome: Final = PkceBrowserSignIn(find_browser=no_browser, environ={}).unavailable_because()
|
||||
|
||||
assert outcome == PkceFailure("no web browser is available on this machine")
|
||||
assert requests_mock.call_count == 0
|
||||
assert echoed == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize("ci", ["true", "1", "True"])
|
||||
|
|
@ -174,13 +175,10 @@ def test_the_browser_sign_in_never_looks_for_a_browser_in_ci(requests_mock: Mock
|
|||
def browser_lookup() -> webbrowser.BaseBrowser:
|
||||
raise AssertionError("looked for a browser while CI is set")
|
||||
|
||||
echoed: list[str] = []
|
||||
|
||||
outcome: Final = pkce_browser_sign_in(BASE_URL, echoed.append, find_browser=browser_lookup, environ={"CI": ci})
|
||||
outcome: Final = PkceBrowserSignIn(find_browser=browser_lookup, environ={"CI": ci}).unavailable_because()
|
||||
|
||||
assert outcome == PkceFailure("the CI environment variable is set, so lite does not open a browser")
|
||||
assert requests_mock.call_count == 0
|
||||
assert echoed == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
|
@ -198,15 +196,12 @@ def test_the_browser_sign_in_never_looks_for_a_browser_in_ci(requests_mock: Mock
|
|||
def test_a_text_mode_browser_fails_at_once_even_when_browser_picks_it(
|
||||
requests_mock: Mocker, name: str, environ: dict[str, str]
|
||||
) -> None:
|
||||
echoed: list[str] = []
|
||||
|
||||
outcome: Final = pkce_browser_sign_in(
|
||||
BASE_URL, echoed.append, find_browser=lambda: webbrowser.GenericBrowser(name), environ=environ
|
||||
)
|
||||
outcome: Final = PkceBrowserSignIn(
|
||||
find_browser=lambda: webbrowser.GenericBrowser(name), environ=environ
|
||||
).unavailable_because()
|
||||
|
||||
assert outcome == PkceFailure("the web browser is a text-mode one, which lite cannot show while a command waits")
|
||||
assert requests_mock.call_count == 0
|
||||
assert echoed == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize("ssh_variable", ["SSH_CONNECTION", "SSH_TTY"])
|
||||
|
|
@ -216,17 +211,14 @@ def test_an_ssh_session_without_a_display_never_opens_a_browser_on_the_remote_sc
|
|||
def browser_lookup() -> webbrowser.BaseBrowser:
|
||||
raise AssertionError("looked for a browser in an SSH session without a display")
|
||||
|
||||
echoed: list[str] = []
|
||||
|
||||
outcome: Final = pkce_browser_sign_in(
|
||||
BASE_URL, echoed.append, find_browser=browser_lookup, environ={ssh_variable: "203.0.113.7 52144 10.0.0.5 22"}
|
||||
)
|
||||
outcome: Final = PkceBrowserSignIn(
|
||||
find_browser=browser_lookup, environ={ssh_variable: "203.0.113.7 52144 10.0.0.5 22"}
|
||||
).unavailable_because()
|
||||
|
||||
assert outcome == PkceFailure(
|
||||
"this is an SSH session without a display, so a browser would open where you cannot see it"
|
||||
)
|
||||
assert requests_mock.call_count == 0
|
||||
assert echoed == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
|
@ -245,10 +237,10 @@ 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)
|
||||
|
||||
outcome: Final = pkce_browser_sign_in(
|
||||
BASE_URL, echoed.append, find_browser=lambda: webbrowser.GenericBrowser("firefox"), environ=environ
|
||||
)
|
||||
assert sign_in.unavailable_because() is None
|
||||
outcome: Final = sign_in(BASE_URL, echoed.append)
|
||||
|
||||
assert isinstance(outcome, PkceFailure)
|
||||
assert discovery.call_count == 1
|
||||
|
|
@ -412,10 +404,9 @@ def test_a_proxy_without_browser_sign_in_is_named_as_too_old_instead_of_pointing
|
|||
requests_mock.post(f"{BASE_URL}/key/generate", status_code=401, json=REJECTED)
|
||||
requests_mock.get(f"{BASE_URL}/.well-known/litellm-cli-auth", status_code=discovery_status)
|
||||
|
||||
def sign_in_with_a_picked_browser(base_url: str, echo: Callable[[str], None]) -> PkceCredential | PkceFailure:
|
||||
return pkce_browser_sign_in(
|
||||
base_url, echo, find_browser=lambda: webbrowser.GenericBrowser("firefox"), environ={"BROWSER": "firefox"}
|
||||
)
|
||||
sign_in_with_a_picked_browser: Final = PkceBrowserSignIn(
|
||||
find_browser=lambda: webbrowser.GenericBrowser("firefox"), environ={"BROWSER": "firefox"}
|
||||
)
|
||||
|
||||
result: Final = CliRunner().invoke(
|
||||
cli,
|
||||
|
|
@ -520,20 +511,46 @@ def test_a_stored_login_the_proxy_accepts_but_a_command_refuses_reports_a_missin
|
|||
assert "sk-member" not in result.output
|
||||
|
||||
|
||||
class SiblingSignsInWhileThisCommandWaits:
|
||||
def __init__(self, vault) -> None:
|
||||
self.vault = vault
|
||||
class AnotherCommandHoldsTheSignInLock:
|
||||
def __init__(self, when_it_lets_go: Callable[[], None], lets_go: bool = True) -> None:
|
||||
self.when_it_lets_go = when_it_lets_go
|
||||
self.lets_go = lets_go
|
||||
self.attempts = 0
|
||||
self.held = False
|
||||
|
||||
def __call__(self) -> "SiblingSignsInWhileThisCommandWaits":
|
||||
def __call__(self) -> "AnotherCommandHoldsTheSignInLock":
|
||||
return self
|
||||
|
||||
def __enter__(self) -> None:
|
||||
save_token(_classic_login("sk-sibling"), vault=self.vault)
|
||||
def acquire(self, timeout: float) -> object:
|
||||
self.attempts += 1
|
||||
if self.attempts == 1 or not self.lets_go:
|
||||
raise Timeout("sign-in.lock")
|
||||
self.when_it_lets_go()
|
||||
self.held = True
|
||||
return self
|
||||
|
||||
def __exit__(self, *exc_info: object) -> None:
|
||||
def release(self) -> object:
|
||||
self.held = False
|
||||
return None
|
||||
|
||||
|
||||
class UnwritableLoginDirectory:
|
||||
def __call__(self) -> "UnwritableLoginDirectory":
|
||||
return self
|
||||
|
||||
def acquire(self, timeout: float) -> object:
|
||||
raise PermissionError(13, "Permission denied", "/home/agent/.litellm")
|
||||
|
||||
def release(self) -> object:
|
||||
raise AssertionError("released a sign-in lock it never held")
|
||||
|
||||
|
||||
def _lite_with_lock(vault: object, sign_in: FakeBrowserSignIn, lock: object) -> Result:
|
||||
return CliRunner().invoke(
|
||||
cli,
|
||||
["--base-url", BASE_URL, "keys", "generate"],
|
||||
obj={"secret_vault": vault, "browser_sign_in": sign_in, "sign_in_lock": lock},
|
||||
)
|
||||
|
||||
|
||||
def test_a_command_that_waits_for_another_commands_sign_in_uses_its_login_instead_of_a_second_browser(
|
||||
|
|
@ -544,19 +561,99 @@ def test_a_command_that_waits_for_another_commands_sign_in_uses_its_login_instea
|
|||
f"{BASE_URL}/key/generate", [{"status_code": 401, "json": REJECTED}, {"status_code": 200, "json": GENERATED}]
|
||||
)
|
||||
sign_in: Final = FakeBrowserSignIn(_credential())
|
||||
sibling: Final = SiblingSignsInWhileThisCommandWaits(vault)
|
||||
other_command: Final = AnotherCommandHoldsTheSignInLock(
|
||||
lambda: save_token(_classic_login("sk-sibling"), vault=vault)
|
||||
)
|
||||
|
||||
result: Final = _lite_with_lock(vault, sign_in, other_command)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert json.loads(result.stdout) == GENERATED
|
||||
assert "Another lite command is waiting for a browser sign-in" in result.stderr
|
||||
assert sign_in.calls == []
|
||||
assert _bearers(route) == [None, "Bearer sk-sibling"]
|
||||
assert other_command.held is False
|
||||
|
||||
|
||||
def test_a_command_that_waited_on_a_sign_in_that_ended_without_a_login_opens_no_second_browser(
|
||||
requests_mock: Mocker, secret_vault_factory: type[FakeSecretVault]
|
||||
) -> None:
|
||||
route: Final = requests_mock.post(f"{BASE_URL}/key/generate", status_code=401, json=REJECTED)
|
||||
sign_in: Final = FakeBrowserSignIn(_credential())
|
||||
other_command: Final = AnotherCommandHoldsTheSignInLock(lambda: None)
|
||||
|
||||
result: Final = _lite_with_lock(secret_vault_factory(), sign_in, other_command)
|
||||
|
||||
assert result.exit_code == 1
|
||||
assert isinstance(result.exception, SystemExit)
|
||||
assert "so lite did not open a second one. Run the command again to sign in." in result.stderr
|
||||
assert sign_in.calls == []
|
||||
assert route.call_count == 1
|
||||
assert other_command.held is False
|
||||
|
||||
|
||||
def test_a_command_stops_waiting_for_a_sign_in_another_command_never_finishes(
|
||||
requests_mock: Mocker, secret_vault_factory: type[FakeSecretVault]
|
||||
) -> None:
|
||||
route: Final = requests_mock.post(f"{BASE_URL}/key/generate", status_code=401, json=REJECTED)
|
||||
sign_in: Final = FakeBrowserSignIn(_credential())
|
||||
other_command: Final = AnotherCommandHoldsTheSignInLock(lambda: None, lets_go=False)
|
||||
|
||||
result: Final = _lite_with_lock(secret_vault_factory(), sign_in, other_command)
|
||||
|
||||
assert result.exit_code == 1
|
||||
assert isinstance(result.exception, SystemExit)
|
||||
assert "Another lite command has been waiting for a browser sign-in for over 6 minutes" in result.stderr
|
||||
assert sign_in.calls == []
|
||||
assert route.call_count == 1
|
||||
|
||||
|
||||
def test_a_login_directory_lite_cannot_write_still_signs_in_and_finishes_the_command(
|
||||
requests_mock: Mocker, secret_vault_factory: type[FakeSecretVault]
|
||||
) -> None:
|
||||
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_with_lock(secret_vault_factory(), sign_in, UnwritableLoginDirectory())
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert json.loads(result.stdout) == GENERATED
|
||||
assert sign_in.calls == [BASE_URL]
|
||||
assert _bearers(route) == [None, f"Bearer {SIGNED_IN_KEY}"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("environ", "reason"),
|
||||
[
|
||||
({"CI": "1"}, "the CI environment variable is set"),
|
||||
({"SSH_CONNECTION": "203.0.113.7 52144 10.0.0.5 22"}, "this is an SSH session without a display"),
|
||||
],
|
||||
ids=["ci", "ssh_without_a_display"],
|
||||
)
|
||||
def test_a_rejected_stored_login_fails_at_once_without_a_key_check_or_lock_where_no_browser_can_open(
|
||||
requests_mock: Mocker, secret_vault_factory: type[FakeSecretVault], environ: dict[str, str], reason: str
|
||||
) -> None:
|
||||
vault: Final = secret_vault_factory()
|
||||
save_token(_classic_login("sk-stale"), vault=vault)
|
||||
key_check: Final = _key_check(requests_mock, 401)
|
||||
route: Final = requests_mock.post(f"{BASE_URL}/key/generate", status_code=401, json=REJECTED)
|
||||
lock: Final = AnotherCommandHoldsTheSignInLock(lambda: None)
|
||||
|
||||
result: Final = CliRunner().invoke(
|
||||
cli,
|
||||
["--base-url", BASE_URL, "keys", "generate"],
|
||||
obj={"secret_vault": vault, "browser_sign_in": sign_in, "sign_in_lock": sibling},
|
||||
obj={"secret_vault": vault, "browser_sign_in": PkceBrowserSignIn(environ=environ), "sign_in_lock": lock},
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert json.loads(result.stdout) == GENERATED
|
||||
assert sign_in.calls == []
|
||||
assert _bearers(route) == [None, "Bearer sk-sibling"]
|
||||
assert sibling.held is False
|
||||
assert result.exit_code == 1
|
||||
assert isinstance(result.exception, SystemExit)
|
||||
assert reason in result.stderr
|
||||
assert "Run `lite login --pkce` from a terminal on a machine with a browser" in result.stderr
|
||||
assert _bearers(route) == ["Bearer sk-stale"]
|
||||
assert key_check.call_count == 0
|
||||
assert lock.attempts == 0
|
||||
|
||||
|
||||
def test_help_never_signs_in(requests_mock: Mocker, secret_vault_factory: type[FakeSecretVault]) -> None:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue