Spaces:
Running
on
CPU Upgrade
Running
on
CPU Upgrade
import { useEffect, useState, useTransition } from "react" | |
import { OAuthResult, oauthHandleRedirectIfPresent, oauthLoginUrl } from "@huggingface/hub" | |
import { useLocalStorage } from "usehooks-ts" | |
import { UserInfo } from "@/types/general" | |
import { useStore } from "./useStore" | |
import { localStorageKeys } from "./localStorageKeys" | |
import { defaultSettings } from "./defaultSettings" | |
import { getCurrentUser } from "../api/actions/users" | |
export function useCurrentUser({ | |
isLoginRequired = false | |
}: { | |
// set this to true, and the page will automatically redirect to the | |
// HF login page if the session is expired | |
isLoginRequired?: boolean | |
} = {}): { | |
user?: UserInfo | |
login: (redirectUrl?: string) => void | |
checkSession: (isLoginRequired: boolean) => Promise<UserInfo | undefined> | |
apiKey: string | |
oauthResult?: OAuthResult | |
// the long standing API is a temporary solution for "PRO" users of AiTube | |
// (users who use Clap files using external tools, | |
// or want ot use their own HF account to generate videos) | |
longStandingApiKey: string | |
setLongStandingApiKey: (apiKey: string, loginOnFailure: boolean) => void | |
} { | |
const [_pending, startTransition] = useTransition() | |
const user = useStore(s => s.currentUser) | |
const setCurrentUser = useStore(s => s.setCurrentUser) | |
const [oauthResult, setOauthResult] = useState<OAuthResult>() | |
const userId = `${user?.id || ""}` | |
// this is the legacy, long-standing API key | |
// which is still required for long generation of Clap files | |
const [huggingfaceApiKey, setHuggingfaceApiKey] = useLocalStorage<string>( | |
localStorageKeys.huggingfaceApiKey, | |
defaultSettings.huggingfaceApiKey | |
) | |
// this is the new recommended API to use, with short expiration rates | |
// in the future this API key will be enough for all our use cases | |
const [huggingfaceTemporaryApiKey, setHuggingfaceTemporaryApiKey] = useLocalStorage<string>( | |
localStorageKeys.huggingfaceTemporaryApiKey, | |
defaultSettings.huggingfaceTemporaryApiKey | |
) | |
// force the API call | |
const checkSession = async (isLoginRequired: boolean = false): Promise<UserInfo | undefined> => { | |
console.log("useCurrentUser.checkSession()") | |
let huggingfaceTemporaryApiKey = localStorage.getItem(localStorageKeys.huggingfaceTemporaryApiKey) || "" | |
console.log("huggingfaceTemporaryApiKey:", huggingfaceTemporaryApiKey) | |
if (huggingfaceTemporaryApiKey.startsWith('"')) { | |
console.log("the key has been corrupted..") | |
localStorage.setItem(localStorageKeys.huggingfaceTemporaryApiKey, JSON.parse(huggingfaceTemporaryApiKey)) | |
huggingfaceTemporaryApiKey = localStorage.getItem(localStorageKeys.huggingfaceTemporaryApiKey) || "" | |
console.log(`the recovered key is: ${huggingfaceTemporaryApiKey}`) | |
} | |
// new way: try to use the safer temporary key whenever possible | |
if (huggingfaceTemporaryApiKey) { | |
try { | |
console.log(`calling getCurrentUser()`, { huggingfaceTemporaryApiKey }) | |
const user = await getCurrentUser(huggingfaceTemporaryApiKey) | |
setCurrentUser(user) | |
return user // we stop there, no need to try the legacy key | |
} catch (err) { | |
console.error("failed to log in using the temporary key:", err) | |
setCurrentUser(undefined) | |
} | |
} | |
// deprecated: the old static key which is harder to renew | |
if (huggingfaceApiKey) { | |
try { | |
const user = await getCurrentUser(huggingfaceApiKey) | |
setCurrentUser(user) | |
return user | |
} catch (err) { | |
console.error("failed to log in using the static key:", err) | |
setCurrentUser(undefined) | |
} | |
} | |
// when we reach this stage, we know that none of the API tokens were valid | |
// we are given the choice to request a login or not | |
// (depending on if it's a secret page or not) | |
if (isLoginRequired) { | |
// await login("/") | |
} | |
return undefined | |
} | |
// can be called many times, but won't do the API call if not necessary | |
const main = (isLoginRequired: boolean) => { | |
// already logged-in, no need to spend an API call | |
// although it is worth noting that the API token might be expired at this stage | |
if (userId) { | |
console.log("we are already logged-in") | |
return | |
} | |
startTransition(async () => { | |
console.log("useCurrentUser(): yes, we need to call synchronizeSession()") | |
await checkSession(isLoginRequired) | |
}) | |
} | |
useEffect(() => { | |
main(isLoginRequired) | |
}, [isLoginRequired, huggingfaceApiKey, huggingfaceTemporaryApiKey, userId]) | |
useEffect(() => { | |
// DIY | |
try { | |
localStorage.setItem( | |
"huggingface.co:oauth:nonce", | |
localStorage.getItem("aitube.at:oauth:nonce") || "" | |
) | |
// localStorage.removeItem("aitube.at:oauth:nonce") | |
localStorage.setItem( | |
"huggingface.co:oauth:code_verifier", | |
localStorage.getItem("aitube.at:oauth:code_verifier") || "" | |
) | |
// localStorage.removeItem("aitube.at:oauth:code_verifier") | |
} catch (err) { | |
console.log("no pending oauth flow to finish") | |
} | |
console.log("useCurrentUser()") | |
const searchParams = new URLSearchParams(window.location.search); | |
console.log("debug:", { | |
"window.location.search:": window.location.search, | |
searchParams, | |
}) | |
const fn = async () => { | |
try { | |
const res = await oauthHandleRedirectIfPresent() | |
console.log("result of oauthHandleRedirectIfPresent:", res) | |
if (res) { | |
console.log("oauthHandleRedirectIfPresent returned something!", res) | |
setOauthResult(res) | |
console.log("debug:", { accessToken: res.accessToken }) | |
setHuggingfaceTemporaryApiKey(res.accessToken) | |
startTransition(async () => { | |
console.log("TODO julian do something, eg. reload the page, remove the things in the URL etc") | |
// await checkSession(isLoginRequired) | |
}) | |
} | |
} catch (err) { | |
console.error(err) | |
} | |
} | |
fn() | |
}, [isLoginRequired]) | |
const login = async ( | |
// used to redirect the user back to the route they were browsing | |
redirectTo: string = "" | |
) => { | |
const oauthUrl = await oauthLoginUrl({ | |
/** | |
* OAuth client ID. | |
* | |
* For static Spaces, you can omit this and it will be loaded from the Space config, as long as `hf_oauth: true` is present in the README.md's metadata. | |
* For other Spaces, it is available to the backend in the OAUTH_CLIENT_ID environment variable, as long as `hf_oauth: true` is present in the README.md's metadata. | |
* | |
* You can also create a Developer Application at https://huggingface.co/settings/connected-applications and use its client ID. | |
*/ | |
clientId: process.env.NEXT_PUBLIC_AI_TUBE_OAUTH_CLIENT_ID, | |
// hubUrl?: string; | |
/** | |
* OAuth scope, a list of space separate scopes. | |
* | |
* For static Spaces, you can omit this and it will be loaded from the Space config, as long as `hf_oauth: true` is present in the README.md's metadata. | |
* For other Spaces, it is available to the backend in the OAUTH_SCOPES environment variable, as long as `hf_oauth: true` is present in the README.md's metadata. | |
* | |
* Defaults to "openid profile". | |
* | |
* You can also create a Developer Application at https://huggingface.co/settings/connected-applications and use its scopes. | |
* | |
* See https://huggingface.co/docs/hub/oauth for a list of available scopes. | |
*/ | |
scopes: "openid profile", | |
/** | |
* Redirect URI, defaults to the current URL. | |
* | |
* For Spaces, any URL within the Space is allowed. | |
* | |
* For Developer Applications, you can add any URL you want to the list of allowed redirect URIs at https://huggingface.co/settings/connected-applications. | |
*/ | |
redirectUrl: `${process.env.NEXT_PUBLIC_DOMAIN}/api/login`, | |
/** | |
* State to pass to the OAuth provider, which will be returned in the call to `oauthLogin` after the redirect. | |
*/ | |
state: JSON.stringify({ redirectTo }) | |
}) | |
// DIY | |
localStorage.setItem( | |
"aitube.at:oauth:nonce", | |
localStorage.getItem("huggingface.co:oauth:nonce") || "" | |
) | |
localStorage.setItem( | |
"aitube.at:oauth:code_verifier", | |
localStorage.getItem("huggingface.co:oauth:code_verifier") || "" | |
) | |
// should we open this in a new tab? | |
window.location.href = oauthUrl | |
} | |
const setLongStandingApiKey = (apiKey: string, loginOnFailure: boolean) => { | |
(async () => { | |
try { | |
const user = await getCurrentUser(apiKey) | |
setHuggingfaceApiKey(apiKey) | |
setCurrentUser(user) | |
} catch (err) { | |
console.error("failed to log in using the long standing key:", err) | |
setHuggingfaceApiKey("") | |
setCurrentUser(undefined) | |
if (loginOnFailure) { | |
// login() | |
} | |
} | |
})() | |
} | |
// this may correspond to either a short or a long standing api key | |
// in the future it may always be a short api key with auto renewal | |
const apiKey = user?.hfApiToken || "" | |
// for now, we still need to keep track of a logn api, but this is purely optional | |
const longStandingApiKey = huggingfaceApiKey | |
return { | |
user, | |
login, | |
checkSession, | |
oauthResult, | |
apiKey, | |
longStandingApiKey, | |
setLongStandingApiKey, | |
} | |
} |