scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-19 00:01:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 122

address greptile review feedback (greploop iteration 3)

Browse source
This commit is contained in:
Ishaan Jaffer 2026-03-11 18:57:55 -07:00
parent f5eac90665
commit 388508b269
3 changed files with 58 additions and 49 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
litellm/proxy/_experimental/mcp_server/db.py
View file

@ -628,6 +628,24 @@ async def store_user_oauth_credential(
)
def is_oauth_credential_expired(cred: Dict[str, Any]) -> bool:
"""Return True if the OAuth2 credential's access_token has expired.
Checks the ``expires_at`` ISO-format string stored in the credential payload.
Returns False when ``expires_at`` is absent or unparseable (treat as non-expired).
"""
expires_at = cred.get("expires_at")
if not expires_at:
return False
try:
exp_dt = datetime.fromisoformat(expires_at)
if exp_dt.tzinfo is None:
exp_dt = exp_dt.replace(tzinfo=timezone.utc)
return datetime.now(timezone.utc) > exp_dt
except (ValueError, TypeError):
return False
async def get_user_oauth_credential(
prisma_client: PrismaClient,
user_id: str,

52
litellm/proxy/_experimental/mcp_server/rest_endpoints.py
View file

@ -69,29 +69,6 @@ if MCP_AVAILABLE:
return server_auth
return mcp_auth_header
def _check_credential_expiry(
cred: Dict[str, Any],
user_id: str,
server_id: str,
) -> bool:
"""Return True if the credential is expired, False otherwise."""
expires_at = cred.get("expires_at")
if not expires_at:
return False
try:
exp_dt = datetime.fromisoformat(expires_at)
if exp_dt.tzinfo is None:
exp_dt = exp_dt.replace(tzinfo=timezone.utc)
if datetime.now(timezone.utc) > exp_dt:
verbose_logger.debug(
f"_get_user_oauth_extra_headers: token expired for "
f"user={user_id} server={server_id}"
)
return True
except (ValueError, TypeError):
pass
return False
async def _get_user_oauth_extra_headers(
server,
user_api_key_dict: UserAPIKeyAuth,
@ -114,12 +91,14 @@ if MCP_AVAILABLE:
if not user_id or not server_id:
return None
try:
from litellm.proxy._experimental.mcp_server.db import (
get_user_oauth_credential,
is_oauth_credential_expired,
)
if prefetched_creds is not None:
cred = prefetched_creds.get(server_id)
else:
from litellm.proxy._experimental.mcp_server.db import (
get_user_oauth_credential,
)
from litellm.proxy.utils import get_prisma_client_or_throw
prisma_client = get_prisma_client_or_throw(
@ -127,7 +106,11 @@ if MCP_AVAILABLE:
)
cred = await get_user_oauth_credential(prisma_client, user_id, server_id)
if cred and cred.get("access_token"):
if _check_credential_expiry(cred, user_id, server_id):
if is_oauth_credential_expired(cred):
verbose_logger.debug(
f"_get_user_oauth_extra_headers: token expired for "
f"user={user_id} server={server_id}"
)
return None
return {"Authorization": f"Bearer {cred['access_token']}"}
except Exception as e:
@ -517,8 +500,19 @@ if MCP_AVAILABLE:
},
)
# Pre-fetch all OAuth credentials for this user once to avoid N+1 DB queries.
prefetched_oauth_creds = await _prefetch_user_oauth_creds(user_api_key_dict)
# Pre-fetch OAuth credentials only when at least one allowed server uses OAuth2,
# to avoid an unnecessary DB round-trip on requests with no OAuth2 MCP servers.
_oauth2_server_ids = {
sid for sid in allowed_server_ids
if getattr(
global_mcp_server_manager.get_mcp_server_by_id(sid), "auth_type", None
) == MCPAuth.oauth2
}
prefetched_oauth_creds = (
await _prefetch_user_oauth_creds(user_api_key_dict)
if _oauth2_server_ids
else {}
)
# Query all servers the user has access to
errors = []

37
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -871,19 +871,6 @@ if MCP_AVAILABLE:
return allowed_mcp_servers
def _is_oauth_cred_expired(cred: Dict[str, Any]) -> bool:
"""Return True if the OAuth2 credential's access_token has expired."""
expires_at = cred.get("expires_at")
if not expires_at:
return False
try:
exp_dt = datetime.fromisoformat(expires_at)
if exp_dt.tzinfo is None:
exp_dt = exp_dt.replace(tzinfo=timezone.utc)
return datetime.now(timezone.utc) > exp_dt
except (ValueError, TypeError):
return False
async def _get_user_oauth_extra_headers_from_db(
server: MCPServer,
user_api_key_auth: Optional[UserAPIKeyAuth],
@ -904,12 +891,14 @@ if MCP_AVAILABLE:
if not user_id or not server_id:
return None
try:
from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415
get_user_oauth_credential,
is_oauth_credential_expired,
)
if prefetched_creds is not None:
cred = prefetched_creds.get(server_id)
else:
from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415
get_user_oauth_credential,
)
from litellm.proxy.utils import ( # noqa: PLC0415
get_prisma_client_or_throw,
)
@ -919,7 +908,7 @@ if MCP_AVAILABLE:
)
cred = await get_user_oauth_credential(prisma_client, user_id, server_id)
if cred and cred.get("access_token"):
if _is_oauth_cred_expired(cred):
if is_oauth_credential_expired(cred):
verbose_logger.debug(
f"_get_user_oauth_extra_headers_from_db: token expired for "
f"user={user_id} server={server_id}"
@ -1104,9 +1093,17 @@ if MCP_AVAILABLE:
mcp_servers=mcp_servers,
)
# Pre-fetch all OAuth credentials for this user once to avoid N+1 DB queries
# inside the asyncio.gather loop below.
_prefetched_oauth_creds = await _prefetch_oauth_creds_for_user(user_api_key_auth)
# Pre-fetch OAuth credentials only when at least one server uses OAuth2,
# to avoid an unnecessary DB round-trip on requests with no OAuth2 MCP servers.
_has_oauth2_server = any(
getattr(s, "auth_type", None) == MCPAuth.oauth2
for s in allowed_mcp_servers
)
_prefetched_oauth_creds = (
await _prefetch_oauth_creds_for_user(user_api_key_auth)
if _has_oauth2_server
else {}
)
async def _fetch_and_filter_server_tools(
server: MCPServer,