From c53db59174071b7d4ee303b6a3674ec39eda53e1 Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 21 Oct 2026 03:04:53 +0000 Subject: [PATCH] feat(cli): add lite credentials connect for per-user GitHub Copilot Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/github_copilot/per_user_auth.py | 9 +- .../proxy/client/cli/commands/credentials.py | 360 +++++++++++- litellm/proxy/client/credentials.py | 83 ++- .../llms/github_copilot/test_per_user_auth.py | 30 +- .../client/cli/test_credentials_commands.py | 526 +++++++++++++++++- tests/unit/proxy/client/test_credentials.py | 135 ++++- tests/unit/test_litellm_cli.py | 2 + 7 files changed, 1109 insertions(+), 36 deletions(-) diff --git a/litellm/llms/github_copilot/per_user_auth.py b/litellm/llms/github_copilot/per_user_auth.py index 75a10f2e671..801781dd3e6 100644 --- a/litellm/llms/github_copilot/per_user_auth.py +++ b/litellm/llms/github_copilot/per_user_auth.py @@ -6,6 +6,7 @@ touched on the per-user path. import asyncio import hmac +import shlex import threading import time from collections.abc import Callable, Mapping @@ -244,7 +245,8 @@ def _reject_connection(credential_name: str) -> CallerCredentialAuthenticationEr return CallerCredentialAuthenticationError( message=( f"GitHub rejected the GitHub Copilot connection for credential '{credential_name}'. " - f"Reconnect GitHub Copilot for credential '{credential_name}' in the LiteLLM UI (LLM Credentials)" + f"Reconnect GitHub Copilot for credential '{credential_name}' in the LiteLLM UI (LLM Credentials) " + f"or run `lite credentials connect {shlex.quote(credential_name)}`" ), llm_provider=_LLM_PROVIDER, model="", @@ -601,7 +603,10 @@ def github_copilot_auth_mode(litellm_credential_name: object, auth_type_value: o def _connect_error(credential_name: str) -> CallerCredentialAuthenticationError: return CallerCredentialAuthenticationError( - message=(f"Connect GitHub Copilot for credential '{credential_name}' in the LiteLLM UI (LLM Credentials)"), + message=( + f"Connect GitHub Copilot for credential '{credential_name}' in the LiteLLM UI (LLM Credentials) " + f"or run `lite credentials connect {shlex.quote(credential_name)}`" + ), llm_provider=_LLM_PROVIDER, model="", ) diff --git a/litellm/proxy/client/cli/commands/credentials.py b/litellm/proxy/client/cli/commands/credentials.py index 2028bdb5599..c628d25bf5a 100644 --- a/litellm/proxy/client/cli/commands/credentials.py +++ b/litellm/proxy/client/cli/commands/credentials.py @@ -1,16 +1,33 @@ import json -from collections.abc import Sequence -from typing import Final, Literal +import re +import shlex +import threading +import time +import webbrowser +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from typing import Final, Literal, NoReturn, Protocol +from urllib.parse import urlparse import click import requests import rich +from pydantic import BaseModel, ConfigDict from rich.table import Table from typing_extensions import NotRequired, ReadOnly, TypedDict -from ...credentials import CredentialsManagementClient +from ...credentials import ( + CredentialsManagementClient, + UserConnection, + UserConnectionPoll, + UserConnectionStart, +) +from ...exceptions import UnauthorizedError from ._cli_context import cli_context_values +SLOW_DOWN_INCREMENT_SECONDS: Final = 5 +GITHUB_DEVICE_LOGIN_URL: Final = "https://github.com/login/device" + class _CredentialInfo(TypedDict): custom_llm_provider: ReadOnly[NotRequired[str]] @@ -35,15 +52,320 @@ class _JsonBodyView(TypedDict): body: ReadOnly[object] +class _ProxyErrorMessage(BaseModel): + model_config = ConfigDict(frozen=True) + + message: str + type: str | None = None + param: str | None = None + code: str | None = None + + +class _FastAPIErrorDetail(BaseModel): + model_config = ConfigDict(frozen=True) + + error: str + + +class _ProxyErrorBody(BaseModel): + model_config = ConfigDict(frozen=True) + + detail: str | _FastAPIErrorDetail | None = None + error: _ProxyErrorMessage | None = None + + +class UserConnectionClient(Protocol): + def start_user_connection(self, credential_name: str) -> UserConnectionStart: ... + + def poll_user_connection(self, credential_name: str, flow_handle: str) -> UserConnectionPoll: ... + + +@dataclass(frozen=True, slots=True) +class UserConnectionConnected: + status: Literal["connected"] + github_login: str | None + + +@dataclass(frozen=True, slots=True) +class UserConnectionFailed: + status: Literal["failed"] + reason: Literal["expired", "denied", "no_copilot_seat"] + message: str + + +def _safe_verification_uri(uri: str) -> str: + try: + parsed: Final = urlparse(uri) + host: Final = parsed.hostname + has_credentials: Final = parsed.username is not None or parsed.password is not None + has_no_port: Final = parsed.port is None + host_is_allowed: Final = host == "github.com" or ( + host is not None and re.fullmatch(r"[a-z0-9-]+\.ghe\.com", host) is not None + ) + if parsed.scheme == "https" and not has_credentials and has_no_port and host_is_allowed: + return uri + except ValueError: + pass + return GITHUB_DEVICE_LOGIN_URL + + +def _user_connection_failure( + reason: Literal["expired", "denied", "no_copilot_seat"], + credential_name: str, +) -> UserConnectionFailed: + command: Final = f"lite credentials connect {shlex.quote(credential_name)}" + match reason: + case "expired": + return UserConnectionFailed( + status="failed", + reason=reason, + message=f"The device code expired before it was approved. Run `{command}` again to get a new code", + ) + case "denied": + return UserConnectionFailed( + status="failed", + reason=reason, + message=f"The request was denied on GitHub. Run `{command}` again if this was a mistake", + ) + case "no_copilot_seat": + return UserConnectionFailed( + status="failed", + reason=reason, + message="This GitHub account does not have GitHub Copilot access, so nothing was saved", + ) + + +class _ConnectionPoller: + def __init__( + self, + client: UserConnectionClient, + credential_name: str, + flow_handle: str, + interval: int, + deadline: float, + sleep: Callable[[float], None], + clock: Callable[[], float], + ) -> None: + self._client = client + self._credential_name = credential_name + self._flow_handle = flow_handle + self._interval = interval + self._deadline = deadline + self._sleep = sleep + self._clock = clock + + def poll(self) -> UserConnectionPoll | None: + if self._clock() >= self._deadline: + return None + self._sleep(self._interval) + if self._clock() >= self._deadline: + return None + response: Final = self._client.poll_user_connection(self._credential_name, self._flow_handle) + if response.status == "slow_down": + self._interval = ( + response.interval if response.interval is not None else self._interval + SLOW_DOWN_INCREMENT_SECONDS + ) + return response + + +def run_user_connection( + client: UserConnectionClient, + credential_name: str, + *, + open_browser: bool, + launch_browser: Callable[[str], object], + sleep: Callable[[float], None], + clock: Callable[[], float], + echo: Callable[[str], None], +) -> UserConnectionConnected | UserConnectionFailed: + start: Final = client.start_user_connection(credential_name) + verification_uri: Final = _safe_verification_uri(start.verification_uri) + echo( + f"To connect GitHub Copilot for credential {credential_name!r}, open {verification_uri} " + f"and enter code {start.user_code}" + ) + minutes: Final = max(1, round(start.expires_in / 60)) + minute_unit: Final = "minute" if minutes == 1 else "minutes" + echo(f"Waiting for approval. The code expires in {minutes} {minute_unit}") + if open_browser: + launch_browser(verification_uri) + poller: Final = _ConnectionPoller( + client, + credential_name, + start.flow_handle, + start.interval, + clock() + start.expires_in, + sleep, + clock, + ) + for response in iter(poller.poll, None): + match response.status: + case "pending" | "slow_down": + continue + case "connected": + if response.github_login is not None: + echo(f"Connected GitHub Copilot for credential {credential_name!r} as {response.github_login}") + else: + echo(f"Connected GitHub Copilot for credential {credential_name!r}") + return UserConnectionConnected(status="connected", github_login=response.github_login) + case "expired" | "denied" | "no_copilot_seat": + return _user_connection_failure(response.status, credential_name) + return _user_connection_failure("expired", credential_name) + + +def _launch_browser(uri: str) -> object: + browser_thread: Final = threading.Thread(target=webbrowser.open, args=(uri,), daemon=True) + browser_thread.start() + return browser_thread + + +def _http_error_detail(error: requests.exceptions.HTTPError) -> str: + response: Final = error.response + if response is None: + return str(error)[:500] + try: + body: Final = _ProxyErrorBody.model_validate(response.json()) + except ValueError: + return response.text[:500] + if body.error is not None: + return body.error.message[:500] + detail: Final = body.detail + if isinstance(detail, str): + return detail[:500] + if detail is not None: + return detail.error[:500] + return response.text[:500] + + +def _raise_http_error(error: requests.exceptions.HTTPError) -> NoReturn: + response: Final = error.response + status: Final = response.status_code if response is not None else "unknown" + raise click.ClickException(f"HTTP {status}: {_http_error_detail(error)}") from error + + +def _raise_proxy_error( + error: requests.exceptions.HTTPError | requests.exceptions.ConnectionError | UnauthorizedError, +) -> NoReturn: + if isinstance(error, UnauthorizedError): + original: Final = error.orig_exception + if isinstance(original, requests.exceptions.HTTPError): + _raise_http_error(original) + raise click.ClickException(str(error)) from error + if isinstance(error, requests.exceptions.HTTPError): + _raise_http_error(error) + raise click.ClickException(f"Could not connect to the LiteLLM proxy: {error}") from error + + def _print_json(data: object) -> None: rich.print_json(data=data) +def _render_table( + title: str, + columns: tuple[tuple[str, str | None], ...], + rows: Sequence[tuple[str, ...]], +) -> None: + table: Final = Table(title=title) + for column_name, style in columns: + if style is None: + table.add_column(column_name) + else: + table.add_column(column_name, style=style) + for row in rows: + table.add_row(*row) + rich.print(table) + + +def _credential_table_row(credential: _CredentialItem) -> tuple[str, str, str, str]: + info: Final = credential.get("credential_info", {}) + return ( + str(credential.get("credential_name", "")), + credential.get("display_name") or "", + str(credential.get("source", "")), + str(info.get("custom_llm_provider", "")), + ) + + +def _connection_table_row(connection: UserConnection) -> tuple[str, str, str, str]: + return ( + connection.credential_name, + "yes" if connection.connected else "no", + connection.github_login or "", + connection.connected_at or "", + ) + + @click.group() def credentials(): """Manage credentials for the LiteLLM proxy server""" +@credentials.command() +@click.argument("credential_name") +@click.option("--no-browser", is_flag=True, help="Do not open the verification page in a browser") +@click.pass_context +def connect(ctx: click.Context, credential_name: str, no_browser: bool) -> None: + """Connect a GitHub Copilot account to a per-user credential""" + context: Final = cli_context_values(ctx) + client: Final = CredentialsManagementClient(context["base_url"], context["api_key"]) + try: + result: Final = run_user_connection( + client, + credential_name, + open_browser=not no_browser, + launch_browser=_launch_browser, + sleep=time.sleep, + clock=time.monotonic, + echo=click.echo, + ) + except (requests.exceptions.HTTPError, requests.exceptions.ConnectionError, UnauthorizedError) as error: + _raise_proxy_error(error) + if isinstance(result, UserConnectionFailed): + raise click.ClickException(result.message) + + +@credentials.command() +@click.option( + "--format", + "output_format", + type=click.Choice(["table", "json"]), + default="table", + help="Output format (table or json)", +) +@click.pass_context +def connections(ctx: click.Context, output_format: Literal["table", "json"]) -> None: + """List GitHub Copilot connections for the authenticated user""" + context: Final = cli_context_values(ctx) + client: Final = CredentialsManagementClient(context["base_url"], context["api_key"]) + try: + response: Final = client.list_user_connections() + except (requests.exceptions.HTTPError, requests.exceptions.ConnectionError, UnauthorizedError) as error: + _raise_proxy_error(error) + if output_format == "json": + _print_json(response.model_dump(mode="json")) + return + rows: Final = tuple(_connection_table_row(connection) for connection in response.connections) + _render_table( + "GitHub Copilot connections", + (("Credential Name", "cyan"), ("Connected", None), ("GitHub Login", None), ("Connected At", None)), + rows, + ) + + +@credentials.command() +@click.argument("credential_name") +@click.pass_context +def disconnect(ctx: click.Context, credential_name: str) -> None: + """Disconnect a GitHub Copilot account from a per-user credential""" + context: Final = cli_context_values(ctx) + client: Final = CredentialsManagementClient(context["base_url"], context["api_key"]) + try: + client.delete_user_connection(credential_name) + except (requests.exceptions.HTTPError, requests.exceptions.ConnectionError, UnauthorizedError) as error: + _raise_proxy_error(error) + click.echo(f"Disconnected GitHub Copilot for credential {credential_name!r}") + + @credentials.command() @click.option( "--format", @@ -62,27 +384,19 @@ def list(ctx: click.Context, output_format: Literal["table", "json"]): if output_format == "json": _print_json(response) - else: # table format - table: Final = Table(title="Credentials") - - # Add columns - table.add_column("Credential Name", style="cyan") - table.add_column("Display Name") - table.add_column("Source") - table.add_column("Custom LLM Provider", style="green") - - # Add rows + else: listed: Final[_CredentialsListView] = {"credentials": response.get("credentials", [])} - for cred in listed["credentials"]: - info = cred.get("credential_info", {}) - table.add_row( - str(cred.get("credential_name", "")), - cred.get("display_name") or "", - str(cred.get("source", "")), - str(info.get("custom_llm_provider", "")), - ) - - rich.print(table) + rows: Final = tuple(_credential_table_row(credential) for credential in listed["credentials"]) + _render_table( + "Credentials", + ( + ("Credential Name", "cyan"), + ("Display Name", None), + ("Source", None), + ("Custom LLM Provider", "green"), + ), + rows, + ) @credentials.command() diff --git a/litellm/proxy/client/credentials.py b/litellm/proxy/client/credentials.py index 649fe529936..f5727d806b5 100644 --- a/litellm/proxy/client/credentials.py +++ b/litellm/proxy/client/credentials.py @@ -1,12 +1,53 @@ from collections.abc import Mapping -from typing import Any, Final +from typing import Any, Final, Literal from urllib.parse import quote import requests +from pydantic import BaseModel, ConfigDict, Field from .exceptions import UnauthorizedError +class UserConnectionStart(BaseModel): + model_config = ConfigDict(frozen=True) + + user_code: str + verification_uri: str + expires_in: int + interval: int + flow_handle: str = Field(repr=False) + + +class UserConnectionPoll(BaseModel): + model_config = ConfigDict(frozen=True) + + status: Literal["pending", "slow_down", "expired", "denied", "connected", "no_copilot_seat"] + interval: int | None = None + github_login: str | None = None + + +class UserConnection(BaseModel): + model_config = ConfigDict(frozen=True) + + credential_name: str + provider: str + connected: bool + github_login: str | None + connected_at: str | None + + +class UserConnections(BaseModel): + model_config = ConfigDict(frozen=True) + + connections: tuple[UserConnection, ...] + + +class UserConnectionDelete(BaseModel): + model_config = ConfigDict(frozen=True) + + status: Literal["disconnected"] + + class CredentialsManagementClient: def __init__(self, base_url: str, api_key: str | None = None, timeout: int = 30): """ @@ -214,3 +255,43 @@ class CredentialsManagementClient: if e.response.status_code == 401: raise UnauthorizedError(e) raise + + def _send_user_connection_request(self, request: requests.Request) -> requests.Response: + with requests.Session() as session: + response: Final = session.send(request.prepare(), timeout=self._timeout) + try: + response.raise_for_status() + except requests.exceptions.HTTPError as error: + if error.response is not None and error.response.status_code == 401: + raise UnauthorizedError(error) from error + raise + return response + + def start_user_connection(self, credential_name: str) -> UserConnectionStart: + url: Final = f"{self._base_url}/credentials/{quote(credential_name, safe='')}/user_connection/start" + request: Final = requests.Request("POST", url, headers=self._get_headers()) + response: Final = self._send_user_connection_request(request) + return UserConnectionStart.model_validate(response.json()) + + def poll_user_connection(self, credential_name: str, flow_handle: str) -> UserConnectionPoll: + url: Final = f"{self._base_url}/credentials/{quote(credential_name, safe='')}/user_connection/poll" + request: Final = requests.Request( + "POST", + url, + headers=self._get_headers(), + json={"flow_handle": flow_handle}, + ) + response: Final = self._send_user_connection_request(request) + return UserConnectionPoll.model_validate(response.json()) + + def list_user_connections(self) -> UserConnections: + url: Final = f"{self._base_url}/credentials/user_connections" + request: Final = requests.Request("GET", url, headers=self._get_headers()) + response: Final = self._send_user_connection_request(request) + return UserConnections.model_validate(response.json()) + + def delete_user_connection(self, credential_name: str) -> UserConnectionDelete: + url: Final = f"{self._base_url}/credentials/{quote(credential_name, safe='')}/user_connection" + request: Final = requests.Request("DELETE", url, headers=self._get_headers()) + response: Final = self._send_user_connection_request(request) + return UserConnectionDelete.model_validate(response.json()) diff --git a/tests/unit/llms/github_copilot/test_per_user_auth.py b/tests/unit/llms/github_copilot/test_per_user_auth.py index cf6fb9fa71f..31c3fa5cf83 100644 --- a/tests/unit/llms/github_copilot/test_per_user_auth.py +++ b/tests/unit/llms/github_copilot/test_per_user_auth.py @@ -254,8 +254,10 @@ def test_exchange_auth_failures_raise_caller_auth_error_and_evict(status): with github_http(_exchange_responder(status=status)): with pytest.raises(CallerCredentialAuthenticationError) as exc: exchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) - assert "Reconnect GitHub Copilot for credential 'copilot-cred' in the LiteLLM UI (LLM Credentials)" in str( - exc.value + assert str(exc.value) == ( + "litellm.AuthenticationError: GitHub rejected the GitHub Copilot connection for credential 'copilot-cred'. " + "Reconnect GitHub Copilot for credential 'copilot-cred' in the LiteLLM UI (LLM Credentials) " + "or run `lite credentials connect copilot-cred`" ) assert ( _SESSION_CACHE.get_cache(_session_cache_key("user-a", "copilot-cred", GITHUB_COPILOT_DEFAULT_HOST)) is None @@ -346,7 +348,11 @@ def test_attach_per_user_without_connection_raises_connect_401(per_user_credenti kwargs = {"litellm_credential_name": "copilot-cred", "secret_fields": {}} with pytest.raises(CallerCredentialAuthenticationError) as exc: attach_github_copilot_user_session(kwargs) - assert "Connect GitHub Copilot for credential 'copilot-cred' in the LiteLLM UI (LLM Credentials)" in str(exc.value) + assert str(exc.value) == ( + "litellm.AuthenticationError: Connect GitHub Copilot for credential 'copilot-cred' " + "in the LiteLLM UI (LLM Credentials) " + "or run `lite credentials connect copilot-cred`" + ) @pytest.mark.parametrize( @@ -362,10 +368,26 @@ def test_attach_rejects_connection_from_a_different_github_host( with github_http(lambda request: pytest.fail(f"unexpected token exchange at {request.url}")) as requests: with pytest.raises(CallerCredentialAuthenticationError) as exc: attach_github_copilot_user_session(kwargs) - assert "Connect GitHub Copilot for credential 'copilot-cred'" in str(exc.value) + assert str(exc.value) == ( + "litellm.AuthenticationError: Connect GitHub Copilot for credential 'copilot-cred' " + "in the LiteLLM UI (LLM Credentials) " + "or run `lite credentials connect copilot-cred`" + ) assert not requests +def test_exchange_auth_failure_quotes_credential_name_in_cli_hint(): + with github_http(_exchange_responder(status=401)): + with pytest.raises(CallerCredentialAuthenticationError) as exc: + exchange_github_token("user-a", "gh-token", "my cred", _GITHUB_ENDPOINTS) + + assert str(exc.value) == ( + "litellm.AuthenticationError: GitHub rejected the GitHub Copilot connection for credential 'my cred'. " + "Reconnect GitHub Copilot for credential 'my cred' in the LiteLLM UI (LLM Credentials) " + "or run `lite credentials connect 'my cred'`" + ) + + def test_legacy_connection_without_host_uses_github_com(per_user_credential): with github_http(_exchange_responder()) as requests: kwargs: Final = _kwargs_with_connection(github_host=None) diff --git a/tests/unit/proxy/client/cli/test_credentials_commands.py b/tests/unit/proxy/client/cli/test_credentials_commands.py index cdbd73057a0..b1c9746f2c7 100644 --- a/tests/unit/proxy/client/cli/test_credentials_commands.py +++ b/tests/unit/proxy/client/cli/test_credentials_commands.py @@ -1,13 +1,20 @@ import json +from collections.abc import Iterator +from typing import Final from unittest.mock import MagicMock import pytest import requests +import responses from click.testing import CliRunner - - +from litellm.proxy.client.cli.commands.credentials import ( + UserConnectionConnected, + UserConnectionFailed, + run_user_connection, +) from litellm.proxy.client.cli.main import cli +from litellm.proxy.client.credentials import UserConnectionPoll, UserConnectionStart @pytest.fixture @@ -264,3 +271,518 @@ def test_update_credential_requires_exactly_one_display_name_flag(cli_runner, mo assert result.exit_code != 0 assert "exactly one" in result.output mock_credentials_client.return_value.update_display_name.assert_not_called() + + +class _SequenceUserConnectionClient: + def __init__( + self, + start: UserConnectionStart, + poll_responses: tuple[UserConnectionPoll, ...], + ) -> None: + self._start = start + self._poll_responses: Iterator[UserConnectionPoll] = iter(poll_responses) + + def start_user_connection(self, credential_name: str) -> UserConnectionStart: + return self._start + + def poll_user_connection(self, credential_name: str, flow_handle: str) -> UserConnectionPoll: + return next(self._poll_responses) + + +class _StringRecorder: + def __init__(self) -> None: + self.values: tuple[str, ...] = () + + def __call__(self, value: str) -> None: + self.values = (*self.values, value) + + +class _SleepRecorder: + def __init__(self) -> None: + self.values: tuple[float, ...] = () + + def __call__(self, value: float) -> None: + self.values = (*self.values, value) + + +class _AdvancingClock: + def __init__(self) -> None: + self.current: float = 0 + self.sleeps: tuple[float, ...] = () + + def __call__(self) -> float: + return self.current + + def sleep(self, seconds: float) -> None: + self.current += seconds + self.sleeps = (*self.sleeps, seconds) + + def advance(self, seconds: float) -> None: + self.current += seconds + + +class _LateConnectedUserConnectionClient: + def __init__(self, clock: _AdvancingClock) -> None: + self._clock = clock + + def start_user_connection(self, _credential_name: str) -> UserConnectionStart: + return _start_response(interval=0, expires_in=3) + + def poll_user_connection(self, _credential_name: str, _flow_handle: str) -> UserConnectionPoll: + self._clock.advance(4) + return UserConnectionPoll(status="connected", github_login="octocat") + + +def _start_response( + verification_uri: str = "https://github.com/login/device", + interval: int = 5, + expires_in: int = 900, +) -> UserConnectionStart: + return UserConnectionStart( + user_code="ABCD-1234", + verification_uri=verification_uri, + expires_in=expires_in, + interval=interval, + flow_handle="flow-handle", + ) + + +def _cli_args(*args: str) -> tuple[str, ...]: + return ("--base-url", "http://fake-proxy", "--api-key", "sk-user", *args) + + +def _register_start_response(credential_name: str = "copilot", expires_in: int = 900) -> None: + responses.add( + responses.POST, + f"http://fake-proxy/credentials/{credential_name}/user_connection/start", + json={ + "user_code": "ABCD-1234", + "verification_uri": "https://github.com/login/device", + "expires_in": expires_in, + "interval": 0, + "flow_handle": "flow-handle", + }, + status=200, + ) + + +@pytest.mark.parametrize( + ("expires_in", "expected_expiry"), + [(899, "15 minutes"), (60, "1 minute"), (20, "1 minute")], +) +@responses.activate +def test_connect_command_polls_until_connected( + cli_runner: CliRunner, + expires_in: int, + expected_expiry: str, +) -> None: + _register_start_response(expires_in=expires_in) + responses.add( + responses.POST, + "http://fake-proxy/credentials/copilot/user_connection/poll", + json={"status": "pending"}, + status=200, + ) + responses.add( + responses.POST, + "http://fake-proxy/credentials/copilot/user_connection/poll", + json={"status": "slow_down", "interval": 0}, + status=200, + ) + responses.add( + responses.POST, + "http://fake-proxy/credentials/copilot/user_connection/poll", + json={"status": "connected", "github_login": "octocat"}, + status=200, + ) + + result: Final = cli_runner.invoke(cli, _cli_args("credentials", "connect", "copilot", "--no-browser")) + + assert result.exit_code == 0, result.output + assert "ABCD-1234" in result.output + assert "https://github.com/login/device" in result.output + assert f"The code expires in {expected_expiry}" in result.output + assert "Connected GitHub Copilot for credential 'copilot' as octocat" in result.output + + +@pytest.mark.parametrize( + ("status", "message"), + [ + ( + "expired", + "The device code expired before it was approved. " + "Run `lite credentials connect copilot` again to get a new code", + ), + ( + "denied", + "The request was denied on GitHub. Run `lite credentials connect copilot` again if this was a mistake", + ), + ( + "no_copilot_seat", + "This GitHub account does not have GitHub Copilot access, so nothing was saved", + ), + ], +) +@responses.activate +def test_connect_command_reports_terminal_poll_statuses( + cli_runner: CliRunner, + status: str, + message: str, +) -> None: + _register_start_response() + responses.add( + responses.POST, + "http://fake-proxy/credentials/copilot/user_connection/poll", + json={"status": status}, + status=200, + ) + + result: Final = cli_runner.invoke(cli, _cli_args("credentials", "connect", "copilot", "--no-browser")) + + assert result.exit_code == 1 + assert result.output.splitlines()[-1] == f"Error: {message}" + + +@responses.activate +def test_connect_command_reports_proxy_connection_error(cli_runner: CliRunner) -> None: + responses.add( + responses.POST, + "http://fake-proxy/credentials/copilot/user_connection/start", + body=requests.exceptions.ConnectionError("connection refused"), + ) + + result: Final = cli_runner.invoke(cli, _cli_args("credentials", "connect", "copilot", "--no-browser")) + + assert result.exit_code == 1 + assert "Could not connect to the LiteLLM proxy" in result.output + + +@responses.activate +def test_connect_command_prints_unauthorized_detail(cli_runner: CliRunner) -> None: + responses.add( + responses.POST, + "http://fake-proxy/credentials/copilot/user_connection/start", + json={"detail": "An authenticated user is required"}, + status=401, + ) + + result: Final = cli_runner.invoke(cli, _cli_args("credentials", "connect", "copilot", "--no-browser")) + + assert result.exit_code == 1 + assert "Error: HTTP 401: An authenticated user is required" in result.output + + +@responses.activate +def test_connect_command_prints_proxy_404_error_message(cli_runner: CliRunner) -> None: + responses.add( + responses.POST, + "http://fake-proxy/credentials/copilot/user_connection/start", + json={ + "error": { + "message": "Credential not found", + "type": "internal_server_error", + "param": None, + "code": "404", + } + }, + status=404, + ) + + result: Final = cli_runner.invoke(cli, _cli_args("credentials", "connect", "copilot", "--no-browser")) + + assert result.exit_code == 1 + assert result.output.splitlines()[-1] == "Error: HTTP 404: Credential not found" + + +@responses.activate +def test_connect_command_prints_proxy_401_error_message(cli_runner: CliRunner) -> None: + responses.add( + responses.POST, + "http://fake-proxy/credentials/copilot/user_connection/start", + json={ + "error": { + "message": "Authentication Error, Invalid API key provided.", + "type": "token_not_found_in_db", + "param": "key", + "code": "401", + } + }, + status=401, + ) + + result: Final = cli_runner.invoke(cli, _cli_args("credentials", "connect", "copilot", "--no-browser")) + + assert result.exit_code == 1 + assert result.output.splitlines()[-1] == "Error: HTTP 401: Authentication Error, Invalid API key provided." + + +@responses.activate +def test_connect_command_prints_nested_fastapi_error_detail(cli_runner: CliRunner) -> None: + responses.add( + responses.POST, + "http://fake-proxy/credentials/copilot/user_connection/start", + json={"detail": {"error": "Invalid flow details"}}, + status=400, + ) + + result: Final = cli_runner.invoke(cli, _cli_args("credentials", "connect", "copilot", "--no-browser")) + + assert result.exit_code == 1 + assert result.output.splitlines()[-1] == "Error: HTTP 400: Invalid flow details" + + +@responses.activate +def test_connect_command_reports_invalid_flow_handle(cli_runner: CliRunner) -> None: + _register_start_response() + responses.add( + responses.POST, + "http://fake-proxy/credentials/copilot/user_connection/poll", + json={"detail": "Invalid flow handle"}, + status=400, + ) + + result: Final = cli_runner.invoke(cli, _cli_args("credentials", "connect", "copilot", "--no-browser")) + + assert result.exit_code == 1 + assert "Error: HTTP 400: Invalid flow handle" in result.output + + +@pytest.mark.parametrize("output_format", ["table", "json"]) +@responses.activate +def test_connections_command_lists_per_user_connections(cli_runner: CliRunner, output_format: str) -> None: + response: Final = { + "connections": [ + { + "credential_name": "copilot-prod", + "provider": "github_copilot", + "connected": True, + "github_login": "octocat", + "connected_at": "2025-01-19T12:00:00+00:00", + }, + { + "credential_name": "copilot-dev", + "provider": "github_copilot", + "connected": False, + "github_login": None, + "connected_at": None, + }, + ] + } + responses.add( + responses.GET, + "http://fake-proxy/credentials/user_connections", + json=response, + status=200, + ) + + result: Final = cli_runner.invoke( + cli, + _cli_args("credentials", "connections", "--format", output_format), + env={"COLUMNS": "200"}, + ) + + assert result.exit_code == 0, result.output + if output_format == "json": + assert json.loads(result.output) == response + else: + assert "GitHub Copilot connections" in result.output + assert "Credential Name" in result.output + assert "Connected" in result.output + assert "GitHub Login" in result.output + assert "Connected At" in result.output + prod_row: Final = next((line for line in result.output.splitlines() if "copilot-prod" in line), "") + dev_row: Final = next((line for line in result.output.splitlines() if "copilot-dev" in line), "") + assert "yes" in prod_row + assert "octocat" in prod_row + assert "2025-01-19T12:00:00+00:00" in prod_row + assert "no" in dev_row + + +@responses.activate +def test_disconnect_command_encodes_credential_name(cli_runner: CliRunner) -> None: + credential_name: Final = "my/ credential" + responses.add( + responses.DELETE, + "http://fake-proxy/credentials/my%2F%20credential/user_connection", + json={"status": "disconnected"}, + status=200, + ) + + result: Final = cli_runner.invoke( + cli, + _cli_args("credentials", "disconnect", credential_name), + ) + + assert result.exit_code == 0, result.output + assert "Disconnected GitHub Copilot for credential 'my/ credential'" in result.output + assert responses.calls[0].request.url == "http://fake-proxy/credentials/my%2F%20credential/user_connection" + + +def test_run_user_connection_uses_response_intervals_and_slow_down_fallback() -> None: + client: Final = _SequenceUserConnectionClient( + _start_response(interval=5), + ( + UserConnectionPoll(status="pending"), + UserConnectionPoll(status="slow_down", interval=10), + UserConnectionPoll(status="slow_down"), + UserConnectionPoll(status="connected", github_login="octocat"), + ), + ) + sleep_calls: Final = _SleepRecorder() + output: Final = _StringRecorder() + + result: Final = run_user_connection( + client, + "copilot", + open_browser=False, + launch_browser=lambda uri: None, + sleep=sleep_calls, + clock=lambda: 0, + echo=output, + ) + + assert result == UserConnectionConnected(status="connected", github_login="octocat") + assert sleep_calls.values == (5, 5, 10, 15) + assert output.values[-1] == "Connected GitHub Copilot for credential 'copilot' as octocat" + + +def test_run_user_connection_honors_server_poll_interval() -> None: + sleep_calls: Final = _SleepRecorder() + client: Final = _SequenceUserConnectionClient( + _start_response(), + ( + UserConnectionPoll(status="slow_down", interval=9), + UserConnectionPoll(status="connected", github_login="octocat"), + ), + ) + + result: Final = run_user_connection( + client, + "copilot", + open_browser=False, + launch_browser=lambda uri: None, + sleep=sleep_calls, + clock=lambda: 0, + echo=_StringRecorder(), + ) + + assert result == UserConnectionConnected(status="connected", github_login="octocat") + assert sleep_calls.values == (5, 9) + + +def test_run_user_connection_stops_at_local_expiry_deadline() -> None: + clock: Final = _AdvancingClock() + client: Final = _SequenceUserConnectionClient( + _start_response(interval=2, expires_in=3), + (UserConnectionPoll(status="pending"), UserConnectionPoll(status="pending")), + ) + + result: Final = run_user_connection( + client, + "copilot", + open_browser=False, + launch_browser=lambda uri: None, + sleep=clock.sleep, + clock=clock, + echo=lambda message: None, + ) + + assert result == UserConnectionFailed( + status="failed", + reason="expired", + message=( + "The device code expired before it was approved. " + "Run `lite credentials connect copilot` again to get a new code" + ), + ) + assert clock.sleeps == (2, 2) + + +def test_run_user_connection_accepts_connected_poll_received_after_deadline() -> None: + clock: Final = _AdvancingClock() + client: Final = _LateConnectedUserConnectionClient(clock) + + result: Final = run_user_connection( + client, + "copilot", + open_browser=False, + launch_browser=lambda uri: None, + sleep=clock.sleep, + clock=clock, + echo=lambda message: None, + ) + + assert result == UserConnectionConnected(status="connected", github_login="octocat") + + +def test_run_user_connection_opens_safe_url_only_when_requested() -> None: + launches: Final = _StringRecorder() + client: Final = _SequenceUserConnectionClient( + _start_response(), + (UserConnectionPoll(status="connected", github_login="octocat"),), + ) + + result: Final = run_user_connection( + client, + "copilot", + open_browser=True, + launch_browser=launches, + sleep=lambda seconds: None, + clock=lambda: 0, + echo=lambda message: None, + ) + + assert result == UserConnectionConnected(status="connected", github_login="octocat") + assert launches.values == ("https://github.com/login/device",) + + +def test_run_user_connection_does_not_launch_browser_when_disabled() -> None: + launches: Final = _StringRecorder() + client: Final = _SequenceUserConnectionClient( + _start_response(), + (UserConnectionPoll(status="connected", github_login="octocat"),), + ) + + run_user_connection( + client, + "copilot", + open_browser=False, + launch_browser=launches, + sleep=lambda seconds: None, + clock=lambda: 0, + echo=lambda message: None, + ) + + assert launches.values == () + + +@pytest.mark.parametrize( + "verification_uri", + [ + "http://github.com/login/device", + "https://github.com:443/login/device", + "https://user@github.com/login/device", + "https://github.com.evil.com/login/device", + ], +) +def test_run_user_connection_replaces_unsafe_verification_uri(verification_uri: str) -> None: + launches: Final = _StringRecorder() + output: Final = _StringRecorder() + client: Final = _SequenceUserConnectionClient( + _start_response(verification_uri=verification_uri, interval=0), + (UserConnectionPoll(status="connected", github_login="octocat"),), + ) + + run_user_connection( + client, + "copilot", + open_browser=True, + launch_browser=launches, + sleep=lambda seconds: None, + clock=lambda: 0, + echo=output, + ) + + assert launches.values == ("https://github.com/login/device",) + assert "https://github.com/login/device" in output.values[0] + assert verification_uri not in output.values[0] diff --git a/tests/unit/proxy/client/test_credentials.py b/tests/unit/proxy/client/test_credentials.py index 9de990da487..c4c0771508b 100644 --- a/tests/unit/proxy/client/test_credentials.py +++ b/tests/unit/proxy/client/test_credentials.py @@ -1,18 +1,40 @@ import time +from typing import Final + import pytest import requests - - - import responses +from pydantic import BaseModel, ConfigDict, TypeAdapter -from litellm.proxy.client.credentials import CredentialsManagementClient +from litellm.proxy.client.credentials import ( + CredentialsManagementClient, + UserConnection, + UserConnectionDelete, + UserConnectionPoll, + UserConnections, + UserConnectionStart, +) from litellm.proxy.client.exceptions import UnauthorizedError from litellm.proxy.credential_endpoints.endpoints import CredentialHelperUtils from litellm.types.utils import CredentialItem +class _CapturedRequest(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + headers: dict[str, str] + body: bytes | str | None + method: str | None + url: str | None + + +class _CapturedCall(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + request: _CapturedRequest + + @pytest.fixture def base_url(): return "http://localhost:8000" @@ -316,3 +338,108 @@ def test_list_gives_up_at_the_timeout_instead_of_hanging(hanging_server): client.list() assert time.monotonic() - started < 10 + + +@responses.activate +def test_user_connection_methods_encode_names_send_auth_and_validate_responses( + client: CredentialsManagementClient, + base_url: str, + api_key: str, +) -> None: + credential_name: Final = "my/ credential" + credential_path: Final = "my%2F%20credential" + start_payload: Final = { + "user_code": "ABCD-1234", + "verification_uri": "https://github.com/login/device", + "expires_in": 900, + "interval": 5, + "flow_handle": "encrypted-handle", + } + poll_payload: Final = {"status": "slow_down", "interval": 10, "github_login": None} + connections_payload: Final = { + "connections": [ + { + "credential_name": credential_name, + "provider": "github_copilot", + "connected": True, + "github_login": "octocat", + "connected_at": "2025-01-19T12:00:00+00:00", + } + ] + } + delete_payload: Final = {"status": "disconnected"} + responses.add( + responses.POST, + f"{base_url}/credentials/{credential_path}/user_connection/start", + json=start_payload, + status=200, + ) + responses.add( + responses.POST, + f"{base_url}/credentials/{credential_path}/user_connection/poll", + json=poll_payload, + status=200, + ) + responses.add( + responses.GET, + f"{base_url}/credentials/user_connections", + json=connections_payload, + status=200, + ) + responses.add( + responses.DELETE, + f"{base_url}/credentials/{credential_path}/user_connection", + json=delete_payload, + status=200, + ) + + start: Final = client.start_user_connection(credential_name) + poll: Final = client.poll_user_connection(credential_name, "encrypted-handle") + connections: Final = client.list_user_connections() + deleted: Final = client.delete_user_connection(credential_name) + calls: Final = TypeAdapter(tuple[_CapturedCall, ...]).validate_python(responses.calls) + + assert start == UserConnectionStart.model_validate(start_payload) + assert poll == UserConnectionPoll.model_validate(poll_payload) + assert connections == UserConnections.model_validate(connections_payload) + assert connections.connections == ( + UserConnection.model_validate(connections_payload["connections"][0]), + ) + assert deleted == UserConnectionDelete.model_validate(delete_payload) + assert tuple(call.request.headers["Authorization"] for call in calls) == (f"Bearer {api_key}",) * 4 + assert calls[1].request.body == b'{"flow_handle": "encrypted-handle"}' + assert tuple(call.request.method for call in calls) == ("POST", "POST", "GET", "DELETE") + assert tuple(call.request.url for call in calls) == ( + f"{base_url}/credentials/{credential_path}/user_connection/start", + f"{base_url}/credentials/{credential_path}/user_connection/poll", + f"{base_url}/credentials/user_connections", + f"{base_url}/credentials/{credential_path}/user_connection", + ) + + +@responses.activate +def test_start_user_connection_raises_unauthorized_error( + client: CredentialsManagementClient, + base_url: str, +) -> None: + responses.add( + responses.POST, + f"{base_url}/credentials/copilot/user_connection/start", + json={"detail": "An authenticated user is required"}, + status=401, + ) + + with pytest.raises(UnauthorizedError): + client.start_user_connection("copilot") + + +def test_user_connection_start_repr_omits_flow_handle() -> None: + start: Final = UserConnectionStart( + user_code="ABCD-1234", + verification_uri="https://github.com/login/device", + expires_in=900, + interval=5, + flow_handle="secret-handle", + ) + + assert "secret-handle" not in repr(start) diff --git a/tests/unit/test_litellm_cli.py b/tests/unit/test_litellm_cli.py index 5d1e510742a..a986b9aee42 100644 --- a/tests/unit/test_litellm_cli.py +++ b/tests/unit/test_litellm_cli.py @@ -21,6 +21,8 @@ _HELP_ARGV: Final = ( ("codex", "--help"), ("pi", "--help"), ("keys", "list", "--help"), + ("credentials", "connect", "--help"), + ("credentials", "connections", "--help"), ) _COMMAND_ARGV: Final = (("logout",), ("whoami",))