From db2039b11b3deb902377e47fcc3e1074b996a722 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Tue, 20 Oct 2026 14:16:35 -0700 Subject: [PATCH] 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> --- litellm/constants.py | 2 + .../litellm_core_utils/get_litellm_params.py | 1 + litellm/llms/github_copilot/github_host.py | 78 +++ litellm/llms/github_copilot/per_user_auth.py | 211 +++++--- litellm/main.py | 1 + litellm/models/credentials.py | 2 + litellm/proxy/auth/auth_utils.py | 1 + .../proxy/credential_endpoints/endpoints.py | 108 +++- .../user_provider_credentials.py | 30 +- litellm/proxy/litellm_pre_call_utils.py | 6 +- .../provider_create_fields.json | 10 + litellm/types/litellm_params.py | 1 + litellm/types/proxy/litellm_pre_call_utils.py | 2 +- litellm/types/utils.py | 2 +- .../test_get_litellm_params.py | 4 + .../llms/github_copilot/test_github_host.py | 59 +++ .../llms/github_copilot/test_per_user_auth.py | 237 +++++++-- tests/unit/proxy/auth/test_auth_utils.py | 2 +- .../credential_endpoints/test_endpoints.py | 469 +++++++++++++++++- .../test_user_provider_credentials.py | 118 +++-- .../unit/proxy/test_litellm_pre_call_utils.py | 29 +- tests/unit/types/test_litellm_params.py | 1 + tests/unit/types/test_router.py | 2 +- .../add_model/provider_auth_types.test.ts | 6 +- .../add_model/provider_auth_types.ts | 3 +- .../CredentialModal.integration.test.tsx | 12 +- ...ubCopilotConnectModal.integration.test.tsx | 12 + .../github_copilot_connect_flow.test.ts | 20 +- .../model_add/github_copilot_connect_flow.ts | 18 +- 29 files changed, 1233 insertions(+), 214 deletions(-) create mode 100644 litellm/llms/github_copilot/github_host.py create mode 100644 tests/unit/llms/github_copilot/test_github_host.py diff --git a/litellm/constants.py b/litellm/constants.py index cb50d74f0cf..919122d6e4d 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -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" diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index 657fd3abfbc..85557090ae9 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -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", diff --git a/litellm/llms/github_copilot/github_host.py b/litellm/llms/github_copilot/github_host.py new file mode 100644 index 00000000000..4fc921d490d --- /dev/null +++ b/litellm/llms/github_copilot/github_host.py @@ -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}", + ) diff --git a/litellm/llms/github_copilot/per_user_auth.py b/litellm/llms/github_copilot/per_user_auth.py index 3dc8c12fe06..75a10f2e671 100644 --- a/litellm/llms/github_copilot/per_user_auth.py +++ b/litellm/llms/github_copilot/per_user_auth.py @@ -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) diff --git a/litellm/main.py b/litellm/main.py index 63375292077..d647cab24e3 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -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", diff --git a/litellm/models/credentials.py b/litellm/models/credentials.py index d569ea12617..e6faf461104 100644 --- a/litellm/models/credentials.py +++ b/litellm/models/credentials.py @@ -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 diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 996c28bebaa..68cd8420166 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -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 diff --git a/litellm/proxy/credential_endpoints/endpoints.py b/litellm/proxy/credential_endpoints/endpoints.py index 953c944bfb2..d20b5602186 100644 --- a/litellm/proxy/credential_endpoints/endpoints.py +++ b/litellm/proxy/credential_endpoints/endpoints.py @@ -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"} diff --git a/litellm/proxy/credential_endpoints/user_provider_credentials.py b/litellm/proxy/credential_endpoints/user_provider_credentials.py index 4253483a228..22ee271410e 100644 --- a/litellm/proxy/credential_endpoints/user_provider_credentials.py +++ b/litellm/proxy/credential_endpoints/user_provider_credentials.py @@ -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} diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index d69778136ee..f8e6498502b 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -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 diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index a794854ce19..4bec42a7868 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -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" diff --git a/litellm/types/litellm_params.py b/litellm/types/litellm_params.py index 6b08920e042..5f318cc3c69 100644 --- a/litellm/types/litellm_params.py +++ b/litellm/types/litellm_params.py @@ -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 diff --git a/litellm/types/proxy/litellm_pre_call_utils.py b/litellm/types/proxy/litellm_pre_call_utils.py index 0b980d38320..f0dfd50bfbe 100644 --- a/litellm/types/proxy/litellm_pre_call_utils.py +++ b/litellm/types/proxy/litellm_pre_call_utils.py @@ -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]] diff --git a/litellm/types/utils.py b/litellm/types/utils.py index e268cc90ffd..61a179b8b65 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -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 " diff --git a/tests/unit/litellm_core_utils/test_get_litellm_params.py b/tests/unit/litellm_core_utils/test_get_litellm_params.py index 251bfa624a5..89173956dc6 100644 --- a/tests/unit/litellm_core_utils/test_get_litellm_params.py +++ b/tests/unit/litellm_core_utils/test_get_litellm_params.py @@ -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.""" diff --git a/tests/unit/llms/github_copilot/test_github_host.py b/tests/unit/llms/github_copilot/test_github_host.py new file mode 100644 index 00000000000..df2efe4669b --- /dev/null +++ b/tests/unit/llms/github_copilot/test_github_host.py @@ -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" diff --git a/tests/unit/llms/github_copilot/test_per_user_auth.py b/tests/unit/llms/github_copilot/test_per_user_auth.py index 9b16d3c0aa7..cf6fb9fa71f 100644 --- a/tests/unit/llms/github_copilot/test_per_user_auth.py +++ b/tests/unit/llms/github_copilot/test_per_user_auth.py @@ -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 diff --git a/tests/unit/proxy/auth/test_auth_utils.py b/tests/unit/proxy/auth/test_auth_utils.py index 67ff7dd9344..7597847ef6d 100644 --- a/tests/unit/proxy/auth/test_auth_utils.py +++ b/tests/unit/proxy/auth/test_auth_utils.py @@ -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 diff --git a/tests/unit/proxy/credential_endpoints/test_endpoints.py b/tests/unit/proxy/credential_endpoints/test_endpoints.py index 6aeceaf4356..c905d38dd41 100644 --- a/tests/unit/proxy/credential_endpoints/test_endpoints.py +++ b/tests/unit/proxy/credential_endpoints/test_endpoints.py @@ -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.""" diff --git a/tests/unit/proxy/credential_endpoints/test_user_provider_credentials.py b/tests/unit/proxy/credential_endpoints/test_user_provider_credentials.py index b549ecc67a9..dae3a0459e2 100644 --- a/tests/unit/proxy/credential_endpoints/test_user_provider_credentials.py +++ b/tests/unit/proxy/credential_endpoints/test_user_provider_credentials.py @@ -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() diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 246496ec48f..c2a5b01fe4e 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -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 diff --git a/tests/unit/types/test_litellm_params.py b/tests/unit/types/test_litellm_params.py index 84ba85f4a68..788d405f99f 100644 --- a/tests/unit/types/test_litellm_params.py +++ b/tests/unit/types/test_litellm_params.py @@ -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", diff --git a/tests/unit/types/test_router.py b/tests/unit/types/test_router.py index e15e95f9d55..2892a829eab 100644 --- a/tests/unit/types/test_router.py +++ b/tests/unit/types/test_router.py @@ -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 diff --git a/ui/litellm-dashboard/src/components/add_model/provider_auth_types.test.ts b/ui/litellm-dashboard/src/components/add_model/provider_auth_types.test.ts index 5467ff0085a..f0b739b877d 100644 --- a/ui/litellm-dashboard/src/components/add_model/provider_auth_types.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/provider_auth_types.test.ts @@ -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", + ]); }); }); diff --git a/ui/litellm-dashboard/src/components/add_model/provider_auth_types.ts b/ui/litellm-dashboard/src/components/add_model/provider_auth_types.ts index 6fe54955133..0bf98407bfa 100644 --- a/ui/litellm-dashboard/src/components/add_model/provider_auth_types.ts +++ b/ui/litellm-dashboard/src/components/add_model/provider_auth_types.ts @@ -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> = { @@ -31,7 +32,7 @@ export const PROVIDER_AUTH_TYPES: Partial { 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"]); }); }); diff --git a/ui/litellm-dashboard/src/components/model_add/GithubCopilotConnectModal.integration.test.tsx b/ui/litellm-dashboard/src/components/model_add/GithubCopilotConnectModal.integration.test.tsx index 963ebab81d3..42f30db08a8 100644 --- a/ui/litellm-dashboard/src/components/model_add/GithubCopilotConnectModal.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/model_add/GithubCopilotConnectModal.integration.test.tsx @@ -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" }); diff --git a/ui/litellm-dashboard/src/components/model_add/github_copilot_connect_flow.test.ts b/ui/litellm-dashboard/src/components/model_add/github_copilot_connect_flow.test.ts index 72c2e783adf..9645839bbb4 100644 --- a/ui/litellm-dashboard/src/components/model_add/github_copilot_connect_flow.test.ts +++ b/ui/litellm-dashboard/src/components/model_add/github_copilot_connect_flow.test.ts @@ -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 }); diff --git a/ui/litellm-dashboard/src/components/model_add/github_copilot_connect_flow.ts b/ui/litellm-dashboard/src/components/model_add/github_copilot_connect_flow.ts index 0e38fc91bb1..6f0e76f9aaf 100644 --- a/ui/litellm-dashboard/src/components/model_add/github_copilot_connect_flow.ts +++ b/ui/litellm-dashboard/src/components/model_add/github_copilot_connect_flow.ts @@ -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,