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

@@ -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<HTMLTextAreaElement>(null)
   // store the input selection for replacing inputValue
   const prevInputSelectionEnd = React.useRef<number>()
@@ -56,14 +56,20 @@ function PromptFormRenderer(
     Record<string, ISearchHit>
   >({})
 
-  useSWR<SearchReponse>(queryCompletionUrl, fetcher, {
-    revalidateOnFocus: false,
-    dedupingInterval: 0,
-    onSuccess: (data, key) => {
-      if (key !== latestFetchKey.current) return
-      setOptions(data?.hits ?? [])
+  const { data } = useSession()
+  const { data: completionData } = useSWR<SearchReponse>(
+    [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([])

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<HealthInfo> {
-  return useSWRImmutable('/v1/health', fetcher)
+  const { data } = useSession()
+  return useSWR(['/v1/health', data?.accessToken], fetcher)
 }

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')

ee/tabby-ui/lib/tabby/auth.tsx

@@ -210,6 +210,7 @@ function useSignOut(): () => Promise<void> {
 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
     }

ee/tabby-ui/lib/tabby/fetcher.ts

@@ -1,9 +1,14 @@
-export default function fetcher(url: string): Promise<any> {
-  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<any> {
+  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())
 }

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<TResult>
 ): SWRResponse<TResult> {
+  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
   )
 }