@datalayer/core
Version:
[](https://datalayer.io)
317 lines (316 loc) • 12 kB
JavaScript
/*
* Copyright (c) 2023-2025 Datalayer, Inc.
* Distributed under the terms of the Modified BSD License.
*/
import { Poll } from '@lumino/polling';
import { useStore } from 'zustand';
import { createStore } from 'zustand/vanilla';
import { ANONYMOUS_USER_TOKEN, ANONYMOUS_USER, asUser, IAMProvidersSpecs, } from '../../models';
import { getStoredToken, getStoredUser, loadRefreshTokenFromCookie, storeToken, storeUser, } from '../storage';
import { requestDatalayerAPI, } from '../../api/DatalayerApi';
import { getCookie, setCookie, deleteCookie } from '../../utils';
import { coreStore } from './CoreState';
/**
* Limit to warn about low credits in milliseconds.
*/
export const RESERVATION_WARNING_TIME_MS = 5 * 60_000;
export const iamStore = createStore((set, get) => {
return {
credits: undefined,
creditsReservations: [],
externalToken: loadRefreshTokenFromCookie() ?? undefined,
user: getStoredUser(),
token: getStoredToken(),
iamRunUrl: coreStore.getState().configuration?.iamRunUrl,
iamProvidersAuthorizationURL: {},
version: '',
isLoginInProgress: false,
addIAMProviderAuthorizationURL: (provider, authorizationURL) => {
set(state => ({
iamProvidersAuthorizationURL: {
...state.iamProvidersAuthorizationURL,
[provider]: authorizationURL,
},
}));
},
login: async (token) => {
const { refreshUserByToken, iamRunUrl, logout } = get();
// Set flag to prevent interference from automatic token refresh
set({ isLoginInProgress: true });
try {
const resp = await requestDatalayerAPI({
url: `${iamRunUrl}/api/iam/v1/login`,
method: 'POST',
body: { token },
});
if (resp.success && resp.token) {
await refreshUserByToken(resp.token);
}
else {
throw new Error('Invalid Token.');
}
}
catch (error) {
console.debug('Failed to login.', error);
if (error.name === 'RunResponseError' &&
error.response.status === 401) {
console.debug('Received 401 error - Logging out.');
logout();
}
throw error;
}
finally {
set({ isLoginInProgress: false });
}
},
logout: () => {
storeUser();
storeToken();
set({
credits: undefined,
creditsReservations: [],
user: ANONYMOUS_USER,
token: ANONYMOUS_USER_TOKEN,
// externalToken: ANONYMOUS_EXTERNAL_TOKEN, // !!! Do not do that... otherwise the automatic login with refresh_token cookie will break!!!
});
},
checkIAMToken: async (token) => {
if (get().token !== token) {
await get().refreshUserByToken(token);
}
},
setIAMProviderAccessToken: (provider, accessToken) => {
const { user } = get();
if (!user) {
throw Error('You need to be authenticated to set an Access Token for an IAM provider.');
}
const iamProvider = IAMProvidersSpecs.getProvider(provider);
if (accessToken) {
setCookie(iamProvider.accessTokenCookieName(user), accessToken);
}
else {
deleteCookie(iamProvider.accessTokenCookieName(user));
}
},
// TODO passing the user as param for now, could/should be changed? If so, check the profile pages are still working on refresh...
getIAMProviderAccessToken: (user, provider) => {
const iamProvider = IAMProvidersSpecs.getProvider(provider);
const cookieName = iamProvider.accessTokenCookieName(user);
const accessToken = getCookie(cookieName);
return accessToken;
},
updateUser: (user) => set((state) => {
const updatedState = {
user: {
...state.user,
...user,
},
};
/*
if (state.user?.email && !updatedState.user.email) {
updatedState.user.email = state.user.email;
}
*/
return updatedState;
}),
refreshCredits: async () => {
const { externalToken, token, iamRunUrl, logout } = get();
if (token) {
try {
const creditsRaw = await requestDatalayerAPI({
url: `${iamRunUrl}/api/iam/v1/usage/credits`,
token,
headers: externalToken
? {
'X-External-Token': externalToken,
}
: undefined,
});
const { credits, reservations: creditsReservations = [] } = creditsRaw;
let available = credits.quota !== null
? credits.quota - credits.credits
: credits.credits;
available -= creditsReservations.reduce((consumed, reservation) => consumed + reservation.credits, 0);
set({
credits: { ...credits, available: Math.max(0, available) },
creditsReservations: creditsReservations,
});
}
catch (error) {
console.error('Failed to refresh user credits.', error);
if (error.name === 'RunResponseError' &&
error.response.status === 401) {
console.error('Received 401, logging out.');
logout();
}
throw error;
}
}
else {
set({ credits: undefined });
}
},
refreshUser: async () => {
const { token, refreshUserByToken, logout } = get();
if (token) {
await refreshUserByToken(token);
}
else {
logout();
}
},
refreshUserByTokenStored: async () => {
const token = getStoredToken();
if (token) {
await get().refreshUserByToken(token);
}
else {
get().logout();
}
},
refreshUserByToken: async (token) => {
const { iamRunUrl, logout, isLoginInProgress } = get();
try {
const data = await requestDatalayerAPI({
url: `${iamRunUrl}/api/iam/v1/whoami`,
token,
});
const user = asUser(data.profile);
storeUser(user);
storeToken(token);
set(() => ({ user, token }));
// TODO Centralize User Setting Management.
const aiagentsRunUrl = user.settings?.aiAgentsUrl;
if (aiagentsRunUrl) {
coreStore.getState().setConfiguration({
aiagentsRunUrl,
});
}
}
catch (error) {
if (error.name === 'RunResponseError' &&
error.response.status === 401) {
console.debug('Invalid token - Received 401 error.');
}
else {
console.debug('Failed to fetch user identity.', error);
}
// Only logout if we're not in the middle of a login attempt
if (!isLoginInProgress) {
logout();
}
}
},
setExternalToken: (externalToken) => set((state) => {
return { externalToken };
}),
setLogin: (user, token) => set((state) => {
storeUser(user);
storeToken(token);
return {
user,
token,
};
}),
setVersion: version => {
if (version && !get().version) {
set(state => ({ version }));
}
},
};
});
export function useIAMStore(selector) {
return useStore(iamStore, selector);
}
// Poll user credits.
const creditsPoll = new Poll({
name: '@datalayer/ui:credits-refresh',
factory: () => iamStore.getState().refreshCredits(),
auto: false,
frequency: {
interval: 60 * 1000,
backoff: true,
max: 600 * 1000,
},
standby: () => (iamStore.getState().user?.id ? 'when-hidden' : true),
});
// Initialize the IAM store with the stored token if it is valid.
iamStore
.getState()
.refreshUserByTokenStored()
.catch(reason => {
console.error('Failed to refresh to validate the stored token.', reason);
})
.finally(() => {
const { externalToken, iamRunUrl, checkIAMToken, token } = iamStore.getState();
// If the stored token is invalid and an external token exists, try authenticating with it.
if (!token && externalToken) {
console.debug('Can not login with token - Trying with the external token.');
requestDatalayerAPI({
url: `${iamRunUrl}/api/iam/v1/login`,
method: 'POST',
body: { token: externalToken },
})
.then(response => {
if (response.token) {
checkIAMToken(response.token);
}
})
.catch(reason => {
console.debug('Can not login with token.', token, reason);
});
}
if (token) {
console.log('Logged in with token and external token.');
}
else {
console.debug('Failed to login with token and no external token available.');
}
// Start the credits poll in any case after trying to validate the user token.
creditsPoll.start();
// Force a refresh when the user comes back to the application tab.
// Useful for checkout platform redirecting to another tab to add credits.
const maybeSetupVisibilityListener = () => {
if (typeof window === 'undefined') {
return;
}
const doc = window.document;
if (!doc || typeof doc.addEventListener !== 'function') {
return;
}
const onVisibilityChange = () => {
if (!doc.hidden && iamStore.getState().user?.id) {
creditsPoll.refresh();
}
};
doc.addEventListener('visibilitychange', onVisibilityChange);
};
maybeSetupVisibilityListener();
});
// Connect the core store with the iam store.
coreStore.subscribe((state, prevState) => {
if (state.configuration?.iamRunUrl &&
state.configuration.iamRunUrl !== prevState.configuration?.iamRunUrl) {
const iamRunUrl = state.configuration.iamRunUrl;
console.log('Updating iamRunUrl with new value', iamRunUrl);
iamStore.setState({ iamRunUrl });
// Check the token is valid with the new server.
if (iamStore.getState().externalToken) {
iamStore
.getState()
.login(iamStore.getState().externalToken)
.catch(reason => {
console.error('Failed to refresh the user after updating the IAM RUN URL.', reason);
});
}
else {
iamStore
.getState()
.refreshUser()
.catch(reason => {
console.error('Failed to refresh the user after updating the IAM server URL.', reason);
});
}
}
});
export default useIAMStore;