mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-27 06:13:07 +00:00
feat(github_copilot): support a GHE.com host on per-user Copilot credentials (#47610)
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Yassin Kortam
GitHub
parent
2a268951c4
commit
db2039b11b
29 changed files with 1233 additions and 214 deletions
|
|
@ -25,6 +25,8 @@ CLAUDE_CODE_CLIENT: Final = "claude-code"
|
|||
MICROSOFT_365_COPILOT_DEFAULT_TIME_ZONE: Final = "UTC"
|
||||
GITHUB_COPILOT_PER_USER_AUTH_TYPE: Final = "per_user_oauth"
|
||||
GITHUB_COPILOT_AUTH_TYPE_KEY: Final = "github_copilot_auth_type"
|
||||
GITHUB_COPILOT_HOST_KEY: Final = "github_copilot_host"
|
||||
GITHUB_COPILOT_DEFAULT_HOST: Final = "github.com"
|
||||
GITHUB_COPILOT_USER_TOKEN_SAFETY_MARGIN_SECONDS: Final = 60
|
||||
GITHUB_COPILOT_USER_CREDENTIAL_CACHE_TTL_SECONDS: Final = 60
|
||||
GITHUB_COPILOT_DEVICE_FLOW_CACHE_PREFIX: Final = "github_copilot_device_flow"
|
||||
|
|
|
|||
|
|
@ -77,6 +77,7 @@ OPTIONAL_KWARGS_KEYS: Final = (
|
|||
"otpm",
|
||||
"use_xai_oauth",
|
||||
"github_copilot_auth_type",
|
||||
"github_copilot_host",
|
||||
"github_copilot_user_session",
|
||||
"fireworks_forward_user_id",
|
||||
"responses_api_version",
|
||||
|
|
|
|||
78
litellm/llms/github_copilot/github_host.py
Normal file
78
litellm/llms/github_copilot/github_host.py
Normal file
|
|
@ -0,0 +1,78 @@
|
|||
import os
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
from litellm.constants import GITHUB_COPILOT_DEFAULT_HOST
|
||||
|
||||
from .authenticator import (
|
||||
DEFAULT_GITHUB_ACCESS_TOKEN_URL,
|
||||
DEFAULT_GITHUB_API_KEY_URL,
|
||||
DEFAULT_GITHUB_CLIENT_ID,
|
||||
DEFAULT_GITHUB_DEVICE_CODE_URL,
|
||||
)
|
||||
|
||||
_GHE_SUBDOMAIN_PATTERN: Final = re.compile(r"[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?")
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class GithubHostEndpoints:
|
||||
host: str
|
||||
device_code_url: str
|
||||
access_token_url: str
|
||||
api_key_url: str
|
||||
user_url: str
|
||||
client_id: str
|
||||
copilot_api_hosts: frozenset[str]
|
||||
copilot_api_host_suffix: str | None
|
||||
default_copilot_api_base: str
|
||||
|
||||
|
||||
def normalized_github_host(raw: object) -> str | None:
|
||||
if raw is None or raw == "":
|
||||
return GITHUB_COPILOT_DEFAULT_HOST
|
||||
if not isinstance(raw, str):
|
||||
return None
|
||||
|
||||
host: Final = raw.lower()
|
||||
if host == GITHUB_COPILOT_DEFAULT_HOST:
|
||||
return host
|
||||
if not host.endswith(".ghe.com"):
|
||||
return None
|
||||
|
||||
subdomain: Final = host.removesuffix(".ghe.com")
|
||||
return host if _GHE_SUBDOMAIN_PATTERN.fullmatch(subdomain) is not None else None
|
||||
|
||||
|
||||
def github_host_endpoints(host: str) -> GithubHostEndpoints:
|
||||
normalized: Final = normalized_github_host(host)
|
||||
if normalized is None:
|
||||
raise ValueError("host must be github.com or a single-subdomain GHE.com hostname")
|
||||
client_id: Final = os.getenv("GITHUB_COPILOT_CLIENT_ID", DEFAULT_GITHUB_CLIENT_ID)
|
||||
if normalized == GITHUB_COPILOT_DEFAULT_HOST:
|
||||
return GithubHostEndpoints(
|
||||
host=normalized,
|
||||
device_code_url=os.getenv("GITHUB_COPILOT_DEVICE_CODE_URL", DEFAULT_GITHUB_DEVICE_CODE_URL),
|
||||
access_token_url=os.getenv("GITHUB_COPILOT_ACCESS_TOKEN_URL", DEFAULT_GITHUB_ACCESS_TOKEN_URL),
|
||||
api_key_url=os.getenv("GITHUB_COPILOT_API_KEY_URL", DEFAULT_GITHUB_API_KEY_URL),
|
||||
user_url="https://api.github.com/user",
|
||||
client_id=client_id,
|
||||
copilot_api_hosts=frozenset({"githubcopilot.com"}),
|
||||
copilot_api_host_suffix=".githubcopilot.com",
|
||||
default_copilot_api_base="https://api.githubcopilot.com",
|
||||
)
|
||||
|
||||
subdomain: Final = normalized.removesuffix(".ghe.com")
|
||||
api_host: Final = f"api.{subdomain}.ghe.com"
|
||||
copilot_api_host: Final = f"copilot-api.{subdomain}.ghe.com"
|
||||
return GithubHostEndpoints(
|
||||
host=normalized,
|
||||
device_code_url=f"https://{normalized}/login/device/code",
|
||||
access_token_url=f"https://{normalized}/login/oauth/access_token",
|
||||
api_key_url=f"https://{api_host}/copilot_internal/v2/token",
|
||||
user_url=f"https://{api_host}/user",
|
||||
client_id=client_id,
|
||||
copilot_api_hosts=frozenset({copilot_api_host}),
|
||||
copilot_api_host_suffix=None,
|
||||
default_copilot_api_base=f"https://{copilot_api_host}",
|
||||
)
|
||||
|
|
@ -5,8 +5,7 @@ touched on the per-user path.
|
|||
"""
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import os
|
||||
import hmac
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Callable, Mapping
|
||||
|
|
@ -27,6 +26,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict
|
|||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.constants import (
|
||||
GITHUB_COPILOT_AUTH_TYPE_KEY,
|
||||
GITHUB_COPILOT_HOST_KEY,
|
||||
GITHUB_COPILOT_PER_USER_AUTH_TYPE,
|
||||
GITHUB_COPILOT_USER_TOKEN_SAFETY_MARGIN_SECONDS,
|
||||
)
|
||||
|
|
@ -38,14 +38,8 @@ from litellm.exceptions import (
|
|||
ServiceUnavailableError,
|
||||
)
|
||||
|
||||
from .authenticator import (
|
||||
DEFAULT_GITHUB_ACCESS_TOKEN_URL,
|
||||
DEFAULT_GITHUB_API_KEY_URL,
|
||||
DEFAULT_GITHUB_CLIENT_ID,
|
||||
DEFAULT_GITHUB_DEVICE_CODE_URL,
|
||||
github_api_headers,
|
||||
)
|
||||
from .common_utils import DEFAULT_GITHUB_COPILOT_API_BASE
|
||||
from .authenticator import github_api_headers
|
||||
from .github_host import GithubHostEndpoints, github_host_endpoints, normalized_github_host
|
||||
|
||||
GITHUB_COPILOT_USER_SESSION_KWARG_KEY: Final = "github_copilot_user_session"
|
||||
_LLM_PROVIDER: Final = "github_copilot"
|
||||
|
|
@ -98,7 +92,7 @@ _SESSION_CACHE: Final = cast( # cast-ok: pins the untyped module cache to _Sess
|
|||
)
|
||||
_EXCHANGE_LOCKS: Final = tuple(threading.Lock() for _ in range(_SYNC_LOCK_COUNT))
|
||||
_IN_FLIGHT: Final[ # mutable-ok: single-flight dedup registry mutated under its own locking
|
||||
dict[str, "asyncio.Future[GithubCopilotUserSession]"]
|
||||
dict[str, tuple[str, "asyncio.Future[GithubCopilotUserSession]"]]
|
||||
] = {}
|
||||
|
||||
|
||||
|
|
@ -141,7 +135,7 @@ _DEVICE_CODE: Final = TypeAdapter(_DeviceCodePayload)
|
|||
_ACCESS_TOKEN_POLL: Final = TypeAdapter(_AccessTokenPollPayload)
|
||||
_GITHUB_USER: Final = TypeAdapter(_GithubUserPayload)
|
||||
_COPILOT_ENDPOINTS: Final = TypeAdapter(_CopilotEndpointsPayload)
|
||||
_CACHED_SESSION_ENTRY: Final = TypeAdapter(tuple[str, str, float])
|
||||
_CACHED_SESSION_ENTRY: Final = TypeAdapter(tuple[str, str, str, float])
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True, repr=False)
|
||||
|
|
@ -178,53 +172,72 @@ class GithubCopilotDeviceFlowPoll:
|
|||
access_token: str | None = field(default=None, repr=False)
|
||||
|
||||
|
||||
def validated_copilot_api_base(endpoints_api: object) -> str:
|
||||
"""Only a genuine *.githubcopilot.com HTTPS endpoint may receive a Copilot
|
||||
token; anything else falls back to the default base."""
|
||||
def validated_copilot_api_base(endpoints_api: object, endpoints: GithubHostEndpoints) -> str:
|
||||
if not isinstance(endpoints_api, str):
|
||||
return DEFAULT_GITHUB_COPILOT_API_BASE
|
||||
parsed: Final = urlparse(endpoints_api)
|
||||
return endpoints.default_copilot_api_base
|
||||
try:
|
||||
parsed: Final = urlparse(endpoints_api)
|
||||
port: Final = parsed.port
|
||||
except ValueError:
|
||||
return endpoints.default_copilot_api_base
|
||||
hostname: Final = (parsed.hostname or "").lower()
|
||||
if (
|
||||
parsed.scheme.lower() != "https"
|
||||
or not (hostname == "githubcopilot.com" or hostname.endswith(".githubcopilot.com"))
|
||||
or not (
|
||||
hostname in endpoints.copilot_api_hosts
|
||||
or (endpoints.copilot_api_host_suffix is not None and hostname.endswith(endpoints.copilot_api_host_suffix))
|
||||
)
|
||||
or parsed.username is not None
|
||||
or parsed.password is not None
|
||||
or port not in (None, 443)
|
||||
):
|
||||
return DEFAULT_GITHUB_COPILOT_API_BASE
|
||||
try:
|
||||
port: Final = parsed.port
|
||||
except ValueError:
|
||||
return DEFAULT_GITHUB_COPILOT_API_BASE
|
||||
if port not in (None, 443):
|
||||
return DEFAULT_GITHUB_COPILOT_API_BASE
|
||||
return endpoints.default_copilot_api_base
|
||||
return endpoints_api.rstrip("/")
|
||||
|
||||
|
||||
def _session_cache_key(user_id: str, github_token: str) -> str:
|
||||
return hashlib.sha256(f"{user_id}\x00{github_token}".encode()).hexdigest()
|
||||
def github_copilot_host_endpoints(credential_name: str) -> GithubHostEndpoints:
|
||||
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
|
||||
|
||||
credential: Final = CredentialAccessor.find_credential(credential_name)
|
||||
host_raw: Final[object] = (
|
||||
credential.credential_values.get(GITHUB_COPILOT_HOST_KEY) if credential is not None else None
|
||||
)
|
||||
host: Final = normalized_github_host(host_raw)
|
||||
if host is None:
|
||||
raise BadRequestError(
|
||||
message=f"Credential '{credential_name}' has an invalid {GITHUB_COPILOT_HOST_KEY}",
|
||||
llm_provider=_LLM_PROVIDER,
|
||||
model="",
|
||||
)
|
||||
return github_host_endpoints(host)
|
||||
|
||||
|
||||
def evict_copilot_user_session(user_id: str, github_token: str) -> None:
|
||||
_SESSION_CACHE.delete_cache(_session_cache_key(user_id, github_token))
|
||||
def _session_cache_key(user_id: str, credential_name: str, github_host: str) -> str:
|
||||
return f"{user_id}\x00{credential_name}\x00{github_host}"
|
||||
|
||||
|
||||
def _cached_session(cache_key: str) -> GithubCopilotUserSession | None:
|
||||
def evict_copilot_user_session(user_id: str, credential_name: str, github_host: str) -> None:
|
||||
_SESSION_CACHE.delete_cache(_session_cache_key(user_id, credential_name, github_host))
|
||||
|
||||
|
||||
def _cached_session(cache_key: str, github_token: str) -> GithubCopilotUserSession | None:
|
||||
cached: Final[object] = _SESSION_CACHE.get_cache(cache_key)
|
||||
try:
|
||||
token, api_base, expires_at = _CACHED_SESSION_ENTRY.validate_python(cached)
|
||||
cached_github_token, token, api_base, expires_at = _CACHED_SESSION_ENTRY.validate_python(cached)
|
||||
except ValidationError:
|
||||
return None
|
||||
if not hmac.compare_digest(cached_github_token, github_token):
|
||||
return None
|
||||
if expires_at - GITHUB_COPILOT_USER_TOKEN_SAFETY_MARGIN_SECONDS <= time.time():
|
||||
return None
|
||||
return GithubCopilotUserSession(token=token, api_base=api_base)
|
||||
|
||||
|
||||
def _cache_session(cache_key: str, session: GithubCopilotUserSession, expires_at: float) -> None:
|
||||
def _cache_session(cache_key: str, github_token: str, session: GithubCopilotUserSession, expires_at: float) -> None:
|
||||
ttl: Final = expires_at - GITHUB_COPILOT_USER_TOKEN_SAFETY_MARGIN_SECONDS - time.time()
|
||||
if ttl <= 0:
|
||||
return
|
||||
_SESSION_CACHE.set_cache(cache_key, (session.token, session.api_base, expires_at), ttl=ttl)
|
||||
_SESSION_CACHE.set_cache(cache_key, (github_token, session.token, session.api_base, expires_at), ttl=ttl)
|
||||
|
||||
|
||||
def _reject_connection(credential_name: str) -> CallerCredentialAuthenticationError:
|
||||
|
|
@ -238,7 +251,13 @@ def _reject_connection(credential_name: str) -> CallerCredentialAuthenticationEr
|
|||
)
|
||||
|
||||
|
||||
def _session_from_response(response: httpx.Response, credential_name: str, cache_key: str) -> GithubCopilotUserSession:
|
||||
def _session_from_response(
|
||||
response: httpx.Response,
|
||||
github_token: str,
|
||||
credential_name: str,
|
||||
cache_key: str,
|
||||
endpoints: GithubHostEndpoints,
|
||||
) -> GithubCopilotUserSession:
|
||||
if response.status_code in (401, 403, 404):
|
||||
_SESSION_CACHE.delete_cache(cache_key)
|
||||
raise _reject_connection(credential_name)
|
||||
|
|
@ -263,15 +282,17 @@ def _session_from_response(response: httpx.Response, credential_name: str, cache
|
|||
model="",
|
||||
) from e
|
||||
|
||||
endpoints: Final[object] = payload.get("endpoints")
|
||||
payload_endpoints: Final[object] = payload.get("endpoints")
|
||||
endpoints_api: Final[object] = (
|
||||
_COPILOT_ENDPOINTS.validate_python(endpoints).get("api") if isinstance(endpoints, Mapping) else None
|
||||
_COPILOT_ENDPOINTS.validate_python(payload_endpoints).get("api")
|
||||
if isinstance(payload_endpoints, Mapping)
|
||||
else None
|
||||
)
|
||||
api_base: Final = validated_copilot_api_base(endpoints_api)
|
||||
api_base: Final = validated_copilot_api_base(endpoints_api, endpoints)
|
||||
session: Final = GithubCopilotUserSession(token=payload["token"], api_base=api_base)
|
||||
expires_at: Final = _expires_at_seconds(payload.get("expires_at"))
|
||||
if expires_at > 0:
|
||||
_cache_session(cache_key, session, expires_at)
|
||||
_cache_session(cache_key, github_token, session, expires_at)
|
||||
return session
|
||||
|
||||
|
||||
|
|
@ -303,7 +324,12 @@ def _expires_at_seconds(raw: object) -> float:
|
|||
return 0.0
|
||||
|
||||
|
||||
def _do_exchange(github_token: str, credential_name: str, cache_key: str) -> GithubCopilotUserSession:
|
||||
def _do_exchange(
|
||||
github_token: str,
|
||||
credential_name: str,
|
||||
cache_key: str,
|
||||
endpoints: GithubHostEndpoints,
|
||||
) -> GithubCopilotUserSession:
|
||||
import litellm
|
||||
|
||||
try:
|
||||
|
|
@ -311,7 +337,7 @@ def _do_exchange(github_token: str, credential_name: str, cache_key: str) -> Git
|
|||
_SyncGetClient, litellm.module_level_client
|
||||
)
|
||||
response: Final = client.get(
|
||||
DEFAULT_GITHUB_API_KEY_URL,
|
||||
endpoints.api_key_url,
|
||||
headers=github_api_headers(github_token),
|
||||
)
|
||||
except Exception as e:
|
||||
|
|
@ -320,7 +346,7 @@ def _do_exchange(github_token: str, credential_name: str, cache_key: str) -> Git
|
|||
llm_provider=_LLM_PROVIDER,
|
||||
model="",
|
||||
) from e
|
||||
return _session_from_response(response, credential_name, cache_key)
|
||||
return _session_from_response(response, github_token, credential_name, cache_key, endpoints)
|
||||
|
||||
|
||||
def _new_async_github_client() -> _AsyncGitHubClient:
|
||||
|
|
@ -332,11 +358,16 @@ def _new_async_github_client() -> _AsyncGitHubClient:
|
|||
)(_LLM_PROVIDER)
|
||||
|
||||
|
||||
async def _do_aexchange(github_token: str, credential_name: str, cache_key: str) -> GithubCopilotUserSession:
|
||||
async def _do_aexchange(
|
||||
github_token: str,
|
||||
credential_name: str,
|
||||
cache_key: str,
|
||||
endpoints: GithubHostEndpoints,
|
||||
) -> GithubCopilotUserSession:
|
||||
client: Final = _new_async_github_client()
|
||||
try:
|
||||
response: Final = await client.get(
|
||||
DEFAULT_GITHUB_API_KEY_URL,
|
||||
endpoints.api_key_url,
|
||||
headers=github_api_headers(github_token),
|
||||
)
|
||||
except Exception as e:
|
||||
|
|
@ -345,43 +376,50 @@ async def _do_aexchange(github_token: str, credential_name: str, cache_key: str)
|
|||
llm_provider=_LLM_PROVIDER,
|
||||
model="",
|
||||
) from e
|
||||
return _session_from_response(response, credential_name, cache_key)
|
||||
return _session_from_response(response, github_token, credential_name, cache_key, endpoints)
|
||||
|
||||
|
||||
def exchange_github_token(
|
||||
user_id: str,
|
||||
github_token: str,
|
||||
credential_name: str,
|
||||
endpoints: GithubHostEndpoints,
|
||||
) -> GithubCopilotUserSession:
|
||||
cache_key: Final = _session_cache_key(user_id, github_token)
|
||||
cached: Final = _cached_session(cache_key)
|
||||
cache_key: Final = _session_cache_key(user_id, credential_name, endpoints.host)
|
||||
cached: Final = _cached_session(cache_key, github_token)
|
||||
if cached is not None:
|
||||
return cached
|
||||
lock: Final = _EXCHANGE_LOCKS[int(cache_key, 16) % _SYNC_LOCK_COUNT]
|
||||
lock: Final = _EXCHANGE_LOCKS[hash(cache_key) % _SYNC_LOCK_COUNT]
|
||||
with lock:
|
||||
rechecked: Final = _cached_session(cache_key)
|
||||
rechecked: Final = _cached_session(cache_key, github_token)
|
||||
if rechecked is not None:
|
||||
return rechecked
|
||||
return _do_exchange(github_token, credential_name, cache_key)
|
||||
return _do_exchange(github_token, credential_name, cache_key, endpoints)
|
||||
|
||||
|
||||
async def aexchange_github_token(
|
||||
user_id: str,
|
||||
github_token: str,
|
||||
credential_name: str,
|
||||
endpoints: GithubHostEndpoints,
|
||||
) -> GithubCopilotUserSession:
|
||||
cache_key: Final = _session_cache_key(user_id, github_token)
|
||||
cached: Final = _cached_session(cache_key)
|
||||
cache_key: Final = _session_cache_key(user_id, credential_name, endpoints.host)
|
||||
cached: Final = _cached_session(cache_key, github_token)
|
||||
if cached is not None:
|
||||
return cached
|
||||
loop: Final = asyncio.get_running_loop()
|
||||
in_flight: Final = _IN_FLIGHT.get(cache_key)
|
||||
if in_flight is not None and not in_flight.done() and in_flight.get_loop() is loop:
|
||||
return await asyncio.shield(in_flight)
|
||||
in_flight_entry: Final = _IN_FLIGHT.get(cache_key)
|
||||
if in_flight_entry is not None:
|
||||
in_flight_token, in_flight = in_flight_entry
|
||||
if not hmac.compare_digest(in_flight_token, github_token):
|
||||
return await _do_aexchange(github_token, credential_name, cache_key, endpoints)
|
||||
if not in_flight.done() and in_flight.get_loop() is loop:
|
||||
return await asyncio.shield(in_flight)
|
||||
future: Final = loop.create_future()
|
||||
_IN_FLIGHT[cache_key] = future
|
||||
entry: Final = (github_token, future)
|
||||
_IN_FLIGHT[cache_key] = entry
|
||||
try:
|
||||
session: Final = await _do_aexchange(github_token, credential_name, cache_key)
|
||||
session: Final = await _do_aexchange(github_token, credential_name, cache_key, endpoints)
|
||||
except BaseException as e:
|
||||
if not future.done():
|
||||
future.set_exception(e)
|
||||
|
|
@ -392,19 +430,17 @@ async def aexchange_github_token(
|
|||
future.set_result(session)
|
||||
return session
|
||||
finally:
|
||||
if _IN_FLIGHT.get(cache_key) is future:
|
||||
if _IN_FLIGHT.get(cache_key) is entry:
|
||||
_IN_FLIGHT.pop(cache_key, None)
|
||||
|
||||
|
||||
async def astart_device_flow() -> GithubCopilotDeviceFlowStart:
|
||||
async def astart_device_flow(endpoints: GithubHostEndpoints) -> GithubCopilotDeviceFlowStart:
|
||||
client: Final = _new_async_github_client()
|
||||
device_code_url: Final = os.getenv("GITHUB_COPILOT_DEVICE_CODE_URL", DEFAULT_GITHUB_DEVICE_CODE_URL)
|
||||
client_id: Final = os.getenv("GITHUB_COPILOT_CLIENT_ID", DEFAULT_GITHUB_CLIENT_ID)
|
||||
try:
|
||||
response: Final = await client.post(
|
||||
device_code_url,
|
||||
endpoints.device_code_url,
|
||||
headers=github_api_headers(),
|
||||
json={"client_id": client_id, "scope": _DEVICE_FLOW_SCOPE},
|
||||
json={"client_id": endpoints.client_id, "scope": _DEVICE_FLOW_SCOPE},
|
||||
)
|
||||
except httpx.HTTPStatusError as e:
|
||||
raise ServiceUnavailableError(
|
||||
|
|
@ -435,16 +471,14 @@ async def astart_device_flow() -> GithubCopilotDeviceFlowStart:
|
|||
)
|
||||
|
||||
|
||||
async def apoll_device_flow(device_code: str) -> GithubCopilotDeviceFlowPoll:
|
||||
async def apoll_device_flow(device_code: str, endpoints: GithubHostEndpoints) -> GithubCopilotDeviceFlowPoll:
|
||||
client: Final = _new_async_github_client()
|
||||
access_token_url: Final = os.getenv("GITHUB_COPILOT_ACCESS_TOKEN_URL", DEFAULT_GITHUB_ACCESS_TOKEN_URL)
|
||||
client_id: Final = os.getenv("GITHUB_COPILOT_CLIENT_ID", DEFAULT_GITHUB_CLIENT_ID)
|
||||
try:
|
||||
response: Final = await client.post(
|
||||
access_token_url,
|
||||
endpoints.access_token_url,
|
||||
headers=github_api_headers(),
|
||||
json={
|
||||
"client_id": client_id,
|
||||
"client_id": endpoints.client_id,
|
||||
"device_code": device_code,
|
||||
"grant_type": "urn:ietf:params:oauth:grant-type:device_code",
|
||||
},
|
||||
|
|
@ -500,11 +534,11 @@ def _poll_payload_from_response(response: httpx.Response) -> GithubCopilotDevice
|
|||
)
|
||||
|
||||
|
||||
async def afetch_github_login(github_token: str) -> str:
|
||||
async def afetch_github_login(github_token: str, endpoints: GithubHostEndpoints) -> str:
|
||||
client: Final = _new_async_github_client()
|
||||
try:
|
||||
response: Final = await client.get(
|
||||
"https://api.github.com/user",
|
||||
endpoints.user_url,
|
||||
headers=github_api_headers(github_token),
|
||||
)
|
||||
except Exception as e:
|
||||
|
|
@ -530,11 +564,17 @@ async def afetch_github_login(github_token: str) -> str:
|
|||
return payload["login"]
|
||||
|
||||
|
||||
async def acheck_copilot_seat(user_id: str, github_token: str, credential_name: str) -> GithubCopilotUserSession:
|
||||
async def acheck_copilot_seat(
|
||||
user_id: str,
|
||||
github_token: str,
|
||||
credential_name: str,
|
||||
endpoints: GithubHostEndpoints,
|
||||
) -> GithubCopilotUserSession:
|
||||
return await aexchange_github_token(
|
||||
user_id=user_id,
|
||||
github_token=github_token,
|
||||
credential_name=credential_name,
|
||||
endpoints=endpoints,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -567,7 +607,11 @@ def _connect_error(credential_name: str) -> CallerCredentialAuthenticationError:
|
|||
)
|
||||
|
||||
|
||||
def _caller_connection(kwargs: Mapping[str, object], credential_name: str) -> tuple[str, str]:
|
||||
def _caller_connection(
|
||||
kwargs: Mapping[str, object],
|
||||
credential_name: str,
|
||||
endpoints: GithubHostEndpoints,
|
||||
) -> tuple[str, str]:
|
||||
secret_fields_raw: Final[object] = kwargs.get("secret_fields")
|
||||
secret_fields: Final = (
|
||||
cast(Mapping[object, object], secret_fields_raw) # cast-ok: value is Mapping-checked or dict-shaped at runtime
|
||||
|
|
@ -585,8 +629,21 @@ def _caller_connection(kwargs: Mapping[str, object], credential_name: str) -> tu
|
|||
if isinstance(credentials_raw, Mapping)
|
||||
else None
|
||||
)
|
||||
token: Final[object] = credentials.get(credential_name) if credentials is not None else None
|
||||
if not isinstance(user_id, str) or not user_id or not isinstance(token, str) or not token:
|
||||
connection_raw: Final[object] = credentials.get(credential_name) if credentials is not None else None
|
||||
connection: Final = (
|
||||
cast(Mapping[object, object], connection_raw) # cast-ok: runtime Mapping check validates the stored connection
|
||||
if isinstance(connection_raw, Mapping)
|
||||
else None
|
||||
)
|
||||
token: Final[object] = connection.get("access_token") if connection is not None else None
|
||||
github_host: Final = normalized_github_host(connection.get("github_host") if connection is not None else None)
|
||||
if (
|
||||
not isinstance(user_id, str)
|
||||
or not user_id
|
||||
or not isinstance(token, str)
|
||||
or not token
|
||||
or github_host != endpoints.host
|
||||
):
|
||||
raise _connect_error(credential_name)
|
||||
return user_id, token
|
||||
|
||||
|
|
@ -705,11 +762,13 @@ def attach_github_copilot_user_session(
|
|||
if isinstance(kwargs.get(GITHUB_COPILOT_USER_SESSION_KWARG_KEY), GithubCopilotUserSession):
|
||||
return
|
||||
credential_name: Final = credential_name_obj if isinstance(credential_name_obj, str) else ""
|
||||
user_id, github_token = _caller_connection(kwargs, credential_name)
|
||||
endpoints: Final = github_copilot_host_endpoints(credential_name)
|
||||
user_id, github_token = _caller_connection(kwargs, credential_name, endpoints)
|
||||
session: Final = exchange_github_token(
|
||||
user_id=user_id,
|
||||
github_token=github_token,
|
||||
credential_name=credential_name,
|
||||
endpoints=endpoints,
|
||||
)
|
||||
kwargs[GITHUB_COPILOT_USER_SESSION_KWARG_KEY] = session # rebind-ok: kwargs is the request mutation channel
|
||||
_strip_caller_authorization(kwargs)
|
||||
|
|
@ -725,11 +784,13 @@ async def aattach_github_copilot_user_session(
|
|||
if isinstance(kwargs.get(GITHUB_COPILOT_USER_SESSION_KWARG_KEY), GithubCopilotUserSession):
|
||||
return
|
||||
credential_name: Final = credential_name_obj if isinstance(credential_name_obj, str) else ""
|
||||
user_id, github_token = _caller_connection(kwargs, credential_name)
|
||||
endpoints: Final = github_copilot_host_endpoints(credential_name)
|
||||
user_id, github_token = _caller_connection(kwargs, credential_name, endpoints)
|
||||
session: Final = await aexchange_github_token(
|
||||
user_id=user_id,
|
||||
github_token=github_token,
|
||||
credential_name=credential_name,
|
||||
endpoints=endpoints,
|
||||
)
|
||||
kwargs[GITHUB_COPILOT_USER_SESSION_KWARG_KEY] = session # rebind-ok: kwargs is the request mutation channel
|
||||
_strip_caller_authorization(kwargs)
|
||||
|
|
|
|||
|
|
@ -5988,6 +5988,7 @@ def completion(
|
|||
*OPENAI_WIF_KWARGS_KEYS,
|
||||
*OAUTH_TOKEN_EXCHANGE_KWARGS_KEYS,
|
||||
"github_copilot_auth_type",
|
||||
"github_copilot_host",
|
||||
"github_copilot_user_session",
|
||||
PROVIDER_AFFINITY_HEADER_KWARG_KEY,
|
||||
"fireworks_forward_user_id",
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ from typing import Literal, TypeAlias
|
|||
|
||||
from pydantic import ConfigDict, Field, model_validator
|
||||
|
||||
from litellm.constants import GITHUB_COPILOT_DEFAULT_HOST
|
||||
from litellm.types.llms.base import LiteLLMBaseModel
|
||||
|
||||
CredentialSource: TypeAlias = Literal["db", "config"]
|
||||
|
|
@ -101,6 +102,7 @@ class UserConnectionFlowHandle(LiteLLMBaseModel):
|
|||
|
||||
user_id: str
|
||||
credential_name: str
|
||||
github_host: str = GITHUB_COPILOT_DEFAULT_HOST
|
||||
device_code: str = Field(repr=False)
|
||||
interval: int
|
||||
expires_at: float
|
||||
|
|
|
|||
|
|
@ -462,6 +462,7 @@ _BANNED_REQUEST_BODY_PARAMS: Final[tuple[str, ...]] = (
|
|||
# stored credential's values and the caller's connection by the proxy's own
|
||||
# secret_fields, so a body-supplied value could only spoof either.
|
||||
"github_copilot_auth_type",
|
||||
"github_copilot_host",
|
||||
"user_provider_credentials",
|
||||
"github_copilot_user_session",
|
||||
# Observability credentials, hosts, and project identifiers: derived
|
||||
|
|
|
|||
|
|
@ -19,6 +19,7 @@ from litellm._internal_context import with_service_target
|
|||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import (
|
||||
GITHUB_COPILOT_AUTH_TYPE_KEY,
|
||||
GITHUB_COPILOT_HOST_KEY,
|
||||
GITHUB_COPILOT_PER_USER_AUTH_TYPE,
|
||||
)
|
||||
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
|
||||
|
|
@ -84,6 +85,22 @@ _DISPLAY_NAME_MAX_LENGTH: Final = 255
|
|||
_DISPLAY_NAME_FORBIDDEN_CATEGORIES: Final = frozenset({"Cc", "Cf", "Cs", "Zl", "Zp"})
|
||||
|
||||
|
||||
def _normalized_github_copilot_values(
|
||||
credential_values: Mapping[str, object],
|
||||
credential_info: Mapping[str, object],
|
||||
) -> Mapping[str, object]:
|
||||
provider: Final = stored_credential_provider(credential_info.get("custom_llm_provider"))
|
||||
if provider != _GITHUB_COPILOT_PROVIDER or GITHUB_COPILOT_HOST_KEY not in credential_values:
|
||||
return credential_values
|
||||
|
||||
from litellm.llms.github_copilot.github_host import normalized_github_host
|
||||
|
||||
host: Final = normalized_github_host(credential_values[GITHUB_COPILOT_HOST_KEY])
|
||||
if host is None:
|
||||
raise HTTPException(status_code=400, detail=f"Invalid {GITHUB_COPILOT_HOST_KEY}")
|
||||
return {**credential_values, GITHUB_COPILOT_HOST_KEY: host}
|
||||
|
||||
|
||||
def _reject_non_admin_wif_fields(
|
||||
wif_fields: tuple[str, ...],
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
|
|
@ -318,7 +335,14 @@ async def create_credential(
|
|||
status_code=400,
|
||||
detail="Credential values are required. Unable to infer credential values from model ID.",
|
||||
)
|
||||
_reject_non_admin_wif_fields(server_owned_wif_fields_written(credential_values), user_api_key_dict)
|
||||
credential_info: Final = _CREDENTIAL_DICT_ADAPTER.validate_python(
|
||||
credential.credential_info # pyright: ignore[reportUnknownMemberType] # untyped credential field
|
||||
)
|
||||
normalized_values: Final = _normalized_github_copilot_values(
|
||||
_CREDENTIAL_DICT_ADAPTER.validate_python(credential_values),
|
||||
credential_info,
|
||||
)
|
||||
_reject_non_admin_wif_fields(server_owned_wif_fields_written(normalized_values), user_api_key_dict)
|
||||
_reject_non_admin_wif_fields(
|
||||
await named_credential_wif_fields(credential.credential_name, prisma_client), user_api_key_dict
|
||||
)
|
||||
|
|
@ -326,8 +350,8 @@ async def create_credential(
|
|||
processed_credential: Final = CredentialItem(
|
||||
credential_name=credential.credential_name,
|
||||
display_name=_normalized_display_name(credential.display_name),
|
||||
credential_values=_without_null_values(_CREDENTIAL_DICT_ADAPTER.validate_python(credential_values)),
|
||||
credential_info=credential.credential_info,
|
||||
credential_values=_without_null_values(normalized_values),
|
||||
credential_info=credential_info,
|
||||
)
|
||||
encrypted_credential: Final = CredentialHelperUtils.encrypt_credential_values(processed_credential)
|
||||
credentials_dict: Final = encrypted_credential.model_dump(exclude_none=True)
|
||||
|
|
@ -614,22 +638,24 @@ async def start_user_connection(
|
|||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), # noqa: B008 # FastAPI resolves the dependency from the default
|
||||
) -> UserConnectionStartResponse:
|
||||
"""Begin a GitHub device flow for the calling user's connection to a per-user credential."""
|
||||
from litellm.llms.github_copilot.per_user_auth import astart_device_flow
|
||||
from litellm.llms.github_copilot.per_user_auth import astart_device_flow, github_copilot_host_endpoints
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
try:
|
||||
user_id: Final = _authenticated_user_id(user_api_key_dict)
|
||||
_per_user_credential_or_404(credential_name)
|
||||
endpoints: Final = github_copilot_host_endpoints(credential_name)
|
||||
if prisma_client is None:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail={"error": CommonProxyErrors.db_not_connected_error.value},
|
||||
)
|
||||
flow: Final = await astart_device_flow()
|
||||
flow: Final = await astart_device_flow(endpoints)
|
||||
handle: Final = encrypt_value_helper(
|
||||
UserConnectionFlowHandle(
|
||||
user_id=user_id,
|
||||
credential_name=credential_name,
|
||||
github_host=endpoints.host,
|
||||
device_code=flow.device_code,
|
||||
interval=flow.interval,
|
||||
expires_at=time.time() + flow.expires_in,
|
||||
|
|
@ -665,12 +691,14 @@ async def poll_user_connection(
|
|||
acheck_copilot_seat,
|
||||
afetch_github_login,
|
||||
apoll_device_flow,
|
||||
github_copilot_host_endpoints,
|
||||
)
|
||||
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
|
||||
|
||||
try:
|
||||
user_id: Final = _authenticated_user_id(user_api_key_dict)
|
||||
_per_user_credential_or_404(credential_name)
|
||||
endpoints: Final = github_copilot_host_endpoints(credential_name)
|
||||
if prisma_client is None:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
|
|
@ -681,11 +709,12 @@ async def poll_user_connection(
|
|||
handle is None
|
||||
or handle.user_id != user_id
|
||||
or handle.credential_name != credential_name
|
||||
or handle.github_host != endpoints.host
|
||||
or handle.expires_at <= time.time()
|
||||
):
|
||||
raise HTTPException(status_code=400, detail="invalid or expired flow_handle")
|
||||
|
||||
poll: Final = await apoll_device_flow(handle.device_code)
|
||||
poll: Final = await apoll_device_flow(handle.device_code, endpoints)
|
||||
if poll.status in ("expired", "denied", "pending", "slow_down"):
|
||||
return UserConnectionPollResponse(status=poll.status, interval=poll.interval)
|
||||
|
||||
|
|
@ -693,12 +722,19 @@ async def poll_user_connection(
|
|||
if not github_token:
|
||||
raise HTTPException(status_code=502, detail="GitHub device flow completed without a token")
|
||||
try:
|
||||
await acheck_copilot_seat(user_id=user_id, github_token=github_token, credential_name=credential_name)
|
||||
await acheck_copilot_seat(
|
||||
user_id=user_id,
|
||||
github_token=github_token,
|
||||
credential_name=credential_name,
|
||||
endpoints=endpoints,
|
||||
)
|
||||
except litellm.CallerCredentialAuthenticationError:
|
||||
return UserConnectionPollResponse(status="no_copilot_seat")
|
||||
github_login: Final = await afetch_github_login(github_token)
|
||||
github_login: Final = await afetch_github_login(github_token, endpoints)
|
||||
connection_payload: Final = GithubCopilotUserConnectionPayload(
|
||||
access_token=github_token, github_login=github_login
|
||||
access_token=github_token,
|
||||
github_login=github_login,
|
||||
github_host=handle.github_host,
|
||||
)
|
||||
await upsert_user_provider_credential(
|
||||
prisma_client,
|
||||
|
|
@ -760,7 +796,7 @@ async def delete_user_connection(
|
|||
)
|
||||
prior: Final = await delete_user_provider_credential(prisma_client, user_id, credential_name)
|
||||
if prior is not None:
|
||||
evict_copilot_user_session(user_id, prior.access_token)
|
||||
evict_copilot_user_session(user_id, credential_name, prior.github_host)
|
||||
return UserConnectionDeleteResponse(status="disconnected")
|
||||
except Exception as e: # noqa: BLE001 # endpoint boundary: every failure becomes the proxy error contract
|
||||
raise handle_exception_on_proxy(e)
|
||||
|
|
@ -785,16 +821,15 @@ async def _purge_user_connections_for_credential(credential_name: str) -> None:
|
|||
"_purge_user_connections_for_credential: failed to drop user connections for %s", credential_name
|
||||
)
|
||||
return
|
||||
access_tokens: Final[dict[str, str]] = {} # mutable-ok: accumulates one entry per deleted row
|
||||
for row in prior_rows:
|
||||
decoded = decode_user_provider_credential(row.credential_b64)
|
||||
if decoded is not None:
|
||||
access_tokens[row.user_id] = decoded.access_token
|
||||
github_hosts: Final = {
|
||||
row.user_id: decoded.github_host
|
||||
for row in prior_rows
|
||||
if (decoded := decode_user_provider_credential(row.credential_b64)) is not None
|
||||
}
|
||||
for user_id in user_ids:
|
||||
await invalidate_user_provider_credential_cache(user_api_key_cache, user_id, credential_name)
|
||||
token = access_tokens.get(user_id)
|
||||
if token:
|
||||
evict_copilot_user_session(user_id, token)
|
||||
if user_id in github_hosts:
|
||||
evict_copilot_user_session(user_id, credential_name, github_hosts[user_id])
|
||||
|
||||
|
||||
@router.delete(
|
||||
|
|
@ -926,7 +961,25 @@ async def update_credential(
|
|||
credential_values=incoming_values,
|
||||
credential_values_to_delete=credential.credential_values_to_delete,
|
||||
)
|
||||
merged_credential: Final = update_db_credential(db_credential, patch)
|
||||
decrypted_stored_credential: Final = decrypted_credential(db_credential)
|
||||
decrypted_values: Final = _CREDENTIAL_DICT_ADAPTER.validate_python(
|
||||
decrypted_stored_credential.credential_values
|
||||
)
|
||||
stored_info: Final = _CREDENTIAL_DICT_ADAPTER.validate_python(
|
||||
db_credential.credential_info # pyright: ignore[reportUnknownMemberType] # untyped credential field
|
||||
)
|
||||
patch_info: Final = _CREDENTIAL_DICT_ADAPTER.validate_python(
|
||||
patch.credential_info # pyright: ignore[reportUnknownMemberType] # untyped credential field
|
||||
)
|
||||
effective_info: Final = {**stored_info, **patch_info}
|
||||
normalized_patch: Final = (
|
||||
patch.model_copy(
|
||||
update={"credential_values": _normalized_github_copilot_values(incoming_values, effective_info)}
|
||||
)
|
||||
if incoming_values.get(GITHUB_COPILOT_HOST_KEY) is not None
|
||||
else patch
|
||||
)
|
||||
merged_credential: Final = update_db_credential(db_credential, normalized_patch)
|
||||
credential_object_jsonified: Final = cast( # cast-ok: deep-copies a model_dump, so keys are str
|
||||
"dict[str, object]", jsonify_object(merged_credential.model_dump(exclude_none=True))
|
||||
)
|
||||
|
|
@ -940,13 +993,20 @@ async def update_credential(
|
|||
)
|
||||
|
||||
# Sync in-memory credential_list (skip if not in memory - e.g., proxy restarted)
|
||||
_sync_in_memory_credential(patch, credential_name)
|
||||
_sync_in_memory_credential(normalized_patch, credential_name)
|
||||
|
||||
merged_values: Final = cast( # cast-ok: credential_values is a plain dict at runtime
|
||||
Mapping[object, object], merged_credential.credential_values
|
||||
from litellm.llms.github_copilot.github_host import normalized_github_host
|
||||
|
||||
new_decrypted_values: Final = _CREDENTIAL_DICT_ADAPTER.validate_python(
|
||||
decrypted_credential(merged_credential).credential_values
|
||||
)
|
||||
still_per_user: Final = merged_values.get(GITHUB_COPILOT_AUTH_TYPE_KEY) == GITHUB_COPILOT_PER_USER_AUTH_TYPE
|
||||
if not still_per_user:
|
||||
still_per_user: Final = (
|
||||
new_decrypted_values.get(GITHUB_COPILOT_AUTH_TYPE_KEY) == GITHUB_COPILOT_PER_USER_AUTH_TYPE
|
||||
)
|
||||
old_host: Final = normalized_github_host(decrypted_values.get(GITHUB_COPILOT_HOST_KEY))
|
||||
new_host: Final = normalized_github_host(new_decrypted_values.get(GITHUB_COPILOT_HOST_KEY))
|
||||
host_changed: Final = old_host is None or old_host != new_host
|
||||
if not still_per_user or host_changed:
|
||||
await _purge_user_connections_for_credential(credential_name)
|
||||
|
||||
return {"success": True, "message": "Credential updated successfully"}
|
||||
|
|
|
|||
|
|
@ -25,6 +25,7 @@ from litellm.caching.dual_cache import DualCache
|
|||
if TYPE_CHECKING:
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
from litellm.constants import (
|
||||
GITHUB_COPILOT_DEFAULT_HOST,
|
||||
GITHUB_COPILOT_USER_CREDENTIAL_CACHE_TTL_SECONDS,
|
||||
USER_PROVIDER_CREDENTIAL_CACHE_PREFIX,
|
||||
USER_PROVIDER_CREDENTIAL_NOT_CONNECTED,
|
||||
|
|
@ -50,6 +51,7 @@ class GithubCopilotUserConnectionPayload(BaseModel):
|
|||
|
||||
access_token: str = Field(repr=False)
|
||||
github_login: str
|
||||
github_host: str = GITHUB_COPILOT_DEFAULT_HOST
|
||||
|
||||
|
||||
class _UserProviderCredentialsRepository(PrismaTableRepository["prisma_models.LiteLLM_UserProviderCredentials"]):
|
||||
|
|
@ -186,7 +188,7 @@ async def _try_cache_get(token_cache: "RedisCache", key: str) -> object:
|
|||
try:
|
||||
return await _string_cache(token_cache).async_get_cache(key)
|
||||
except Exception: # noqa: BLE001 # a Redis outage must fall back to the database, never reject a connected user
|
||||
verbose_proxy_logger.warning("aget_user_provider_tokens: Redis get failed; falling back to the database")
|
||||
verbose_proxy_logger.warning("aget_user_provider_credentials: Redis get failed; falling back to the database")
|
||||
return None
|
||||
|
||||
|
||||
|
|
@ -196,7 +198,7 @@ async def _try_cache_set(token_cache: "RedisCache", key: str, value: str, nx: bo
|
|||
key, value, nx=nx, ttl=GITHUB_COPILOT_USER_CREDENTIAL_CACHE_TTL_SECONDS
|
||||
)
|
||||
except Exception: # noqa: BLE001 # caching is best-effort; a Redis outage is not worth failing the request
|
||||
verbose_proxy_logger.warning("aget_user_provider_tokens: Redis set failed; skipping the cache write")
|
||||
verbose_proxy_logger.warning("aget_user_provider_credentials: Redis set failed; skipping the cache write")
|
||||
|
||||
|
||||
@with_service_target("user_provider_connections")
|
||||
|
|
@ -292,16 +294,16 @@ async def set_user_provider_credential_cache(
|
|||
|
||||
|
||||
@with_service_target("user_provider_connections")
|
||||
async def aget_user_provider_tokens(
|
||||
async def aget_user_provider_credentials(
|
||||
prisma_client: "PrismaClient",
|
||||
cache: DualCache,
|
||||
user_id: str,
|
||||
credential_names: Sequence[str],
|
||||
) -> Mapping[str, str]:
|
||||
"""Map each per-user credential name to the caller's stored GitHub token.
|
||||
) -> Mapping[str, Mapping[str, str]]:
|
||||
"""Map each per-user credential name to the caller's stored GitHub connection.
|
||||
|
||||
Cached values are the stored ciphertext or the ``_NOT_CONNECTED`` marker;
|
||||
plaintext tokens only ever live in the returned dict. Redis is the only
|
||||
plaintext connection data only ever lives in the returned mapping. Redis is the only
|
||||
cache layer used: without it there is no shared invalidation, so a worker
|
||||
serving a stale local entry after a disconnect is worse than a DB read."""
|
||||
token_cache: Final = cache.redis_cache
|
||||
|
|
@ -320,9 +322,9 @@ async def aget_user_provider_tokens(
|
|||
misses.append(name)
|
||||
if not misses:
|
||||
return {
|
||||
name: token
|
||||
name: connection
|
||||
for name, ciphertext in cached.items()
|
||||
if (token := _token_from_ciphertext(ciphertext)) is not None
|
||||
if (connection := _connection_from_ciphertext(ciphertext)) is not None
|
||||
}
|
||||
|
||||
try:
|
||||
|
|
@ -333,7 +335,7 @@ async def aget_user_provider_tokens(
|
|||
where={"user_id": user_id},
|
||||
)
|
||||
except Exception: # noqa: BLE001 # a DB outage reads as not connected so the caller gets a 401, never a 500
|
||||
verbose_proxy_logger.exception("aget_user_provider_tokens: DB read failed for user_id=%s", user_id)
|
||||
verbose_proxy_logger.exception("aget_user_provider_credentials: DB read failed")
|
||||
return {}
|
||||
found: Final = {row.credential_name: row.credential_b64 for row in rows}
|
||||
for name in misses:
|
||||
|
|
@ -342,10 +344,14 @@ async def aget_user_provider_tokens(
|
|||
if name in found:
|
||||
cached[name] = found[name]
|
||||
return {
|
||||
name: token for name, ciphertext in cached.items() if (token := _token_from_ciphertext(ciphertext)) is not None
|
||||
name: connection
|
||||
for name, ciphertext in cached.items()
|
||||
if (connection := _connection_from_ciphertext(ciphertext)) is not None
|
||||
}
|
||||
|
||||
|
||||
def _token_from_ciphertext(ciphertext: str) -> str | None:
|
||||
def _connection_from_ciphertext(ciphertext: str) -> Mapping[str, str] | None:
|
||||
payload: Final = decode_user_provider_credential(ciphertext)
|
||||
return payload.access_token if payload is not None else None
|
||||
if payload is None:
|
||||
return None
|
||||
return {"access_token": payload.access_token, "github_host": payload.github_host}
|
||||
|
|
|
|||
|
|
@ -3321,12 +3321,12 @@ async def _resolve_user_provider_credentials_for_request(
|
|||
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
|
||||
from litellm.types.proxy.litellm_pre_call_utils import RedactedDict
|
||||
|
||||
from .credential_endpoints.user_provider_credentials import aget_user_provider_tokens
|
||||
from .credential_endpoints.user_provider_credentials import aget_user_provider_credentials
|
||||
|
||||
if prisma_client is None:
|
||||
return
|
||||
try:
|
||||
tokens: Final = await aget_user_provider_tokens(
|
||||
connections: Final = await aget_user_provider_credentials(
|
||||
prisma_client=prisma_client,
|
||||
cache=user_api_key_cache,
|
||||
user_id=user_id,
|
||||
|
|
@ -3341,7 +3341,7 @@ async def _resolve_user_provider_credentials_for_request(
|
|||
secret_fields: Final = data.get("secret_fields")
|
||||
if not isinstance(secret_fields, dict):
|
||||
return
|
||||
secret_fields["user_provider_credentials"] = RedactedDict(dict(tokens))
|
||||
secret_fields["user_provider_credentials"] = RedactedDict(dict(connections))
|
||||
secret_fields["user_provider_credentials_user_id"] = user_id
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1731,6 +1731,16 @@
|
|||
"field_type": "password",
|
||||
"options": null,
|
||||
"default_value": null
|
||||
},
|
||||
{
|
||||
"key": "github_copilot_host",
|
||||
"label": "GitHub host",
|
||||
"placeholder": "github.com",
|
||||
"tooltip": "Leave empty for github.com. For GHE.com data residency enter your SUBDOMAIN.ghe.com host",
|
||||
"required": false,
|
||||
"field_type": "text",
|
||||
"options": null,
|
||||
"default_value": null
|
||||
}
|
||||
],
|
||||
"default_model_placeholder": "gpt-3.5-turbo"
|
||||
|
|
|
|||
|
|
@ -91,6 +91,7 @@ class ProviderConnection:
|
|||
configurable_clientside_auth_params: "Sequence[str | ConfigurableClientsideParamsCustomAuth] | None" = None
|
||||
use_xai_oauth: bool | None = None
|
||||
github_copilot_auth_type: str | None = None
|
||||
github_copilot_host: str | None = None
|
||||
github_copilot_user_session: object | None = None
|
||||
token_exchange_endpoint: str | None = None
|
||||
token_exchange_profile: str | None = None
|
||||
|
|
|
|||
|
|
@ -24,5 +24,5 @@ class SecretFields(TypedDict):
|
|||
"""
|
||||
|
||||
raw_headers: dict
|
||||
user_provider_credentials: NotRequired[ReadOnly[Mapping[str, str]]]
|
||||
user_provider_credentials: NotRequired[ReadOnly[Mapping[str, Mapping[str, str]]]]
|
||||
user_provider_credentials_user_id: NotRequired[ReadOnly[str]]
|
||||
|
|
|
|||
|
|
@ -4211,7 +4211,7 @@ oauth_token_exchange_litellm_params: Final = (
|
|||
"token_exchange_profile",
|
||||
"token_exchange_scope",
|
||||
)
|
||||
github_copilot_oauth_litellm_params: Final = ("github_copilot_auth_type",)
|
||||
github_copilot_oauth_litellm_params: Final = ("github_copilot_auth_type", "github_copilot_host")
|
||||
OIDC_SECRET_REFERENCE_PREFIX: Final = "oidc/"
|
||||
API_KEY_SECRET_REFERENCE_REJECTION: Final = (
|
||||
"Rejected Request: api_key holds an oidc/ secret reference, which only a deployment or credential a proxy "
|
||||
|
|
|
|||
|
|
@ -317,6 +317,10 @@ class TestGetLitellmParamsExplicitFields:
|
|||
result = get_litellm_params(no_log=True)
|
||||
assert result["no-log"] is True
|
||||
|
||||
def test_github_copilot_host_is_registered_as_a_litellm_param(self):
|
||||
result: Final = get_litellm_params(github_copilot_host="octo.ghe.com")
|
||||
assert result["github_copilot_host"] == "octo.ghe.com"
|
||||
|
||||
|
||||
class TestGetLitellmParamsDataResidency:
|
||||
"""Verify that data_residency is inferred from OpenAI regional api_base."""
|
||||
|
|
|
|||
59
tests/unit/llms/github_copilot/test_github_host.py
Normal file
59
tests/unit/llms/github_copilot/test_github_host.py
Normal file
|
|
@ -0,0 +1,59 @@
|
|||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.llms.github_copilot.github_host import (
|
||||
github_host_endpoints,
|
||||
normalized_github_host,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("raw", "expected"),
|
||||
[
|
||||
(None, "github.com"),
|
||||
("", "github.com"),
|
||||
("GitHub.com", "github.com"),
|
||||
("OCTO.ghe.com", "octo.ghe.com"),
|
||||
("octo-1.ghe.com", "octo-1.ghe.com"),
|
||||
("https://octo.ghe.com", None),
|
||||
("ghes.example.com", None),
|
||||
("octo.ghe.com:443", None),
|
||||
("a.b.ghe.com", None),
|
||||
("octo.ghe.com.", None),
|
||||
(" api.octo.ghe.com", None),
|
||||
("octo.ghe.com/path", None),
|
||||
],
|
||||
)
|
||||
def test_normalized_github_host_accepts_only_bare_supported_hosts(raw: object, expected: str | None) -> None:
|
||||
assert normalized_github_host(raw) == expected
|
||||
|
||||
|
||||
def test_github_com_endpoints_read_url_overrides_at_call_time(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("GITHUB_COPILOT_DEVICE_CODE_URL", "https://device.example/code")
|
||||
monkeypatch.setenv("GITHUB_COPILOT_ACCESS_TOKEN_URL", "https://device.example/token")
|
||||
monkeypatch.setenv("GITHUB_COPILOT_API_KEY_URL", "https://api.example/copilot")
|
||||
monkeypatch.setenv("GITHUB_COPILOT_CLIENT_ID", "custom-client")
|
||||
|
||||
endpoints: Final = github_host_endpoints("github.com")
|
||||
|
||||
assert endpoints.device_code_url == "https://device.example/code"
|
||||
assert endpoints.access_token_url == "https://device.example/token"
|
||||
assert endpoints.api_key_url == "https://api.example/copilot"
|
||||
assert endpoints.client_id == "custom-client"
|
||||
|
||||
|
||||
def test_ghe_endpoints_ignore_github_com_url_overrides(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("GITHUB_COPILOT_DEVICE_CODE_URL", "https://device.example/code")
|
||||
monkeypatch.setenv("GITHUB_COPILOT_ACCESS_TOKEN_URL", "https://device.example/token")
|
||||
monkeypatch.setenv("GITHUB_COPILOT_API_KEY_URL", "https://api.example/copilot")
|
||||
monkeypatch.setenv("GITHUB_COPILOT_CLIENT_ID", "custom-client")
|
||||
|
||||
endpoints: Final = github_host_endpoints("octo.ghe.com")
|
||||
|
||||
assert endpoints.device_code_url == "https://octo.ghe.com/login/device/code"
|
||||
assert endpoints.access_token_url == "https://octo.ghe.com/login/oauth/access_token"
|
||||
assert endpoints.api_key_url == "https://api.octo.ghe.com/copilot_internal/v2/token"
|
||||
assert endpoints.user_url == "https://api.octo.ghe.com/user"
|
||||
assert endpoints.default_copilot_api_base == "https://copilot-api.octo.ghe.com"
|
||||
assert endpoints.client_id == "custom-client"
|
||||
|
|
@ -12,6 +12,7 @@ import httpx
|
|||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.constants import GITHUB_COPILOT_DEFAULT_HOST
|
||||
from litellm.exceptions import (
|
||||
AuthenticationError,
|
||||
BadRequestError,
|
||||
|
|
@ -32,20 +33,33 @@ from litellm.llms.github_copilot.per_user_auth import (
|
|||
attach_github_copilot_user_session,
|
||||
evict_copilot_user_session,
|
||||
exchange_github_token,
|
||||
github_copilot_host_endpoints,
|
||||
github_copilot_auth_mode,
|
||||
github_copilot_user_session_from,
|
||||
validated_copilot_api_base,
|
||||
)
|
||||
from litellm.llms.github_copilot.github_host import github_host_endpoints
|
||||
from litellm.types.utils import CredentialItem, ModelResponse
|
||||
|
||||
EXCHANGE_URL = "https://api.github.com/copilot_internal/v2/token"
|
||||
DEVICE_CODE_URL = "https://github.com/login/device/code"
|
||||
ACCESS_TOKEN_URL = "https://github.com/login/oauth/access_token"
|
||||
_GITHUB_ENDPOINTS: Final = github_host_endpoints("github.com")
|
||||
|
||||
|
||||
def _credential(auth_type="per_user_oauth"):
|
||||
values = {} if auth_type is None else {"github_copilot_auth_type": auth_type}
|
||||
return CredentialItem(credential_name="copilot-cred", credential_values=values, credential_info={})
|
||||
def _credential(
|
||||
auth_type: str | None = "per_user_oauth",
|
||||
github_host: str | None = None,
|
||||
) -> CredentialItem:
|
||||
values: Final = {} if auth_type is None else {"github_copilot_auth_type": auth_type}
|
||||
return CredentialItem(
|
||||
credential_name="copilot-cred",
|
||||
credential_values={
|
||||
**values,
|
||||
**({"github_copilot_host": github_host} if github_host is not None else {}),
|
||||
},
|
||||
credential_info={},
|
||||
)
|
||||
|
||||
|
||||
def _exchange_payload(token="copilot-token", expires_at=None, api_base="https://api.githubcopilot.com"):
|
||||
|
|
@ -143,7 +157,61 @@ def test_unknown_named_credential_falls_back_to_kwargs():
|
|||
],
|
||||
)
|
||||
def test_validated_copilot_api_base(raw, expected):
|
||||
assert validated_copilot_api_base(raw) == expected
|
||||
assert validated_copilot_api_base(raw, github_host_endpoints("github.com")) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"raw",
|
||||
[
|
||||
"https://copilot-api.evil.ghe.com",
|
||||
"https://api.githubcopilot.com",
|
||||
"http://copilot-api.octo.ghe.com",
|
||||
"https://user@copilot-api.octo.ghe.com",
|
||||
"https://copilot-api.octo.ghe.com:8443",
|
||||
],
|
||||
)
|
||||
def test_ghe_exchange_response_must_stay_on_its_tenant_host(raw: object) -> None:
|
||||
endpoints: Final = github_host_endpoints("octo.ghe.com")
|
||||
assert validated_copilot_api_base(raw, endpoints) == "https://copilot-api.octo.ghe.com"
|
||||
|
||||
|
||||
def test_github_com_exchange_rejects_a_ghe_copilot_host() -> None:
|
||||
endpoints: Final = github_host_endpoints("github.com")
|
||||
assert validated_copilot_api_base("https://copilot-api.octo.ghe.com", endpoints) == "https://api.githubcopilot.com"
|
||||
|
||||
|
||||
def test_ghe_exchange_uses_stored_host_not_request_kwargs():
|
||||
credential: Final = CredentialItem(
|
||||
credential_name="copilot-cred",
|
||||
credential_values={"github_copilot_auth_type": "per_user_oauth", "github_copilot_host": "octo.ghe.com"},
|
||||
credential_info={},
|
||||
)
|
||||
|
||||
def respond(request: httpx.Request) -> httpx.Response:
|
||||
assert str(request.url) == "https://api.octo.ghe.com/copilot_internal/v2/token"
|
||||
return httpx.Response(200, json=_exchange_payload(api_base="https://copilot-api.octo.ghe.com"), request=request)
|
||||
|
||||
with patch.object(litellm, "credential_list", [credential]), github_http(respond) as requests:
|
||||
kwargs: Final = _kwargs_with_connection(github_host="octo.ghe.com")
|
||||
kwargs["github_copilot_host"] = "evil.ghe.com"
|
||||
attach_github_copilot_user_session(kwargs)
|
||||
|
||||
session: Final = github_copilot_user_session_from(kwargs)
|
||||
assert session is not None and session.api_base == "https://copilot-api.octo.ghe.com"
|
||||
assert len(requests) == 1
|
||||
|
||||
|
||||
def test_invalid_stored_host_fails_before_exchange():
|
||||
credential: Final = CredentialItem(
|
||||
credential_name="copilot-cred",
|
||||
credential_values={"github_copilot_auth_type": "per_user_oauth", "github_copilot_host": "https://octo.ghe.com"},
|
||||
credential_info={},
|
||||
)
|
||||
with patch.object(litellm, "credential_list", [credential]):
|
||||
with github_http(lambda request: pytest.fail(f"unexpected request to {request.url}")) as requests:
|
||||
with pytest.raises(BadRequestError):
|
||||
attach_github_copilot_user_session(_kwargs_with_connection())
|
||||
assert not requests
|
||||
|
||||
|
||||
def test_session_repr_and_str_hide_the_token(caplog):
|
||||
|
|
@ -159,49 +227,59 @@ def test_session_repr_and_str_hide_the_token(caplog):
|
|||
|
||||
def test_exchange_returns_session_and_caches_per_user():
|
||||
with github_http(_exchange_responder()) as requests:
|
||||
s1 = exchange_github_token("user-a", "gh-token", "copilot-cred")
|
||||
s1 = exchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS)
|
||||
assert len(requests) == 1
|
||||
assert requests[0].headers["Authorization"] == "token gh-token"
|
||||
s2 = exchange_github_token("user-a", "gh-token", "copilot-cred")
|
||||
s2 = exchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS)
|
||||
assert len(requests) == 1
|
||||
assert s2.token == s1.token == "copilot-token"
|
||||
# different user -> separate exchange
|
||||
exchange_github_token("user-b", "gh-token", "copilot-cred")
|
||||
exchange_github_token("user-b", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS)
|
||||
assert len(requests) == 2
|
||||
# different github token -> separate exchange
|
||||
exchange_github_token("user-a", "gh-token-2", "copilot-cred")
|
||||
exchange_github_token("user-a", "gh-token-2", "copilot-cred", _GITHUB_ENDPOINTS)
|
||||
assert len(requests) == 3
|
||||
|
||||
|
||||
def test_cached_session_is_not_reused_after_github_token_changes():
|
||||
with github_http(_exchange_responder()) as requests:
|
||||
exchange_github_token("user-a", "gh-token-a", "copilot-cred", _GITHUB_ENDPOINTS)
|
||||
exchange_github_token("user-a", "gh-token-b", "copilot-cred", _GITHUB_ENDPOINTS)
|
||||
|
||||
assert [request.headers["Authorization"] for request in requests] == ["token gh-token-a", "token gh-token-b"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("status", [401, 403, 404])
|
||||
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")
|
||||
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 _SESSION_CACHE.get_cache(_session_cache_key("user-a", "gh-token")) is None
|
||||
assert (
|
||||
_SESSION_CACHE.get_cache(_session_cache_key("user-a", "copilot-cred", GITHUB_COPILOT_DEFAULT_HOST)) is None
|
||||
)
|
||||
|
||||
|
||||
def test_exchange_429_raises_caller_rate_limit():
|
||||
with github_http(_exchange_responder(status=429)):
|
||||
with pytest.raises(CallerCredentialRateLimitError):
|
||||
exchange_github_token("user-a", "gh-token", "copilot-cred")
|
||||
exchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS)
|
||||
|
||||
|
||||
def test_exchange_other_failure_raises_service_unavailable_without_body():
|
||||
with github_http(_exchange_responder(status=500, payload={"secret": "body"})):
|
||||
with pytest.raises(ServiceUnavailableError) as exc:
|
||||
exchange_github_token("user-a", "gh-token", "copilot-cred")
|
||||
exchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS)
|
||||
assert "body" not in str(exc.value)
|
||||
|
||||
|
||||
def test_exchange_expiry_margin_refetches_near_expiry():
|
||||
with github_http(_exchange_responder(_exchange_payload(expires_at=int(time.time()) + 30))) as requests:
|
||||
exchange_github_token("user-a", "gh-token", "copilot-cred")
|
||||
exchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS)
|
||||
# ttl under the 60s safety margin -> not cached, second call refetches
|
||||
exchange_github_token("user-a", "gh-token", "copilot-cred")
|
||||
exchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS)
|
||||
assert len(requests) == 2
|
||||
|
||||
|
||||
|
|
@ -212,8 +290,8 @@ async def test_aexchange_single_flight_one_http_call():
|
|||
|
||||
with github_http(respond) as requests:
|
||||
s1, s2 = await asyncio.gather(
|
||||
aexchange_github_token("user-a", "gh-token", "copilot-cred"),
|
||||
aexchange_github_token("user-a", "gh-token", "copilot-cred"),
|
||||
aexchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS),
|
||||
aexchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS),
|
||||
)
|
||||
assert len(requests) == 1
|
||||
assert s1.token == s2.token == "copilot-token"
|
||||
|
|
@ -223,16 +301,26 @@ async def test_aexchange_single_flight_one_http_call():
|
|||
async def test_aexchange_401_raises_and_evicts():
|
||||
with github_http(_exchange_responder(status=401)):
|
||||
with pytest.raises(CallerCredentialAuthenticationError):
|
||||
await aexchange_github_token("user-a", "gh-token", "copilot-cred")
|
||||
assert _SESSION_CACHE.get_cache(_session_cache_key("user-a", "gh-token")) is None
|
||||
await aexchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS)
|
||||
assert (
|
||||
_SESSION_CACHE.get_cache(_session_cache_key("user-a", "copilot-cred", GITHUB_COPILOT_DEFAULT_HOST)) is None
|
||||
)
|
||||
|
||||
|
||||
def _kwargs_with_connection(token="gho_user_token"):
|
||||
def _kwargs_with_connection(
|
||||
token: str = "gho_user_token",
|
||||
github_host: str | None = "github.com",
|
||||
) -> dict[str, object]:
|
||||
return {
|
||||
"litellm_credential_name": "copilot-cred",
|
||||
"secret_fields": {
|
||||
"user_provider_credentials_user_id": "user-a",
|
||||
"user_provider_credentials": {"copilot-cred": token},
|
||||
"user_provider_credentials": {
|
||||
"copilot-cred": {
|
||||
"access_token": token,
|
||||
**({"github_host": github_host} if github_host is not None else {}),
|
||||
}
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
|
@ -261,6 +349,31 @@ def test_attach_per_user_without_connection_raises_connect_401(per_user_credenti
|
|||
assert "Connect GitHub Copilot for credential 'copilot-cred' in the LiteLLM UI (LLM Credentials)" in str(exc.value)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("credential_host", "connection_host"),
|
||||
[("github.com", "octo.ghe.com"), ("octo.ghe.com", "github.com")],
|
||||
)
|
||||
def test_attach_rejects_connection_from_a_different_github_host(
|
||||
credential_host: str,
|
||||
connection_host: str,
|
||||
):
|
||||
with patch.object(litellm, "credential_list", [_credential(github_host=credential_host)]):
|
||||
kwargs: Final = _kwargs_with_connection(github_host=connection_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 not requests
|
||||
|
||||
|
||||
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)
|
||||
attach_github_copilot_user_session(kwargs)
|
||||
assert len(requests) == 1
|
||||
assert str(requests[0].url) == EXCHANGE_URL
|
||||
|
||||
|
||||
def test_attach_shared_mode_does_nothing_without_connection():
|
||||
with patch.object(litellm, "credential_list", [_credential(auth_type=None)]):
|
||||
kwargs = {"litellm_credential_name": "copilot-cred"}
|
||||
|
|
@ -270,9 +383,9 @@ def test_attach_shared_mode_does_nothing_without_connection():
|
|||
|
||||
def test_evict_drops_cached_session():
|
||||
with github_http(_exchange_responder()) as requests:
|
||||
exchange_github_token("user-a", "gh-token", "copilot-cred")
|
||||
evict_copilot_user_session("user-a", "gh-token")
|
||||
exchange_github_token("user-a", "gh-token", "copilot-cred")
|
||||
exchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS)
|
||||
evict_copilot_user_session("user-a", "copilot-cred", _GITHUB_ENDPOINTS.host)
|
||||
exchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS)
|
||||
assert len(requests) == 2
|
||||
|
||||
|
||||
|
|
@ -293,7 +406,7 @@ async def test_astart_device_flow_shape():
|
|||
)
|
||||
|
||||
with github_http(respond):
|
||||
start = await astart_device_flow()
|
||||
start = await astart_device_flow(_GITHUB_ENDPOINTS)
|
||||
assert (start.device_code, start.user_code, start.verification_uri, start.expires_in, start.interval) == (
|
||||
"dc",
|
||||
"UC-123",
|
||||
|
|
@ -319,7 +432,7 @@ async def test_apoll_device_flow_error_mapping(github_error, expected_status):
|
|||
return httpx.Response(200, json={"error": github_error}, request=request)
|
||||
|
||||
with github_http(respond):
|
||||
poll = await apoll_device_flow("dc")
|
||||
poll = await apoll_device_flow("dc", _GITHUB_ENDPOINTS)
|
||||
assert poll.status == expected_status
|
||||
|
||||
|
||||
|
|
@ -330,7 +443,7 @@ async def test_apoll_device_flow_connected():
|
|||
return httpx.Response(200, json={"access_token": "gho_x"}, request=request)
|
||||
|
||||
with github_http(respond):
|
||||
poll = await apoll_device_flow("dc")
|
||||
poll = await apoll_device_flow("dc", _GITHUB_ENDPOINTS)
|
||||
assert poll.status == "connected" and poll.access_token == "gho_x"
|
||||
|
||||
|
||||
|
|
@ -363,7 +476,12 @@ def _chat_completion_response():
|
|||
def _secret_fields(user_id: str, github_token: str):
|
||||
return {
|
||||
"user_provider_credentials_user_id": user_id,
|
||||
"user_provider_credentials": {"copilot-cred": github_token},
|
||||
"user_provider_credentials": {
|
||||
"copilot-cred": {
|
||||
"access_token": github_token,
|
||||
"github_host": GITHUB_COPILOT_DEFAULT_HOST,
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
|
|
@ -842,7 +960,9 @@ def test_sync_exchange_single_flight_one_http_call():
|
|||
with github_http(slow_respond) as requests:
|
||||
threads = [
|
||||
threading.Thread(
|
||||
target=lambda: sessions.append(exchange_github_token("user-a", "gh-token", "copilot-cred"))
|
||||
target=lambda: sessions.append(
|
||||
exchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS)
|
||||
)
|
||||
)
|
||||
for _ in range(6)
|
||||
]
|
||||
|
|
@ -856,12 +976,34 @@ def test_sync_exchange_single_flight_one_http_call():
|
|||
|
||||
def test_two_users_same_github_token_do_not_share_session():
|
||||
with github_http(_exchange_responder()):
|
||||
s1 = exchange_github_token("user-a", "gh-same", "copilot-cred")
|
||||
s2 = exchange_github_token("user-b", "gh-same", "copilot-cred")
|
||||
assert _session_cache_key("user-a", "gh-same") != _session_cache_key("user-b", "gh-same")
|
||||
s1: Final = exchange_github_token("user-a", "gh-same", "copilot-cred", _GITHUB_ENDPOINTS)
|
||||
s2: Final = exchange_github_token("user-b", "gh-same", "copilot-cred", _GITHUB_ENDPOINTS)
|
||||
user_a_key: Final = _session_cache_key("user-a", "copilot-cred", _GITHUB_ENDPOINTS.host)
|
||||
user_b_key: Final = _session_cache_key("user-b", "copilot-cred", _GITHUB_ENDPOINTS.host)
|
||||
other_credential_key: Final = _session_cache_key("user-a", "another-credential", _GITHUB_ENDPOINTS.host)
|
||||
other_host_key: Final = _session_cache_key("user-a", "copilot-cred", "octo.ghe.com")
|
||||
assert user_a_key != user_b_key
|
||||
assert user_a_key != other_credential_key
|
||||
assert user_a_key != other_host_key
|
||||
assert s1.token == s2.token == "copilot-token"
|
||||
|
||||
|
||||
def test_same_user_and_token_do_not_share_sessions_across_hosts():
|
||||
from litellm.llms.github_copilot.github_host import github_host_endpoints
|
||||
|
||||
ghe_endpoints: Final = github_host_endpoints("octo.ghe.com")
|
||||
|
||||
def respond(request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(200, json=_exchange_payload(), request=request)
|
||||
|
||||
with github_http(respond) as requests:
|
||||
github_session: Final = exchange_github_token("user-a", "gh-same", "copilot-cred", _GITHUB_ENDPOINTS)
|
||||
ghe_session: Final = exchange_github_token("user-a", "gh-same", "copilot-cred", ghe_endpoints)
|
||||
assert len(requests) == 2
|
||||
assert github_session.api_base == "https://api.githubcopilot.com"
|
||||
assert ghe_session.api_base == "https://copilot-api.octo.ghe.com"
|
||||
|
||||
|
||||
def test_router_constructs_with_per_user_credential_deployment(per_user_credential):
|
||||
from litellm.llms.github_copilot.authenticator import Authenticator
|
||||
|
||||
|
|
@ -928,9 +1070,9 @@ async def test_aexchange_cancelled_waiter_does_not_cancel_the_shared_exchange(pe
|
|||
with patch( # test-quality-ok: the gate holds the HTTP edge open so two waiters overlap
|
||||
"litellm.llms.custom_httpx.http_handler.get_async_httpx_client", return_value=_GatedClient()
|
||||
):
|
||||
first = asyncio.ensure_future(aexchange_github_token("user-a", "gh-shared", "copilot-cred"))
|
||||
first = asyncio.ensure_future(aexchange_github_token("user-a", "gh-shared", "copilot-cred", _GITHUB_ENDPOINTS))
|
||||
await asyncio.sleep(0)
|
||||
second = asyncio.ensure_future(aexchange_github_token("user-a", "gh-shared", "copilot-cred"))
|
||||
second = asyncio.ensure_future(aexchange_github_token("user-a", "gh-shared", "copilot-cred", _GITHUB_ENDPOINTS))
|
||||
await asyncio.sleep(0)
|
||||
second.cancel()
|
||||
gate.set()
|
||||
|
|
@ -940,6 +1082,33 @@ async def test_aexchange_cancelled_waiter_does_not_cancel_the_shared_exchange(pe
|
|||
await second
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aexchange_does_not_share_an_in_flight_session_after_token_changes(per_user_credential):
|
||||
gate = asyncio.Event()
|
||||
auth_headers: asyncio.Queue[str] = asyncio.Queue()
|
||||
|
||||
class _GatedClient:
|
||||
async def get(self, url, headers=None):
|
||||
await auth_headers.put(headers["authorization"])
|
||||
await gate.wait()
|
||||
token = "copilot-token-a" if headers["authorization"] == "token gh-a" else "copilot-token-b"
|
||||
return httpx.Response(
|
||||
200,
|
||||
json=_exchange_payload(token=token),
|
||||
request=httpx.Request("GET", url),
|
||||
)
|
||||
|
||||
with patch("litellm.llms.custom_httpx.http_handler.get_async_httpx_client", return_value=_GatedClient()):
|
||||
first = asyncio.ensure_future(aexchange_github_token("user-a", "gh-a", "copilot-cred", _GITHUB_ENDPOINTS))
|
||||
assert await asyncio.wait_for(auth_headers.get(), timeout=1) == "token gh-a"
|
||||
second = asyncio.ensure_future(aexchange_github_token("user-a", "gh-b", "copilot-cred", _GITHUB_ENDPOINTS))
|
||||
assert await asyncio.wait_for(auth_headers.get(), timeout=1) == "token gh-b"
|
||||
gate.set()
|
||||
sessions = await asyncio.gather(first, second)
|
||||
|
||||
assert [session.token for session in sessions] == ["copilot-token-a", "copilot-token-b"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aexchange_failure_with_a_single_caller_emits_no_unretrieved_warning(per_user_credential):
|
||||
"""The leader removes the failed future before anyone awaits it; the stored
|
||||
|
|
@ -955,7 +1124,7 @@ async def test_aexchange_failure_with_a_single_caller_emits_no_unretrieved_warni
|
|||
github_http(_exchange_responder(status=401)),
|
||||
pytest.raises(CallerCredentialAuthenticationError),
|
||||
):
|
||||
await aexchange_github_token("user-a", "gh-fail", "copilot-cred")
|
||||
await aexchange_github_token("user-a", "gh-fail", "copilot-cred", _GITHUB_ENDPOINTS)
|
||||
_IN_FLIGHT_KEYS = list(pua._IN_FLIGHT)
|
||||
import gc
|
||||
|
||||
|
|
|
|||
|
|
@ -5088,7 +5088,7 @@ def test_get_model_from_request_vertex_ai_passthrough(request_data, route, expec
|
|||
|
||||
@pytest.mark.parametrize(
|
||||
"field",
|
||||
["github_copilot_auth_type", "user_provider_credentials", "github_copilot_user_session"],
|
||||
["github_copilot_auth_type", "github_copilot_host", "user_provider_credentials", "github_copilot_user_session"],
|
||||
)
|
||||
def test_per_user_credential_slots_in_request_body_are_rejected(field: str) -> None:
|
||||
"""The per-user Copilot mode is decided by the stored credential's values and the caller's
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ import json
|
|||
from contextlib import contextmanager
|
||||
from types import SimpleNamespace
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
from unittest.mock import AsyncMock, MagicMock, call, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
|
@ -14,11 +14,13 @@ from fastapi import HTTPException, Request
|
|||
from fastapi.testclient import TestClient
|
||||
|
||||
import litellm
|
||||
from litellm.models.credentials import CredentialSource
|
||||
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
|
||||
from litellm.llms.github_copilot.per_user_auth import _SESSION_CACHE, _session_cache_key
|
||||
from litellm.models.credentials import CredentialSource, UserConnectionFlowHandle
|
||||
from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper
|
||||
from litellm.proxy.credential_endpoints.endpoints import get_llm_router
|
||||
from litellm.proxy.credential_endpoints.endpoints import _decode_flow_handle, get_llm_router
|
||||
from litellm.proxy.proxy_server import app
|
||||
from litellm.types.utils import CredentialItem
|
||||
|
||||
|
|
@ -1428,6 +1430,298 @@ def test_create_credential_still_answers_200_for_a_name_that_is_free(credential_
|
|||
find_by_name.assert_not_awaited(), "the unique index is the guard; create must not add a lookup"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"host",
|
||||
["https://x.ghe.com", "ghes.example.com", "x.ghe.com:443", "a.b.ghe.com"],
|
||||
)
|
||||
def test_create_credential_rejects_invalid_github_copilot_hosts(credential_store, host):
|
||||
create: Final = AsyncMock()
|
||||
credential_store(create=create)
|
||||
|
||||
response: Final = _create_credential(
|
||||
{
|
||||
"credential_name": "copilot",
|
||||
"credential_values": {"github_copilot_host": host},
|
||||
"credential_info": {"custom_llm_provider": "github_copilot"},
|
||||
}
|
||||
)
|
||||
|
||||
assert response.status_code == 400
|
||||
create.assert_not_awaited()
|
||||
|
||||
|
||||
def test_create_credential_normalizes_github_copilot_host(credential_store):
|
||||
credential_store(create=AsyncMock())
|
||||
|
||||
response = _create_credential(
|
||||
{
|
||||
"credential_name": "copilot",
|
||||
"credential_values": {"github_copilot_host": "OCTO.ghe.com"},
|
||||
"credential_info": {"custom_llm_provider": "github_copilot"},
|
||||
}
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert CredentialAccessor.find_credential("copilot").credential_values["github_copilot_host"] == "octo.ghe.com"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"host",
|
||||
["https://x.ghe.com", "ghes.example.com", "x.ghe.com:443", "a.b.ghe.com"],
|
||||
)
|
||||
def test_update_credential_rejects_invalid_github_copilot_hosts(credential_store, host):
|
||||
stored: Final = CredentialItem(
|
||||
credential_name="copilot",
|
||||
credential_values={"github_copilot_host": "octo.ghe.com"},
|
||||
credential_info={"custom_llm_provider": "github_copilot"},
|
||||
)
|
||||
update: Final = AsyncMock()
|
||||
credential_store(in_memory=(stored,), find_by_name=AsyncMock(return_value=stored), update_by_name=update)
|
||||
|
||||
response = _patch_credential(
|
||||
"copilot",
|
||||
{
|
||||
"credential_name": "copilot",
|
||||
"credential_values": {"github_copilot_host": host},
|
||||
"credential_info": {"custom_llm_provider": "github_copilot"},
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 400
|
||||
update.assert_not_awaited()
|
||||
|
||||
|
||||
def test_update_credential_normalizes_github_copilot_host(credential_store):
|
||||
stored = CredentialItem(
|
||||
credential_name="copilot",
|
||||
credential_values={},
|
||||
credential_info={"custom_llm_provider": "github_copilot"},
|
||||
)
|
||||
credential_store(
|
||||
in_memory=(stored,),
|
||||
find_by_name=AsyncMock(return_value=stored),
|
||||
update_by_name=AsyncMock(),
|
||||
)
|
||||
|
||||
response = _patch_credential(
|
||||
"copilot",
|
||||
{
|
||||
"credential_name": "copilot",
|
||||
"credential_values": {"github_copilot_host": "OCTO.ghe.com"},
|
||||
"credential_info": {"custom_llm_provider": "github_copilot"},
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert CredentialAccessor.find_credential("copilot").credential_values["github_copilot_host"] == "octo.ghe.com"
|
||||
|
||||
|
||||
def _encrypted_per_user_copilot_credential(host: str) -> CredentialItem:
|
||||
return CredentialItem(
|
||||
credential_name="copilot",
|
||||
credential_values={
|
||||
"github_copilot_auth_type": encrypt_value_helper("per_user_oauth"),
|
||||
"github_copilot_host": encrypt_value_helper(host),
|
||||
},
|
||||
credential_info={"custom_llm_provider": "github_copilot"},
|
||||
)
|
||||
|
||||
|
||||
def test_update_credential_purges_connections_when_github_host_changes(monkeypatch, restore_credential_list):
|
||||
stored = _encrypted_per_user_copilot_credential("github.com")
|
||||
with _repository_holding(stored):
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
table = MagicMock()
|
||||
table.find_many = AsyncMock(return_value=[_connection_row(credential_name="copilot")])
|
||||
table.delete_many = AsyncMock(return_value=None)
|
||||
proxy_server.prisma_client.db.litellm_userprovidercredentials = table
|
||||
proxy_server.prisma_client.writer_db = proxy_server.prisma_client.db
|
||||
|
||||
cached_values: Final[dict[str, object]] = {} # mutable-ok: fake Redis stores invalidation writes
|
||||
|
||||
async def set_cache(key: str, value: str, **kwargs: object) -> None:
|
||||
cached_values[key] = value
|
||||
|
||||
async def get_cache(key: str, **kwargs: object) -> object:
|
||||
return cached_values.get(key)
|
||||
|
||||
redis_cache: Final = MagicMock()
|
||||
redis_cache.async_set_cache = AsyncMock(side_effect=set_cache)
|
||||
redis_cache.async_get_cache = AsyncMock(side_effect=get_cache)
|
||||
monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache(redis_cache=redis_cache))
|
||||
|
||||
session_key: Final = _session_cache_key("user-a", "copilot", "github.com")
|
||||
_SESSION_CACHE.set_cache(
|
||||
session_key,
|
||||
("gho_secret", "copilot-token", "https://api.githubcopilot.com", 4102444800),
|
||||
ttl=3600,
|
||||
)
|
||||
response = _patch_credential(
|
||||
"copilot",
|
||||
{
|
||||
"credential_name": "copilot",
|
||||
"credential_values": {"github_copilot_host": "octo.ghe.com"},
|
||||
"credential_info": {},
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
table.find_many.assert_has_awaits(
|
||||
[
|
||||
call(where={"credential_name": "copilot"}),
|
||||
call(where={"credential_name": "copilot"}),
|
||||
]
|
||||
)
|
||||
table.delete_many.assert_awaited_once_with(where={"credential_name": "copilot"})
|
||||
assert cached_values[upc._cache_key("user-a", "copilot")] == upc._NOT_CONNECTED
|
||||
assert _SESSION_CACHE.get_cache(session_key) is None
|
||||
|
||||
|
||||
def test_update_credential_does_not_purge_connections_for_case_only_host_change(restore_credential_list):
|
||||
stored = _encrypted_per_user_copilot_credential("octo.ghe.com")
|
||||
with _repository_holding(stored):
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
table = MagicMock()
|
||||
table.find_many = AsyncMock(
|
||||
return_value=[_connection_row(credential_name="copilot", github_host="octo.ghe.com")]
|
||||
)
|
||||
table.delete_many = AsyncMock(return_value=None)
|
||||
proxy_server.prisma_client.db.litellm_userprovidercredentials = table
|
||||
proxy_server.prisma_client.writer_db = proxy_server.prisma_client.db
|
||||
|
||||
session_key: Final = _session_cache_key("user-a", "copilot", "octo.ghe.com")
|
||||
_SESSION_CACHE.set_cache(
|
||||
session_key,
|
||||
("gho_secret", "copilot-token", "https://copilot-api.octo.ghe.com", 4102444800),
|
||||
ttl=3600,
|
||||
)
|
||||
response = _patch_credential(
|
||||
"copilot",
|
||||
{
|
||||
"credential_name": "copilot",
|
||||
"credential_values": {"github_copilot_host": "OCTO.ghe.com"},
|
||||
"credential_info": {},
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
table.find_many.assert_not_awaited()
|
||||
table.delete_many.assert_not_awaited()
|
||||
assert _SESSION_CACHE.get_cache(session_key) is not None
|
||||
_SESSION_CACHE.delete_cache(session_key)
|
||||
|
||||
|
||||
def test_update_credential_does_not_purge_connections_for_legacy_host_case_difference(restore_credential_list):
|
||||
stored = _encrypted_per_user_copilot_credential("OCTO.ghe.com")
|
||||
with _repository_holding(stored):
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
table = MagicMock()
|
||||
table.find_many = AsyncMock(
|
||||
return_value=[_connection_row(credential_name="copilot", github_host="octo.ghe.com")]
|
||||
)
|
||||
table.delete_many = AsyncMock(return_value=None)
|
||||
proxy_server.prisma_client.db.litellm_userprovidercredentials = table
|
||||
proxy_server.prisma_client.writer_db = proxy_server.prisma_client.db
|
||||
|
||||
session_key: Final = _session_cache_key("user-a", "copilot", "octo.ghe.com")
|
||||
_SESSION_CACHE.set_cache(
|
||||
session_key,
|
||||
("gho_secret", "copilot-token", "https://copilot-api.octo.ghe.com", 4102444800),
|
||||
ttl=3600,
|
||||
)
|
||||
response = _patch_credential(
|
||||
"copilot",
|
||||
{
|
||||
"credential_name": "copilot",
|
||||
"credential_values": {"github_copilot_host": "octo.ghe.com"},
|
||||
"credential_info": {},
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
table.find_many.assert_not_awaited()
|
||||
table.delete_many.assert_not_awaited()
|
||||
assert _SESSION_CACHE.get_cache(session_key) is not None
|
||||
_SESSION_CACHE.delete_cache(session_key)
|
||||
|
||||
|
||||
def test_update_credential_purges_connections_when_github_host_is_deleted(restore_credential_list):
|
||||
stored = _encrypted_per_user_copilot_credential("octo.ghe.com")
|
||||
with _repository_holding(stored):
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
table = MagicMock()
|
||||
table.find_many = AsyncMock(return_value=[_connection_row(credential_name="copilot")])
|
||||
table.delete_many = AsyncMock(return_value=None)
|
||||
proxy_server.prisma_client.db.litellm_userprovidercredentials = table
|
||||
proxy_server.prisma_client.writer_db = proxy_server.prisma_client.db
|
||||
|
||||
response = _patch_credential(
|
||||
"copilot",
|
||||
{
|
||||
"credential_name": "copilot",
|
||||
"credential_values": {},
|
||||
"credential_values_to_delete": ["github_copilot_host"],
|
||||
"credential_info": {},
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
table.delete_many.assert_awaited_once_with(where={"credential_name": "copilot"})
|
||||
|
||||
|
||||
def test_update_credential_purges_connections_for_invalid_legacy_github_host(restore_credential_list):
|
||||
stored = _encrypted_per_user_copilot_credential("legacy.invalid")
|
||||
with _repository_holding(stored):
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
table = MagicMock()
|
||||
table.find_many = AsyncMock(return_value=[_connection_row(credential_name="copilot")])
|
||||
table.delete_many = AsyncMock(return_value=None)
|
||||
proxy_server.prisma_client.db.litellm_userprovidercredentials = table
|
||||
proxy_server.prisma_client.writer_db = proxy_server.prisma_client.db
|
||||
|
||||
response = _patch_credential(
|
||||
"copilot",
|
||||
{
|
||||
"credential_name": "copilot",
|
||||
"credential_values": {"api_key": "sk-updated"},
|
||||
"credential_info": {},
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
table.delete_many.assert_awaited_once_with(where={"credential_name": "copilot"})
|
||||
|
||||
|
||||
def test_update_credential_does_not_purge_connections_for_unrelated_value_update(restore_credential_list):
|
||||
stored = _encrypted_per_user_copilot_credential("octo.ghe.com")
|
||||
with _repository_holding(stored):
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
table = MagicMock()
|
||||
table.find_many = AsyncMock(return_value=[_connection_row(credential_name="copilot")])
|
||||
table.delete_many = AsyncMock(return_value=None)
|
||||
proxy_server.prisma_client.db.litellm_userprovidercredentials = table
|
||||
proxy_server.prisma_client.writer_db = proxy_server.prisma_client.db
|
||||
|
||||
response = _patch_credential(
|
||||
"copilot",
|
||||
{
|
||||
"credential_name": "copilot",
|
||||
"credential_values": {"api_key": "sk-updated"},
|
||||
"credential_info": {},
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
table.find_many.assert_not_awaited()
|
||||
table.delete_many.assert_not_awaited()
|
||||
|
||||
|
||||
def test_update_credential_resolves_credential_values_from_model_id_like_create(credential_store):
|
||||
"""Regression: PATCH dropped ``model_id`` from the body, so an update that named a
|
||||
deployment instead of raw values wrote whatever the caller sent, or nothing."""
|
||||
|
|
@ -1522,24 +1816,38 @@ def _as_user(user_id="user-a"):
|
|||
return lambda: UserAPIKeyAuth(api_key="test-key", user_role="internal_user", user_id=user_id)
|
||||
|
||||
|
||||
def _per_user_credential(name="copilot-cred"):
|
||||
def _per_user_credential(
|
||||
name: str = "copilot-cred",
|
||||
host: str | None = None,
|
||||
) -> CredentialItem:
|
||||
return CredentialItem(
|
||||
credential_name=name,
|
||||
credential_values={"github_copilot_auth_type": "per_user_oauth"},
|
||||
credential_values={
|
||||
"github_copilot_auth_type": "per_user_oauth",
|
||||
**({"github_copilot_host": host} if host is not None else {}),
|
||||
},
|
||||
credential_info={},
|
||||
)
|
||||
|
||||
|
||||
def _connection_row(user_id="user-a", credential_name="copilot-cred", login="octo"):
|
||||
def _connection_row(
|
||||
user_id: str = "user-a",
|
||||
credential_name: str = "copilot-cred",
|
||||
login: str = "octo",
|
||||
github_host: str = "github.com",
|
||||
) -> SimpleNamespace:
|
||||
from datetime import datetime, timezone
|
||||
from types import SimpleNamespace
|
||||
|
||||
return SimpleNamespace(
|
||||
user_id=user_id,
|
||||
credential_name=credential_name,
|
||||
provider="github_copilot",
|
||||
credential_b64=upc._encode(
|
||||
upc.GithubCopilotUserConnectionPayload(access_token="gho_secret", github_login=login)
|
||||
upc.GithubCopilotUserConnectionPayload(
|
||||
access_token="gho_secret",
|
||||
github_login=login,
|
||||
github_host=github_host,
|
||||
)
|
||||
),
|
||||
created_at=datetime(2026, 1, 1, tzinfo=timezone.utc),
|
||||
updated_at=datetime(2026, 1, 1, tzinfo=timezone.utc),
|
||||
|
|
@ -1635,6 +1943,20 @@ def test_user_connection_start_returns_device_flow_fields(monkeypatch):
|
|||
}
|
||||
handle = response.json()["flow_handle"]
|
||||
assert isinstance(handle, str) and "device_code" not in handle and _DEVICE_FLOW_START["device_code"] not in handle
|
||||
decoded_handle: Final = _decode_flow_handle(handle)
|
||||
assert decoded_handle is not None
|
||||
assert decoded_handle.github_host == "github.com"
|
||||
|
||||
|
||||
def test_legacy_user_connection_flow_handle_defaults_to_github_com():
|
||||
handle: Final = UserConnectionFlowHandle(
|
||||
user_id="user-a",
|
||||
credential_name="copilot-cred",
|
||||
device_code="DC-1",
|
||||
interval=5,
|
||||
expires_at=4102444800,
|
||||
)
|
||||
assert handle.github_host == "github.com"
|
||||
|
||||
|
||||
def test_user_connection_start_404s_for_shared_credential(monkeypatch):
|
||||
|
|
@ -1657,7 +1979,7 @@ def test_user_connection_poll_rejects_an_undecryptable_flow_handle(monkeypatch):
|
|||
|
||||
|
||||
def test_user_connection_poll_connected_persists_and_reports_login(monkeypatch):
|
||||
table = _patch_user_connection_env(monkeypatch, [_per_user_credential()])
|
||||
table: Final = _patch_user_connection_env(monkeypatch, [_per_user_credential()])
|
||||
with _github_http(
|
||||
{
|
||||
_DEVICE_CODE_URL: _DEVICE_FLOW_START,
|
||||
|
|
@ -1666,9 +1988,13 @@ def test_user_connection_poll_connected_persists_and_reports_login(monkeypatch):
|
|||
_GITHUB_USER_URL: {"login": "octo"},
|
||||
}
|
||||
):
|
||||
start = _call_as("POST", "/credentials/copilot-cred/user_connection/start", auth=_as_user())
|
||||
start: Final = _call_as("POST", "/credentials/copilot-cred/user_connection/start", auth=_as_user())
|
||||
assert start.status_code == 200, start.text
|
||||
poll = _call_as(
|
||||
encrypted_handle: Final = start.json()["flow_handle"]
|
||||
decoded_handle: Final = _decode_flow_handle(encrypted_handle)
|
||||
assert decoded_handle is not None
|
||||
assert decoded_handle.github_host == "github.com"
|
||||
poll: Final = _call_as(
|
||||
"POST",
|
||||
"/credentials/copilot-cred/user_connection/poll",
|
||||
json_body={"flow_handle": start.json()["flow_handle"]},
|
||||
|
|
@ -1677,6 +2003,110 @@ def test_user_connection_poll_connected_persists_and_reports_login(monkeypatch):
|
|||
assert poll.status_code == 200, poll.text
|
||||
assert poll.json() == {"status": "connected", "interval": None, "github_login": "octo"}
|
||||
table.upsert.assert_awaited_once()
|
||||
stored_connection: Final = upc.decode_user_provider_credential(
|
||||
table.upsert.await_args.kwargs["data"]["create"]["credential_b64"]
|
||||
)
|
||||
assert stored_connection is not None
|
||||
assert stored_connection.github_host == "github.com"
|
||||
|
||||
|
||||
def test_poll_host_change_race_requires_reconnect_on_the_next_request(monkeypatch):
|
||||
from litellm.llms.github_copilot.per_user_auth import attach_github_copilot_user_session
|
||||
|
||||
table: Final = _patch_user_connection_env(monkeypatch, [_per_user_credential(host="octo.ghe.com")])
|
||||
|
||||
def change_host_after_token_exchange(request: httpx.Request) -> dict[str, str]:
|
||||
monkeypatch.setattr(litellm, "credential_list", [_per_user_credential()])
|
||||
return {"access_token": "gho_1"}
|
||||
|
||||
with _github_http(
|
||||
{
|
||||
"https://octo.ghe.com/login/device/code": _DEVICE_FLOW_START,
|
||||
"https://octo.ghe.com/login/oauth/access_token": change_host_after_token_exchange,
|
||||
"https://api.octo.ghe.com/copilot_internal/v2/token": {
|
||||
"token": "copilot-tok",
|
||||
"expires_at": 4102444800,
|
||||
"endpoints": {"api": "https://copilot-api.octo.ghe.com"},
|
||||
},
|
||||
"https://api.octo.ghe.com/user": {"login": "octo"},
|
||||
}
|
||||
):
|
||||
start: Final = _call_as("POST", "/credentials/copilot-cred/user_connection/start", auth=_as_user())
|
||||
poll: Final = _call_as(
|
||||
"POST",
|
||||
"/credentials/copilot-cred/user_connection/poll",
|
||||
json_body={"flow_handle": start.json()["flow_handle"]},
|
||||
auth=_as_user(),
|
||||
)
|
||||
assert poll.status_code == 200, poll.text
|
||||
stored_connection: Final = upc.decode_user_provider_credential(
|
||||
table.upsert.await_args.kwargs["data"]["create"]["credential_b64"]
|
||||
)
|
||||
assert stored_connection is not None
|
||||
|
||||
kwargs: Final[dict[str, object]] = {
|
||||
"litellm_credential_name": "copilot-cred",
|
||||
"secret_fields": {
|
||||
"user_provider_credentials_user_id": "user-a",
|
||||
"user_provider_credentials": {
|
||||
"copilot-cred": {
|
||||
"access_token": stored_connection.access_token,
|
||||
"github_host": stored_connection.github_host,
|
||||
}
|
||||
},
|
||||
},
|
||||
}
|
||||
module_level_client: Final = MagicMock()
|
||||
module_level_client.get.side_effect = AssertionError("unexpected token exchange")
|
||||
with _github_http({}), patch.object(litellm, "module_level_client", module_level_client):
|
||||
with pytest.raises(litellm.CallerCredentialAuthenticationError) as exc:
|
||||
attach_github_copilot_user_session(kwargs)
|
||||
assert "Connect GitHub Copilot for credential 'copilot-cred'" in str(exc.value)
|
||||
assert stored_connection.github_host == "octo.ghe.com"
|
||||
|
||||
|
||||
def test_ghe_user_connection_uses_tenant_urls_for_device_flow_exchange_and_login(monkeypatch):
|
||||
table: Final = _patch_user_connection_env(monkeypatch, [_per_user_credential(host="octo.ghe.com")])
|
||||
with _github_http(
|
||||
{
|
||||
"https://octo.ghe.com/login/device/code": _DEVICE_FLOW_START,
|
||||
"https://octo.ghe.com/login/oauth/access_token": {"access_token": "gho_1"},
|
||||
"https://api.octo.ghe.com/copilot_internal/v2/token": {
|
||||
"token": "copilot-tok",
|
||||
"expires_at": 4102444800,
|
||||
"endpoints": {"api": "https://copilot-api.octo.ghe.com"},
|
||||
},
|
||||
"https://api.octo.ghe.com/user": {"login": "octo"},
|
||||
}
|
||||
):
|
||||
start: Final = _call_as("POST", "/credentials/copilot-cred/user_connection/start", auth=_as_user())
|
||||
assert start.status_code == 200, start.text
|
||||
poll: Final = _call_as(
|
||||
"POST",
|
||||
"/credentials/copilot-cred/user_connection/poll",
|
||||
json_body={"flow_handle": start.json()["flow_handle"]},
|
||||
auth=_as_user(),
|
||||
)
|
||||
assert poll.status_code == 200, poll.text
|
||||
assert poll.json() == {"status": "connected", "interval": None, "github_login": "octo"}
|
||||
table.upsert.assert_awaited_once()
|
||||
stored_connection: Final = upc.decode_user_provider_credential(
|
||||
table.upsert.await_args.kwargs["data"]["create"]["credential_b64"]
|
||||
)
|
||||
assert stored_connection is not None
|
||||
assert stored_connection.github_host == "octo.ghe.com"
|
||||
|
||||
|
||||
def test_user_connection_poll_rejects_handle_from_different_host(monkeypatch):
|
||||
_patch_user_connection_env(monkeypatch, [_per_user_credential(host="octo.ghe.com")])
|
||||
response = _call_as(
|
||||
"POST",
|
||||
"/credentials/copilot-cred/user_connection/poll",
|
||||
json_body={"flow_handle": _flow_handle_for()},
|
||||
auth=_as_user(),
|
||||
)
|
||||
assert response.status_code == 400
|
||||
assert response.json()["error"]["message"] == "invalid or expired flow_handle"
|
||||
|
||||
|
||||
def test_user_connection_poll_pending(monkeypatch):
|
||||
|
|
@ -2038,6 +2468,23 @@ def test_pending_device_flow_survives_disconnect_via_stateless_handle(monkeypatc
|
|||
assert delete.json() == {"status": "disconnected"}
|
||||
|
||||
|
||||
def test_delete_user_connection_evicts_copilot_session(monkeypatch):
|
||||
table: Final = _patch_user_connection_env(monkeypatch, [_per_user_credential()], rows=[_connection_row()])
|
||||
session_key: Final = _session_cache_key("user-a", "copilot-cred", "github.com")
|
||||
_SESSION_CACHE.set_cache(
|
||||
session_key,
|
||||
("gho_secret", "copilot-token", "https://api.githubcopilot.com", 4102444800),
|
||||
ttl=3600,
|
||||
)
|
||||
|
||||
response: Final = _call_as("DELETE", "/credentials/copilot-cred/user_connection", auth=_as_user())
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json() == {"status": "disconnected"}
|
||||
assert _SESSION_CACHE.get_cache(session_key) is None
|
||||
table.delete_many.assert_awaited_once_with(where={"user_id": "user-a", "credential_name": "copilot-cred"})
|
||||
|
||||
|
||||
def test_user_connection_cache_keys_cannot_collide_on_colons():
|
||||
"""user_id/credential_name are joined losslessly, so ('a:b','c') and ('a','b:c')
|
||||
can never share a cache or device-flow entry."""
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
"""Tests for the per-user provider credentials DB/cache module."""
|
||||
|
||||
import json
|
||||
from types import SimpleNamespace
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
|
@ -7,9 +8,11 @@ from unittest.mock import AsyncMock, MagicMock
|
|||
import pytest
|
||||
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.constants import GITHUB_COPILOT_DEFAULT_HOST
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper
|
||||
from litellm.proxy.credential_endpoints.user_provider_credentials import (
|
||||
GithubCopilotUserConnectionPayload,
|
||||
aget_user_provider_tokens,
|
||||
aget_user_provider_credentials,
|
||||
decode_user_provider_credential,
|
||||
delete_user_provider_credential,
|
||||
delete_user_provider_credentials_for_credential,
|
||||
|
|
@ -110,7 +113,7 @@ async def test_delete_for_credential_returns_affected_user_ids():
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aget_tokens_returns_plaintext_and_caches_ciphertext(monkeypatch):
|
||||
async def test_aget_credentials_returns_plaintext_and_caches_ciphertext(monkeypatch):
|
||||
payload = GithubCopilotUserConnectionPayload(access_token="gho_secret", github_login="octo")
|
||||
from litellm.proxy.credential_endpoints import user_provider_credentials as upc
|
||||
|
||||
|
|
@ -121,12 +124,12 @@ async def test_aget_tokens_returns_plaintext_and_caches_ciphertext(monkeypatch):
|
|||
redis = _fake_redis(monkeypatch)
|
||||
cache = DualCache(redis_cache=redis)
|
||||
|
||||
tokens = await aget_user_provider_tokens(prisma_client, cache, "user-a", ["copilot-cred"])
|
||||
assert tokens == {"copilot-cred": "gho_secret"}
|
||||
connections = await aget_user_provider_credentials(prisma_client, cache, "user-a", ["copilot-cred"])
|
||||
assert connections == {"copilot-cred": {"access_token": "gho_secret", "github_host": "github.com"}}
|
||||
|
||||
# second call served from cache: no new DB read, and cache holds ciphertext not plaintext
|
||||
tokens2 = await aget_user_provider_tokens(prisma_client, cache, "user-a", ["copilot-cred"])
|
||||
assert tokens2 == {"copilot-cred": "gho_secret"}
|
||||
connections2 = await aget_user_provider_credentials(prisma_client, cache, "user-a", ["copilot-cred"])
|
||||
assert connections2 == {"copilot-cred": {"access_token": "gho_secret", "github_host": "github.com"}}
|
||||
assert table.find_many.await_count == 1
|
||||
cached = await redis.async_get_cache(upc._cache_key("user-a", "copilot-cred"))
|
||||
assert cached == ciphertext
|
||||
|
|
@ -135,14 +138,59 @@ async def test_aget_tokens_returns_plaintext_and_caches_ciphertext(monkeypatch):
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aget_tokens_negative_caches_missing_connection(monkeypatch):
|
||||
async def test_aget_connections_preserves_github_host(monkeypatch):
|
||||
payload: Final = GithubCopilotUserConnectionPayload(
|
||||
access_token="gho_secret",
|
||||
github_login="octo",
|
||||
github_host="octo.ghe.com",
|
||||
)
|
||||
from litellm.proxy.credential_endpoints import user_provider_credentials as upc
|
||||
|
||||
table: Final = MagicMock()
|
||||
table.find_many = AsyncMock(return_value=[_row(credential_b64=upc._encode(payload))])
|
||||
prisma_client: Final = _prisma(table)
|
||||
cache: Final = DualCache(redis_cache=_fake_redis(monkeypatch))
|
||||
|
||||
connections: Final = await aget_user_provider_credentials(
|
||||
prisma_client,
|
||||
cache,
|
||||
"user-a",
|
||||
["copilot-cred"],
|
||||
)
|
||||
assert connections == {"copilot-cred": {"access_token": "gho_secret", "github_host": "octo.ghe.com"}}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_legacy_connection_without_github_host_defaults_to_github_com(monkeypatch):
|
||||
ciphertext: Final = encrypt_value_helper(json.dumps({"access_token": "gho_secret", "github_login": "octo"}))
|
||||
table: Final = MagicMock()
|
||||
table.find_many = AsyncMock(return_value=[_row(credential_b64=ciphertext)])
|
||||
prisma_client: Final = _prisma(table)
|
||||
cache: Final = DualCache(redis_cache=_fake_redis(monkeypatch))
|
||||
|
||||
connections: Final = await aget_user_provider_credentials(
|
||||
prisma_client,
|
||||
cache,
|
||||
"user-a",
|
||||
["copilot-cred"],
|
||||
)
|
||||
assert connections == {
|
||||
"copilot-cred": {
|
||||
"access_token": "gho_secret",
|
||||
"github_host": GITHUB_COPILOT_DEFAULT_HOST,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aget_credentials_negative_caches_missing_connection(monkeypatch):
|
||||
table = MagicMock()
|
||||
table.find_many = AsyncMock(return_value=[])
|
||||
prisma_client = _prisma(table)
|
||||
cache = DualCache(redis_cache=_fake_redis(monkeypatch))
|
||||
|
||||
assert await aget_user_provider_tokens(prisma_client, cache, "user-a", ["copilot-cred"]) == {}
|
||||
assert await aget_user_provider_tokens(prisma_client, cache, "user-a", ["copilot-cred"]) == {}
|
||||
assert await aget_user_provider_credentials(prisma_client, cache, "user-a", ["copilot-cred"]) == {}
|
||||
assert await aget_user_provider_credentials(prisma_client, cache, "user-a", ["copilot-cred"]) == {}
|
||||
assert table.find_many.await_count == 1
|
||||
|
||||
|
||||
|
|
@ -156,10 +204,10 @@ async def test_invalidate_cache_forces_db_refetch(monkeypatch):
|
|||
prisma_client = _prisma(table)
|
||||
cache = DualCache(redis_cache=_fake_redis(monkeypatch))
|
||||
|
||||
await aget_user_provider_tokens(prisma_client, cache, "user-a", ["copilot-cred"])
|
||||
await aget_user_provider_credentials(prisma_client, cache, "user-a", ["copilot-cred"])
|
||||
await drop_user_provider_credential_cache(cache, "user-a", "copilot-cred")
|
||||
tokens = await aget_user_provider_tokens(prisma_client, cache, "user-a", ["copilot-cred"])
|
||||
assert tokens == {"copilot-cred": "gho_new"}
|
||||
connections = await aget_user_provider_credentials(prisma_client, cache, "user-a", ["copilot-cred"])
|
||||
assert connections == {"copilot-cred": {"access_token": "gho_new", "github_host": "github.com"}}
|
||||
assert table.find_many.await_count == 2
|
||||
|
||||
|
||||
|
|
@ -184,19 +232,19 @@ async def test_disconnect_on_one_worker_invalidates_other_workers_via_redis(monk
|
|||
prisma_client = _prisma(table)
|
||||
|
||||
# worker A "connects" (DB read + cache write)
|
||||
tokens_a = await aget_user_provider_tokens(prisma_client, cache_a, "user-a", ["copilot-cred"])
|
||||
assert tokens_a == {"copilot-cred": "gho_secret"}
|
||||
connections_a = await aget_user_provider_credentials(prisma_client, cache_a, "user-a", ["copilot-cred"])
|
||||
assert connections_a == {"copilot-cred": {"access_token": "gho_secret", "github_host": "github.com"}}
|
||||
|
||||
# worker B reads the same connection straight from Redis (its in-memory layer is empty)
|
||||
tokens_b = await aget_user_provider_tokens(prisma_client, cache_b, "user-a", ["copilot-cred"])
|
||||
assert tokens_b == {"copilot-cred": "gho_secret"}
|
||||
connections_b = await aget_user_provider_credentials(prisma_client, cache_b, "user-a", ["copilot-cred"])
|
||||
assert connections_b == {"copilot-cred": {"access_token": "gho_secret", "github_host": "github.com"}}
|
||||
|
||||
# worker A disconnects: the key is overwritten with the not-connected tombstone
|
||||
await invalidate_user_provider_credential_cache(cache_a, "user-a", "copilot-cred")
|
||||
|
||||
# worker B sees the tombstone from Redis, no worker-local staleness, no DB read needed
|
||||
tokens_b_after = await aget_user_provider_tokens(prisma_client, cache_b, "user-a", ["copilot-cred"])
|
||||
assert tokens_b_after == {}
|
||||
connections_b_after = await aget_user_provider_credentials(prisma_client, cache_b, "user-a", ["copilot-cred"])
|
||||
assert connections_b_after == {}
|
||||
assert table.find_many.await_count == 1
|
||||
|
||||
|
||||
|
|
@ -222,8 +270,8 @@ async def test_stale_fill_cannot_overwrite_a_disconnect_tombstone(monkeypatch):
|
|||
table.find_many = AsyncMock(side_effect=_read_while_disconnecting)
|
||||
prisma_client = _prisma(table)
|
||||
|
||||
await aget_user_provider_tokens(prisma_client, cache, "user-a", ["copilot-cred"])
|
||||
assert await aget_user_provider_tokens(prisma_client, cache, "user-a", ["copilot-cred"]) == {}
|
||||
await aget_user_provider_credentials(prisma_client, cache, "user-a", ["copilot-cred"])
|
||||
assert await aget_user_provider_credentials(prisma_client, cache, "user-a", ["copilot-cred"]) == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -240,8 +288,8 @@ async def test_redis_read_failure_falls_back_to_the_database(monkeypatch):
|
|||
table.find_many = AsyncMock(return_value=[_row(credential_b64=upc._encode(payload))])
|
||||
prisma_client = _prisma(table)
|
||||
|
||||
assert await aget_user_provider_tokens(prisma_client, cache, "user-a", ["copilot-cred"]) == {
|
||||
"copilot-cred": "gho_secret"
|
||||
assert await aget_user_provider_credentials(prisma_client, cache, "user-a", ["copilot-cred"]) == {
|
||||
"copilot-cred": {"access_token": "gho_secret", "github_host": "github.com"}
|
||||
}
|
||||
|
||||
|
||||
|
|
@ -259,8 +307,8 @@ async def test_redis_write_and_delete_failures_do_not_fail_the_request(monkeypat
|
|||
table.find_many = AsyncMock(return_value=[_row(credential_b64=upc._encode(payload))])
|
||||
prisma_client = _prisma(table)
|
||||
|
||||
assert await aget_user_provider_tokens(prisma_client, cache, "user-a", ["copilot-cred"]) == {
|
||||
"copilot-cred": "gho_secret"
|
||||
assert await aget_user_provider_credentials(prisma_client, cache, "user-a", ["copilot-cred"]) == {
|
||||
"copilot-cred": {"access_token": "gho_secret", "github_host": "github.com"}
|
||||
}
|
||||
await invalidate_user_provider_credential_cache(cache, "user-a", "copilot-cred")
|
||||
|
||||
|
|
@ -279,14 +327,14 @@ async def test_without_redis_no_worker_local_entries_are_kept():
|
|||
cache_a = DualCache()
|
||||
cache_b = DualCache()
|
||||
|
||||
assert await aget_user_provider_tokens(prisma_client, cache_a, "user-a", ["copilot-cred"]) == {
|
||||
"copilot-cred": "gho_secret"
|
||||
assert await aget_user_provider_credentials(prisma_client, cache_a, "user-a", ["copilot-cred"]) == {
|
||||
"copilot-cred": {"access_token": "gho_secret", "github_host": "github.com"}
|
||||
}
|
||||
assert await aget_user_provider_tokens(prisma_client, cache_b, "user-a", ["copilot-cred"]) == {
|
||||
"copilot-cred": "gho_secret"
|
||||
assert await aget_user_provider_credentials(prisma_client, cache_b, "user-a", ["copilot-cred"]) == {
|
||||
"copilot-cred": {"access_token": "gho_secret", "github_host": "github.com"}
|
||||
}
|
||||
await invalidate_user_provider_credential_cache(cache_b, "user-a", "copilot-cred")
|
||||
assert await aget_user_provider_tokens(prisma_client, cache_a, "user-a", ["copilot-cred"]) == {}
|
||||
assert await aget_user_provider_credentials(prisma_client, cache_a, "user-a", ["copilot-cred"]) == {}
|
||||
assert table.find_many.await_count == 3
|
||||
|
||||
|
||||
|
|
@ -306,8 +354,8 @@ async def test_connect_overwrites_a_not_connected_tombstone(monkeypatch):
|
|||
table = MagicMock()
|
||||
table.find_many = AsyncMock(return_value=[])
|
||||
prisma_client = _prisma(table)
|
||||
assert await aget_user_provider_tokens(prisma_client, cache, "user-a", ["copilot-cred"]) == {
|
||||
"copilot-cred": "gho_secret"
|
||||
assert await aget_user_provider_credentials(prisma_client, cache, "user-a", ["copilot-cred"]) == {
|
||||
"copilot-cred": {"access_token": "gho_secret", "github_host": "github.com"}
|
||||
}
|
||||
table.find_many.assert_not_awaited()
|
||||
|
||||
|
|
@ -332,9 +380,9 @@ async def test_stale_not_connected_fill_cannot_overwrite_a_connect_write(monkeyp
|
|||
table.find_many = AsyncMock(side_effect=_read_while_connecting)
|
||||
prisma_client = _prisma(table)
|
||||
|
||||
await aget_user_provider_tokens(prisma_client, cache, "user-a", ["copilot-cred"])
|
||||
assert await aget_user_provider_tokens(prisma_client, cache, "user-a", ["copilot-cred"]) == {
|
||||
"copilot-cred": "gho_secret"
|
||||
await aget_user_provider_credentials(prisma_client, cache, "user-a", ["copilot-cred"])
|
||||
assert await aget_user_provider_credentials(prisma_client, cache, "user-a", ["copilot-cred"]) == {
|
||||
"copilot-cred": {"access_token": "gho_secret", "github_host": "github.com"}
|
||||
}
|
||||
|
||||
|
||||
|
|
@ -391,7 +439,7 @@ async def test_reads_hit_the_writer_engine_not_a_stale_replica():
|
|||
prisma_client.db.litellm_userprovidercredentials = reader_table
|
||||
prisma_client.writer_db.litellm_userprovidercredentials = writer_table
|
||||
|
||||
await aget_user_provider_tokens(prisma_client, DualCache(), "user-a", ["copilot-cred"])
|
||||
await aget_user_provider_credentials(prisma_client, DualCache(), "user-a", ["copilot-cred"])
|
||||
|
||||
writer_table.find_many.assert_awaited_once()
|
||||
reader_table.find_many.assert_not_awaited()
|
||||
|
|
|
|||
|
|
@ -9347,7 +9347,9 @@ class TestResolveUserProviderCredentials:
|
|||
llm_router=self._router(),
|
||||
)
|
||||
|
||||
assert data["secret_fields"]["user_provider_credentials"] == {"copilot-cred": "gho_secret"}
|
||||
assert data["secret_fields"]["user_provider_credentials"] == {
|
||||
"copilot-cred": {"access_token": "gho_secret", "github_host": "github.com"}
|
||||
}
|
||||
assert data["secret_fields"]["user_provider_credentials_user_id"] == "user-a"
|
||||
# the shared fallback credential is never looked up
|
||||
call_where = table.find_many.await_args.kwargs["where"]
|
||||
|
|
@ -9804,7 +9806,9 @@ class TestResolveUserProviderCredentials:
|
|||
team_id=None,
|
||||
llm_router=router,
|
||||
)
|
||||
assert data["secret_fields"]["user_provider_credentials"] == {"copilot-cred": "gho_secret"}
|
||||
assert data["secret_fields"]["user_provider_credentials"] == {
|
||||
"copilot-cred": {"access_token": "gho_secret", "github_host": "github.com"}
|
||||
}
|
||||
call_where: Final = table.find_many.await_args.kwargs["where"]
|
||||
assert call_where == {"AND": ({"user_id": "user-a"}, {"credential_name": {"in": ["copilot-cred"]}})}
|
||||
|
||||
|
|
@ -10290,7 +10294,8 @@ async def test_secret_fields_user_provider_credentials_repr_hides_tokens(monkeyp
|
|||
assert isinstance(credentials, RedactedDict)
|
||||
assert "gho_secret" not in str(data)
|
||||
assert "gho_secret" not in repr(credentials)
|
||||
assert credentials["copilot-cred"] == "gho_secret", "provider reads still get the token"
|
||||
assert credentials["copilot-cred"]["access_token"] == "gho_secret", "provider reads still get the token"
|
||||
assert credentials["copilot-cred"]["github_host"] == "github.com"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -10304,8 +10309,10 @@ async def test_header_mapped_user_id_cannot_load_another_users_connection(monkey
|
|||
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
|
||||
|
||||
_per_user_env(monkeypatch)
|
||||
load = AsyncMock(return_value={"copilot-cred": "gho_user_a_token"})
|
||||
monkeypatch.setattr(upc, "aget_user_provider_tokens", load)
|
||||
load = AsyncMock(
|
||||
return_value={"copilot-cred": {"access_token": "gho_user_a_token", "github_host": "github.com"}}
|
||||
)
|
||||
monkeypatch.setattr(upc, "aget_user_provider_credentials", load)
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
|
|
@ -10344,7 +10351,8 @@ async def test_header_mapped_user_id_cannot_load_another_users_connection(monkey
|
|||
called_user_ids = [call.kwargs["user_id"] for call in load.await_args_list]
|
||||
assert called_user_ids == ["user-a"], f"header-mapped user-b must never reach the token loader: {called_user_ids}"
|
||||
credentials = updated["secret_fields"]["user_provider_credentials"]
|
||||
assert credentials["copilot-cred"] == "gho_user_a_token"
|
||||
assert credentials["copilot-cred"]["access_token"] == "gho_user_a_token"
|
||||
assert credentials["copilot-cred"]["github_host"] == "github.com"
|
||||
assert updated["secret_fields"]["user_provider_credentials_user_id"] == "user-a"
|
||||
|
||||
|
||||
|
|
@ -10357,8 +10365,10 @@ async def test_per_user_token_resolution_uses_key_identity_without_header_mappin
|
|||
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
|
||||
|
||||
_per_user_env(monkeypatch)
|
||||
load = AsyncMock(return_value={"copilot-cred": "gho_user_a_token"})
|
||||
monkeypatch.setattr(upc, "aget_user_provider_tokens", load)
|
||||
load = AsyncMock(
|
||||
return_value={"copilot-cred": {"access_token": "gho_user_a_token", "github_host": "github.com"}}
|
||||
)
|
||||
monkeypatch.setattr(upc, "aget_user_provider_credentials", load)
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
|
|
@ -10394,7 +10404,8 @@ async def test_per_user_token_resolution_uses_key_identity_without_header_mappin
|
|||
|
||||
called_user_ids = [call.kwargs["user_id"] for call in load.await_args_list]
|
||||
assert called_user_ids == ["user-a"]
|
||||
assert updated["secret_fields"]["user_provider_credentials"]["copilot-cred"] == "gho_user_a_token"
|
||||
assert updated["secret_fields"]["user_provider_credentials"]["copilot-cred"]["access_token"] == "gho_user_a_token"
|
||||
assert updated["secret_fields"]["user_provider_credentials"]["copilot-cred"]["github_host"] == "github.com"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -91,6 +91,7 @@ CONNECTION_NAMES: Final = (
|
|||
"azure_ad_token_provider",
|
||||
"litellm_credential_name",
|
||||
"github_copilot_auth_type",
|
||||
"github_copilot_host",
|
||||
"configurable_clientside_auth_params",
|
||||
"use_xai_oauth",
|
||||
"fireworks_forward_user_id",
|
||||
|
|
|
|||
|
|
@ -361,7 +361,7 @@ def test_server_owned_registry_includes_anthropic_openai_and_oauth_token_exchang
|
|||
"token_exchange_profile",
|
||||
"token_exchange_scope",
|
||||
)
|
||||
assert github_copilot_oauth_litellm_params == ("github_copilot_auth_type",)
|
||||
assert github_copilot_oauth_litellm_params == ("github_copilot_auth_type", "github_copilot_host")
|
||||
assert server_owned_wif_litellm_params == (
|
||||
anthropic_wif_litellm_params
|
||||
+ openai_wif_litellm_params
|
||||
|
|
|
|||
|
|
@ -28,6 +28,7 @@ describe("provider auth types", () => {
|
|||
"token_exchange_audience",
|
||||
"api_key",
|
||||
]);
|
||||
expect(authTypeFieldKeys("GITHUB_COPILOT")).toContain("github_copilot_host");
|
||||
});
|
||||
|
||||
it("infers auth types from required fields and falls back when only optional defaults are present", () => {
|
||||
|
|
@ -85,6 +86,9 @@ describe("provider auth types", () => {
|
|||
expect(inferAuthTypeId(authTypes, { api_key: "" })).toBe("shared_device_login");
|
||||
expect(inferAuthTypeId(authTypes, { github_copilot_auth_type: "per_user_oauth" })).toBe("per_user_oauth");
|
||||
expect(hiddenAuthFieldKeys(authTypes, "per_user_oauth")).toEqual(["api_base", "api_key"]);
|
||||
expect(hiddenAuthFieldKeys(authTypes, "shared_device_login")).toEqual(["github_copilot_auth_type"]);
|
||||
expect(hiddenAuthFieldKeys(authTypes, "shared_device_login")).toEqual([
|
||||
"github_copilot_auth_type",
|
||||
"github_copilot_host",
|
||||
]);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ export interface ProviderAuthType {
|
|||
const EMPTY_PROVIDER_AUTH_TYPES: readonly ProviderAuthType[] = [];
|
||||
|
||||
export const GITHUB_COPILOT_AUTH_TYPE_KEY = "github_copilot_auth_type";
|
||||
export const GITHUB_COPILOT_HOST_KEY = "github_copilot_host";
|
||||
export const GITHUB_COPILOT_PER_USER_AUTH_TYPE = "per_user_oauth";
|
||||
|
||||
export const PROVIDER_AUTH_TYPES: Partial<Record<keyof typeof Providers, readonly ProviderAuthType[]>> = {
|
||||
|
|
@ -31,7 +32,7 @@ export const PROVIDER_AUTH_TYPES: Partial<Record<keyof typeof Providers, readonl
|
|||
credentialOnly: true,
|
||||
description:
|
||||
"Each LiteLLM user connects their own GitHub account from LLM Credentials, and their requests use their own GitHub Copilot access. Users who have not connected get a 401.",
|
||||
fieldKeys: [GITHUB_COPILOT_AUTH_TYPE_KEY],
|
||||
fieldKeys: [GITHUB_COPILOT_AUTH_TYPE_KEY, GITHUB_COPILOT_HOST_KEY],
|
||||
requiredFieldKeys: [GITHUB_COPILOT_AUTH_TYPE_KEY],
|
||||
fixedValues: { [GITHUB_COPILOT_AUTH_TYPE_KEY]: GITHUB_COPILOT_PER_USER_AUTH_TYPE },
|
||||
},
|
||||
|
|
|
|||
|
|
@ -67,6 +67,13 @@ vi.mock("../networking", async () => {
|
|||
credential_fields: [
|
||||
{ key: "api_base", label: "API Base", field_type: "text" },
|
||||
{ key: "api_key", label: "API Key", field_type: "password" },
|
||||
{
|
||||
key: "github_copilot_host",
|
||||
label: "GitHub host",
|
||||
placeholder: "github.com",
|
||||
tooltip: "Leave empty for github.com. For GHE.com data residency enter your SUBDOMAIN.ghe.com host",
|
||||
field_type: "text",
|
||||
},
|
||||
],
|
||||
},
|
||||
]),
|
||||
|
|
@ -532,7 +539,7 @@ describe("CredentialModal with Anthropic workload identity federation", () => {
|
|||
describe("CredentialModal with GitHub Copilot auth types", () => {
|
||||
const perUserCopilotCredential: CredentialItem = {
|
||||
credential_name: "copilot-per-user",
|
||||
credential_values: { github_copilot_auth_type: "per_user_oauth" },
|
||||
credential_values: { github_copilot_auth_type: "per_user_oauth", github_copilot_host: "octo.ghe.com" },
|
||||
credential_info: { custom_llm_provider: "GITHUB_COPILOT" },
|
||||
};
|
||||
|
||||
|
|
@ -577,6 +584,7 @@ describe("CredentialModal with GitHub Copilot auth types", () => {
|
|||
const onSubmit = renderModal({ mode: "edit", existingCredential: perUserCopilotCredential });
|
||||
|
||||
expect(await screen.findByRole("combobox", { name: "Auth Type:" })).toHaveTextContent("Per-user GitHub OAuth");
|
||||
expect(screen.getByLabelText("GitHub host")).toHaveValue("octo.ghe.com");
|
||||
await chooseOption(user, /^Auth Type:/, "Shared device login");
|
||||
await user.click(screen.getByRole("button", { name: "Update Credential" }));
|
||||
|
||||
|
|
@ -586,7 +594,7 @@ describe("CredentialModal with GitHub Copilot auth types", () => {
|
|||
custom_llm_provider: "GITHUB_COPILOT",
|
||||
credential_values: {},
|
||||
});
|
||||
expect(valuesToDelete).toEqual(["github_copilot_auth_type"]);
|
||||
expect([...valuesToDelete].sort()).toEqual(["github_copilot_auth_type", "github_copilot_host"]);
|
||||
});
|
||||
});
|
||||
|
||||
|
|
|
|||
|
|
@ -79,6 +79,18 @@ describe("GithubCopilotConnectModal", () => {
|
|||
expect(props.onDisconnect).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it("uses the GHE verification URI returned by the start response", async () => {
|
||||
const verificationUri = "https://octo.ghe.com/login/device";
|
||||
networking.userConnectionStartCall.mockResolvedValue({
|
||||
...START_RESPONSE,
|
||||
verification_uri: verificationUri,
|
||||
});
|
||||
renderModal();
|
||||
await advance(0);
|
||||
|
||||
expect(screen.getByRole("link", { name: verificationUri })).toHaveAttribute("href", verificationUri);
|
||||
});
|
||||
|
||||
it("stops polling and offers a restart when the device code expires", async () => {
|
||||
networking.userConnectionStartCall.mockResolvedValue({ ...START_RESPONSE, expires_in: 12 });
|
||||
networking.userConnectionPollCall.mockResolvedValue({ status: "pending" });
|
||||
|
|
|
|||
|
|
@ -30,11 +30,29 @@ describe("connectFlowReducer", () => {
|
|||
expect(started()).toEqual(expected);
|
||||
});
|
||||
|
||||
it("never links anywhere but github.com", () => {
|
||||
it("falls back for an untrusted verification host", () => {
|
||||
const state = started("https://evil.example/login/device");
|
||||
expect(state.kind === "awaiting" && state.verificationUri).toBe("https://github.com/login/device");
|
||||
});
|
||||
|
||||
it("preserves a single-label GHE.com verification URL", () => {
|
||||
const uri = "https://octo.ghe.com/login/device";
|
||||
const state = started(uri);
|
||||
expect(state.kind === "awaiting" && state.verificationUri).toBe(uri);
|
||||
});
|
||||
|
||||
it.each([
|
||||
"http://octo.ghe.com/login/device",
|
||||
"https://a.b.ghe.com/login/device",
|
||||
"https://octo.ghe.com.evil.example/login/device",
|
||||
"https://user@octo.ghe.com/login/device",
|
||||
"https://octo.ghe.com:8443/login/device",
|
||||
"not a url",
|
||||
])("falls back for unsafe verification URL %s", (uri) => {
|
||||
const state = started(uri);
|
||||
expect(state.kind === "awaiting" && state.verificationUri).toBe("https://github.com/login/device");
|
||||
});
|
||||
|
||||
it("keeps waiting on pending and slows down by the server interval or by five seconds", () => {
|
||||
const pending = connectFlowReducer(started(), { type: "polled", response: { status: "pending" } });
|
||||
expect(pending).toMatchObject({ kind: "awaiting", pollCount: 1, intervalSeconds: 5 });
|
||||
|
|
|
|||
|
|
@ -34,8 +34,22 @@ export type ConnectFlowEvent =
|
|||
|
||||
export const initialConnectFlowState: ConnectFlowState = { kind: "starting", attempt: 0 };
|
||||
|
||||
const safeVerificationUri = (uri: string): string =>
|
||||
uri.startsWith("https://github.com/") ? uri : GITHUB_DEVICE_LOGIN_URL;
|
||||
const safeVerificationUri = (uri: string): string => {
|
||||
try {
|
||||
const url = new URL(uri);
|
||||
const usesHttps = url.protocol === "https:";
|
||||
const hasCredentials = url.username !== "" || url.password !== "";
|
||||
const hasNoPort = url.port === "";
|
||||
const hostIsAllowed =
|
||||
url.hostname === "github.com" || /^[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?\.ghe\.com$/.test(url.hostname);
|
||||
const urlHasAllowedAuthority = !hasCredentials && hasNoPort && hostIsAllowed;
|
||||
const isSafeUrl = usesHttps && urlHasAllowedAuthority;
|
||||
|
||||
return isSafeUrl ? uri : GITHUB_DEVICE_LOGIN_URL;
|
||||
} catch {
|
||||
return GITHUB_DEVICE_LOGIN_URL;
|
||||
}
|
||||
};
|
||||
|
||||
const afterPoll = (
|
||||
state: Extract<ConnectFlowState, { kind: "awaiting" }>,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue