Make OAuth start call modular

This commit is contained in:
Allen Zhou
2025-02-03 20:06:21 -08:00
parent 4c89aed4d9
commit 7957d9f577

View File

@@ -121,6 +121,14 @@ export function useConnection({
} }
}; };
const initiateOAuthFlow = async () => {
sessionStorage.removeItem(SESSION_KEYS.ACCESS_TOKEN);
sessionStorage.removeItem(SESSION_KEYS.REFRESH_TOKEN);
sessionStorage.setItem(SESSION_KEYS.SERVER_URL, sseUrl);
const redirectUrl = await startOAuthFlow(sseUrl);
window.location.href = redirectUrl;
};
const handleTokenRefresh = async () => { const handleTokenRefresh = async () => {
try { try {
const tokens = await refreshAccessToken(sseUrl); const tokens = await refreshAccessToken(sseUrl);
@@ -131,13 +139,7 @@ export function useConnection({
return tokens.access_token; return tokens.access_token;
} catch (error) { } catch (error) {
console.error("Token refresh failed:", error); console.error("Token refresh failed:", error);
// If refresh token is expired/invalid (401) or any other error, await initiateOAuthFlow();
// clear tokens and redirect to home to trigger re-authentication
sessionStorage.removeItem(SESSION_KEYS.ACCESS_TOKEN);
sessionStorage.removeItem(SESSION_KEYS.REFRESH_TOKEN);
sessionStorage.setItem(SESSION_KEYS.SERVER_URL, sseUrl);
const redirectUrl = await startOAuthFlow(sseUrl);
window.location.href = redirectUrl;
throw error; throw error;
} }
}; };
@@ -182,25 +184,19 @@ export function useConnection({
const response = await fetch(url, { ...init, headers }); const response = await fetch(url, { ...init, headers });
if (response.status === 401) { if (response.status === 401) {
// First try to refresh if we have a refresh token
if (sessionStorage.getItem(SESSION_KEYS.REFRESH_TOKEN)) { if (sessionStorage.getItem(SESSION_KEYS.REFRESH_TOKEN)) {
try { try {
const newAccessToken = await handleTokenRefresh(); const newAccessToken = await handleTokenRefresh();
headers["Authorization"] = `Bearer ${newAccessToken}`; headers["Authorization"] = `Bearer ${newAccessToken}`;
// Retry the request with new token
return fetch(url, { ...init, headers }); return fetch(url, { ...init, headers });
} catch (error) { } catch (error) {
console.error("Token refresh failed:", error); console.error("Token refresh failed:", error);
} }
} }
// If we have an access token but refresh failed or wasn't available,
// we need to re-authenticate since the token is invalid
if (sessionStorage.getItem(SESSION_KEYS.ACCESS_TOKEN)) { if (sessionStorage.getItem(SESSION_KEYS.ACCESS_TOKEN)) {
sessionStorage.setItem(SESSION_KEYS.SERVER_URL, sseUrl); await initiateOAuthFlow();
const redirectUrl = await startOAuthFlow(sseUrl); return new Response();
window.location.href = redirectUrl;
return new Response(); // This won't actually be used due to redirect
} }
} }
@@ -231,13 +227,9 @@ export function useConnection({
} catch (error) { } catch (error) {
console.error("Failed to connect to MCP server:", error); console.error("Failed to connect to MCP server:", error);
if (error instanceof SseError && error.code === 401) { if (error instanceof SseError && error.code === 401) {
// Store the server URL for the callback handler await initiateOAuthFlow();
sessionStorage.setItem(SESSION_KEYS.SERVER_URL, sseUrl);
const redirectUrl = await startOAuthFlow(sseUrl);
window.location.href = redirectUrl;
return; return;
} }
throw error; throw error;
} }