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:
Ishaan Jaffer 2026-03-11 18:20:29 -07:00
parent f8137174a4
commit ba3e243e1d
7 changed files with 97 additions and 7 deletions

View file

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

View file

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

View file

@ -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": []
},
{

View file

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

View file

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

View file

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

View file

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