mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
feat(cli): connect a per-user GitHub Copilot credential from the command line (#50863)
* feat(cli): connect a per-user GitHub Copilot credential from the command line lite credentials connect runs the GitHub device login with the caller's own key, so users can connect Copilot without opening the LiteLLM UI. * fix(cli): stop connect before the GitHub code expires and explain lost connections The connect command kept polling until the proxy rejected the expired code with a raw HTTP 400, so it now stops one interval before expiry and prints its own expired message. A dropped connection or timeout while waiting now prints a short rerun message instead of a traceback. * fix(cli): keep checking for approval until the GitHub code actually expires Connect stopped one interval early, so an approval in the last seconds was missed. The countdown now starts when the code is requested, which is never later than the proxy's own cutoff, and polling continues while another poll fits.
This commit is contained in:
parent
0fd6dd4168
commit
e95bf288f3
4 changed files with 297 additions and 9 deletions
|
|
@ -260,6 +260,16 @@ Delete a credential:
|
|||
lite credentials delete <credential-name>
|
||||
```
|
||||
|
||||
#### Connect a GitHub Copilot Credential
|
||||
|
||||
Connect your own GitHub Copilot account to a per-user GitHub Copilot credential, the same as the Connect button in the UI. Run it with your own virtual key, since the connection belongs to the user that key resolves to:
|
||||
|
||||
```bash
|
||||
lite --api-key sk-your-key credentials connect <credential-name>
|
||||
```
|
||||
|
||||
It prints a code and a GitHub link, then waits. Open the link, enter the code, and approve. The command exits 0 once GitHub confirms the login, and exits 1 if the code expires, the login is denied, or the GitHub account has no Copilot seat.
|
||||
|
||||
### Keys Management
|
||||
|
||||
The CLI provides commands for managing API keys on your LiteLLM proxy server:
|
||||
|
|
|
|||
|
|
@ -1,5 +1,8 @@
|
|||
import json
|
||||
from collections.abc import Sequence
|
||||
import time
|
||||
from collections.abc import Callable, Sequence
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal
|
||||
|
||||
import click
|
||||
|
|
@ -8,9 +11,30 @@ import rich
|
|||
from rich.table import Table
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
from litellm.models.credentials import UserConnectionPollResponse, UserConnectionPollStatus
|
||||
|
||||
from ...credentials import CredentialsManagementClient
|
||||
from ...exceptions import UnauthorizedError
|
||||
from ._cli_context import cli_context_values
|
||||
|
||||
_SLOW_DOWN_STEP_SECONDS: Final = 5
|
||||
_CONNECT_FAILURES: Final[MappingProxyType[UserConnectionPollStatus, str]] = MappingProxyType(
|
||||
{
|
||||
"expired": "The code expired before it was approved. Run the command again for a new code",
|
||||
"denied": "The GitHub login was denied",
|
||||
"no_copilot_seat": "That GitHub account has no GitHub Copilot seat",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PollClock:
|
||||
now: Callable[[], float]
|
||||
sleep: Callable[[float], None]
|
||||
|
||||
|
||||
SYSTEM_CLOCK: Final = PollClock(now=time.monotonic, sleep=time.sleep)
|
||||
|
||||
|
||||
class _CredentialInfo(TypedDict):
|
||||
custom_llm_provider: ReadOnly[NotRequired[str]]
|
||||
|
|
@ -179,3 +203,80 @@ def get(ctx: click.Context, credential_name: str):
|
|||
client: Final = CredentialsManagementClient(context["base_url"], context["api_key"])
|
||||
response: Final = client.get(credential_name)
|
||||
_print_json(response)
|
||||
|
||||
|
||||
def _next_poll_interval(poll: UserConnectionPollResponse, current: int) -> int:
|
||||
if poll.interval is not None:
|
||||
return poll.interval
|
||||
return current + _SLOW_DOWN_STEP_SECONDS if poll.status == "slow_down" else current
|
||||
|
||||
|
||||
def _poll_until_settled(
|
||||
client: CredentialsManagementClient,
|
||||
credential_name: str,
|
||||
flow_handle: str,
|
||||
first_interval: int,
|
||||
expires_at: float,
|
||||
clock: PollClock,
|
||||
) -> UserConnectionPollResponse:
|
||||
interval = first_interval # rebind-ok: GitHub can ask for a longer wait on any poll
|
||||
while clock.now() + interval < expires_at:
|
||||
clock.sleep(interval)
|
||||
if (poll := client.poll_user_connection(credential_name, flow_handle)).status not in ("pending", "slow_down"):
|
||||
return poll
|
||||
interval = _next_poll_interval(poll, interval)
|
||||
return UserConnectionPollResponse(status="expired")
|
||||
|
||||
|
||||
def connect_user_credential(
|
||||
client: CredentialsManagementClient,
|
||||
credential_name: str,
|
||||
clock: PollClock,
|
||||
) -> UserConnectionPollResponse:
|
||||
"""Run the GitHub device login for the calling key's user and return the final poll result.
|
||||
|
||||
Polling stops once the next poll would land after the code expires, since the proxy rejects polls on an expired
|
||||
code. The countdown starts before the code is requested, so it never runs past the proxy's own."""
|
||||
requested_at: Final = clock.now()
|
||||
started: Final = client.start_user_connection(credential_name)
|
||||
click.echo(f"Go to {started.verification_uri} and enter code {started.user_code}")
|
||||
click.echo("Waiting for approval...")
|
||||
expires_at: Final = requested_at + started.expires_in
|
||||
return _poll_until_settled(client, credential_name, started.flow_handle, started.interval, expires_at, clock)
|
||||
|
||||
|
||||
def _report_http_error(response: requests.Response) -> None:
|
||||
click.echo(f"Error: HTTP {response.status_code}", err=True)
|
||||
try:
|
||||
error_body: Final[_JsonBodyView] = {"body": response.json()}
|
||||
_print_json(error_body["body"])
|
||||
except json.JSONDecodeError:
|
||||
click.echo(response.text, err=True)
|
||||
|
||||
|
||||
@credentials.command()
|
||||
@click.argument("credential_name")
|
||||
@click.pass_context
|
||||
def connect(ctx: click.Context, credential_name: str) -> None:
|
||||
"""Connect your own GitHub Copilot account to a per-user credential, without the UI"""
|
||||
context: Final = cli_context_values(ctx)
|
||||
client: Final = CredentialsManagementClient(context["base_url"], context["api_key"])
|
||||
try:
|
||||
result: Final = connect_user_credential(client, credential_name, SYSTEM_CLOCK)
|
||||
except UnauthorizedError as e:
|
||||
if isinstance(e.orig_exception, requests.exceptions.HTTPError) and e.orig_exception.response is not None:
|
||||
_report_http_error(e.orig_exception.response)
|
||||
raise click.Abort()
|
||||
except requests.exceptions.HTTPError as e:
|
||||
_report_http_error(e.response)
|
||||
raise click.Abort()
|
||||
except requests.exceptions.RequestException as e:
|
||||
click.echo(
|
||||
f"Error: could not reach the proxy at {context['base_url']} ({type(e).__name__}). Run the command again",
|
||||
err=True,
|
||||
)
|
||||
raise click.Abort()
|
||||
if result.status != "connected":
|
||||
click.echo(f"Error: {_CONNECT_FAILURES.get(result.status, result.status)}", err=True)
|
||||
raise click.exceptions.Exit(1)
|
||||
click.echo(f"Connected as {result.github_login} for credential '{credential_name}'")
|
||||
|
|
|
|||
|
|
@ -5,6 +5,8 @@ from urllib.parse import quote
|
|||
import requests
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm.models.credentials import UserConnectionPollResponse, UserConnectionStartResponse
|
||||
|
||||
from .exceptions import UnauthorizedError
|
||||
|
||||
_DECODED_JSON: Final = TypeAdapter(object)
|
||||
|
|
@ -217,3 +219,33 @@ class CredentialsManagementClient:
|
|||
if e.response.status_code == 401:
|
||||
raise UnauthorizedError(e)
|
||||
raise
|
||||
|
||||
def _send_json(self, request: requests.Request) -> object:
|
||||
try:
|
||||
response: Final = requests.Session().send(request.prepare(), timeout=self._timeout)
|
||||
response.raise_for_status()
|
||||
return _DECODED_JSON.validate_python(response.json())
|
||||
except requests.exceptions.HTTPError as e:
|
||||
if e.response.status_code == 401:
|
||||
raise UnauthorizedError(e)
|
||||
raise
|
||||
|
||||
def _user_connection_url(self, credential_name: str, action: str) -> str:
|
||||
return f"{self._base_url}/credentials/{quote(credential_name, safe='')}/user_connection/{action}"
|
||||
|
||||
def start_user_connection(self, credential_name: str) -> UserConnectionStartResponse:
|
||||
"""Start the GitHub device login that connects the calling key's user to a per-user credential."""
|
||||
request: Final = requests.Request(
|
||||
"POST", self._user_connection_url(credential_name, "start"), headers=self._get_headers()
|
||||
)
|
||||
return UserConnectionStartResponse.model_validate(self._send_json(request))
|
||||
|
||||
def poll_user_connection(self, credential_name: str, flow_handle: str) -> UserConnectionPollResponse:
|
||||
"""Check once whether the user approved the device login that ``start_user_connection`` began."""
|
||||
request: Final = requests.Request(
|
||||
"POST",
|
||||
self._user_connection_url(credential_name, "poll"),
|
||||
headers=self._get_headers(),
|
||||
json={"flow_handle": flow_handle},
|
||||
)
|
||||
return UserConnectionPollResponse.model_validate(self._send_json(request))
|
||||
|
|
|
|||
|
|
@ -3,11 +3,12 @@ from unittest.mock import MagicMock
|
|||
|
||||
import pytest
|
||||
import requests
|
||||
import responses
|
||||
from click.testing import CliRunner
|
||||
|
||||
|
||||
|
||||
from litellm.proxy.client.cli.commands.credentials import PollClock, connect_user_credential
|
||||
from litellm.proxy.client.cli.main import cli
|
||||
from litellm.proxy.client.credentials import CredentialsManagementClient
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -136,9 +137,7 @@ def test_acreate_credential_http_error(cli_runner, mock_credentials_client):
|
|||
mock_error_response = MagicMock()
|
||||
mock_error_response.status_code = 400
|
||||
mock_error_response.json.return_value = {"error": "Invalid request"}
|
||||
mock_instance.create.side_effect = requests.exceptions.HTTPError(
|
||||
response=mock_error_response
|
||||
)
|
||||
mock_instance.create.side_effect = requests.exceptions.HTTPError(response=mock_error_response)
|
||||
|
||||
# Run command
|
||||
result = cli_runner.invoke(
|
||||
|
|
@ -182,9 +181,7 @@ def test_adelete_credential_http_error(cli_runner, mock_credentials_client):
|
|||
mock_error_response = MagicMock()
|
||||
mock_error_response.status_code = 404
|
||||
mock_error_response.json.return_value = {"error": "Credential not found"}
|
||||
mock_instance.delete.side_effect = requests.exceptions.HTTPError(
|
||||
response=mock_error_response
|
||||
)
|
||||
mock_instance.delete.side_effect = requests.exceptions.HTTPError(response=mock_error_response)
|
||||
|
||||
# Run command
|
||||
result = cli_runner.invoke(cli, ["credentials", "delete", "test-cred"])
|
||||
|
|
@ -264,3 +261,151 @@ 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()
|
||||
|
||||
|
||||
PROXY = "http://proxy.test"
|
||||
CONNECT_URL = f"{PROXY}/credentials/team%2Fcopilot/user_connection"
|
||||
|
||||
|
||||
def _start_body(interval: int = 0, expires_in: int = 900) -> dict[str, object]:
|
||||
return {
|
||||
"user_code": "ABCD-1234",
|
||||
"verification_uri": "https://github.com/login/device",
|
||||
"expires_in": expires_in,
|
||||
"interval": interval,
|
||||
"flow_handle": "sealed-handle",
|
||||
}
|
||||
|
||||
|
||||
class _FakeClock:
|
||||
def __init__(self) -> None:
|
||||
self.elapsed = 0.0
|
||||
self.waits: tuple[float, ...] = ()
|
||||
|
||||
def sleep(self, seconds: float) -> None:
|
||||
self.elapsed += seconds
|
||||
self.waits = (*self.waits, seconds)
|
||||
|
||||
def as_poll_clock(self) -> PollClock:
|
||||
return PollClock(now=lambda: self.elapsed, sleep=self.sleep)
|
||||
|
||||
|
||||
def _run_connect(cli_runner):
|
||||
return cli_runner.invoke(
|
||||
cli, ["--base-url", PROXY, "--api-key", "sk-user", "credentials", "connect", "team/copilot"]
|
||||
)
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_connect_shows_the_github_code_and_reports_the_connected_login(cli_runner):
|
||||
responses.add(responses.POST, f"{CONNECT_URL}/start", json=_start_body())
|
||||
responses.add(responses.POST, f"{CONNECT_URL}/poll", json={"status": "pending"})
|
||||
responses.add(responses.POST, f"{CONNECT_URL}/poll", json={"status": "connected", "github_login": "octocat"})
|
||||
|
||||
result = _run_connect(cli_runner)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "https://github.com/login/device" in result.output
|
||||
assert "ABCD-1234" in result.output
|
||||
assert "Connected as octocat" in result.output
|
||||
polls = [call for call in responses.calls if call.request.url.endswith("/poll")]
|
||||
assert len(polls) == 2
|
||||
assert all(json.loads(call.request.body) == {"flow_handle": "sealed-handle"} for call in polls)
|
||||
assert all(call.request.headers["Authorization"] == "Bearer sk-user" for call in responses.calls)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("status", "message"),
|
||||
[("expired", "expired"), ("denied", "denied"), ("no_copilot_seat", "no GitHub Copilot seat")],
|
||||
)
|
||||
@responses.activate
|
||||
def test_connect_exits_nonzero_when_the_login_does_not_complete(cli_runner, status, message):
|
||||
responses.add(responses.POST, f"{CONNECT_URL}/start", json=_start_body())
|
||||
responses.add(responses.POST, f"{CONNECT_URL}/poll", json={"status": status})
|
||||
|
||||
result = _run_connect(cli_runner)
|
||||
|
||||
assert result.exit_code == 1
|
||||
assert message in result.output
|
||||
assert "Connected as" not in result.output
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("status_code", "detail"),
|
||||
[(401, "An authenticated user is required"), (404, "Credential not found")],
|
||||
)
|
||||
@responses.activate
|
||||
def test_connect_surfaces_the_proxy_error_and_never_polls(cli_runner, status_code, detail):
|
||||
responses.add(responses.POST, f"{CONNECT_URL}/start", status=status_code, json={"detail": detail})
|
||||
|
||||
result = _run_connect(cli_runner)
|
||||
|
||||
assert result.exit_code != 0
|
||||
assert f"HTTP {status_code}" in result.output
|
||||
assert detail in result.output
|
||||
assert not any(call.request.url.endswith("/poll") for call in responses.calls)
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_connect_waits_the_interval_github_asks_for_and_backs_off_on_slow_down():
|
||||
responses.add(responses.POST, f"{CONNECT_URL}/start", json=_start_body(interval=5))
|
||||
for poll in (
|
||||
{"status": "slow_down"},
|
||||
{"status": "pending"},
|
||||
{"status": "slow_down", "interval": 20},
|
||||
{"status": "connected", "github_login": "octocat"},
|
||||
):
|
||||
responses.add(responses.POST, f"{CONNECT_URL}/poll", json=poll)
|
||||
clock = _FakeClock()
|
||||
|
||||
result = connect_user_credential(
|
||||
CredentialsManagementClient(PROXY, "sk-user"), "team/copilot", clock.as_poll_clock()
|
||||
)
|
||||
|
||||
assert result.status == "connected"
|
||||
assert clock.waits == (5, 10, 10, 20)
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_connect_stops_polling_once_the_next_poll_would_land_after_the_code_expires():
|
||||
responses.add(responses.POST, f"{CONNECT_URL}/start", json=_start_body(interval=5, expires_in=19))
|
||||
for _ in range(3):
|
||||
responses.add(responses.POST, f"{CONNECT_URL}/poll", json={"status": "pending"})
|
||||
responses.add(responses.POST, f"{CONNECT_URL}/poll", status=400, json={"detail": "invalid or expired flow_handle"})
|
||||
clock = _FakeClock()
|
||||
|
||||
result = connect_user_credential(
|
||||
CredentialsManagementClient(PROXY, "sk-user"), "team/copilot", clock.as_poll_clock()
|
||||
)
|
||||
|
||||
assert result.status == "expired"
|
||||
assert clock.waits == (5, 5, 5)
|
||||
assert sum(call.request.url.endswith("/poll") for call in responses.calls) == 3
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_connect_still_collects_an_approval_on_the_last_poll_before_the_code_expires():
|
||||
responses.add(responses.POST, f"{CONNECT_URL}/start", json=_start_body(interval=5, expires_in=19))
|
||||
responses.add(responses.POST, f"{CONNECT_URL}/poll", json={"status": "pending"})
|
||||
responses.add(responses.POST, f"{CONNECT_URL}/poll", json={"status": "pending"})
|
||||
responses.add(responses.POST, f"{CONNECT_URL}/poll", json={"status": "connected", "github_login": "octocat"})
|
||||
clock = _FakeClock()
|
||||
|
||||
result = connect_user_credential(
|
||||
CredentialsManagementClient(PROXY, "sk-user"), "team/copilot", clock.as_poll_clock()
|
||||
)
|
||||
|
||||
assert result.status == "connected"
|
||||
assert clock.elapsed == 15
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_connect_reports_a_lost_connection_while_waiting_and_asks_to_rerun(cli_runner):
|
||||
responses.add(responses.POST, f"{CONNECT_URL}/start", json=_start_body())
|
||||
responses.add(responses.POST, f"{CONNECT_URL}/poll", body=requests.exceptions.ConnectionError("connection reset"))
|
||||
|
||||
result = _run_connect(cli_runner)
|
||||
|
||||
assert result.exit_code == 1
|
||||
assert "ABCD-1234" in result.output
|
||||
assert f"could not reach the proxy at {PROXY} (ConnectionError). Run the command again" in result.output
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue