From 8336ae3604d09a9e4610697e5e81dd795d0cc936 Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 21 Oct 2026 00:19:17 +0000 Subject: [PATCH] feat(github_copilot): refresh expiring per-user GitHub tokens Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 4 + litellm/llms/github_copilot/per_user_auth.py | 110 ++- .../proxy/credential_endpoints/endpoints.py | 8 + .../github_copilot_token_refresh.py | 357 +++++++++ .../user_provider_credentials.py | 83 +- litellm/proxy/litellm_pre_call_utils.py | 40 +- .../redis_refresh_coordinator.py | 18 +- .../llms/github_copilot/test_per_user_auth.py | 139 +++- .../credential_endpoints/test_endpoints.py | 48 +- .../test_github_copilot_token_refresh.py | 753 ++++++++++++++++++ .../test_user_provider_credentials.py | 106 +++ .../unit/proxy/test_litellm_pre_call_utils.py | 20 +- 12 files changed, 1644 insertions(+), 42 deletions(-) create mode 100644 litellm/proxy/credential_endpoints/github_copilot_token_refresh.py create mode 100644 tests/unit/proxy/credential_endpoints/test_github_copilot_token_refresh.py diff --git a/litellm/constants.py b/litellm/constants.py index 8e80c07b4ed..56e3c90849a 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -30,6 +30,10 @@ 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_GITHUB_TOKEN_REFRESH_MARGIN_SECONDS: Final = 300 +GITHUB_COPILOT_TOKEN_REFRESH_LOCK_PREFIX: Final = "github_copilot_token_refresh_lock:" +GITHUB_COPILOT_TOKEN_REFRESH_REREAD_SECONDS: Final = 2.0 +GITHUB_COPILOT_TOKEN_REFRESH_REREAD_INTERVAL_SECONDS: Final = 0.1 GITHUB_COPILOT_DEVICE_FLOW_CACHE_PREFIX: Final = "github_copilot_device_flow" USER_PROVIDER_CREDENTIAL_CACHE_PREFIX: Final = "user_provider_credential" USER_PROVIDER_CREDENTIAL_NOT_CONNECTED: Final = "__not_connected__" diff --git a/litellm/llms/github_copilot/per_user_auth.py b/litellm/llms/github_copilot/per_user_auth.py index 75a10f2e671..67b3aab2f7c 100644 --- a/litellm/llms/github_copilot/per_user_auth.py +++ b/litellm/llms/github_copilot/per_user_auth.py @@ -6,6 +6,7 @@ touched on the per-user path. import asyncio import hmac +import json import threading import time from collections.abc import Callable, Mapping @@ -48,6 +49,10 @@ _SESSION_CACHE_MAX_SIZE: Final = 1000 _SYNC_LOCK_COUNT: Final = 64 +class GithubCopilotTokenRejectedError(CallerCredentialAuthenticationError): + pass + + class _SessionCache(Protocol): def get_cache(self, key: str) -> object: ... @@ -115,6 +120,9 @@ class _DeviceCodePayload(TypedDict): @with_config(ConfigDict(extra="allow", strict=True)) class _AccessTokenPollPayload(TypedDict): access_token: ReadOnly[NotRequired[str]] + refresh_token: ReadOnly[NotRequired[str]] + expires_in: ReadOnly[NotRequired[object]] + refresh_token_expires_in: ReadOnly[NotRequired[object]] error: ReadOnly[NotRequired[object]] error_description: ReadOnly[NotRequired[object]] interval: ReadOnly[NotRequired[object]] @@ -170,6 +178,30 @@ class GithubCopilotDeviceFlowPoll: status: DeviceFlowPollStatus interval: int | None = None access_token: str | None = field(default=None, repr=False) + refresh_token: str | None = field(default=None, repr=False) + expires_in: int | None = None + refresh_token_expires_in: int | None = None + + +@dataclass(frozen=True, slots=True) +class GithubTokenRefreshed: + access_token: str = field(repr=False) + refresh_token: str | None = field(repr=False) + expires_in: int | None + refresh_token_expires_in: int | None + + +@dataclass(frozen=True, slots=True) +class GithubTokenRefreshRejected: + pass + + +@dataclass(frozen=True, slots=True) +class GithubTokenRefreshUnavailable: + pass + + +GithubTokenRefreshResult: TypeAlias = GithubTokenRefreshed | GithubTokenRefreshRejected | GithubTokenRefreshUnavailable def validated_copilot_api_base(endpoints_api: object, endpoints: GithubHostEndpoints) -> str: @@ -240,8 +272,11 @@ def _cache_session(cache_key: str, github_token: str, session: GithubCopilotUser _SESSION_CACHE.set_cache(cache_key, (github_token, session.token, session.api_base, expires_at), ttl=ttl) -def _reject_connection(credential_name: str) -> CallerCredentialAuthenticationError: - return CallerCredentialAuthenticationError( +def _reject_connection( + credential_name: str, + exception_type: type[CallerCredentialAuthenticationError] = CallerCredentialAuthenticationError, +) -> CallerCredentialAuthenticationError: + return exception_type( message=( f"GitHub rejected the GitHub Copilot connection for credential '{credential_name}'. " f"Reconnect GitHub Copilot for credential '{credential_name}' in the LiteLLM UI (LLM Credentials)" @@ -260,7 +295,10 @@ def _session_from_response( ) -> GithubCopilotUserSession: if response.status_code in (401, 403, 404): _SESSION_CACHE.delete_cache(cache_key) - raise _reject_connection(credential_name) + raise _reject_connection( + credential_name, + GithubCopilotTokenRejectedError if response.status_code == 401 else CallerCredentialAuthenticationError, + ) if response.status_code == 429: raise CallerCredentialRateLimitError( message="GitHub rate limited the GitHub Copilot token exchange", @@ -311,6 +349,12 @@ def _to_int(raw: object) -> int: return 0 +def _positive_int(raw: object) -> int | None: + if isinstance(raw, int) and not isinstance(raw, bool) and raw > 0: + return raw + return None + + def _expires_at_seconds(raw: object) -> float: if isinstance(raw, bool): return 0.0 @@ -512,7 +556,14 @@ def _poll_payload_from_response(response: httpx.Response) -> GithubCopilotDevice access_token: Final = payload.get("access_token") if isinstance(access_token, str) and access_token: - return GithubCopilotDeviceFlowPoll(status="connected", access_token=access_token) + raw_refresh_token: Final = payload.get("refresh_token") + return GithubCopilotDeviceFlowPoll( + status="connected", + access_token=access_token, + refresh_token=raw_refresh_token if isinstance(raw_refresh_token, str) and raw_refresh_token else None, + expires_in=_positive_int(payload.get("expires_in")), + refresh_token_expires_in=_positive_int(payload.get("refresh_token_expires_in")), + ) error: Final = payload.get("error") raw_interval: Final = payload.get("interval") @@ -534,6 +585,57 @@ def _poll_payload_from_response(response: httpx.Response) -> GithubCopilotDevice ) +async def arefresh_github_token( + refresh_token: str, + endpoints: GithubHostEndpoints, +) -> GithubTokenRefreshResult: + client: Final = _new_async_github_client() + try: + response: Final = await client.post( + endpoints.access_token_url, + headers=github_api_headers(), + json={ + "client_id": endpoints.client_id, + "grant_type": "refresh_token", + "refresh_token": refresh_token, + }, + ) + except httpx.HTTPStatusError as exc: + status: Final = exc.response.status_code + if status == 429 or status >= 500: + return GithubTokenRefreshUnavailable() + if 400 <= status < 500: + return GithubTokenRefreshRejected() + return GithubTokenRefreshUnavailable() + except (httpx.RequestError, APIConnectionError): + return GithubTokenRefreshUnavailable() + + if response.status_code == 429 or response.status_code >= 500: + return GithubTokenRefreshUnavailable() + if 400 <= response.status_code < 500: + return GithubTokenRefreshRejected() + if not response.is_success: + return GithubTokenRefreshUnavailable() + try: + payload: Final = _ACCESS_TOKEN_POLL.validate_python(response.json()) + except json.JSONDecodeError: + return GithubTokenRefreshUnavailable() + except ValidationError: + return GithubTokenRefreshUnavailable() + if payload.get("error") is not None: + return GithubTokenRefreshRejected() + access_token: Final = payload.get("access_token") + if not isinstance(access_token, str) or not access_token: + return GithubTokenRefreshRejected() + raw_refresh_token: Final = payload.get("refresh_token") + return GithubTokenRefreshed( + access_token=access_token, + refresh_token=raw_refresh_token if isinstance(raw_refresh_token, str) and raw_refresh_token else None, + expires_in=_positive_int(payload.get("expires_in")), + refresh_token_expires_in=_positive_int(payload.get("refresh_token_expires_in")), + ) + + async def afetch_github_login(github_token: str, endpoints: GithubHostEndpoints) -> str: client: Final = _new_async_github_client() try: diff --git a/litellm/proxy/credential_endpoints/endpoints.py b/litellm/proxy/credential_endpoints/endpoints.py index 5b3507f9b6d..d22df5268f7 100644 --- a/litellm/proxy/credential_endpoints/endpoints.py +++ b/litellm/proxy/credential_endpoints/endpoints.py @@ -750,10 +750,14 @@ async def poll_user_connection( except litellm.CallerCredentialAuthenticationError: return UserConnectionPollResponse(status="no_copilot_seat") github_login: Final = await afetch_github_login(github_token, endpoints) + now: Final = time.time() connection_payload: Final = GithubCopilotUserConnectionPayload( access_token=github_token, github_login=github_login, github_host=handle.github_host, + refresh_token=poll.refresh_token, + access_token_expires_at=now + poll.expires_in if poll.expires_in else None, + refresh_token_expires_at=(now + poll.refresh_token_expires_in if poll.refresh_token_expires_in else None), ) await upsert_user_provider_credential( prisma_client, @@ -814,6 +818,10 @@ async def delete_user_connection( detail={"error": "could not revoke cached connection, retry"}, ) prior: Final = await delete_user_provider_credential(prisma_client, user_id, credential_name) + if not await invalidate_user_provider_credential_cache(user_api_key_cache, user_id, credential_name): + verbose_proxy_logger.warning( + "delete_user_connection: failed to refresh the cache tombstone after deleting the connection" + ) if prior is not None: evict_copilot_user_session(user_id, credential_name, prior.github_host) return UserConnectionDeleteResponse(status="disconnected") diff --git a/litellm/proxy/credential_endpoints/github_copilot_token_refresh.py b/litellm/proxy/credential_endpoints/github_copilot_token_refresh.py new file mode 100644 index 00000000000..d1acb1f75fe --- /dev/null +++ b/litellm/proxy/credential_endpoints/github_copilot_token_refresh.py @@ -0,0 +1,357 @@ +import asyncio +import time +from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass, field +from typing import Final, Protocol, TypeAlias, cast # noqa: TID251 # RedisCache is narrowed to its async Protocol + +from litellm.caching.dual_cache import DualCache +from litellm.constants import ( + GITHUB_COPILOT_GITHUB_TOKEN_REFRESH_MARGIN_SECONDS, + GITHUB_COPILOT_TOKEN_REFRESH_LOCK_PREFIX, + GITHUB_COPILOT_TOKEN_REFRESH_REREAD_INTERVAL_SECONDS, + GITHUB_COPILOT_TOKEN_REFRESH_REREAD_SECONDS, +) +from litellm.exceptions import BadRequestError +from litellm.llms.github_copilot.github_host import GithubHostEndpoints +from litellm.llms.github_copilot.per_user_auth import ( + GithubCopilotTokenRejectedError, + GithubTokenRefreshed, + GithubTokenRefreshRejected, + GithubTokenRefreshUnavailable, + aexchange_github_token, + arefresh_github_token, + github_copilot_host_endpoints, +) +from litellm.proxy.credential_endpoints.user_provider_credentials import ( + GithubCopilotUserConnectionPayload, + compare_and_set_user_provider_credential, + drop_user_provider_credential_cache, + get_user_provider_credential_row, + invalidate_user_provider_credential_cache, + set_user_provider_credential_cache, +) +from litellm.proxy.mcp_server.outbound_credentials.redis_distributed_lock import ( + RedisCommands, + RedisDistributedLock, +) +from litellm.proxy.mcp_server.outbound_credentials.redis_refresh_coordinator import ( + RedisRefreshCoordinator, +) +from litellm.repositories.prisma_protocols import DatabaseClient + +ConnectionPayload: TypeAlias = GithubCopilotUserConnectionPayload +_InFlightTasks: TypeAlias = dict[ + tuple[str, str], + asyncio.Task[ConnectionPayload | None], +] + + +class _RedisCache(Protocol): + def init_async_client(self) -> RedisCommands: ... + + def check_and_fix_namespace(self, key: str) -> str: ... + + +@dataclass(slots=True) +class GithubCopilotTokenRefreshSingleFlight: + _lock: asyncio.Lock = field(default_factory=asyncio.Lock) + _in_flight: _InFlightTasks = field( + default_factory=dict, + repr=False, + ) + + async def run( + self, + key: tuple[str, str], + refresh: Callable[[], Awaitable[ConnectionPayload | None]], + ) -> ConnectionPayload | None: + task: Final = await self._get_or_create_task(key, refresh) + return await asyncio.shield(task) + + async def _get_or_create_task( + self, + key: tuple[str, str], + refresh: Callable[[], Awaitable[ConnectionPayload | None]], + ) -> asyncio.Task[ConnectionPayload | None]: + async with self._lock: + task: Final = self._in_flight.get(key) + if task is not None and not task.done(): + return task + + async def execute_refresh() -> ConnectionPayload | None: + return await refresh() + + created_task: Final = asyncio.create_task(execute_refresh()) + self._in_flight[key] = created_task + + def remove_completed_task(completed_task: asyncio.Task[ConnectionPayload | None]) -> None: + if self._in_flight.get(key) is completed_task: + self._in_flight.pop(key, None) + + created_task.add_done_callback(remove_completed_task) + return created_task + + +_SINGLE_FLIGHT: Final = GithubCopilotTokenRefreshSingleFlight() + + +async def arefresh_expiring_connections( + prisma_client: DatabaseClient, + cache: DualCache, + user_id: str, + connections: Mapping[str, ConnectionPayload], + *, + single_flight: GithubCopilotTokenRefreshSingleFlight | None = None, + sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, + clock: Callable[[], float] = time.monotonic, +) -> Mapping[str, ConnectionPayload]: + refresher: Final = single_flight if single_flight is not None else _SINGLE_FLIGHT + items: Final = tuple(connections.items()) + refreshed: Final = await asyncio.gather( + *( + _refresh_connection( + prisma_client=prisma_client, + cache=cache, + user_id=user_id, + credential_name=name, + stale=payload, + single_flight=refresher, + sleep=sleep, + clock=clock, + ) + for name, payload in items + ) + ) + return {name: payload for (name, _), payload in zip(items, refreshed) if payload is not None} + + +async def _refresh_connection( + prisma_client: DatabaseClient, + cache: DualCache, + user_id: str, + credential_name: str, + stale: ConnectionPayload, + single_flight: GithubCopilotTokenRefreshSingleFlight, + sleep: Callable[[float], Awaitable[None]], + clock: Callable[[], float], +) -> ConnectionPayload | None: + now: Final = time.time() + if stale.refresh_token is None: + return stale + if stale.refresh_token_expires_at is not None and stale.refresh_token_expires_at <= now: + return stale + try: + endpoints: Final = github_copilot_host_endpoints(credential_name) + except BadRequestError: + return stale + if stale.github_host != endpoints.host: + return stale + if stale.access_token_expires_at is None or ( + stale.access_token_expires_at - GITHUB_COPILOT_GITHUB_TOKEN_REFRESH_MARGIN_SECONDS > now + ): + try: + await aexchange_github_token( + user_id=user_id, + github_token=stale.access_token, + credential_name=credential_name, + endpoints=endpoints, + ) + except GithubCopilotTokenRejectedError: + pass + except Exception: # noqa: BLE001 # other Copilot exchange errors stay on the existing request path + return stale + else: + return stale + key: Final = (user_id, credential_name) + return await single_flight.run( + key, + lambda: _coordinate_refresh( + prisma_client=prisma_client, + cache=cache, + user_id=user_id, + credential_name=credential_name, + stale_access_token=stale.access_token, + endpoints=endpoints, + sleep=sleep, + clock=clock, + ), + ) + + +async def _coordinate_refresh( + prisma_client: DatabaseClient, + cache: DualCache, + user_id: str, + credential_name: str, + stale_access_token: str, + endpoints: GithubHostEndpoints, + sleep: Callable[[float], Awaitable[None]], + clock: Callable[[], float], +) -> ConnectionPayload | None: + async def refresh() -> ConnectionPayload | None: + return await _refresh_elected( + prisma_client=prisma_client, + cache=cache, + user_id=user_id, + credential_name=credential_name, + stale_access_token=stale_access_token, + endpoints=endpoints, + sleep=sleep, + clock=clock, + ) + + redis_cache_raw: Final = cache.redis_cache + if redis_cache_raw is None: + return await refresh() + try: + redis_cache: Final = cast( # cast-ok: runtime RedisCache exposes async methods omitted from its protocol + _RedisCache, redis_cache_raw + ) + lock: Final = RedisDistributedLock( + redis_cache.init_async_client(), + namespace_key=redis_cache.check_and_fix_namespace, + ) + coordinator: Final = RedisRefreshCoordinator( + lock, + key_prefix=GITHUB_COPILOT_TOKEN_REFRESH_LOCK_PREFIX, + ) + except Exception: # noqa: BLE001 # Redis setup failures fail open to an uncoordinated refresh + return await refresh() + return await coordinator.run( + user_id=user_id, + server_id=credential_name, + refresh=refresh, + reread=lambda: _read_connection(prisma_client, user_id, credential_name), + ) + + +async def _refresh_elected( + prisma_client: DatabaseClient, + cache: DualCache, + user_id: str, + credential_name: str, + stale_access_token: str, + endpoints: GithubHostEndpoints, + sleep: Callable[[float], Awaitable[None]], + clock: Callable[[], float], +) -> ConnectionPayload | None: + row: Final = await get_user_provider_credential_row(prisma_client, user_id, credential_name) + if row is None: + return None + ciphertext, payload = row + if payload.access_token != stale_access_token or payload.refresh_token is None: + return payload + result: Final = await arefresh_github_token(payload.refresh_token, endpoints) + match result: + case GithubTokenRefreshRejected(): + return await _reread_after_rejection( + prisma_client=prisma_client, + cache=cache, + user_id=user_id, + credential_name=credential_name, + ciphertext=ciphertext, + stale_access_token=stale_access_token, + stale=payload, + sleep=sleep, + clock=clock, + ) + case GithubTokenRefreshUnavailable(): + return payload + case GithubTokenRefreshed(): + now: Final = time.time() + refreshed: Final = payload.model_copy( + update={ + "access_token": result.access_token, + "refresh_token": result.refresh_token, + "access_token_expires_at": (now + result.expires_in if result.expires_in is not None else None), + "refresh_token_expires_at": ( + now + result.refresh_token_expires_in if result.refresh_token_expires_in is not None else None + ), + } + ) + return await _persist_connection( + prisma_client=prisma_client, + cache=cache, + user_id=user_id, + credential_name=credential_name, + expected_ciphertext=ciphertext, + payload=refreshed, + ) + + +async def _persist_connection( + prisma_client: DatabaseClient, + cache: DualCache, + user_id: str, + credential_name: str, + expected_ciphertext: str, + payload: ConnectionPayload, +) -> ConnectionPayload | None: + if not await compare_and_set_user_provider_credential( + prisma_client=prisma_client, + user_id=user_id, + credential_name=credential_name, + expected_ciphertext=expected_ciphertext, + payload=payload, + ): + return await _read_connection(prisma_client, user_id, credential_name) + await set_user_provider_credential_cache(cache, user_id, credential_name, payload) + current: Final = await _read_connection(prisma_client, user_id, credential_name) + if current is None: + await invalidate_user_provider_credential_cache(cache, user_id, credential_name) + return None + if current != payload: + await drop_user_provider_credential_cache(cache, user_id, credential_name) + return current + return payload + + +async def _reread_after_rejection( + prisma_client: DatabaseClient, + cache: DualCache, + user_id: str, + credential_name: str, + ciphertext: str, + stale_access_token: str, + stale: ConnectionPayload, + sleep: Callable[[float], Awaitable[None]], + clock: Callable[[], float], +) -> ConnectionPayload | None: + deadline: Final = clock() + GITHUB_COPILOT_TOKEN_REFRESH_REREAD_SECONDS + + async def read_changed_connection() -> tuple[bool, ConnectionPayload | None]: + current: Final = await _read_connection(prisma_client, user_id, credential_name) + if current is None: + return True, None + if current.access_token != stale_access_token: + return True, current + return False, None + + while True: + is_changed, current = await read_changed_connection() + if is_changed: + return current + if deadline <= clock(): + return await _persist_connection( + prisma_client=prisma_client, + cache=cache, + user_id=user_id, + credential_name=credential_name, + expected_ciphertext=ciphertext, + payload=stale.model_copy(update={"refresh_token": None, "refresh_token_expires_at": None}), + ) + await sleep( + min( + GITHUB_COPILOT_TOKEN_REFRESH_REREAD_INTERVAL_SECONDS, + max(0.0, deadline - clock()), + ) + ) + + +async def _read_connection( + prisma_client: DatabaseClient, + user_id: str, + credential_name: str, +) -> ConnectionPayload | None: + row: Final = await get_user_provider_credential_row(prisma_client, user_id, credential_name) + return row[1] if row is not None else None diff --git a/litellm/proxy/credential_endpoints/user_provider_credentials.py b/litellm/proxy/credential_endpoints/user_provider_credentials.py index 22ee271410e..0cd8c11eede 100644 --- a/litellm/proxy/credential_endpoints/user_provider_credentials.py +++ b/litellm/proxy/credential_endpoints/user_provider_credentials.py @@ -35,7 +35,7 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( encrypt_value_helper, ) from litellm.repositories.chunked_in import find_many_in -from litellm.repositories.prisma_protocols import TableActions +from litellm.repositories.prisma_protocols import DatabaseClient, TableActions from litellm.repositories.table_repositories import PrismaTableRepository if TYPE_CHECKING: @@ -52,6 +52,9 @@ class GithubCopilotUserConnectionPayload(BaseModel): access_token: str = Field(repr=False) github_login: str github_host: str = GITHUB_COPILOT_DEFAULT_HOST + refresh_token: str | None = Field(default=None, repr=False) + access_token_expires_at: float | None = None + refresh_token_expires_at: float | None = None class _UserProviderCredentialsRepository(PrismaTableRepository["prisma_models.LiteLLM_UserProviderCredentials"]): @@ -59,9 +62,12 @@ class _UserProviderCredentialsRepository(PrismaTableRepository["prisma_models.Li def _table( - prisma_client: "PrismaClient", + prisma_client: DatabaseClient, ) -> TableActions["prisma_models.LiteLLM_UserProviderCredentials"]: - return _UserProviderCredentialsRepository(prisma_client, use_writer=True).table + repository_client: Final = cast( # cast-ok: the repository constructor retains its legacy client type + "PrismaClient", prisma_client + ) + return _UserProviderCredentialsRepository(repository_client, use_writer=True).table def _cache_key(user_id: str, credential_name: str) -> str: @@ -122,6 +128,38 @@ async def get_user_provider_credential( return decode_user_provider_credential(row.credential_b64) +async def get_user_provider_credential_row( + prisma_client: DatabaseClient, + user_id: str, + credential_name: str, +) -> tuple[str, GithubCopilotUserConnectionPayload] | None: + row: Final = await _table(prisma_client).find_unique( + where={"user_id_credential_name": {"user_id": user_id, "credential_name": credential_name}} + ) + if row is None: + return None + payload: Final = decode_user_provider_credential(row.credential_b64) + return (row.credential_b64, payload) if payload is not None else None + + +async def compare_and_set_user_provider_credential( + prisma_client: DatabaseClient, + user_id: str, + credential_name: str, + expected_ciphertext: str, + payload: GithubCopilotUserConnectionPayload, +) -> bool: + updated_count: Final = await _table(prisma_client).update_many( + where={ + "user_id": user_id, + "credential_name": credential_name, + "credential_b64": expected_ciphertext, + }, + data={"credential_b64": _encode(payload)}, + ) + return updated_count == 1 + + async def delete_user_provider_credential( prisma_client: "PrismaClient", user_id: str, @@ -294,13 +332,13 @@ async def set_user_provider_credential_cache( @with_service_target("user_provider_connections") -async def aget_user_provider_credentials( - prisma_client: "PrismaClient", +async def aget_user_provider_connections( + prisma_client: DatabaseClient, cache: DualCache, user_id: str, credential_names: Sequence[str], -) -> Mapping[str, Mapping[str, str]]: - """Map each per-user credential name to the caller's stored GitHub connection. +) -> Mapping[str, GithubCopilotUserConnectionPayload]: + """Map each per-user credential name to the caller's stored GitHub connection payload. Cached values are the stored ciphertext or the ``_NOT_CONNECTED`` marker; plaintext connection data only ever lives in the returned mapping. Redis is the only @@ -321,11 +359,7 @@ async def aget_user_provider_credentials( else: misses.append(name) if not misses: - return { - name: connection - for name, ciphertext in cached.items() - if (connection := _connection_from_ciphertext(ciphertext)) is not None - } + return dict(item for item in map(_decode_connection_item, cached.items()) if item is not None) try: rows: Final = await find_many_in( @@ -343,15 +377,30 @@ async def aget_user_provider_credentials( await _try_cache_set(token_cache, _cache_key(user_id, name), found.get(name, _NOT_CONNECTED), nx=True) if name in found: cached[name] = found[name] + return dict(item for item in map(_decode_connection_item, cached.items()) if item is not None) + + +async def aget_user_provider_credentials( + prisma_client: "PrismaClient", + cache: DualCache, + user_id: str, + credential_names: Sequence[str], +) -> Mapping[str, Mapping[str, str]]: + connections: Final = await aget_user_provider_connections( + prisma_client=prisma_client, + cache=cache, + user_id=user_id, + credential_names=credential_names, + ) return { - name: connection - for name, ciphertext in cached.items() - if (connection := _connection_from_ciphertext(ciphertext)) is not None + name: {"access_token": payload.access_token, "github_host": payload.github_host} + for name, payload in connections.items() } -def _connection_from_ciphertext(ciphertext: str) -> Mapping[str, str] | None: +def _decode_connection_item(item: tuple[str, str]) -> tuple[str, GithubCopilotUserConnectionPayload] | None: + name, ciphertext = item payload: Final = decode_user_provider_credential(ciphertext) if payload is None: return None - return {"access_token": payload.access_token, "github_host": payload.github_host} + return name, payload diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index c1e541d20a9..6c7bb1a5f60 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -342,12 +342,15 @@ _ENABLE_TEAM_STALE_ALIAS_BYPASS: bool | None = None if TYPE_CHECKING: from opentelemetry.trace import Span as OtelSpan + from litellm.caching.dual_cache import DualCache from litellm.integrations.otel.model.destination import OtelDestination from litellm.models.tag import LiteLLM_TagTable from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.credential_endpoints.user_provider_credentials import GithubCopilotUserConnectionPayload from litellm.proxy.policy_engine.attachment_registry import AttachmentRegistry from litellm.proxy.proxy_server import ProxyConfig as _ProxyConfig from litellm.proxy.utils import PrismaClient + from litellm.repositories.prisma_protocols import DatabaseClient from litellm.types.proxy.policy_engine import Policy, PolicyMatchContext ProxyConfig = _ProxyConfig @@ -3502,12 +3505,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_credentials + from .credential_endpoints.user_provider_credentials import aget_user_provider_connections if prisma_client is None: return try: - connections: Final = await aget_user_provider_credentials( + connections: Final = await aget_user_provider_connections( prisma_client=prisma_client, cache=user_api_key_cache, user_id=user_id, @@ -3518,14 +3521,45 @@ async def _resolve_user_provider_credentials_for_request( "_resolve_user_provider_credentials_for_request: failed to load user provider credentials" ) return + refreshed_connections: Final = await _refresh_user_provider_connections( + prisma_client=prisma_client, + cache=user_api_key_cache, + user_id=user_id, + connections=connections, + ) secret_fields: Final = data.get("secret_fields") if not isinstance(secret_fields, dict): return - secret_fields["user_provider_credentials"] = RedactedDict(dict(connections)) + secret_fields["user_provider_credentials"] = RedactedDict( + { + name: {"access_token": payload.access_token, "github_host": payload.github_host} + for name, payload in refreshed_connections.items() + } + ) secret_fields["user_provider_credentials_user_id"] = user_id +async def _refresh_user_provider_connections( + prisma_client: "DatabaseClient", + cache: "DualCache", + user_id: str, + connections: "Mapping[str, GithubCopilotUserConnectionPayload]", +) -> "Mapping[str, GithubCopilotUserConnectionPayload]": + from .credential_endpoints.github_copilot_token_refresh import arefresh_expiring_connections + + try: + return await arefresh_expiring_connections( + prisma_client=prisma_client, + cache=cache, + user_id=user_id, + connections=connections, + ) + except Exception: # noqa: BLE001 # refresh errors must not block the existing credential path + verbose_proxy_logger.exception("_refresh_user_provider_connections: refresh failed; using stored connections") + return connections + + def _warn_stale_team_alias_once(warning_key: str, message: str, *args: str) -> None: if warning_key in _STALE_TEAM_ALIAS_WARNING_KEYS: return diff --git a/litellm/proxy/mcp_server/outbound_credentials/redis_refresh_coordinator.py b/litellm/proxy/mcp_server/outbound_credentials/redis_refresh_coordinator.py index 98446f2b66e..b1e7e0df6f6 100644 --- a/litellm/proxy/mcp_server/outbound_credentials/redis_refresh_coordinator.py +++ b/litellm/proxy/mcp_server/outbound_credentials/redis_refresh_coordinator.py @@ -26,11 +26,9 @@ from collections.abc import Awaitable, Callable from contextlib import suppress from dataclasses import KW_ONLY, dataclass from enum import Enum -from typing import Final, Protocol +from typing import Final, Protocol, TypeVar -from litellm.proxy.mcp_server.outbound_credentials.oauth_token_store import ( - OAuthToken, -) +RefreshValue: Final = TypeVar("RefreshValue") class LockAcquisition(Enum): @@ -86,9 +84,9 @@ class RedisRefreshCoordinator: self, user_id: str, server_id: str, - refresh: Callable[[], Awaitable[OAuthToken | None]], - reread: Callable[[], Awaitable[OAuthToken | None]], - ) -> OAuthToken | None: + refresh: Callable[[], Awaitable[RefreshValue | None]], + reread: Callable[[], Awaitable[RefreshValue | None]], + ) -> RefreshValue | None: key: Final = self._key(user_id, server_id) token: Final = self.new_token() match await self.lock.acquire(key, token, self.lock_ttl_seconds): @@ -111,8 +109,8 @@ class RedisRefreshCoordinator: self, key: str, token: str, - refresh: Callable[[], Awaitable[OAuthToken | None]], - ) -> OAuthToken | None: + refresh: Callable[[], Awaitable[RefreshValue | None]], + ) -> RefreshValue | None: refresh_task: Final = asyncio.ensure_future(refresh()) renewal_task: Final = asyncio.create_task(self._renew_lease_until_done(key, token, refresh_task)) try: @@ -127,7 +125,7 @@ class RedisRefreshCoordinator: self, key: str, token: str, - refresh_task: asyncio.Future[OAuthToken | None], + refresh_task: asyncio.Future[RefreshValue | None], ) -> None: budget_deadline: Final = self.clock() + self.refresh_budget_seconds while not refresh_task.done() and self.clock() < budget_deadline: diff --git a/tests/unit/llms/github_copilot/test_per_user_auth.py b/tests/unit/llms/github_copilot/test_per_user_auth.py index cf6fb9fa71f..227d18e9518 100644 --- a/tests/unit/llms/github_copilot/test_per_user_auth.py +++ b/tests/unit/llms/github_copilot/test_per_user_auth.py @@ -2,6 +2,7 @@ import asyncio import datetime +import json import time from collections.abc import Callable from contextlib import contextmanager @@ -24,11 +25,16 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.github_copilot.per_user_auth import ( _SESSION_CACHE, GITHUB_COPILOT_USER_SESSION_KWARG_KEY, + GithubCopilotTokenRejectedError, + GithubTokenRefreshed, + GithubTokenRefreshRejected, + GithubTokenRefreshUnavailable, GithubCopilotUserSession, _session_cache_key, aattach_github_copilot_user_session, aexchange_github_token, apoll_device_flow, + arefresh_github_token, astart_device_flow, attach_github_copilot_user_session, evict_copilot_user_session, @@ -250,10 +256,11 @@ def test_cached_session_is_not_reused_after_github_token_changes(): @pytest.mark.parametrize("status", [401, 403, 404]) -def test_exchange_auth_failures_raise_caller_auth_error_and_evict(status): +def test_exchange_auth_failures_raise_caller_auth_error_and_evict(status: int) -> None: with github_http(_exchange_responder(status=status)): with pytest.raises(CallerCredentialAuthenticationError) as exc: exchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) + assert isinstance(exc.value, GithubCopilotTokenRejectedError) == (status == 401) assert "Reconnect GitHub Copilot for credential 'copilot-cred' in the LiteLLM UI (LLM Credentials)" in str( exc.value ) @@ -440,11 +447,137 @@ async def test_apoll_device_flow_error_mapping(github_error, expected_status): async def test_apoll_device_flow_connected(): def respond(request: httpx.Request) -> httpx.Response: assert str(request.url) == ACCESS_TOKEN_URL - return httpx.Response(200, json={"access_token": "gho_x"}, request=request) + return httpx.Response( + 200, + json={ + "access_token": "gho_x", + "refresh_token": "ghr_x", + "expires_in": 28800, + "refresh_token_expires_in": 15897600, + }, + request=request, + ) with github_http(respond): poll = await apoll_device_flow("dc", _GITHUB_ENDPOINTS) - assert poll.status == "connected" and poll.access_token == "gho_x" + assert ( + poll.status, + poll.access_token, + poll.refresh_token, + poll.expires_in, + poll.refresh_token_expires_in, + ) == ("connected", "gho_x", "ghr_x", 28800, 15897600) + + +@pytest.mark.asyncio +async def test_device_flow_omits_empty_refresh_token_and_nonpositive_expiries(): + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json={ + "access_token": "gho_x", + "refresh_token": "", + "expires_in": 0, + "refresh_token_expires_in": -1, + }, + request=request, + ) + + with github_http(respond): + poll = await apoll_device_flow("dc", _GITHUB_ENDPOINTS) + assert poll.status == "connected" and poll.access_token == "gho_x" + assert (poll.refresh_token, poll.expires_in, poll.refresh_token_expires_in) == (None, None, None) + + +@pytest.mark.asyncio +async def test_refresh_grant_uses_configured_ghe_endpoint_and_rotates_pair(): + endpoints: Final = github_host_endpoints("octo.ghe.com") + seen: Final[list[httpx.Request]] = [] + + def respond(request: httpx.Request) -> httpx.Response: + seen.append(request) + return httpx.Response( + 200, + json={ + "access_token": "gho_rotated", + "refresh_token": "ghr_rotated", + "expires_in": 28800, + "refresh_token_expires_in": 15897600, + }, + request=request, + ) + + with github_http(respond): + result: Final = await arefresh_github_token("ghr_old", endpoints) + assert result == GithubTokenRefreshed( + access_token="gho_rotated", + refresh_token="ghr_rotated", + expires_in=28800, + refresh_token_expires_in=15897600, + ) + assert str(seen[0].url) == "https://octo.ghe.com/login/oauth/access_token" + assert json.loads(seen[0].content) == { + "client_id": endpoints.client_id, + "grant_type": "refresh_token", + "refresh_token": "ghr_old", + } + assert "client_secret" not in json.loads(seen[0].content) + assert "Authorization" not in seen[0].headers + + +@pytest.mark.asyncio +async def test_refresh_grant_classifies_rejections_and_unavailability(): + responses: Final = ( + (200, {"error": "invalid_grant"}, GithubTokenRefreshRejected), + (400, {"error": "invalid_grant"}, GithubTokenRefreshRejected), + (429, {"error": "rate_limited"}, GithubTokenRefreshUnavailable), + (500, {"error": "server_error"}, GithubTokenRefreshUnavailable), + ) + for status, payload, expected_type in responses: + with github_http( + lambda request, response_status=status, response_payload=payload: httpx.Response( + response_status, + json=response_payload, + request=request, + ) + ): + result: Final = await arefresh_github_token("ghr_old", _GITHUB_ENDPOINTS) + assert isinstance(result, expected_type) + + +@pytest.mark.asyncio +async def test_refresh_grant_transport_and_invalid_json_are_unavailable(): + def disconnected(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("offline", request=request) + + with github_http(disconnected): + disconnected_result: Final = await arefresh_github_token("ghr_old", _GITHUB_ENDPOINTS) + with github_http(lambda request: httpx.Response(200, content=b"{", request=request)): + invalid_json_result: Final = await arefresh_github_token("ghr_old", _GITHUB_ENDPOINTS) + assert isinstance(disconnected_result, GithubTokenRefreshUnavailable) + assert isinstance(invalid_json_result, GithubTokenRefreshUnavailable) + + +@pytest.mark.asyncio +async def test_github_auth_endpoint_overrides_accept_plain_http_localhost(monkeypatch, per_user_credential): + token_url: Final = "http://127.0.0.1:8123/login/oauth/access_token" + api_key_url: Final = "http://127.0.0.1:8123/copilot_internal/v2/token" + monkeypatch.setenv("GITHUB_COPILOT_ACCESS_TOKEN_URL", token_url) + monkeypatch.setenv("GITHUB_COPILOT_API_KEY_URL", api_key_url) + endpoints: Final = github_host_endpoints("github.com") + + def respond(request: httpx.Request) -> httpx.Response: + if str(request.url) == token_url: + return httpx.Response(200, json={"access_token": "gho_next", "refresh_token": "ghr_next"}, request=request) + if str(request.url) == api_key_url: + return httpx.Response(200, json=_exchange_payload(), request=request) + pytest.fail(f"unexpected request to {request.url}") + + with github_http(respond) as requests: + refreshed: Final = await arefresh_github_token("ghr_old", endpoints) + await aexchange_github_token("user-a", "gho_access", "copilot-cred", endpoints) + assert isinstance(refreshed, GithubTokenRefreshed) + assert [str(request.url) for request in requests] == [token_url, api_key_url] def test_is_github_copilot_per_user_request(): diff --git a/tests/unit/proxy/credential_endpoints/test_endpoints.py b/tests/unit/proxy/credential_endpoints/test_endpoints.py index 47dd6d561b3..483328d40f4 100644 --- a/tests/unit/proxy/credential_endpoints/test_endpoints.py +++ b/tests/unit/proxy/credential_endpoints/test_endpoints.py @@ -2241,10 +2241,16 @@ def test_user_connection_poll_rejects_an_undecryptable_flow_handle(monkeypatch): def test_user_connection_poll_connected_persists_and_reports_login(monkeypatch): table: Final = _patch_user_connection_env(monkeypatch, [_per_user_credential()]) + monkeypatch.setattr("litellm.proxy.credential_endpoints.endpoints.time.time", lambda: 1000.0) with _github_http( { _DEVICE_CODE_URL: _DEVICE_FLOW_START, - _ACCESS_TOKEN_URL: {"access_token": "gho_1"}, + _ACCESS_TOKEN_URL: { + "access_token": "gho_1", + "refresh_token": "ghr_1", + "expires_in": 28800, + "refresh_token_expires_in": 15897600, + }, _COPILOT_TOKEN_URL: {"token": "copilot-tok", "expires_at": 4102444800}, _GITHUB_USER_URL: {"login": "octo"}, } @@ -2269,6 +2275,9 @@ def test_user_connection_poll_connected_persists_and_reports_login(monkeypatch): ) assert stored_connection is not None assert stored_connection.github_host == "github.com" + assert stored_connection.refresh_token == "ghr_1" + assert stored_connection.access_token_expires_at == 29800.0 + assert stored_connection.refresh_token_expires_at == 15898600.0 def test_poll_host_change_race_requires_reconnect_on_the_next_request(monkeypatch): @@ -2746,6 +2755,43 @@ def test_delete_user_connection_evicts_copilot_session(monkeypatch): table.delete_many.assert_awaited_once_with(where={"user_id": "user-a", "credential_name": "copilot-cred"}) +def test_delete_user_connection_rewrites_tombstone_after_a_racing_cache_set(monkeypatch): + from types import SimpleNamespace + from litellm.caching.dual_cache import DualCache + from litellm.proxy import proxy_server + + table: Final = _patch_user_connection_env(monkeypatch, [_per_user_credential()], rows=[_connection_row()]) + cached_values: Final[dict[str, str]] = {} # mutable-ok: fake Redis records cache 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 = SimpleNamespace( + async_set_cache=AsyncMock(side_effect=set_cache), + async_get_cache=AsyncMock(side_effect=get_cache), + async_delete_cache=AsyncMock(), + ) + monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache(redis_cache=redis_cache)) + ciphertext: Final = upc._encode( + upc.GithubCopilotUserConnectionPayload(access_token="gho_refreshed", github_login="octo") + ) + + async def set_refresh_between_tombstone_and_delete(*, where: dict[str, str]) -> None: + await redis_cache.async_set_cache(upc._cache_key("user-a", "copilot-cred"), ciphertext, ttl=60) + table.find_unique.return_value = None + + table.delete_many = AsyncMock(side_effect=set_refresh_between_tombstone_and_delete) + 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 cached_values[upc._cache_key("user-a", "copilot-cred")] == upc._NOT_CONNECTED + 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_github_copilot_token_refresh.py b/tests/unit/proxy/credential_endpoints/test_github_copilot_token_refresh.py new file mode 100644 index 00000000000..b1790e08a09 --- /dev/null +++ b/tests/unit/proxy/credential_endpoints/test_github_copilot_token_refresh.py @@ -0,0 +1,753 @@ +import asyncio +import json +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from types import SimpleNamespace +from typing import Final +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest + +import litellm +from litellm.caching.dual_cache import DualCache +from litellm.caching.redis_cache import RedisCache +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.types.utils import CredentialItem +from litellm.proxy.credential_endpoints import user_provider_credentials as upc +from litellm.proxy.credential_endpoints.github_copilot_token_refresh import ( + GithubCopilotTokenRefreshSingleFlight, + arefresh_expiring_connections, +) +from litellm.proxy.credential_endpoints.user_provider_credentials import GithubCopilotUserConnectionPayload +from litellm.llms.github_copilot.per_user_auth import aattach_github_copilot_user_session + +_ACCESS_TOKEN_URL = "https://github.com/login/oauth/access_token" +_API_KEY_URL = "https://api.github.com/copilot_internal/v2/token" + + +class _Store: + def __init__(self, payload: GithubCopilotUserConnectionPayload): + self.ciphertext: str | None = upc._encode(payload) + self.exists = True + + async def find_unique(self, *, where: object) -> SimpleNamespace | None: + return SimpleNamespace(credential_b64=self.ciphertext) if self.exists and self.ciphertext is not None else None + + async def update_many(self, *, where: dict[str, object], data: dict[str, object]) -> int: + if ( + not self.exists + or self.ciphertext is None + or self.ciphertext != where["credential_b64"] + or where["user_id"] != "user-a" + or where["credential_name"] != "copilot-cred" + ): + return 0 + self.ciphertext = str(data["credential_b64"]) + return 1 + + +class _RacingStore(_Store): + def __init__(self, payload: GithubCopilotUserConnectionPayload): + super().__init__(payload) + self._first_reads = 0 + self._both_read = asyncio.Event() + + async def find_unique(self, *, where: object) -> SimpleNamespace | None: + if self._first_reads < 2: + ciphertext: Final = self.ciphertext + self._first_reads += 1 + if self._first_reads == 2: + self._both_read.set() + await self._both_read.wait() + return SimpleNamespace(credential_b64=ciphertext) + return await super().find_unique(where=where) + + +class _Clock: + def __init__(self) -> None: + self.value = 0.0 + + def monotonic(self) -> float: + return self.value + + async def sleep(self, seconds: float) -> None: + self.value += seconds + await asyncio.sleep(0) + + +class _HeldLockClient: + async def set(self, name: str, value: str, *, nx: bool = False, px: int | None = None) -> object | None: + return None + + async def eval(self, script: str, numkeys: int, *keys_and_args: str) -> object: + return 1 + + async def exists(self, *names: str) -> int: + return 0 + + +class _RedisCache: + def __init__(self) -> None: + self.lock_client = _HeldLockClient() + + def init_async_client(self) -> _HeldLockClient: + return self.lock_client + + def check_and_fix_namespace(self, key: str) -> str: + return key + + +def _cache_with_fake_redis(monkeypatch: pytest.MonkeyPatch) -> tuple[DualCache, RedisCache]: + import fakeredis + + redis_client: Final = fakeredis.FakeAsyncRedis() + redis_cache: Final = RedisCache(host="localhost", port=6379) + monkeypatch.setattr(redis_cache, "init_async_client", lambda: redis_client) + return DualCache(redis_cache=redis_cache), redis_cache + + +def _credential() -> CredentialItem: + return CredentialItem( + credential_name="copilot-cred", + credential_values={"github_copilot_auth_type": "per_user_oauth"}, + credential_info={}, + ) + + +def _payload( + access_token: str = "gho_old", + refresh_token: str | None = "ghr_old", + access_expires_at: float | None = 1.0, + refresh_expires_at: float | None = None, +) -> GithubCopilotUserConnectionPayload: + return GithubCopilotUserConnectionPayload( + access_token=access_token, + github_login="octo", + refresh_token=refresh_token, + access_token_expires_at=access_expires_at, + refresh_token_expires_at=refresh_expires_at, + ) + + +def _prisma(store: _Store) -> tuple[MagicMock, MagicMock]: + table: Final = MagicMock() + table.find_unique = AsyncMock(side_effect=store.find_unique) + table.update_many = AsyncMock(side_effect=store.update_many) + table.upsert = AsyncMock() + prisma_client: Final = MagicMock() + prisma_client.db.litellm_userprovidercredentials = table + prisma_client.writer_db = prisma_client.db + return prisma_client, table + + +@contextmanager +def _github_http( + respond: Callable[[httpx.Request], httpx.Response], +) -> Iterator[list[httpx.Request]]: + requests: Final[list[httpx.Request]] = [] + + def capture(request: httpx.Request) -> httpx.Response: + requests.append(request) + return respond(request) + + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(capture)) + with ( + patch("litellm.llms.custom_httpx.http_handler.get_async_httpx_client", return_value=client), + patch("litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client", return_value=client), + ): + yield requests + + +@pytest.fixture(autouse=True) +def _environment(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + from litellm.llms.github_copilot.per_user_auth import _SESSION_CACHE + + _SESSION_CACHE.flush_cache() + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "sk-test-master") + monkeypatch.setattr(litellm, "credential_list", [_credential()]) + yield + _SESSION_CACHE.flush_cache() + + +@pytest.mark.asyncio +async def test_proactive_refresh_persists_rotated_pair_and_rewrites_redis_cache(monkeypatch): + store: Final = _Store(_payload()) + prisma_client, _ = _prisma(store) + import fakeredis + + from litellm.caching.redis_cache import RedisCache + + redis_client: Final = fakeredis.FakeAsyncRedis() + redis_cache: Final = RedisCache(host="localhost", port=6379) + monkeypatch.setattr(redis_cache, "init_async_client", lambda: redis_client) + cache: Final = DualCache(redis_cache=redis_cache) + + def respond(request: httpx.Request) -> httpx.Response: + assert str(request.url) == _ACCESS_TOKEN_URL + return httpx.Response( + 200, + json={ + "access_token": "gho_rotated", + "refresh_token": "ghr_rotated", + "expires_in": 3600, + "refresh_token_expires_in": 7200, + }, + request=request, + ) + + with _github_http(respond) as requests: + result: Final = await arefresh_expiring_connections( + prisma_client, + cache, + "user-a", + {"copilot-cred": _payload()}, + ) + persisted: Final = await upc.get_user_provider_credential(prisma_client, "user-a", "copilot-cred") + assert result["copilot-cred"].access_token == "gho_rotated" + assert result["copilot-cred"].refresh_token == "ghr_rotated" + assert persisted == result["copilot-cred"] + assert persisted.access_token_expires_at is not None + assert persisted.refresh_token_expires_at is not None + cached_ciphertext: Final = await redis_cache.async_get_cache(upc._cache_key("user-a", "copilot-cred")) + assert upc.decode_user_provider_credential(cached_ciphertext) == persisted + assert len(requests) == 1 + + +@pytest.mark.asyncio +async def test_same_worker_concurrent_refreshes_share_one_refresh_request(): + store: Final = _Store(_payload()) + prisma_client, _ = _prisma(store) + single_flight: Final = GithubCopilotTokenRefreshSingleFlight() + + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json={"access_token": "gho_rotated", "refresh_token": "ghr_rotated"}, + request=request, + ) + + with _github_http(respond) as requests: + results: Final = await asyncio.gather( + arefresh_expiring_connections( + prisma_client, + DualCache(), + "user-a", + {"copilot-cred": _payload()}, + single_flight=single_flight, + ), + arefresh_expiring_connections( + prisma_client, + DualCache(), + "user-a", + {"copilot-cred": _payload()}, + single_flight=single_flight, + ), + ) + assert len(requests) == 1 + assert [result["copilot-cred"].access_token for result in results] == ["gho_rotated", "gho_rotated"] + + +@pytest.mark.asyncio +async def test_single_flight_reuses_inflight_refresh_task(): + payload: Final = _payload() + single_flight: Final = GithubCopilotTokenRefreshSingleFlight() + refresh_started: Final = asyncio.Event() + finish_refresh: Final = asyncio.Event() + + async def wait_for_refresh() -> GithubCopilotUserConnectionPayload: + refresh_started.set() + await finish_refresh.wait() + return payload + + refresh_call: Final = AsyncMock(side_effect=wait_for_refresh) + + async def refresh() -> GithubCopilotUserConnectionPayload | None: + return await refresh_call() + + first: Final = asyncio.create_task(single_flight.run(("user-a", "copilot-cred"), refresh)) + await refresh_started.wait() + second: Final = asyncio.create_task(single_flight.run(("user-a", "copilot-cred"), refresh)) + await asyncio.sleep(0) + assert refresh_call.await_count == 1 + finish_refresh.set() + results: Final = await asyncio.gather(first, second) + assert results[0] is payload + assert results[1] is payload + assert refresh_call.await_count == 1 + + +@pytest.mark.asyncio +async def test_completed_single_flight_runs_second_refresh_with_rotated_token(monkeypatch): + from litellm.proxy.credential_endpoints import github_copilot_token_refresh as refresh_module + + monkeypatch.setattr(refresh_module.time, "time", lambda: 1000.0) + store: Final = _Store(_payload()) + prisma_client, _ = _prisma(store) + single_flight: Final = GithubCopilotTokenRefreshSingleFlight() + refresh_tokens: Final[list[str]] = [] + + def respond(request: httpx.Request) -> httpx.Response: + body: Final = json.loads(request.content) + refresh_tokens.append(body["refresh_token"]) + token_number: Final = len(refresh_tokens) + return httpx.Response( + 200, + json={ + "access_token": f"gho_{token_number}", + "refresh_token": f"ghr_{token_number}", + "expires_in": 100, + "refresh_token_expires_in": 3600, + }, + request=request, + ) + + with _github_http(respond) as requests: + first: Final = await arefresh_expiring_connections( + prisma_client, + DualCache(), + "user-a", + {"copilot-cred": _payload()}, + single_flight=single_flight, + ) + second: Final = await arefresh_expiring_connections( + prisma_client, + DualCache(), + "user-a", + first, + single_flight=single_flight, + ) + assert len(requests) == 2 + assert refresh_tokens == ["ghr_old", "ghr_1"] + assert second["copilot-cred"].access_token == "gho_2" + assert second["copilot-cred"].refresh_token == "ghr_2" + + +@pytest.mark.asyncio +async def test_nonexpiring_connection_and_expired_refresh_token_skip_token_calls(): + store: Final = _Store(_payload()) + prisma_client, _ = _prisma(store) + connections: Final = { + "copilot-cred": _payload(refresh_token=None), + "expired-refresh": _payload(refresh_expires_at=1.0), + } + with _github_http(lambda request: pytest.fail(f"unexpected GitHub request to {request.url}")) as requests: + result: Final = await arefresh_expiring_connections(prisma_client, DualCache(), "user-a", connections) + assert result == connections + assert not requests + + +@pytest.mark.asyncio +async def test_reactive_exchange_401_refreshes_the_github_token(): + store: Final = _Store(_payload(access_expires_at=9999999999.0)) + prisma_client, _ = _prisma(store) + + def respond(request: httpx.Request) -> httpx.Response: + if str(request.url) == _API_KEY_URL: + return httpx.Response(401, json={"message": "expired"}, request=request) + return httpx.Response( + 200, + json={"access_token": "gho_reactive", "refresh_token": "ghr_reactive"}, + request=request, + ) + + with _github_http(respond) as requests: + result: Final = await arefresh_expiring_connections( + prisma_client, + DualCache(), + "user-a", + {"copilot-cred": _payload(access_expires_at=9999999999.0)}, + ) + assert [str(request.url) for request in requests] == [_API_KEY_URL, _ACCESS_TOKEN_URL] + assert result["copilot-cred"].access_token == "gho_reactive" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status", [403, 404]) +async def test_exchange_permission_failures_do_not_refresh_github_token(status: int) -> None: + stale: Final = _payload(access_expires_at=9999999999.0) + store: Final = _Store(stale) + original_ciphertext: Final = store.ciphertext + prisma_client, table = _prisma(store) + + def respond(request: httpx.Request) -> httpx.Response: + if str(request.url) == _API_KEY_URL: + return httpx.Response(status, json={"message": "Permission denied"}, request=request) + return httpx.Response( + 200, + json={"access_token": "gho_rotated", "refresh_token": "ghr_rotated"}, + request=request, + ) + + with _github_http(respond) as requests: + first: Final = await arefresh_expiring_connections( + prisma_client, + DualCache(), + "user-a", + {"copilot-cred": stale}, + ) + second: Final = await arefresh_expiring_connections( + prisma_client, + DualCache(), + "user-a", + {"copilot-cred": first["copilot-cred"]}, + ) + + stored: Final = await upc.get_user_provider_credential(prisma_client, "user-a", "copilot-cred") + assert [str(request.url) for request in requests] == [_API_KEY_URL, _API_KEY_URL] + assert table.update_many.await_count == 0 + assert store.ciphertext == original_ciphertext + assert stored == stale + assert first["copilot-cred"] == stale + assert second["copilot-cred"] == stale + + +@pytest.mark.asyncio +async def test_redis_lock_loser_rereads_the_persisted_connection(): + old: Final = _payload() + rotated: Final = _payload("gho_winner", "ghr_winner", 9999999999.0) + store: Final = _Store(rotated) + prisma_client, _ = _prisma(store) + cache: Final = DualCache(redis_cache=_RedisCache()) + with _github_http(lambda request: pytest.fail(f"lock loser refreshed at {request.url}")) as requests: + result: Final = await arefresh_expiring_connections( + prisma_client, + cache, + "user-a", + {"copilot-cred": old}, + ) + assert result["copilot-cred"] == rotated + assert not requests + + +@pytest.mark.asyncio +async def test_compare_and_set_loss_does_not_resurrect_disconnected_connection(): + store: Final = _Store(_payload()) + prisma_client, table = _prisma(store) + + def respond(request: httpx.Request) -> httpx.Response: + store.exists = False + return httpx.Response( + 200, + json={"access_token": "gho_rotated", "refresh_token": "ghr_rotated"}, + request=request, + ) + + with _github_http(respond): + result: Final = await arefresh_expiring_connections( + prisma_client, + DualCache(), + "user-a", + {"copilot-cred": _payload()}, + ) + assert result == {} + table.update_many.assert_awaited_once() + table.upsert.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_refresh_deleted_after_cache_set_restores_disconnect_tombstone(monkeypatch: pytest.MonkeyPatch): + store: Final = _Store(_payload()) + prisma_client, _ = _prisma(store) + cache, redis_cache = _cache_with_fake_redis(monkeypatch) + await upc.invalidate_user_provider_credential_cache(cache, "user-a", "copilot-cred") + original_set_cache: Final = redis_cache.async_set_cache + + async def set_then_delete_row(key: str | None, value: object, **kwargs: object) -> None: + await original_set_cache(key, value, **kwargs) + if isinstance(value, str) and value != upc._NOT_CONNECTED: + store.exists = False + + monkeypatch.setattr(redis_cache, "async_set_cache", set_then_delete_row) + + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json={"access_token": "gho_rotated", "refresh_token": "ghr_rotated"}, + request=request, + ) + + with _github_http(respond): + result: Final = await arefresh_expiring_connections( + prisma_client, + cache, + "user-a", + {"copilot-cred": _payload()}, + ) + cached: Final = await redis_cache.async_get_cache(upc._cache_key("user-a", "copilot-cred")) + assert result == {} + assert cached == upc._NOT_CONNECTED + + +@pytest.mark.asyncio +async def test_refresh_returns_replacement_written_after_cache_set_and_drops_stale_cache( + monkeypatch: pytest.MonkeyPatch, +): + store: Final = _Store(_payload()) + prisma_client, _ = _prisma(store) + cache, redis_cache = _cache_with_fake_redis(monkeypatch) + replacement: Final = _payload("gho_reconnected", "ghr_reconnected", 9999999999.0) + original_set_cache: Final = redis_cache.async_set_cache + + async def set_then_replace_row(key: str | None, value: object, **kwargs: object) -> None: + await original_set_cache(key, value, **kwargs) + if isinstance(value, str) and value != upc._NOT_CONNECTED: + store.ciphertext = upc._encode(replacement) + + monkeypatch.setattr(redis_cache, "async_set_cache", set_then_replace_row) + + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json={"access_token": "gho_rotated", "refresh_token": "ghr_rotated"}, + request=request, + ) + + with _github_http(respond): + result: Final = await arefresh_expiring_connections( + prisma_client, + cache, + "user-a", + {"copilot-cred": _payload()}, + ) + cached: Final = await redis_cache.async_get_cache(upc._cache_key("user-a", "copilot-cred")) + assert result["copilot-cred"] == replacement + assert cached is None + + +@pytest.mark.asyncio +async def test_rejected_refresh_waits_for_winner_then_returns_winner_token(): + old: Final = _payload() + store: Final = _RacingStore(old) + prisma_client, _ = _prisma(store) + clock: Final = _Clock() + attempts: Final[list[str]] = [] + + def respond(request: httpx.Request) -> httpx.Response: + body: Final = json.loads(request.content) + attempts.append(body["refresh_token"]) + if len(attempts) == 1: + return httpx.Response( + 200, + json={"access_token": "gho_winner", "refresh_token": "ghr_winner"}, + request=request, + ) + return httpx.Response(400, json={"error": "invalid_grant"}, request=request) + + with _github_http(respond) as requests: + results: Final = await asyncio.gather( + arefresh_expiring_connections( + prisma_client, + DualCache(), + "user-a", + {"copilot-cred": old}, + single_flight=GithubCopilotTokenRefreshSingleFlight(), + sleep=clock.sleep, + clock=clock.monotonic, + ), + arefresh_expiring_connections( + prisma_client, + DualCache(), + "user-a", + {"copilot-cred": old}, + single_flight=GithubCopilotTokenRefreshSingleFlight(), + sleep=clock.sleep, + clock=clock.monotonic, + ), + ) + assert attempts == ["ghr_old", "ghr_old"] + assert len(requests) == 2 + assert [result["copilot-cred"].access_token for result in results] == ["gho_winner", "gho_winner"] + + +@pytest.mark.asyncio +async def test_rejected_refresh_times_out_then_exchange_returns_reconnect_401(): + store: Final = _Store(_payload()) + prisma_client, _ = _prisma(store) + clock: Final = _Clock() + + def respond(request: httpx.Request) -> httpx.Response: + if str(request.url) == _ACCESS_TOKEN_URL: + return httpx.Response(400, json={"error": "invalid_grant"}, request=request) + return httpx.Response(401, json={"message": "expired"}, request=request) + + with _github_http(respond): + result: Final = await arefresh_expiring_connections( + prisma_client, + DualCache(), + "user-a", + {"copilot-cred": _payload()}, + sleep=clock.sleep, + clock=clock.monotonic, + ) + 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": result["copilot-cred"].access_token, + "github_host": "github.com", + } + }, + }, + } + with pytest.raises(litellm.CallerCredentialAuthenticationError): + await aattach_github_copilot_user_session(kwargs) + assert clock.value == pytest.approx(2.0) + + +@pytest.mark.asyncio +async def test_unavailable_refresh_returns_stale_without_polling(): + store: Final = _Store(_payload()) + prisma_client, _ = _prisma(store) + clock: Final = _Clock() + + with _github_http(lambda request: httpx.Response(500, request=request)): + result: Final = await arefresh_expiring_connections( + prisma_client, + DualCache(), + "user-a", + {"copilot-cred": _payload()}, + sleep=clock.sleep, + clock=clock.monotonic, + ) + assert result["copilot-cred"] == _payload() + assert clock.value == 0.0 + + +@pytest.mark.asyncio +async def test_rejected_refresh_clears_metadata_for_later_requests(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy.credential_endpoints import github_copilot_token_refresh as refresh_module + + monkeypatch.setattr(refresh_module.time, "time", lambda: 1000.0) + stale: Final = _payload(access_expires_at=900.0, refresh_expires_at=2000.0) + store: Final = _Store(stale) + prisma_client, _ = _prisma(store) + clock: Final = _Clock() + + def reject_refresh(request: httpx.Request) -> httpx.Response: + return httpx.Response(400, json={"error": "bad_refresh_token"}, request=request) + + with _github_http(reject_refresh) as requests: + first: Final = await arefresh_expiring_connections( + prisma_client, + DualCache(), + "user-a", + {"copilot-cred": stale}, + sleep=clock.sleep, + clock=clock.monotonic, + ) + stored: Final = await upc.get_user_provider_credential(prisma_client, "user-a", "copilot-cred") + assert stored is not None + assert stored.access_token == stale.access_token + assert stored.access_token_expires_at == stale.access_token_expires_at + assert stored.refresh_token is None + assert stored.refresh_token_expires_at is None + assert first["copilot-cred"] == stored + before_second: Final = clock.monotonic() + second: Final = await arefresh_expiring_connections( + prisma_client, + DualCache(), + "user-a", + {"copilot-cred": stored}, + sleep=clock.sleep, + clock=clock.monotonic, + ) + assert clock.monotonic() == before_second + + assert second["copilot-cred"] == stored + assert len(requests) == 1 + assert json.loads(requests[0].content)["refresh_token"] == stale.refresh_token + + +@pytest.mark.asyncio +async def test_unavailable_refresh_preserves_metadata_and_retries(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy.credential_endpoints import github_copilot_token_refresh as refresh_module + + monkeypatch.setattr(refresh_module.time, "time", lambda: 1000.0) + stale: Final = _payload(access_expires_at=900.0, refresh_expires_at=2000.0) + store: Final = _Store(stale) + prisma_client, _ = _prisma(store) + clock: Final = _Clock() + + def unavailable(request: httpx.Request) -> httpx.Response: + return httpx.Response(503, json={"error": "temporarily_unavailable"}, request=request) + + with _github_http(unavailable) as requests: + first: Final = await arefresh_expiring_connections( + prisma_client, + DualCache(), + "user-a", + {"copilot-cred": stale}, + sleep=clock.sleep, + clock=clock.monotonic, + ) + stored: Final = await upc.get_user_provider_credential(prisma_client, "user-a", "copilot-cred") + assert stored is not None + assert stored == stale + second: Final = await arefresh_expiring_connections( + prisma_client, + DualCache(), + "user-a", + {"copilot-cred": stored}, + sleep=clock.sleep, + clock=clock.monotonic, + ) + + assert first["copilot-cred"] == stale + assert second["copilot-cred"] == stale + assert [json.loads(request.content)["refresh_token"] for request in requests] == ["ghr_old", "ghr_old"] + assert clock.value == 0.0 + + +@pytest.mark.asyncio +async def test_rejected_refresh_cas_loss_returns_replacement_without_cache_write( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy.credential_endpoints import github_copilot_token_refresh as refresh_module + + monkeypatch.setattr(refresh_module.time, "time", lambda: 1000.0) + stale: Final = _payload(access_expires_at=900.0, refresh_expires_at=2000.0) + replacement: Final = _payload( + access_token="gho_reconnected", + refresh_token="ghr_reconnected", + access_expires_at=3000.0, + refresh_expires_at=4000.0, + ) + store: Final = _Store(stale) + prisma_client, table = _prisma(store) + cache, redis_cache = _cache_with_fake_redis(monkeypatch) + clock: Final = _Clock() + cache_writes: Final[list[object]] = [] + original_update_many: Final = store.update_many + original_set_cache: Final = redis_cache.async_set_cache + + async def replace_before_cas(*, where: dict[str, object], data: dict[str, object]) -> int: + store.ciphertext = upc._encode(replacement) + return await original_update_many(where=where, data=data) + + async def record_cache_write(key: str | None, value: object, **kwargs: object) -> None: + cache_writes.append(value) + await original_set_cache(key, value, **kwargs) + + table.update_many = AsyncMock(side_effect=replace_before_cas) + monkeypatch.setattr(redis_cache, "async_set_cache", record_cache_write) + + def reject_refresh(request: httpx.Request) -> httpx.Response: + return httpx.Response(400, json={"error": "bad_refresh_token"}, request=request) + + with _github_http(reject_refresh): + result: Final = await arefresh_expiring_connections( + prisma_client, + cache, + "user-a", + {"copilot-cred": stale}, + sleep=clock.sleep, + clock=clock.monotonic, + ) + + stored: Final = await upc.get_user_provider_credential(prisma_client, "user-a", "copilot-cred") + assert result["copilot-cred"] == replacement + assert stored == replacement + assert cache_writes == [] 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 dae3a0459e2..2621319a1f1 100644 --- a/tests/unit/proxy/credential_endpoints/test_user_provider_credentials.py +++ b/tests/unit/proxy/credential_endpoints/test_user_provider_credentials.py @@ -1,6 +1,7 @@ """Tests for the per-user provider credentials DB/cache module.""" import json +from collections.abc import Mapping from types import SimpleNamespace from typing import Final from unittest.mock import AsyncMock, MagicMock @@ -12,12 +13,15 @@ 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_connections, aget_user_provider_credentials, + compare_and_set_user_provider_credential, decode_user_provider_credential, delete_user_provider_credential, delete_user_provider_credentials_for_credential, drop_user_provider_credential_cache, get_user_provider_credential, + get_user_provider_credential_row, invalidate_user_provider_credential_cache, set_user_provider_credential_cache, upsert_user_provider_credential, @@ -182,6 +186,108 @@ async def test_legacy_connection_without_github_host_defaults_to_github_com(monk } +@pytest.mark.asyncio +async def test_aget_connections_returns_refresh_metadata_and_legacy_defaults(monkeypatch): + from litellm.proxy.credential_endpoints import user_provider_credentials as upc + + refreshed: Final = GithubCopilotUserConnectionPayload( + access_token="gho_access", + github_login="octo", + refresh_token="ghr_refresh", + access_token_expires_at=1000.0, + refresh_token_expires_at=2000.0, + ) + legacy_ciphertext: Final = encrypt_value_helper(json.dumps({"access_token": "gho_legacy", "github_login": "octo"})) + table: Final = MagicMock() + table.find_many = AsyncMock( + return_value=[ + _row(credential_name="copilot-cred", credential_b64=upc._encode(refreshed)), + _row(credential_name="legacy-cred", credential_b64=legacy_ciphertext), + ] + ) + connections: Final = await aget_user_provider_connections( + _prisma(table), + DualCache(redis_cache=_fake_redis(monkeypatch)), + "user-a", + ["copilot-cred", "legacy-cred"], + ) + assert connections == { + "copilot-cred": refreshed, + "legacy-cred": GithubCopilotUserConnectionPayload( + access_token="gho_legacy", + github_login="octo", + ), + } + assert connections["legacy-cred"].refresh_token is None + assert connections["legacy-cred"].access_token_expires_at is None + + +@pytest.mark.asyncio +async def test_compare_and_set_updates_only_the_expected_ciphertext(): + from litellm.proxy.credential_endpoints import user_provider_credentials as upc + + initial: Final = GithubCopilotUserConnectionPayload(access_token="gho_old", github_login="octo") + initial_ciphertext: Final = upc._encode(initial) + stored: dict[str, SimpleNamespace] = {"row": _row(credential_b64=initial_ciphertext)} + + async def update_many(*, where: Mapping[str, str], data: Mapping[str, str]) -> int: + row = stored.get("row") + if ( + row is None + or row.user_id != where["user_id"] + or row.credential_name != where["credential_name"] + or row.credential_b64 != where["credential_b64"] + ): + return 0 + stored["row"] = _row(credential_b64=data["credential_b64"]) + return 1 + + async def find_unique(*, where: Mapping[str, Mapping[str, str]]) -> SimpleNamespace | None: + return stored.get("row") + + table: Final = MagicMock() + table.update_many = AsyncMock(side_effect=update_many) + table.find_unique = AsyncMock(side_effect=find_unique) + prisma_client: Final = _prisma(table) + rotated: Final = GithubCopilotUserConnectionPayload( + access_token="gho_new", + github_login="octo", + refresh_token="ghr_new", + ) + assert await compare_and_set_user_provider_credential( + prisma_client, + "user-a", + "copilot-cred", + initial_ciphertext, + rotated, + ) + assert (await get_user_provider_credential(prisma_client, "user-a", "copilot-cred")) == rotated + assert table.update_many.await_args.kwargs["where"] == { + "user_id": "user-a", + "credential_name": "copilot-cred", + "credential_b64": initial_ciphertext, + } + + +@pytest.mark.asyncio +async def test_compare_and_set_does_not_resurrect_a_missing_row(): + table: Final = MagicMock() + table.update_many = AsyncMock(return_value=0) + table.find_unique = AsyncMock(return_value=None) + prisma_client: Final = _prisma(table) + payload: Final = GithubCopilotUserConnectionPayload(access_token="gho_new", github_login="octo") + assert not await compare_and_set_user_provider_credential( + prisma_client, + "user-a", + "copilot-cred", + "expected-ciphertext", + payload, + ) + assert await get_user_provider_credential_row(prisma_client, "user-a", "copilot-cred") is None + table.update_many.assert_awaited_once() + table.upsert.assert_not_called() + + @pytest.mark.asyncio async def test_aget_credentials_negative_caches_missing_connection(monkeypatch): table = MagicMock() diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 2edb032b048..0010ddb32d7 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -10948,10 +10948,16 @@ 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": {"access_token": "gho_user_a_token", "github_host": "github.com"}} + return_value={ + "copilot-cred": upc.GithubCopilotUserConnectionPayload( + access_token="gho_user_a_token", + github_login="octo", + ) + } ) - monkeypatch.setattr(upc, "aget_user_provider_credentials", load) + monkeypatch.setattr(upc, "aget_user_provider_connections", load) router = Router( model_list=[ { @@ -11004,10 +11010,16 @@ 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": {"access_token": "gho_user_a_token", "github_host": "github.com"}} + return_value={ + "copilot-cred": upc.GithubCopilotUserConnectionPayload( + access_token="gho_user_a_token", + github_login="octo", + ) + } ) - monkeypatch.setattr(upc, "aget_user_provider_credentials", load) + monkeypatch.setattr(upc, "aget_user_provider_connections", load) router = Router( model_list=[ {