diff --git a/ee/tabby-ui/components/providers.tsx b/ee/tabby-ui/components/providers.tsx
index 8922264..f26ea99 100644
--- a/ee/tabby-ui/components/providers.tsx
+++ b/ee/tabby-ui/components/providers.tsx
@@ -5,11 +5,14 @@ import { ThemeProvider as NextThemesProvider } from 'next-themes'
import { ThemeProviderProps } from 'next-themes/dist/types'
import { TooltipProvider } from '@/components/ui/tooltip'
+import { AuthProvider } from '@/lib/tabby/auth'
export function Providers({ children, ...props }: ThemeProviderProps) {
return (
- {children}
+
+ {children}
+
)
}
diff --git a/ee/tabby-ui/lib/gql/generates/gql.ts b/ee/tabby-ui/lib/gql/generates/gql.ts
index 4c6b1d5..2267b05 100644
--- a/ee/tabby-ui/lib/gql/generates/gql.ts
+++ b/ee/tabby-ui/lib/gql/generates/gql.ts
@@ -16,7 +16,9 @@ const documents = {
'\n query GetRegistrationToken {\n registrationToken\n }\n':
types.GetRegistrationTokenDocument,
'\n query GetWorkers {\n workers {\n kind\n name\n addr\n device\n arch\n cpuInfo\n cpuCount\n cudaDevices\n }\n }\n':
- types.GetWorkersDocument
+ types.GetWorkersDocument,
+ '\n mutation refreshToken($refreshToken: String!) {\n refreshToken(refreshToken: $refreshToken) {\n accessToken\n refreshToken\n }\n }\n':
+ types.RefreshTokenDocument
}
/**
@@ -45,6 +47,12 @@ export function graphql(
export function graphql(
source: '\n query GetWorkers {\n workers {\n kind\n name\n addr\n device\n arch\n cpuInfo\n cpuCount\n cudaDevices\n }\n }\n'
): (typeof documents)['\n query GetWorkers {\n workers {\n kind\n name\n addr\n device\n arch\n cpuInfo\n cpuCount\n cudaDevices\n }\n }\n']
+/**
+ * The graphql function is used to parse GraphQL queries into a document that can be used by GraphQL clients.
+ */
+export function graphql(
+ source: '\n mutation refreshToken($refreshToken: String!) {\n refreshToken(refreshToken: $refreshToken) {\n accessToken\n refreshToken\n }\n }\n'
+): (typeof documents)['\n mutation refreshToken($refreshToken: String!) {\n refreshToken(refreshToken: $refreshToken) {\n accessToken\n refreshToken\n }\n }\n']
export function graphql(source: string) {
return (documents as any)[source] ?? {}
diff --git a/ee/tabby-ui/lib/gql/generates/graphql.ts b/ee/tabby-ui/lib/gql/generates/graphql.ts
index c1233b4..6293c3b 100644
--- a/ee/tabby-ui/lib/gql/generates/graphql.ts
+++ b/ee/tabby-ui/lib/gql/generates/graphql.ts
@@ -163,6 +163,19 @@ export type GetWorkersQuery = {
}>
}
+export type RefreshTokenMutationVariables = Exact<{
+ refreshToken: Scalars['String']['input']
+}>
+
+export type RefreshTokenMutation = {
+ __typename?: 'Mutation'
+ refreshToken: {
+ __typename?: 'RefreshTokenResponse'
+ accessToken: string
+ refreshToken: string
+ }
+}
+
export const GetRegistrationTokenDocument = {
kind: 'Document',
definitions: [
@@ -214,3 +227,55 @@ export const GetWorkersDocument = {
}
]
} as unknown as DocumentNode
+export const RefreshTokenDocument = {
+ kind: 'Document',
+ definitions: [
+ {
+ kind: 'OperationDefinition',
+ operation: 'mutation',
+ name: { kind: 'Name', value: 'refreshToken' },
+ variableDefinitions: [
+ {
+ kind: 'VariableDefinition',
+ variable: {
+ kind: 'Variable',
+ name: { kind: 'Name', value: 'refreshToken' }
+ },
+ type: {
+ kind: 'NonNullType',
+ type: { kind: 'NamedType', name: { kind: 'Name', value: 'String' } }
+ }
+ }
+ ],
+ selectionSet: {
+ kind: 'SelectionSet',
+ selections: [
+ {
+ kind: 'Field',
+ name: { kind: 'Name', value: 'refreshToken' },
+ arguments: [
+ {
+ kind: 'Argument',
+ name: { kind: 'Name', value: 'refreshToken' },
+ value: {
+ kind: 'Variable',
+ name: { kind: 'Name', value: 'refreshToken' }
+ }
+ }
+ ],
+ selectionSet: {
+ kind: 'SelectionSet',
+ selections: [
+ { kind: 'Field', name: { kind: 'Name', value: 'accessToken' } },
+ { kind: 'Field', name: { kind: 'Name', value: 'refreshToken' } }
+ ]
+ }
+ }
+ ]
+ }
+ }
+ ]
+} as unknown as DocumentNode<
+ RefreshTokenMutation,
+ RefreshTokenMutationVariables
+>
diff --git a/ee/tabby-ui/lib/hooks/use-interval.ts b/ee/tabby-ui/lib/hooks/use-interval.ts
new file mode 100644
index 0000000..08f284b
--- /dev/null
+++ b/ee/tabby-ui/lib/hooks/use-interval.ts
@@ -0,0 +1,31 @@
+import * as React from 'react'
+
+export default function useInterval(
+ callback: () => void,
+ delay: number | null
+): React.MutableRefObject {
+ const savedCallback = React.useRef(callback)
+ const intervalRef = React.useRef(null)
+
+ // Remember the latest callback if it changes.
+ React.useEffect(() => {
+ savedCallback.current = callback
+ }, [callback])
+
+ // Set up the interval.
+ React.useEffect(() => {
+ const tick = () => savedCallback.current()
+
+ if (typeof delay === 'number') {
+ intervalRef.current = window.setInterval(tick, delay * 60 * 1000)
+ }
+
+ return () => {
+ if (intervalRef.current) {
+ window.clearTimeout(intervalRef.current)
+ }
+ }
+ }, [delay])
+
+ return intervalRef
+}
diff --git a/ee/tabby-ui/lib/tabby/auth.tsx b/ee/tabby-ui/lib/tabby/auth.tsx
new file mode 100644
index 0000000..16d4321
--- /dev/null
+++ b/ee/tabby-ui/lib/tabby/auth.tsx
@@ -0,0 +1,237 @@
+import * as React from 'react'
+import { graphql } from '@/lib/gql/generates'
+import useInterval from '@/lib/hooks/use-interval'
+import { gqlClient } from '@/lib/tabby-gql-client'
+import { jwtDecode } from 'jwt-decode'
+
+interface AuthData {
+ accessToken: string
+ refreshToken: string
+}
+
+type AuthState =
+ | {
+ status: 'authenticated'
+ data: AuthData
+ }
+ | {
+ status: 'loading' | 'unauthenticated'
+ data: null
+ }
+
+enum AuthActionType {
+ Init,
+ SignIn,
+ SignOut,
+ Refresh
+}
+
+interface InitAction {
+ type: AuthActionType.Init
+ data: AuthData | null
+}
+
+interface SignInAction {
+ type: AuthActionType.SignIn
+ data: AuthData
+}
+
+interface SignOutAction {
+ type: AuthActionType.SignOut
+}
+
+interface RefreshAction {
+ type: AuthActionType.Refresh
+ data: AuthData
+}
+
+type AuthActions = InitAction | SignInAction | SignOutAction | RefreshAction
+
+function authReducer(state: AuthState, action: AuthActions): AuthState {
+ switch (action.type) {
+ case AuthActionType.Init:
+ case AuthActionType.SignIn:
+ case AuthActionType.Refresh:
+ if (action.data) {
+ return {
+ status: 'authenticated',
+ data: action.data
+ }
+ } else {
+ return {
+ status: 'unauthenticated',
+ data: null
+ }
+ }
+ case AuthActionType.SignOut:
+ TokenStorage.reset()
+ return {
+ status: 'unauthenticated',
+ data: null
+ }
+ }
+}
+
+class TokenStorage {
+ static authName = '_tabby_auth'
+
+ initialState(): AuthData | null {
+ const authData = localStorage.getItem(TokenStorage.authName)
+ if (authData) {
+ return JSON.parse(authData)
+ } else {
+ return null
+ }
+ }
+
+ persist(state: AuthData) {
+ localStorage.setItem(TokenStorage.authName, JSON.stringify(state))
+ }
+
+ static reset() {
+ localStorage.removeItem(TokenStorage.authName)
+ }
+}
+
+interface AuthProviderProps {
+ children: React.ReactNode
+}
+
+interface AuthStore {
+ authState: AuthState | null
+ dispatch: React.Dispatch
+}
+
+const AuthContext = React.createContext(null)
+
+const refreshTokenMutation = graphql(/* GraphQL */ `
+ mutation refreshToken($refreshToken: String!) {
+ refreshToken(refreshToken: $refreshToken) {
+ accessToken
+ refreshToken
+ }
+ }
+`)
+
+async function doRefresh(token: string, dispatch: React.Dispatch) {
+ let action: AuthActions
+ try {
+ action = {
+ type: AuthActionType.Refresh,
+ data: (
+ await gqlClient.request(refreshTokenMutation, { refreshToken: token })
+ ).refreshToken
+ }
+ } catch (err) {
+ console.error('Failed to refresh token', err)
+ action = {
+ type: AuthActionType.SignOut
+ }
+ }
+
+ dispatch(action)
+}
+
+const AuthProvider: React.FunctionComponent = ({
+ children
+}) => {
+ const storage = new TokenStorage()
+
+ const [authState, dispatch] = React.useReducer(authReducer, {
+ status: 'loading',
+ data: null
+ })
+
+ React.useEffect(() => {
+ const data = storage.initialState()
+ if (data?.refreshToken) {
+ doRefresh(data.refreshToken, dispatch)
+ } else {
+ dispatch({ type: AuthActionType.Init, data: null })
+ }
+ }, [])
+
+ React.useEffect(() => {
+ authState.data && storage.persist(authState.data)
+ }, [authState])
+
+ useInterval(async () => {
+ if (authState.status !== 'authenticated') {
+ return
+ }
+
+ await doRefresh(authState.data.refreshToken, dispatch)
+ }, 5)
+
+ return (
+
+ {children}
+
+ )
+}
+
+function useAuthStore(): AuthStore {
+ const context = React.useContext(AuthContext)
+
+ if (!context) {
+ throw new Error(
+ 'AuthProvider is missing. Please add the AuthProvider at root level'
+ )
+ }
+
+ return context
+}
+
+function useSignIn(): (params: AuthData) => Promise {
+ const { dispatch } = useAuthStore()
+ return async data => {
+ dispatch({
+ type: AuthActionType.SignIn,
+ data
+ })
+
+ return true
+ }
+}
+
+function useSignOut(): () => Promise {
+ const { dispatch } = useAuthStore()
+ return async () => {
+ dispatch({ type: AuthActionType.SignOut })
+ }
+}
+
+interface User {
+ email: string
+ isAdmin: boolean
+}
+
+type Session =
+ | {
+ data: null
+ status: 'loading' | 'unauthenticated'
+ }
+ | {
+ data: User
+ status: 'authenticated'
+ }
+
+function useSession(): Session {
+ const { authState } = useAuthStore()
+ if (authState?.status == 'authenticated') {
+ const { user } = jwtDecode<{ user: User }>(authState.data.accessToken)
+ return {
+ data: user,
+ status: authState.status
+ }
+ } else {
+ return {
+ status: authState?.status ?? 'loading',
+ data: null
+ }
+ }
+}
+
+export type { AuthStore, User, Session }
+
+export { AuthProvider, useSignIn, useSignOut, useSession }
diff --git a/ee/tabby-ui/package.json b/ee/tabby-ui/package.json
index eacbe8a..50f702f 100644
--- a/ee/tabby-ui/package.json
+++ b/ee/tabby-ui/package.json
@@ -40,6 +40,7 @@
"focus-trap-react": "^10.1.1",
"graphql": "^16.8.1",
"graphql-request": "^6.1.0",
+ "jwt-decode": "^4.0.0",
"lodash-es": "^4.17.21",
"nanoid": "^4.0.2",
"next": "^13.4.7",
diff --git a/ee/tabby-ui/yarn.lock b/ee/tabby-ui/yarn.lock
index 7f4ca9b..7445eb0 100644
--- a/ee/tabby-ui/yarn.lock
+++ b/ee/tabby-ui/yarn.lock
@@ -4422,6 +4422,11 @@ jsonify@^0.0.1:
object.assign "^4.1.4"
object.values "^1.1.6"
+jwt-decode@^4.0.0:
+ version "4.0.0"
+ resolved "https://registry.yarnpkg.com/jwt-decode/-/jwt-decode-4.0.0.tgz#2270352425fd413785b2faf11f6e755c5151bd4b"
+ integrity sha512-+KJGIyHgkGuIq3IEBNftfhW/LfWhXUIY6OmyVWjliu5KH1y0fw7VQ8YndE2O4qZdMSd9SqbnC8GOcZEy0Om7sA==
+
katex@^0.16.0:
version "0.16.8"
resolved "https://registry.yarnpkg.com/katex/-/katex-0.16.8.tgz#89b453f40e8557f423f31a1009e9298dd99d5ceb"