scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-08 08:08:37 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 80

feat(github_copilot): refresh expiring per-user GitHub tokens

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin 2026-10-21 00:19:17 +00:00
parent 4c8b90945f
commit 8336ae3604
12 changed files with 1644 additions and 42 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/constants.py
View file

@ -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__"

110
litellm/llms/github_copilot/per_user_auth.py
View file

@ -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:

8
litellm/proxy/credential_endpoints/endpoints.py
View file

@ -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")

357
litellm/proxy/credential_endpoints/github_copilot_token_refresh.py Normal file
View file

@ -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

83
litellm/proxy/credential_endpoints/user_provider_credentials.py
View file

@ -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

40
litellm/proxy/litellm_pre_call_utils.py
View file

@ -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

18
litellm/proxy/mcp_server/outbound_credentials/redis_refresh_coordinator.py
View file

@ -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:

139
tests/unit/llms/github_copilot/test_per_user_auth.py
View file

@ -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():

48
tests/unit/proxy/credential_endpoints/test_endpoints.py
View file

@ -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."""

753
tests/unit/proxy/credential_endpoints/test_github_copilot_token_refresh.py Normal file
View file

@ -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 == []

106
tests/unit/proxy/credential_endpoints/test_user_provider_credentials.py
View file

@ -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()

20
tests/unit/proxy/test_litellm_pre_call_utils.py
View file

@ -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=[
{