diff --git a/crates/tabby/src/serve.rs b/crates/tabby/src/serve.rs index 5713f3aa6036..9df6c528bbbc 100644 --- a/crates/tabby/src/serve.rs +++ b/crates/tabby/src/serve.rs @@ -92,6 +92,10 @@ pub struct ServeArgs { /// memory requirement e.g., GPU vRAM. #[clap(long, default_value_t = 1)] parallelism: u8, + + #[cfg(feature = "ee")] + #[clap(hide = true, long, default_value_t = false)] + webserver: bool, } pub async fn main(config: &Config, args: &ServeArgs) { @@ -114,7 +118,12 @@ pub async fn main(config: &Config, args: &ServeArgs) { .merge(SwaggerUi::new("/swagger-ui").url("/api-docs/openapi.json", ApiDoc::openapi())); #[cfg(feature = "ee")] - let (api, ui) = tabby_webserver::attach_webserver(api, ui, logger, code).await; + let (api, ui) = if args.webserver { + tabby_webserver::attach_webserver(api, ui, logger, code).await + } else { + let ui = ui.fallback(|| async { axum::response::Redirect::permanent("/swagger-ui") }); + (api, ui) + }; #[cfg(not(feature = "ee"))] let ui = ui.fallback(|| async { axum::response::Redirect::permanent("/swagger-ui") }); diff --git a/ee/tabby-ui/app/(dashboard)/page.tsx b/ee/tabby-ui/app/(dashboard)/page.tsx index 491bbedddb84..298f400680f0 100644 --- a/ee/tabby-ui/app/(dashboard)/page.tsx +++ b/ee/tabby-ui/app/(dashboard)/page.tsx @@ -7,7 +7,7 @@ import { WorkerKind } from '@/lib/gql/generates/graphql' import { useHealth } from '@/lib/hooks/use-health' import { useWorkers } from '@/lib/hooks/use-workers' import { useSession } from '@/lib/tabby/auth' -import { useGraphQLQuery } from '@/lib/tabby/gql' +import { useAuthenticatedGraphQLQuery, useGraphQLQuery } from '@/lib/tabby/gql' import { buttonVariants } from '@/components/ui/button' import { Dialog, @@ -90,7 +90,7 @@ const getRegistrationTokenDocument = graphql(/* GraphQL */ ` function MainPanel() { const { data: healthInfo } = useHealth() const workers = useWorkers(healthInfo) - const { data: registrationTokenRes } = useGraphQLQuery( + const { data: registrationTokenRes } = useAuthenticatedGraphQLQuery( getRegistrationTokenDocument ) diff --git a/ee/tabby-ui/app/auth/signup/components/signup.tsx b/ee/tabby-ui/app/auth/signup/components/signup.tsx index 0e09fa8d6571..4c6340f0c9e5 100644 --- a/ee/tabby-ui/app/auth/signup/components/signup.tsx +++ b/ee/tabby-ui/app/auth/signup/components/signup.tsx @@ -12,7 +12,7 @@ export default function Signup() { const title = isAdmin ? 'Create an admin account' : 'Create an account' const description = isAdmin - ? 'After creating an admin account, your instance is secured, and only registered users can access it.' + ? 'Your instance will be secured, only registered users can access it.' : 'Fill form below to create your account' if (isAdmin || invitationCode) { diff --git a/ee/tabby-ui/components/header.tsx b/ee/tabby-ui/components/header.tsx index 760b1e9d11f0..ad31e5ac3155 100644 --- a/ee/tabby-ui/components/header.tsx +++ b/ee/tabby-ui/components/header.tsx @@ -9,6 +9,7 @@ import { WorkerKind } from '@/lib/gql/generates/graphql' import { useHealth } from '@/lib/hooks/use-health' import { ReleaseInfo, useLatestRelease } from '@/lib/hooks/use-latest-release' import { useWorkers } from '@/lib/hooks/use-workers' +import { useAuthenticatedSession } from '@/lib/tabby/auth' import { cn } from '@/lib/utils' import { buttonVariants } from '@/components/ui/button' import { IconGitHub, IconNotice } from '@/components/ui/icons' @@ -16,6 +17,9 @@ import { IconGitHub, IconNotice } from '@/components/ui/icons' import { ThemeToggle } from './theme-toggle' export function Header() { + // Ensure login status. + useAuthenticatedSession() + const { data } = useHealth() const workers = useWorkers(data) const isChatEnabled = has(workers, WorkerKind.Chat) diff --git a/ee/tabby-ui/components/prompt-form.tsx b/ee/tabby-ui/components/prompt-form.tsx index e5a28162c4b2..8d10c4401f50 100644 --- a/ee/tabby-ui/components/prompt-form.tsx +++ b/ee/tabby-ui/components/prompt-form.tsx @@ -4,7 +4,7 @@ import { debounce, has, isEqual } from 'lodash-es' import useSWR from 'swr' import { useEnterSubmit } from '@/lib/hooks/use-enter-submit' -import { useSession } from '@/lib/tabby/auth' +import { useAuthenticatedApi, useSession } from '@/lib/tabby/auth' import fetcher from '@/lib/tabby/fetcher' import type { ISearchHit, SearchReponse } from '@/lib/types' import { cn } from '@/lib/utils' @@ -56,9 +56,8 @@ function PromptFormRenderer( Record >({}) - const { data } = useSession() const { data: completionData } = useSWR( - [queryCompletionUrl, data?.accessToken], + useAuthenticatedApi(queryCompletionUrl), fetcher, { revalidateOnFocus: false, diff --git a/ee/tabby-ui/components/user-panel.tsx b/ee/tabby-ui/components/user-panel.tsx index a0f2b70b87ce..26b73e2a09d0 100644 --- a/ee/tabby-ui/components/user-panel.tsx +++ b/ee/tabby-ui/components/user-panel.tsx @@ -1,55 +1,23 @@ import React from 'react' -import Link from 'next/link' -import { - useAuthenticatedSession, - useIsAdminInitialized, - useSession, - useSignOut -} from '@/lib/tabby/auth' -import { cn } from '@/lib/utils' +import { useAuthenticatedSession, useSignOut } from '@/lib/tabby/auth' -import { IconLogout, IconUnlock } from './ui/icons' +import { IconLogout } from './ui/icons' export default function UserPanel() { - const isAdminInitialized = useIsAdminInitialized() - - const Component = isAdminInitialized ? UserInfoPanel : EnableAdminPanel - - return ( -
- -
- ) -} - -function UserInfoPanel({ className }: React.ComponentProps<'span'>) { const session = useAuthenticatedSession() const signOut = useSignOut() return ( session && ( - - - +
+ + + + + {session.email} - {session.email} - +
) ) } - -function EnableAdminPanel({ className }: React.ComponentProps<'span'>) { - return ( - - Secure Access - - ) -} diff --git a/ee/tabby-ui/lib/hooks/use-health.tsx b/ee/tabby-ui/lib/hooks/use-health.tsx index 5368731b1cd3..64f204991bce 100644 --- a/ee/tabby-ui/lib/hooks/use-health.tsx +++ b/ee/tabby-ui/lib/hooks/use-health.tsx @@ -4,7 +4,7 @@ import useSWR, { SWRResponse } from 'swr' import fetcher from '@/lib/tabby/fetcher' -import { useSession } from '../tabby/auth' +import { useAuthenticatedApi, useSession } from '../tabby/auth' export interface HealthInfo { device: 'metal' | 'cpu' | 'cuda' @@ -20,6 +20,5 @@ export interface HealthInfo { } export function useHealth(): SWRResponse { - const { data } = useSession() - return useSWR(['/v1/health', data?.accessToken], fetcher) + return useSWR(useAuthenticatedApi('/v1/health'), fetcher) } diff --git a/ee/tabby-ui/lib/hooks/use-workers.ts b/ee/tabby-ui/lib/hooks/use-workers.ts index 42c92cf3c814..b2cf93f56958 100644 --- a/ee/tabby-ui/lib/hooks/use-workers.ts +++ b/ee/tabby-ui/lib/hooks/use-workers.ts @@ -3,7 +3,7 @@ import { findIndex, groupBy, slice } from 'lodash-es' import { graphql } from '@/lib/gql/generates' import { Worker, WorkerKind } from '@/lib/gql/generates/graphql' -import { useGraphQLQuery } from '@/lib/tabby/gql' +import { useAuthenticatedGraphQLQuery, useGraphQLQuery } from '@/lib/tabby/gql' import type { HealthInfo } from './use-health' @@ -44,7 +44,7 @@ export const getAllWorkersDocument = graphql(/* GraphQL */ ` `) function useWorkers(healthInfo?: HealthInfo) { - const { data } = useGraphQLQuery(getAllWorkersDocument) + const { data } = useAuthenticatedGraphQLQuery(getAllWorkersDocument) let workers = data?.workers const groupedWorkers = React.useMemo(() => { diff --git a/ee/tabby-ui/lib/tabby/auth.tsx b/ee/tabby-ui/lib/tabby/auth.tsx index 0b6d3d1bab9d..30083deb74bf 100644 --- a/ee/tabby-ui/lib/tabby/auth.tsx +++ b/ee/tabby-ui/lib/tabby/auth.tsx @@ -251,20 +251,15 @@ export const getIsAdminInitialized = graphql(/* GraphQL */ ` } `) -function useIsAdminInitialized() { - const { data } = useGraphQLQuery(getIsAdminInitialized) - return data?.isAdminInitialized -} - function useAuthenticatedSession() { const { data } = useGraphQLQuery(getIsAdminInitialized) const router = useRouter() const { data: session, status } = useSession() React.useEffect(() => { - if (!data?.isAdminInitialized) return - - if (status === 'unauthenticated') { + if (data?.isAdminInitialized === false) { + router.replace('/auth/signup?isAdmin=true') + } else if (status === 'unauthenticated') { router.replace('/auth/signin') } }, [data, status]) @@ -272,6 +267,11 @@ function useAuthenticatedSession() { return session } +function useAuthenticatedApi(path: string | null): [string, string] | null { + const { data, status } = useSession() + return path && status === 'authenticated' ? [path, data.accessToken] : null +} + export type { AuthStore, User, Session } export { @@ -279,6 +279,6 @@ export { useSignIn, useSignOut, useSession, - useIsAdminInitialized, - useAuthenticatedSession + useAuthenticatedSession, + useAuthenticatedApi } diff --git a/ee/tabby-ui/lib/tabby/fetcher.ts b/ee/tabby-ui/lib/tabby/fetcher.ts index 13b19d548cb2..394214d6fce3 100644 --- a/ee/tabby-ui/lib/tabby/fetcher.ts +++ b/ee/tabby-ui/lib/tabby/fetcher.ts @@ -1,14 +1,13 @@ -export default function tokenFetcher([url, token]: Array< - string | undefined ->): Promise { +export default function tokenFetcher([url, token]: [ + string, + string +]): Promise { const headers = new Headers() - if (token) { - headers.append('authorization', `Bearer ${token}`) - } + headers.append('authorization', `Bearer ${token}`) if (process.env.NODE_ENV !== 'production') { url = `${process.env.NEXT_PUBLIC_TABBY_SERVER_URL}${url}` } - return fetch(url!, { headers }).then(x => x.json()) + return fetch(url, { headers }).then(x => x.json()) } diff --git a/ee/tabby-ui/lib/tabby/gql.ts b/ee/tabby-ui/lib/tabby/gql.ts index 6cb7dda42572..af02901eba9b 100644 --- a/ee/tabby-ui/lib/tabby/gql.ts +++ b/ee/tabby-ui/lib/tabby/gql.ts @@ -73,18 +73,37 @@ export function useGraphQLQuery< variables?: TVariables, swrConfiguration?: SWRConfiguration ): SWRResponse { - const { data } = useSession() return useSWR( - [document, variables, data?.accessToken], + [document, variables], + ([document, variables]) => + gqlClient.request({ + document, + variables + }), + swrConfiguration + ) +} + +export function useAuthenticatedGraphQLQuery< + TResult, + TVariables extends Variables | undefined +>( + document: TypedDocumentNode, + variables?: TVariables, + swrConfiguration?: SWRConfiguration +): SWRResponse { + const { data, status } = useSession() + return useSWR( + status === 'authenticated' + ? [document, variables, data?.accessToken] + : null, ([document, variables, accessToken]) => gqlClient.request({ document, variables, - requestHeaders: accessToken - ? { - authorization: `Bearer ${accessToken}` - } - : undefined + requestHeaders: { + authorization: `Bearer ${accessToken}` + } }), swrConfiguration ) diff --git a/ee/tabby-webserver/src/schema/mod.rs b/ee/tabby-webserver/src/schema/mod.rs index 70b286a45d48..eaad8e974d4f 100644 --- a/ee/tabby-webserver/src/schema/mod.rs +++ b/ee/tabby-webserver/src/schema/mod.rs @@ -72,36 +72,27 @@ pub struct Query; #[graphql_object(context = Context)] impl Query { async fn workers(ctx: &Context) -> Result> { - if ctx.locator.auth().is_admin_initialized().await? { - if let Some(claims) = &ctx.claims { - if claims.user_info().is_admin() { - let workers = ctx.locator.worker().list_workers().await; - return Ok(workers); - } + if let Some(claims) = &ctx.claims { + if claims.user_info().is_admin() { + let workers = ctx.locator.worker().list_workers().await; + return Ok(workers); } - Err(CoreError::Unauthorized( - "Only admin is able to read workers", - )) - } else { - Ok(ctx.locator.worker().list_workers().await) } + Err(CoreError::Unauthorized( + "Only admin is able to read workers", + )) } async fn registration_token(ctx: &Context) -> Result { - if ctx.locator.auth().is_admin_initialized().await? { - if let Some(claims) = &ctx.claims { - if claims.user_info().is_admin() { - let token = ctx.locator.worker().read_registration_token().await?; - return Ok(token); - } + if let Some(claims) = &ctx.claims { + if claims.user_info().is_admin() { + let token = ctx.locator.worker().read_registration_token().await?; + return Ok(token); } - Err(CoreError::Unauthorized( - "Only admin is able to read registeration_token", - )) - } else { - let token = ctx.locator.worker().read_registration_token().await?; - Ok(token) } + Err(CoreError::Unauthorized( + "Only admin is able to read registeration_token", + )) } async fn is_admin_initialized(ctx: &Context) -> Result { diff --git a/ee/tabby-webserver/src/service/mod.rs b/ee/tabby-webserver/src/service/mod.rs index 5e3f94de0d61..ae2d8ef3f699 100644 --- a/ee/tabby-webserver/src/service/mod.rs +++ b/ee/tabby-webserver/src/service/mod.rs @@ -48,10 +48,7 @@ impl ServerContext { async fn authorize_request(&self, request: &Request) -> bool { let path = request.uri().path(); - if (path.starts_with("/v1/") || path.starts_with("/v1beta/")) - // Authorization is enabled - && self.db_conn.is_admin_initialized().await.unwrap_or(false) - { + if path.starts_with("/v1/") || path.starts_with("/v1beta/") { let token = { let authorization = request .headers()