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 88

address greptile review feedback

Browse source
This commit is contained in:
Ishaan Jaffer 2026-03-10 20:28:23 -07:00
parent 2e560e958e
commit 4626862580
7 changed files with 275 additions and 57 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

59
litellm/proxy/management_endpoints/mcp_management_endpoints.py
View file

@ -36,6 +36,11 @@ from fastapi import (
)
from fastapi.responses import JSONResponse
try:
from prisma.errors import RecordNotFoundError
except ImportError:
RecordNotFoundError = Exception # type: ignore
import litellm
from litellm._logging import verbose_logger, verbose_proxy_logger
from litellm._uuid import uuid
@ -84,6 +89,7 @@ if MCP_AVAILABLE:
delete_user_credential,
get_all_mcp_servers_for_user,
get_mcp_server,
get_mcp_servers,
get_mcp_submissions,
get_user_oauth_credential,
list_user_oauth_credentials,
@ -1427,14 +1433,11 @@ if MCP_AVAILABLE:
expires_in=payload.expires_in,
scopes=payload.scopes,
)
from datetime import timedelta
from datetime import timezone as _tz
expires_at: Optional[str] = None
if payload.expires_in is not None:
expires_at = (
datetime.now(_tz.utc) + timedelta(seconds=payload.expires_in)
).isoformat()
# Read back the persisted record so the response reflects the stored
# expires_at rather than recomputing it here (which could diverge by
# milliseconds or if the storage logic ever adds a grace period).
stored = await get_user_oauth_credential(prisma_client, user_id, server_id)
expires_at: Optional[str] = stored.get("expires_at") if stored else None
return MCPOAuthUserCredentialStatus(
server_id=server_id,
has_credential=True,
@ -1463,10 +1466,15 @@ if MCP_AVAILABLE:
status_code=status.HTTP_400_BAD_REQUEST,
detail={"error": "User ID not found in token"},
)
try:
await delete_user_credential(prisma_client, user_id, server_id)
except Exception:
pass # Already gone
# Only delete if the stored credential is actually an OAuth2 token.
# This prevents accidentally deleting a BYOK credential if one exists
# for the same (user_id, server_id) pair.
cred_to_delete = await get_user_oauth_credential(prisma_client, user_id, server_id)
if cred_to_delete is not None:
try:
await delete_user_credential(prisma_client, user_id, server_id)
except RecordNotFoundError:
pass # Already gone — treat as a successful delete
return MCPOAuthUserCredentialStatus(
server_id=server_id,
has_credential=False,
@ -1485,8 +1493,6 @@ if MCP_AVAILABLE:
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""Return credential status (has_credential, expiry) without exposing the token."""
from datetime import timezone as _tz
prisma_client = get_prisma_client_or_throw(
"Database not connected. Connect a database to your proxy"
)
@ -1505,10 +1511,8 @@ if MCP_AVAILABLE:
is_expired = False
if expires_at:
try:
from datetime import datetime as _dt
exp = _dt.fromisoformat(expires_at)
is_expired = exp < _dt.now(_tz.utc)
exp = datetime.fromisoformat(expires_at)
is_expired = exp < datetime.now(timezone.utc)
except Exception:
pass
return MCPOAuthUserCredentialStatus(
@ -1530,8 +1534,6 @@ if MCP_AVAILABLE:
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""Return all servers the calling user has connected via OAuth2."""
from datetime import timezone as _tz
prisma_client = get_prisma_client_or_throw(
"Database not connected. Connect a database to your proxy"
)
@ -1544,13 +1546,12 @@ if MCP_AVAILABLE:
oauth_creds = await list_user_oauth_credentials(prisma_client, user_id)
if not oauth_creds:
return []
# Fetch server metadata for display names
# Fetch server metadata for display names — single batch query instead of N+1.
server_ids = [c["server_id"] for c in oauth_creds]
servers = {}
for sid in server_ids:
srv = await get_mcp_server(prisma_client, sid)
if srv is not None:
servers[sid] = srv
servers = {
srv.server_id: srv
for srv in await get_mcp_servers(prisma_client, server_ids)
}
items: List[MCPUserCredentialListItem] = []
for cred in oauth_creds:
sid = cred["server_id"]
@ -1559,9 +1560,7 @@ if MCP_AVAILABLE:
is_expired = False
if expires_at:
try:
from datetime import datetime as _dt
is_expired = _dt.fromisoformat(expires_at) < _dt.now(_tz.utc)
is_expired = datetime.fromisoformat(expires_at) < datetime.now(timezone.utc)
except Exception:
pass
items.append(
@ -1571,7 +1570,7 @@ if MCP_AVAILABLE:
alias=getattr(srv, "alias", None) if srv else None,
credential_type="oauth2",
has_credential=True,
expires_at=None if is_expired else expires_at,
expires_at=expires_at, # always pass the raw timestamp; client computes expiry state
connected_at=cred.get("connected_at"),
)
)

180
tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py
View file

@ -1854,3 +1854,183 @@ class TestValidateMCPRequiredFields:
_validate_mcp_required_fields(payload)
assert exc_info.value.status_code == 500
assert "source_Url" in str(exc_info.value.detail)
# ── OAuth user credential endpoint unit tests ──────────────────────────────────
def _make_user_auth(user_id: str = "user-abc") -> "UserAPIKeyAuth":
return UserAPIKeyAuth(
api_key="sk-test",
user_id=user_id,
user_role=LitellmUserRoles.INTERNAL_USER,
)
def _make_prisma_client():
"""Return a minimal mock PrismaClient accepted by get_prisma_client_or_throw."""
client = MagicMock()
client.db = MagicMock()
return client
@pytest.mark.asyncio
async def test_store_mcp_oauth_user_credential_returns_status():
"""store_mcp_oauth_user_credential persists the token and echoes back status."""
from litellm.proxy._types import (
MCPOAuthUserCredentialRequest,
MCPOAuthUserCredentialStatus,
)
if not mgmt_endpoints.MCP_AVAILABLE:
pytest.skip("MCP module not installed")
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
store_mcp_oauth_user_credential,
)
server_id = "srv-1"
user_id = "user-123"
stored_payload = {
"type": "oauth2",
"access_token": "tok",
"expires_at": "2099-01-01T00:00:00+00:00",
"connected_at": "2026-01-01T00:00:00+00:00",
"server_id": server_id,
}
mock_prisma = _make_prisma_client()
with (
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
return_value=mock_prisma,
),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server",
new=AsyncMock(return_value=generate_mock_mcp_server_db_record(server_id=server_id)),
),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.store_user_oauth_credential",
new=AsyncMock(return_value=None),
),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_user_oauth_credential",
new=AsyncMock(return_value=stored_payload),
),
):
result = await store_mcp_oauth_user_credential(
server_id=server_id,
payload=MCPOAuthUserCredentialRequest(
access_token="tok",
expires_in=3600,
),
user_api_key_dict=_make_user_auth(user_id),
)
assert isinstance(result, MCPOAuthUserCredentialStatus)
assert result.has_credential is True
assert result.server_id == server_id
# expires_at should come from the stored record, not be recomputed
assert result.expires_at == "2099-01-01T00:00:00+00:00"
@pytest.mark.asyncio
async def test_delete_mcp_oauth_user_credential_only_deletes_oauth():
"""delete_mcp_oauth_user_credential only deletes OAuth2 credentials, not BYOK."""
from litellm.proxy._types import MCPOAuthUserCredentialStatus
if not mgmt_endpoints.MCP_AVAILABLE:
pytest.skip("MCP module not installed")
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
delete_mcp_oauth_user_credential,
)
server_id = "srv-2"
user_id = "user-456"
delete_mock = AsyncMock(return_value=None)
# When get_user_oauth_credential returns None (no OAuth cred), delete should NOT be called.
with (
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
return_value=_make_prisma_client(),
),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_user_oauth_credential",
new=AsyncMock(return_value=None),
),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.delete_user_credential",
new=delete_mock,
),
):
result = await delete_mcp_oauth_user_credential(
server_id=server_id,
user_api_key_dict=_make_user_auth(user_id),
)
delete_mock.assert_not_called()
assert isinstance(result, MCPOAuthUserCredentialStatus)
assert result.has_credential is False
@pytest.mark.asyncio
async def test_list_mcp_user_credentials_batch_server_fetch():
"""list_mcp_user_credentials uses a single batch DB call, not N+1 queries."""
from litellm.proxy._types import MCPUserCredentialListItem
if not mgmt_endpoints.MCP_AVAILABLE:
pytest.skip("MCP module not installed")
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
list_mcp_user_credentials,
)
user_id = "user-789"
server_id = "srv-3"
stored_creds = [
{
"type": "oauth2",
"access_token": "tok",
"expires_at": "2099-01-01T00:00:00+00:00",
"connected_at": "2026-01-01T00:00:00+00:00",
"server_id": server_id,
}
]
mock_server = generate_mock_mcp_server_db_record(server_id=server_id, alias="My Server")
# get_mcp_servers (batch) should be called once; get_mcp_server (single) must not be called.
batch_mock = AsyncMock(return_value=[mock_server])
single_mock = AsyncMock(return_value=mock_server)
with (
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
return_value=_make_prisma_client(),
),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.list_user_oauth_credentials",
new=AsyncMock(return_value=stored_creds),
),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_servers",
new=batch_mock,
),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server",
new=single_mock,
),
):
result = await list_mcp_user_credentials(
user_api_key_dict=_make_user_auth(user_id),
)
batch_mock.assert_called_once()
single_mock.assert_not_called()
assert len(result) == 1
assert isinstance(result[0], MCPUserCredentialListItem)
assert result[0].server_id == server_id
assert result[0].alias == "My Server"
# expires_at should always be the raw timestamp (not set to None when expired)
assert result[0].expires_at == "2099-01-01T00:00:00+00:00"

5
ui/litellm-dashboard/src/components/chat/OAuthConnectModal.tsx
View file

@ -66,11 +66,12 @@ export const OAuthConnectModal: React.FC<OAuthConnectModalProps> = ({
scopes,
onSuccess: () => {
onSuccess(server.server_id);
handleClose();
// handleClose is invoked by the useEffect below to avoid calling it twice.
},
});
// If we return from OAuth callback, close the modal so the user isn't stuck.
// Close the modal whenever the flow reaches "success" — covers both the
// in-page path (status set directly) and the post-redirect resume path.
useEffect(() => {
if (status === "success") {
handleClose();

24
ui/litellm-dashboard/src/components/networking.tsx
View file

@ -9600,7 +9600,17 @@ export const storeMCPOAuthUserCredential = async (
});
if (!response.ok) {
const err = await response.json().catch(() => ({}));
throw new Error((err as { detail?: { error?: string } })?.detail?.error || "Failed to store OAuth credential");
const errObj = err as { detail?: unknown };
const detail = errObj?.detail;
const detailMsg =
Array.isArray(detail)
? detail.map((d: unknown) => (d && typeof d === "object" ? (d as Record<string, unknown>).msg ?? JSON.stringify(d) : String(d))).join("; ")
: typeof detail === "string"
? detail
: detail && typeof (detail as Record<string, unknown>).error === "string"
? (detail as Record<string, unknown>).error as string
: undefined;
throw new Error(detailMsg || "Failed to store OAuth credential");
}
return response.json();
};
@ -9618,7 +9628,17 @@ export const deleteMCPOAuthUserCredential = async (
});
if (!response.ok) {
const err = await response.json().catch(() => ({}));
throw new Error((err as { detail?: { error?: string } })?.detail?.error || "Failed to revoke OAuth credential");
const errObj = err as { detail?: unknown };
const detail = errObj?.detail;
const detailMsg =
Array.isArray(detail)
? detail.map((d: unknown) => (d && typeof d === "object" ? (d as Record<string, unknown>).msg ?? JSON.stringify(d) : String(d))).join("; ")
: typeof detail === "string"
? detail
: detail && typeof (detail as Record<string, unknown>).error === "string"
? (detail as Record<string, unknown>).error as string
: undefined;
throw new Error(detailMsg || "Failed to revoke OAuth credential");
}
return response.json();
};

12
ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.tsx
View file

@ -10,20 +10,10 @@ import {
registerMcpOAuthClient,
serverRootPath,
} from "@/components/networking";
import { extractErrorMessage } from "@/utils/errorUtils";
export type McpOAuthStatus = "idle" | "authorizing" | "exchanging" | "success" | "error";
function extractErrorMessage(err: unknown): string {
if (err instanceof Error) return err.message;
if (err && typeof err === "object") {
const e = err as Record<string, unknown>;
if (typeof e.detail === "string") return e.detail;
if (typeof e.message === "string") return e.message;
return JSON.stringify(err);
}
return String(err);
}
interface UseMcpOAuthFlowOptions {
accessToken: string | null;
getCredentials: () => {

17
ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.tsx
View file

@ -22,20 +22,10 @@ import {
storeMCPOAuthUserCredential,
} from "@/components/networking";
import NotificationsManager from "@/components/molecules/notifications_manager";
import { extractErrorMessage } from "@/utils/errorUtils";
export type UserMcpOAuthStatus = "idle" | "authorizing" | "exchanging" | "success" | "error";
function extractErrorMessage(err: unknown): string {
if (err instanceof Error) return err.message;
if (err && typeof err === "object") {
const e = err as Record<string, unknown>;
if (typeof e.detail === "string") return e.detail;
if (typeof e.message === "string") return e.message;
return JSON.stringify(err);
}
return String(err);
}
interface UseUserMcpOAuthFlowOptions {
accessToken: string;
serverId: string;
@ -89,8 +79,11 @@ const genChallenge = async (verifier: string) => {
const setStorage = (key: string, value: string) => {
try {
// Use sessionStorage only — do not write to localStorage.
// The flow state may contain the LiteLLM access token; writing it to
// localStorage would persist it across browser sessions and make it
// readable by any injected script (XSS).
window.sessionStorage.setItem(key, value);
window.localStorage.setItem(key, value);
} catch (_) {}
};

35
ui/litellm-dashboard/src/utils/errorUtils.ts Normal file
View file

@ -0,0 +1,35 @@
/**
* Shared error-message extraction utility.
*
* Handles the common shapes returned by LiteLLM / FastAPI:
* - Error instances (err.message)
* - { detail: "string" }
* - { detail: [{ msg, loc, type }] } (FastAPI 422)
* - { detail: { error: "string" } }
* - { message: "string" }
* - anything else → JSON.stringify / String()
*/
export function extractErrorMessage(err: unknown): string {
if (err instanceof Error) return err.message;
if (err && typeof err === "object") {
const e = err as Record<string, unknown>;
const detail = e.detail;
if (typeof detail === "string") return detail;
if (Array.isArray(detail)) {
return detail.map((d: unknown) => {
if (d && typeof d === "object") {
const item = d as Record<string, unknown>;
return typeof item.msg === "string" ? item.msg : JSON.stringify(d);
}
return String(d);
}).join("; ");
}
if (detail && typeof detail === "object") {
const detailObj = detail as Record<string, unknown>;
if (typeof detailObj.error === "string") return detailObj.error;
}
if (typeof e.message === "string") return e.message;
return JSON.stringify(err);
}
return String(err);
}