diff --git a/changelog/1340.fixed.md b/changelog/1340.fixed.md new file mode 100644 index 000000000..afddb62fb --- /dev/null +++ b/changelog/1340.fixed.md @@ -0,0 +1 @@ +When an API key is valid but the model license has not been accepted, TabPFN now opens the license acceptance page directly, or links to it when no browser is available. diff --git a/src/tabpfn/browser_auth.py b/src/tabpfn/browser_auth.py index 8588ed32f..d2d624668 100644 --- a/src/tabpfn/browser_auth.py +++ b/src/tabpfn/browser_auth.py @@ -22,8 +22,10 @@ import urllib.parse import urllib.request import webbrowser +from enum import Enum from pathlib import Path from typing import TYPE_CHECKING +from typing_extensions import assert_never from tabpfn.errors import ( TabPFNError, @@ -42,6 +44,43 @@ _accepted_repos: set[str] = set() +class _AuthStep(Enum): + """What the user still has to do in the browser.""" + + # No usable API key: log in (or register), then accept the license. + LOGIN = "login" + # The API key is valid; only the license acceptance is missing. + ACCEPT_LICENSE = "accept_license" + + +class _NoBrowserToken(Enum): + """Why the browser flow produced no API key.""" + + # TABPFN_NO_BROWSER is set, so the browser flow never ran. + DISABLED = "disabled" + # No interactive terminal, or the user aborted the flow. + UNAVAILABLE = "unavailable" + + +def _browser_url( + gui_url: str, + step: _AuthStep, + *, + hf_repo_id: str, + callback_url: str | None = None, +) -> str: + """Return the web app page that completes ``step``.""" + match step: + case _AuthStep.LOGIN: + path = "login" + case _AuthStep.ACCEPT_LICENSE: + path = "accept-license" + case _: + assert_never(step) + url = f"{gui_url}/{path}?hf_repo_id={urllib.parse.quote(hf_repo_id)}" + return f"{url}&callback={callback_url}" if callback_url else url + + # --------------------------------------------------------------------------- # Environment detection # --------------------------------------------------------------------------- @@ -353,9 +392,7 @@ def _poll_for_token( return received_token[0] -def _headless_interactive_login( - gui_url: str, hf_repo_id: str | None = None -) -> str | None: +def _headless_interactive_login(gui_url: str, *, hf_repo_id: str) -> str | None: """Token acquisition for headless but interactive environments (e.g. SSH). Shows the login URL, offers single-keypress clipboard copy via OSC 52, @@ -363,9 +400,7 @@ def _headless_interactive_login( Returns the JWT on success, or ``None`` on abort / EOF. """ - login_url = f"{gui_url}/login" - if hf_repo_id: - login_url += f"?hf_repo_id={urllib.parse.quote(hf_repo_id)}" + login_url = _browser_url(gui_url, _AuthStep.LOGIN, hf_repo_id=hf_repo_id) print( # noqa: T201 "\nTabPFN requires a one-time license acceptance to download" @@ -479,7 +514,12 @@ def _headless_readline_loop(login_url: str) -> str | None: return None -def try_browser_login(gui_url: str, hf_repo_id: str | None = None) -> str | None: +def try_browser_login( + gui_url: str, + *, + hf_repo_id: str, + step: _AuthStep, +) -> str | None: """Obtain a token via browser callback and/or manual paste concurrently. Chooses the right strategy based on the environment: @@ -496,7 +536,15 @@ def try_browser_login(gui_url: str, hf_repo_id: str | None = None) -> str | None return None if not _has_display(): - return _headless_interactive_login(gui_url, hf_repo_id=hf_repo_id) + match step: + case _AuthStep.LOGIN: + return _headless_interactive_login(gui_url, hf_repo_id=hf_repo_id) + case _AuthStep.ACCEPT_LICENSE: + # The key is already set, so there is nothing to paste; the + # caller's error links to the acceptance page instead. + return None + case _: + assert_never(step) auth_event = threading.Event() received_token: list[str | None] = [None] @@ -508,10 +556,12 @@ def try_browser_login(gui_url: str, hf_repo_id: str | None = None) -> str | None logger.debug("Could not create callback server", exc_info=True) return None - callback_url = f"http://localhost:{port}" - login_url = f"{gui_url}/login?callback={callback_url}" - if hf_repo_id: - login_url += f"&hf_repo_id={urllib.parse.quote(hf_repo_id)}" + login_url = _browser_url( + gui_url, + step, + hf_repo_id=hf_repo_id, + callback_url=f"http://localhost:{port}", + ) server_thread = threading.Thread( target=_serve_until_event, args=(httpd, auth_event), daemon=True @@ -522,19 +572,34 @@ def try_browser_login(gui_url: str, hf_repo_id: str | None = None) -> str | None webbrowser.open(login_url) # --- print unified instructions --- + match step: + case _AuthStep.LOGIN: + instructions = ( + "\nOpening your browser to complete login/registration…\n" + f"\n {login_url}\n" + "\nWaiting for login to complete…\n" + "\nHaving trouble? You can also authenticate manually:\n" + f" 1. Open {gui_url}/account in a browser" + " (log in or register if needed)\n" + f" 2. Accept the license at {gui_url}/account/licenses\n" + " 3. Copy your API Key\n" + " 4. Paste the API key below\n" + ) + case _AuthStep.ACCEPT_LICENSE: + instructions = ( + "\nYour API key is valid; only the license acceptance is missing." + "\nOpening your browser to accept it…\n" + f"\n {login_url}\n" + "\nWaiting for acceptance to complete…\n" + "\nHaving trouble? Accept the license at the link above," + " then paste your API key below.\n" + ) + case _: + assert_never(step) print( # noqa: T201 "\nTabPFN requires a one-time license acceptance to download" " model weights for local inference." - "\nOpening your browser to complete login/registration…\n" - f"\n {login_url}\n" - "\nWaiting for login to complete…\n" - "\nHaving trouble? You can also authenticate manually:\n" - f" 1. Open {gui_url}/account in a browser" - " (log in or register if needed)\n" - " 2. Accept the license at" - f" {gui_url}/account/licenses\n" - " 3. Copy your API Key\n" - " 4. Paste the API key below\n" + f"{instructions}" ) # --- main thread: poll stdin while waiting for callback --- @@ -552,7 +617,7 @@ def try_browser_login(gui_url: str, hf_repo_id: str | None = None) -> str | None # --------------------------------------------------------------------------- -def ensure_license_accepted(hf_repo_id: str) -> Literal[True]: # noqa: C901 +def ensure_license_accepted(hf_repo_id: str) -> Literal[True]: # noqa: C901, PLR0912 """Ensure the user has accepted the TabPFN license. Checks for a cached token, verifies it, and falls back to browser login @@ -574,6 +639,7 @@ def ensure_license_accepted(hf_repo_id: str) -> Literal[True]: # noqa: C901 api_url = settings.tabpfn.auth_api_url license_version = _get_license_name(hf_repo_id) + step = _AuthStep.LOGIN token = get_cached_token() if token is not None: @@ -590,8 +656,8 @@ def ensure_license_accepted(hf_repo_id: str) -> Literal[True]: # noqa: C901 "Could not reach the license server to verify acceptance.\n\n" "Please check your internet connection and try again." ) - # license_status is False — license not yet accepted. - # Fall through to browser login so the GUI can show the acceptance form. + # license_status is False: the key works, only acceptance is missing. + step = _AuthStep.ACCEPT_LICENSE logger.info( "Token valid but license not accepted; opening browser for acceptance.", ) @@ -605,31 +671,45 @@ def ensure_license_accepted(hf_repo_id: str) -> Literal[True]: # noqa: C901 logger.info("Cached token is invalid; deleting and re-authenticating.") delete_cached_token() - # No valid cached token — need browser login. no_browser = os.environ.get("TABPFN_NO_BROWSER", "").strip() + browser_result: str | _NoBrowserToken if no_browser and no_browser not in ("0", "false", "no", "off"): - raise TabPFNLicenseError( - "TabPFN requires a one-time license acceptance to download\n" - "model weights for local inference, but browser login is\n" - "disabled (TABPFN_NO_BROWSER is set).\n\n" - "Set the TABPFN_TOKEN environment variable with a valid API key\n" - "obtained from https://ux.priorlabs.ai" + browser_result = _NoBrowserToken.DISABLED + else: + browser_result = ( + try_browser_login(gui_url, hf_repo_id=hf_repo_id, step=step) + or _NoBrowserToken.UNAVAILABLE ) - token = try_browser_login(gui_url, hf_repo_id=hf_repo_id) - if token is None: - raise TabPFNLicenseError( - "TabPFN requires a one-time license acceptance to download\n" - "model weights for local inference, but no interactive terminal\n" - "is available.\n\n" - "To authenticate in a non-interactive environment:\n" - f" 1. Open {gui_url} in a browser and log in (or register)\n" - f" 2. Accept the license on the Licenses tab\n" - f" 3. Copy your API Key from {gui_url}/account\n" - ' 4. Set the environment variable: export TABPFN_TOKEN=""\n' - " or in Python (before calling .fit()):" - ' import os; os.environ["TABPFN_TOKEN"] = ""' - ) + match step, browser_result: + case _, str() as token: + pass + case _AuthStep.ACCEPT_LICENSE, _NoBrowserToken(): + raise _license_not_accepted_error(gui_url, hf_repo_id) + case _AuthStep.LOGIN, _NoBrowserToken.DISABLED: + raise TabPFNLicenseError( + "TabPFN requires a one-time license acceptance to download\n" + "model weights for local inference, but browser login is\n" + "disabled (TABPFN_NO_BROWSER is set).\n\n" + "Set the TABPFN_TOKEN environment variable with a valid API key\n" + "obtained from https://ux.priorlabs.ai" + ) + case _AuthStep.LOGIN, _NoBrowserToken.UNAVAILABLE: + raise TabPFNLicenseError( + "TabPFN requires a one-time license acceptance to download\n" + "model weights for local inference, but no interactive terminal\n" + "is available.\n\n" + "To authenticate in a non-interactive environment:\n" + f" 1. Open {gui_url} in a browser and log in (or register)\n" + f" 2. Accept the license on the Licenses tab\n" + f" 3. Copy your API Key from {gui_url}/account\n" + " 4. Set the environment variable:" + ' export TABPFN_TOKEN=""\n' + " or in Python (before calling .fit()):" + ' import os; os.environ["TABPFN_TOKEN"] = ""' + ) + case _: + assert_never((step, browser_result)) # Verify the token we just received from the browser. status = verify_token(token, api_url) @@ -644,6 +724,20 @@ def ensure_license_accepted(hf_repo_id: str) -> Literal[True]: # noqa: C901 license_status = check_license_accepted(token, api_url, license_version) if license_status is True: + # The browser may be signed in to another account than the valid key + # we started with. A TABPFN_TOKEN key outranks the one just saved, so + # without this check every later run would reopen the browser. Only + # the acceptance step had a verified key; after a login, a leftover + # key is just invalid, not another account. + next_run_token = get_cached_token() + if ( + step is _AuthStep.ACCEPT_LICENSE + and next_run_token is not None + and next_run_token != token + and check_license_accepted(next_run_token, api_url, license_version) + is False + ): + raise _different_account_error(gui_url, hf_repo_id) print("License accepted — API key cached for future sessions.\n") # noqa: T201 _accepted_repos.add(hf_repo_id) return True @@ -653,8 +747,25 @@ def ensure_license_accepted(hf_repo_id: str) -> Literal[True]: # noqa: C901 "Please check your internet connection and try again." ) # license_status is False - encoded = urllib.parse.quote(hf_repo_id) - raise TabPFNLicenseError( - "License not yet accepted. Please complete the acceptance form at\n" - f"{gui_url}/accept-license?hf_repo_id={encoded} and try again." + raise _license_not_accepted_error(gui_url, hf_repo_id) + + +def _license_not_accepted_error(gui_url: str, hf_repo_id: str) -> TabPFNLicenseError: + accept_url = _browser_url(gui_url, _AuthStep.ACCEPT_LICENSE, hf_repo_id=hf_repo_id) + return TabPFNLicenseError( + f"You are logged in, but the license for {hf_repo_id} has not been\n" + "accepted for your account yet. Accept it once at\n\n" + f" {accept_url}\n\n" + "then try again." + ) + + +def _different_account_error(gui_url: str, hf_repo_id: str) -> TabPFNLicenseError: + accept_url = _browser_url(gui_url, _AuthStep.ACCEPT_LICENSE, hf_repo_id=hf_repo_id) + return TabPFNLicenseError( + f"The license for {hf_repo_id} was accepted in the browser for a\n" + "different account than the one your TABPFN_TOKEN key belongs to.\n" + "Sign in to the browser with that key's account and accept it at\n\n" + f" {accept_url}\n\n" + "then try again." ) diff --git a/tests/test_browser_auth.py b/tests/test_browser_auth.py index cf49f19ab..08de93313 100644 --- a/tests/test_browser_auth.py +++ b/tests/test_browser_auth.py @@ -339,7 +339,121 @@ def test_browser_login_returns_none_error_includes_steps(self): ): self._import_ensure()("tabpfn_2_6") - def test_browser_token_rejected_raises(self): + @pytest.mark.parametrize("browser", ["disabled", "unavailable", "logged_in"]) + def test_valid_token_without_license_links_to_acceptance( + self, monkeypatch: pytest.MonkeyPatch, browser: str + ): + """A valid token whose license is not accepted gets a direct link.""" + if browser == "logged_in": + login_result = "browser-tok" + else: + login_result = None + monkeypatch.setenv("TABPFN_TOKEN", "valid-tok") + if browser == "disabled": + monkeypatch.setenv("TABPFN_NO_BROWSER", "1") + with ( + patch("tabpfn.browser_auth.verify_token", return_value=True), + patch( + "tabpfn.browser_auth.check_license_accepted", + return_value=False, + ), + patch( + "tabpfn.browser_auth.try_browser_login", return_value=login_result + ) as login, + pytest.raises( + TabPFNLicenseError, + match=r"(?s)not been\naccepted.*accept-license\?hf_repo_id=tabpfn_2_6", + ), + ): + self._import_ensure()("tabpfn_2_6") + if browser == "disabled": + login.assert_not_called() + + @pytest.mark.parametrize( + ("license_accepted", "step"), + [(False, "accept_license"), (None, "login")], + ids=["valid_key_without_license", "no_cached_key"], + ) + def test_browser_is_asked_for_the_missing_step( + self, + monkeypatch: pytest.MonkeyPatch, + license_accepted: bool | None, + step: str, + ): + """A valid key only needs acceptance; without one the user must log in.""" + from tabpfn.browser_auth import _AuthStep # noqa: PLC0415 + + if license_accepted is not None: + monkeypatch.setenv("TABPFN_TOKEN", "valid-tok") + with ( + patch("tabpfn.browser_auth.verify_token", return_value=True), + patch( + "tabpfn.browser_auth.check_license_accepted", + # Key not accepted, browser key accepted, and the same account's + # key accepted on the re-check. + side_effect=[license_accepted, True, True] + if license_accepted is not None + else [True], + ), + patch( + "tabpfn.browser_auth.try_browser_login", return_value="browser-tok" + ) as login, + ): + assert self._import_ensure()("tabpfn_2_6") is True + login.assert_called_once_with( + "https://ux.priorlabs.ai", hf_repo_id="tabpfn_2_6", step=_AuthStep(step) + ) + + @pytest.mark.parametrize("key_source", ["env", "cache_file"]) + def test_acceptance_in_a_different_browser_account( + self, monkeypatch: pytest.MonkeyPatch, key_source: str + ): + """The browser may be signed in to another account than the key. + + Only a TABPFN_TOKEN key keeps winning on later runs; a cached key is + replaced by the browser's key, so that case just works. + """ + from tabpfn import browser_auth # noqa: PLC0415 + + if key_source == "env": + monkeypatch.setenv("TABPFN_TOKEN", "key-tok") + else: + browser_auth.save_token("key-tok") + accepted = {"key-tok": False, "browser-tok": True} + with ( + patch("tabpfn.browser_auth.verify_token", return_value=True), + patch( + "tabpfn.browser_auth.check_license_accepted", + side_effect=lambda token, *_: accepted[token], + ), + patch("tabpfn.browser_auth.try_browser_login", return_value="browser-tok"), + ): + if key_source == "env": + with pytest.raises(TabPFNLicenseError, match="different account"): + self._import_ensure()("tabpfn_2_6") + else: + assert self._import_ensure()("tabpfn_2_6") is True + assert browser_auth.get_cached_token() == "browser-tok" + + def test_invalid_env_key_then_browser_login_succeeds( + self, monkeypatch: pytest.MonkeyPatch + ): + """A rejected TABPFN_TOKEN is not a different account: the login still works.""" + monkeypatch.setenv("TABPFN_TOKEN", "stale-tok") + with ( + patch( + "tabpfn.browser_auth.verify_token", + side_effect=lambda key, _: key == "browser-tok", + ), + patch( + "tabpfn.browser_auth.check_license_accepted", + side_effect=lambda key, *_: key == "browser-tok", + ), + patch("tabpfn.browser_auth.try_browser_login", return_value="browser-tok"), + ): + assert self._import_ensure()("tabpfn_2_6") is True + + def test_login_result_rejected_raises(self): """Token from browser rejected by server -> error.""" with ( patch( @@ -573,7 +687,7 @@ def block_termios(name, *args, **kwargs): # noqa: ANN202 return_value="jwt-val", ) as mock_readline, ): - result = headless_login("https://ux.priorlabs.ai") + result = headless_login("https://ux.priorlabs.ai", hf_repo_id="tabpfn_2_6") assert result == "jwt-val" mock_readline.assert_called_once() @@ -601,13 +715,32 @@ def _import_try_login(self): # noqa: ANN202 def test_non_interactive_returns_none(self): """Non-TTY stdin → returns None without attempting any login.""" + from tabpfn.browser_auth import _AuthStep # noqa: PLC0415 + try_login = self._import_try_login() with patch("tabpfn.browser_auth.sys.stdin") as mock_stdin: mock_stdin.isatty.return_value = False - assert try_login("https://ux.priorlabs.ai") is None + assert ( + try_login( + "https://ux.priorlabs.ai", + hf_repo_id="tabpfn_2_6", + step=_AuthStep.LOGIN, + ) + is None + ) + + @pytest.mark.parametrize( + ("step", "expected"), + [("login", "headless-jwt"), ("accept_license", None)], + ) + def test_headless_routes_by_step(self, step: str, expected: str | None): + """TTY + no display: log in by pasting a key; acceptance has nothing to paste. + + For acceptance the key is already set, so returning None lets the caller + raise the error that links to the acceptance page. + """ + from tabpfn.browser_auth import _AuthStep # noqa: PLC0415 - def test_headless_routes_to_headless_login(self): - """TTY + no display → delegates to _headless_interactive_login.""" try_login = self._import_try_login() with ( patch("tabpfn.browser_auth.sys.stdin") as mock_stdin, @@ -618,15 +751,45 @@ def test_headless_routes_to_headless_login(self): ) as mock_headless, ): mock_stdin.isatty.return_value = True - result = try_login("https://ux.priorlabs.ai", hf_repo_id="tabpfn_2_6") + result = try_login( + "https://ux.priorlabs.ai", hf_repo_id="tabpfn_2_6", step=_AuthStep(step) + ) - assert result == "headless-jwt" - mock_headless.assert_called_once_with( - "https://ux.priorlabs.ai", hf_repo_id="tabpfn_2_6" + assert result == expected + if expected is None: + mock_headless.assert_not_called() + else: + mock_headless.assert_called_once_with( + "https://ux.priorlabs.ai", hf_repo_id="tabpfn_2_6" + ) + + def test_graphical_accept_step_opens_acceptance_page(self): + from tabpfn.browser_auth import _AuthStep # noqa: PLC0415 + + try_login = self._import_try_login() + with ( + patch("tabpfn.browser_auth.sys.stdin") as mock_stdin, + patch("tabpfn.browser_auth._has_display", return_value=True), + patch("tabpfn.browser_auth.webbrowser.open") as mock_browser, + patch("tabpfn.browser_auth._poll_for_token", return_value="browser-jwt"), + ): + mock_stdin.isatty.return_value = True + try_login( + "https://ux.priorlabs.ai", + hf_repo_id="tabpfn_2_6", + step=_AuthStep.ACCEPT_LICENSE, + ) + + (url,), _ = mock_browser.call_args + assert url.startswith( + "https://ux.priorlabs.ai/accept-license?hf_repo_id=tabpfn_2_6" + "&callback=http://localhost:" ) def test_graphical_opens_browser(self): """TTY + display → opens browser (existing flow).""" + from tabpfn.browser_auth import _AuthStep # noqa: PLC0415 + try_login = self._import_try_login() with ( patch("tabpfn.browser_auth.sys.stdin") as mock_stdin, @@ -635,7 +798,9 @@ def test_graphical_opens_browser(self): patch("tabpfn.browser_auth._poll_for_token", return_value="browser-jwt"), ): mock_stdin.isatty.return_value = True - result = try_login("https://ux.priorlabs.ai") + result = try_login( + "https://ux.priorlabs.ai", hf_repo_id="tabpfn_2_6", step=_AuthStep.LOGIN + ) assert result == "browser-jwt" mock_browser.assert_called_once()