Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/_experimental/mcp_server/proxy_api_credentials.py: 22%

1"""The proxy-API side of the native-client sign-in: turning a consented OAuth grant into 

2the same per-user credential ``lite login`` stores, so the bearer a CLI obtains through 

3the browser flow is accepted on every proxy route with user and team attribution.""" 

4 

5from __future__ import annotations 

6 

7from collections.abc import Sequence 

8from typing import Final 

9 

10from litellm.constants import CLI_JWT_EXPIRATION_HOURS 

11from litellm.proxy._experimental.mcp_server.bridge_token_flow import load_active_user_by_id 

12from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( 

13 ConsentTeam, 

14 MintedProxyCredential, 

15 ProxyCredentialMintFailure, 

16 ReloadUserFailure, 

17) 

18from litellm.proxy._types import LiteLLM_UserTable 

19from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken, effective_user_role 

20from litellm.proxy.management_endpoints.ui_sso import ( 

21 CliSsoTeamDetail, 

22 fetch_cli_sso_team_details, 

23 selected_cli_sso_team_detail, 

24) 

25 

26 

27async def lookup_consent_teams(user_id: str) -> tuple[ConsentTeam, ...] | ReloadUserFailure: 

28 user: Final = await load_active_user_by_id(user_id) 

29 if isinstance(user, str): 

30 return user 

31 details: Final = await _team_details(user.teams) 

32 if details is None: 

33 return "unavailable" 

34 return tuple( 

35 ConsentTeam(team_id=detail.team_id, team_alias=detail.team_alias) 

36 for detail in details 

37 if detail.team_id is not None 

38 ) 

39 

40 

41async def mint_proxy_credential( 

42 user_id: str, team_id: str | None 

43) -> MintedProxyCredential | ProxyCredentialMintFailure: 

44 """Mint the ``lite login`` credential for a consented grant. Membership is checked 

45 live against the database row, so a team the user left between consent and redemption (or between refreshes) 

46 refuses the grant instead of minting a credential attributed to a team they are no 

47 longer on. The team is exactly the one the consent page sealed into the grant; nothing 

48 is picked on the user's behalf here, so a refresh can never move the credential, and a 

49 grant that names no team is refused for a user with a live team to pick from (the same 

50 rule ``lite login`` applies), so a user cannot step outside their teams' attribution by 

51 posting the consent form without one. Memberships whose team rows are gone count as no 

52 team at all, the way ``lite login`` treats them, so they can never lock a user out. The 

53 user row handed to the minter carries no team list, exactly like ``lite login``'s, so 

54 the minter's own first-team fallback stays inert. The credential carries the role the 

55 proxy already enforces for the user on every request, so a row with no role (JWT auth's 

56 upsert writes none) mints as an internal user instead of being refused.""" 

57 user: Final = await load_active_user_by_id(user_id, source="database") 

58 if isinstance(user, str): 

59 return user 

60 if team_id is not None and team_id not in user.teams: 

61 return "not_a_member" 

62 details: Final = await _team_details(user.teams) if user.teams else () 

63 if details is None: 

64 return "unavailable" 

65 if team_id is None and any(detail.team_id is not None for detail in details): 

66 return "team_required" 

67 selected: Final = selected_cli_sso_team_detail(details, team_id) 

68 if selected is None: 

69 return "not_a_member" 

70 key: Final = ExperimentalUIJWTToken.get_cli_jwt_auth_token( 

71 user_info=LiteLLM_UserTable( 

72 user_id=user.user_id, user_role=effective_user_role(user.user_role).value, models=user.models 

73 ), 

74 team_id=team_id, 

75 team_alias=selected.team_alias, 

76 team_models=selected.team_models, 

77 team_model_aliases=selected.team_model_aliases, 

78 ) 

79 return MintedProxyCredential( 

80 key=key, 

81 expires_in=CLI_JWT_EXPIRATION_HOURS * 3600, 

82 user_id=user.user_id, 

83 team_id=team_id, 

84 ) 

85 

86 

87async def _team_details(teams: Sequence[str]) -> tuple[CliSsoTeamDetail, ...] | None: 

88 from litellm.proxy.proxy_server import prisma_client # noqa: PLC0415 # rebound after startup, so read it per call 

89 

90 if prisma_client is None: 

91 return None 

92 return await fetch_cli_sso_team_details(prisma_client, teams)