From 8c02c223739e6854e19c023c99dba85dec6529c8 Mon Sep 17 00:00:00 2001 From: Meng Zhang Date: Sat, 9 Dec 2023 00:15:26 +0800 Subject: [PATCH] feat(tabby-ui): properly set accessToken for api access if user is login (#989) * feat(tabby-ui): properly set accessToken for api access if user is login * [autofix.ci] apply automated fixes * fix: use return data from useSWR (#991) * [autofix.ci] apply automated fixes --------- Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> Co-authored-by: aliang <1098486429@qq.com> --- ee/tabby-ui/components/prompt-form.tsx | 25 ++++++++++++-------- ee/tabby-ui/lib/hooks/use-health.tsx | 8 ++++--- ee/tabby-ui/lib/hooks/use-patch-fetch.ts | 24 +++++++++++++++---- ee/tabby-ui/lib/tabby/auth.tsx | 4 +++- ee/tabby-ui/lib/tabby/fetcher.ts | 19 +++++++++------ ee/tabby-ui/lib/tabby/gql.ts | 30 ++++++++++++++++++++---- 6 files changed, 80 insertions(+), 30 deletions(-) diff --git a/ee/tabby-ui/components/prompt-form.tsx b/ee/tabby-ui/components/prompt-form.tsx index 7b29e11..e5a2816 100644 --- a/ee/tabby-ui/components/prompt-form.tsx +++ b/ee/tabby-ui/components/prompt-form.tsx @@ -1,9 +1,10 @@ import * as React from 'react' import { UseChatHelpers } from 'ai/react' -import { debounce, has } from 'lodash-es' +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 fetcher from '@/lib/tabby/fetcher' import type { ISearchHit, SearchReponse } from '@/lib/types' import { cn } from '@/lib/utils' @@ -45,7 +46,6 @@ function PromptFormRenderer( const [queryCompletionUrl, setQueryCompletionUrl] = React.useState< string | null >(null) - const latestFetchKey = React.useRef('') const inputRef = React.useRef(null) // store the input selection for replacing inputValue const prevInputSelectionEnd = React.useRef() @@ -56,14 +56,20 @@ function PromptFormRenderer( Record >({}) - useSWR(queryCompletionUrl, fetcher, { - revalidateOnFocus: false, - dedupingInterval: 0, - onSuccess: (data, key) => { - if (key !== latestFetchKey.current) return - setOptions(data?.hits ?? []) + const { data } = useSession() + const { data: completionData } = useSWR( + [queryCompletionUrl, data?.accessToken], + fetcher, + { + revalidateOnFocus: false, + dedupingInterval: 0, + errorRetryCount: 0 } - }) + ) + + React.useEffect(() => { + setOptions(completionData?.hits ?? []) + }, [completionData?.hits]) React.useImperativeHandle(ref, () => { return { @@ -102,7 +108,6 @@ function PromptFormRenderer( if (queryName) { const query = encodeURIComponent(`name:${queryName} AND kind:function`) const url = `/v1beta/search?q=${query}` - latestFetchKey.current = url setQueryCompletionUrl(url) } else { setOptions([]) diff --git a/ee/tabby-ui/lib/hooks/use-health.tsx b/ee/tabby-ui/lib/hooks/use-health.tsx index 459dcb6..5368731 100644 --- a/ee/tabby-ui/lib/hooks/use-health.tsx +++ b/ee/tabby-ui/lib/hooks/use-health.tsx @@ -1,10 +1,11 @@ 'use client' -import { SWRResponse } from 'swr' -import useSWRImmutable from 'swr/immutable' +import useSWR, { SWRResponse } from 'swr' import fetcher from '@/lib/tabby/fetcher' +import { useSession } from '../tabby/auth' + export interface HealthInfo { device: 'metal' | 'cpu' | 'cuda' model?: string @@ -19,5 +20,6 @@ export interface HealthInfo { } export function useHealth(): SWRResponse { - return useSWRImmutable('/v1/health', fetcher) + const { data } = useSession() + return useSWR(['/v1/health', data?.accessToken], fetcher) } diff --git a/ee/tabby-ui/lib/hooks/use-patch-fetch.ts b/ee/tabby-ui/lib/hooks/use-patch-fetch.ts index 408f4bc..61ba2ed 100644 --- a/ee/tabby-ui/lib/hooks/use-patch-fetch.ts +++ b/ee/tabby-ui/lib/hooks/use-patch-fetch.ts @@ -6,29 +6,43 @@ import { type AIStreamCallbacksAndOptions } from 'ai' +import { useSession } from '../tabby/auth' + const serverUrl = process.env.NEXT_PUBLIC_TABBY_SERVER_URL || '' export function usePatchFetch() { + const { data } = useSession() + useEffect(() => { - const fetch = window.fetch + if (!(window as any)._originFetch) { + ;(window as any)._originFetch = window.fetch + } + + const fetch = (window as any)._originFetch as typeof window.fetch window.fetch = async function (url, options) { if (url !== '/api/chat') { return fetch(url, options) } + const headers: HeadersInit = { + 'Content-Type': 'application/json' + } + + if (data?.accessToken) { + headers['Authorization'] = `Bearer ${data?.accessToken}` + } + const res = await fetch(`${serverUrl}/v1beta/chat/completions`, { ...options, method: 'POST', - headers: { - 'Content-Type': 'application/json' - } + headers }) const stream = StreamAdapter(res, undefined) return new StreamingTextResponse(stream) } - }, []) + }, [data?.accessToken]) } const utf8Decoder = new TextDecoder('utf-8') diff --git a/ee/tabby-ui/lib/tabby/auth.tsx b/ee/tabby-ui/lib/tabby/auth.tsx index 43261c6..0b6d3d1 100644 --- a/ee/tabby-ui/lib/tabby/auth.tsx +++ b/ee/tabby-ui/lib/tabby/auth.tsx @@ -210,6 +210,7 @@ function useSignOut(): () => Promise { interface User { email: string isAdmin: boolean + accessToken: string } type Session = @@ -231,7 +232,8 @@ function useSession(): Session { return { data: { email: user.email, - isAdmin: user.is_admin + isAdmin: user.is_admin, + accessToken: authState.data.accessToken }, status: authState.status } diff --git a/ee/tabby-ui/lib/tabby/fetcher.ts b/ee/tabby-ui/lib/tabby/fetcher.ts index e1c4671..13b19d5 100644 --- a/ee/tabby-ui/lib/tabby/fetcher.ts +++ b/ee/tabby-ui/lib/tabby/fetcher.ts @@ -1,9 +1,14 @@ -export default function fetcher(url: string): Promise { - if (process.env.NODE_ENV === 'production') { - return fetch(url).then(x => x.json()) - } else { - return fetch(`${process.env.NEXT_PUBLIC_TABBY_SERVER_URL}${url}`).then(x => - x.json() - ) +export default function tokenFetcher([url, token]: Array< + string | undefined +>): Promise { + const headers = new Headers() + if (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()) } diff --git a/ee/tabby-ui/lib/tabby/gql.ts b/ee/tabby-ui/lib/tabby/gql.ts index bfe74a8..6cb7dda 100644 --- a/ee/tabby-ui/lib/tabby/gql.ts +++ b/ee/tabby-ui/lib/tabby/gql.ts @@ -3,6 +3,8 @@ import { GraphQLClient, Variables } from 'graphql-request' import { GraphQLResponse } from 'graphql-request/build/esm/types' import useSWR, { SWRConfiguration, SWRResponse } from 'swr' +import { useSession } from './auth' + export const gqlClient = new GraphQLClient( `${process.env.NEXT_PUBLIC_TABBY_SERVER_URL ?? ''}/graphql` ) @@ -26,10 +28,20 @@ export function useGraphQLForm< onError?: (path: string, message: string) => void } ) { - const onSubmit = async (values: TVariables) => { + const { data } = useSession() + const accessToken = data?.accessToken + const onSubmit = async (variables: TVariables) => { let res try { - res = await gqlClient.request(document, values) + res = await gqlClient.request({ + document, + variables, + requestHeaders: accessToken + ? { + authorization: `Bearer ${accessToken}` + } + : undefined + }) } catch (err) { const { errors = [] } = (err as any).response as GraphQLResponse for (const error of errors) { @@ -61,9 +73,19 @@ export function useGraphQLQuery< variables?: TVariables, swrConfiguration?: SWRConfiguration ): SWRResponse { + const { data } = useSession() return useSWR( - [document, variables], - ([document, variables]) => gqlClient.request(document, variables), + [document, variables, data?.accessToken], + ([document, variables, accessToken]) => + gqlClient.request({ + document, + variables, + requestHeaders: accessToken + ? { + authorization: `Bearer ${accessToken}` + } + : undefined + }), swrConfiguration ) }