mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
fix(mcp): OAuth2 chat connect - tools fetch, auth flow, and status fixes
- schema.prisma: add missing MCP table fields (approval_status, submitted_by, submitted_at, reviewed_at, review_notes) to prevent destructive migrations - rest_endpoints.py: inject user OAuth token via extra_headers for OAuth2 servers so tools list is populated; add server name->UUID resolution so MCPConnectPicker name lookups work - mcp_registry.json: fix Atlassian defaults (transport: http, url: .../v1/mcp) - ChatPage.tsx: read mcpOauthReturn param to init sidebarView="apps" on OAuth return, clean up param after mount - MCPAppsPanel.tsx: auto-add OAuth2 servers to selectedServers when credential detected; onConnect also enables server for chat; disconnect removes from selectedServers - mcp_servers.tsx: sort servers by created_at DESC - useUserMcpOAuthFlow.tsx: append mcpOauthReturn=apps to return URL so Apps panel is mounted on return
This commit is contained in:
parent
f8137174a4
commit
ba3e243e1d
7 changed files with 97 additions and 7 deletions
|
|
@ -315,6 +315,11 @@ model LiteLLM_MCPServerTable {
|
|||
is_byok Boolean @default(false)
|
||||
byok_description String[] @default([])
|
||||
byok_api_key_help_url String?
|
||||
approval_status String @default("approved")
|
||||
submitted_by String?
|
||||
submitted_at DateTime?
|
||||
reviewed_at DateTime?
|
||||
review_notes String?
|
||||
}
|
||||
|
||||
// Per-user BYOK credentials for MCP servers
|
||||
|
|
|
|||
|
|
@ -69,6 +69,40 @@ if MCP_AVAILABLE:
|
|||
return server_auth
|
||||
return mcp_auth_header
|
||||
|
||||
async def _get_user_oauth_extra_headers(
|
||||
server,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> Optional[Dict[str, str]]:
|
||||
"""
|
||||
For OAuth2 servers, look up the user's stored access token and return it
|
||||
as extra_headers {"Authorization": "Bearer <token>"} so that it reaches
|
||||
the MCP server the same way the admin "Add MCP / Authorize and Fetch" flow does.
|
||||
Returns None for non-OAuth2 servers or when no credential is stored.
|
||||
"""
|
||||
from litellm.types.mcp import MCPAuth
|
||||
|
||||
if getattr(server, "auth_type", None) != MCPAuth.oauth2:
|
||||
return None
|
||||
user_id = getattr(user_api_key_dict, "user_id", None)
|
||||
server_id = getattr(server, "server_id", None)
|
||||
if not user_id or not server_id:
|
||||
return None
|
||||
try:
|
||||
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(
|
||||
"Database not connected. Connect a database to use OAuth2 MCP tools."
|
||||
)
|
||||
cred = await get_user_oauth_credential(prisma_client, user_id, server_id)
|
||||
if cred and cred.get("access_token"):
|
||||
return {"Authorization": f"Bearer {cred['access_token']}"}
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
|
||||
def _create_tool_response_objects(tools, server_mcp_info):
|
||||
"""Helper function to create tool response objects."""
|
||||
return [
|
||||
|
|
@ -162,11 +196,13 @@ if MCP_AVAILABLE:
|
|||
server_auth_header,
|
||||
raw_headers: Optional[Dict[str, str]] = None,
|
||||
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
|
||||
extra_headers: Optional[Dict[str, str]] = None,
|
||||
):
|
||||
"""Helper function to get tools for a single server."""
|
||||
tools = await global_mcp_server_manager._get_tools_from_server(
|
||||
server=server,
|
||||
mcp_auth_header=server_auth_header,
|
||||
extra_headers=extra_headers,
|
||||
add_prefix=False,
|
||||
raw_headers=raw_headers,
|
||||
)
|
||||
|
|
@ -294,6 +330,13 @@ if MCP_AVAILABLE:
|
|||
|
||||
# If server_id is specified, only query that specific server
|
||||
if server_id:
|
||||
# Resolve a server name to its UUID if needed (MCPConnectPicker passes
|
||||
# server_name strings, but allowed_server_ids_set contains UUIDs).
|
||||
if server_id not in allowed_server_ids:
|
||||
_resolved = global_mcp_server_manager.get_mcp_server_by_name(server_id)
|
||||
if _resolved is not None and _resolved.server_id in set(allowed_server_ids):
|
||||
server_id = _resolved.server_id
|
||||
|
||||
if server_id not in allowed_server_ids:
|
||||
_server = global_mcp_server_manager.get_mcp_server_by_id(server_id)
|
||||
if (
|
||||
|
|
@ -333,6 +376,7 @@ if MCP_AVAILABLE:
|
|||
server_auth_header = _get_server_auth_header(
|
||||
server, mcp_server_auth_headers, mcp_auth_header
|
||||
)
|
||||
user_oauth_extra_headers = await _get_user_oauth_extra_headers(server, user_api_key_dict)
|
||||
|
||||
try:
|
||||
list_tools_result = await _get_tools_for_single_server(
|
||||
|
|
@ -340,6 +384,7 @@ if MCP_AVAILABLE:
|
|||
server_auth_header,
|
||||
raw_headers_from_request,
|
||||
user_api_key_dict,
|
||||
extra_headers=user_oauth_extra_headers,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.exception(
|
||||
|
|
@ -385,6 +430,7 @@ if MCP_AVAILABLE:
|
|||
server_auth_header = _get_server_auth_header(
|
||||
server, mcp_server_auth_headers, mcp_auth_header
|
||||
)
|
||||
user_oauth_extra_headers = await _get_user_oauth_extra_headers(server, user_api_key_dict)
|
||||
|
||||
try:
|
||||
tools_result = await _get_tools_for_single_server(
|
||||
|
|
@ -392,6 +438,7 @@ if MCP_AVAILABLE:
|
|||
server_auth_header,
|
||||
raw_headers_from_request,
|
||||
user_api_key_dict,
|
||||
extra_headers=user_oauth_extra_headers,
|
||||
)
|
||||
list_tools_result.extend(tools_result)
|
||||
except Exception as e:
|
||||
|
|
|
|||
|
|
@ -33,8 +33,8 @@
|
|||
"icon_url": "https://cdn.simpleicons.org/atlassian",
|
||||
"category": "Developer Tools",
|
||||
"registry_url": "https://registry.modelcontextprotocol.io/servers/com.atlassian%2Fatlassian-mcp-server",
|
||||
"transport": "sse",
|
||||
"url": "https://mcp.atlassian.com/v1/sse",
|
||||
"transport": "http",
|
||||
"url": "https://mcp.atlassian.com/v1/mcp",
|
||||
"env_vars": []
|
||||
},
|
||||
{
|
||||
|
|
|
|||
|
|
@ -144,7 +144,10 @@ const ChatPage: React.FC<ChatPageProps> = ({ accessToken, userRole, userId, user
|
|||
const [inputText, setInputText] = useState("");
|
||||
const [mcpPopoverOpen, setMcpPopoverOpen] = useState(false);
|
||||
const [sidebarCollapsed, setSidebarCollapsed] = useState(false);
|
||||
const [sidebarView, setSidebarView] = useState<"chats" | "apps" | "credentials">("chats");
|
||||
const _oauthReturn = searchParams?.get("mcpOauthReturn");
|
||||
const [sidebarView, setSidebarView] = useState<"chats" | "apps" | "credentials">(
|
||||
_oauthReturn === "apps" ? "apps" : "chats"
|
||||
);
|
||||
const [storageBannerDismissed, setStorageBannerDismissed] = useState(false);
|
||||
|
||||
// Comparison mode state (active when selectedModels.length > 1)
|
||||
|
|
@ -172,6 +175,15 @@ const ChatPage: React.FC<ChatPageProps> = ({ accessToken, userRole, userId, user
|
|||
renameConversation,
|
||||
} = useChatHistory(activeConversationId);
|
||||
|
||||
// Clean up the OAuth return param after it's been consumed
|
||||
useEffect(() => {
|
||||
if (_oauthReturn && typeof window !== "undefined") {
|
||||
const url = new URL(window.location.href);
|
||||
url.searchParams.delete("mcpOauthReturn");
|
||||
window.history.replaceState({}, "", url.toString());
|
||||
}
|
||||
}, []); // eslint-disable-line react-hooks/exhaustive-deps
|
||||
|
||||
// Load models
|
||||
useEffect(() => {
|
||||
if (!accessToken) return;
|
||||
|
|
|
|||
|
|
@ -155,6 +155,18 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange
|
|||
return () => { cancelled = true; };
|
||||
}, [accessToken]);
|
||||
|
||||
// Auto-enable oauth2 servers for the current chat session when a valid
|
||||
// credential is detected (either on mount or after a fresh OAuth sign-in).
|
||||
useEffect(() => {
|
||||
if (oauthConnected.size === 0) return;
|
||||
const namesToAdd = servers
|
||||
.filter((s) => oauthConnected.has(s.server_id) && !selectedServers.includes(nameOf(s)))
|
||||
.map(nameOf);
|
||||
if (namesToAdd.length > 0) {
|
||||
onChange([...selectedServers, ...namesToAdd]);
|
||||
}
|
||||
}, [oauthConnected]); // eslint-disable-line react-hooks/exhaustive-deps
|
||||
|
||||
const handleToggle = async (serverName: string, checked: boolean, serverId?: string) => {
|
||||
if (!checked) {
|
||||
onChange(selectedServers.filter((s) => s !== serverName));
|
||||
|
|
@ -283,6 +295,7 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange
|
|||
// Ignore — credential may already be gone; update UI regardless.
|
||||
}
|
||||
setOauthConnected((prev) => { const n = new Set(prev); n.delete(detailServer.server_id); return n; });
|
||||
onChange(selectedServers.filter((s) => s !== name));
|
||||
}}
|
||||
style={{ borderRadius: 8, fontWeight: 600, height: 38, minWidth: 110 }}
|
||||
>
|
||||
|
|
@ -292,7 +305,10 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange
|
|||
<OAuth2ConnectButton
|
||||
server={detailServer}
|
||||
accessToken={accessToken}
|
||||
onConnect={(id) => setOauthConnected((prev) => new Set(prev).add(id))}
|
||||
onConnect={(id) => {
|
||||
setOauthConnected((prev) => new Set(prev).add(id));
|
||||
handleToggle(name, true, detailServer.server_id);
|
||||
}}
|
||||
variant="button"
|
||||
/>
|
||||
)
|
||||
|
|
@ -517,7 +533,10 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange
|
|||
<OAuth2ConnectButton
|
||||
server={server}
|
||||
accessToken={accessToken}
|
||||
onConnect={(id) => setOauthConnected((prev) => new Set(prev).add(id))}
|
||||
onConnect={(id) => {
|
||||
setOauthConnected((prev) => new Set(prev).add(id));
|
||||
handleToggle(nameOf(server), true, server.server_id);
|
||||
}}
|
||||
variant="badge"
|
||||
/>
|
||||
)
|
||||
|
|
|
|||
|
|
@ -128,7 +128,12 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID })
|
|||
server.mcp_access_groups?.some((g: any) => (typeof g === "string" ? g === group : g && g.name === group)),
|
||||
);
|
||||
}
|
||||
setFilteredServers(filtered);
|
||||
const sorted = [...filtered].sort((a, b) => {
|
||||
if (!a.created_at) return 1;
|
||||
if (!b.created_at) return -1;
|
||||
return new Date(b.created_at).getTime() - new Date(a.created_at).getTime();
|
||||
});
|
||||
setFilteredServers(sorted);
|
||||
}, [serversWithHealth]);
|
||||
|
||||
// Handle team filter change
|
||||
|
|
|
|||
|
|
@ -179,7 +179,9 @@ export const useUserMcpOAuthFlow = ({
|
|||
};
|
||||
|
||||
setStorage(FLOW_STATE_KEY, JSON.stringify(flowState));
|
||||
setStorage(RETURN_URL_KEY, window.location.href);
|
||||
const returnUrl = new URL(window.location.href);
|
||||
returnUrl.searchParams.set("mcpOauthReturn", "apps");
|
||||
setStorage(RETURN_URL_KEY, returnUrl.toString());
|
||||
|
||||
window.location.href = authorizeUrl;
|
||||
} catch (err) {
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue