fix(ui): add token auto-refresh to the chat playground (#1625)

* fix(ui): Add token auto-refresh to the chat playground

* update

ee/tabby-ui/lib/hooks/use-patch-fetch.ts

@@ -1,17 +1,15 @@
 import { useEffect } from 'react'
 import { OpenAIStream, StreamingTextResponse } from 'ai'
 
-import { useSession } from '../tabby/auth'
+import fetcher from '../tabby/fetcher'
 
 export function usePatchFetch() {
-  const { data } = useSession()
-
   useEffect(() => {
-    if (!(window as any)._originFetch) {
-      ;(window as any)._originFetch = window.fetch
+    if (!window._originFetch) {
+      window._originFetch = window.fetch
     }
 
-    const fetch = (window as any)._originFetch as typeof window.fetch
+    const fetch = window._originFetch
 
     window.fetch = async function (url, options) {
       if (url !== '/api/chat') {
@@ -22,18 +20,25 @@ export function usePatchFetch() {
         'Content-Type': 'application/json'
       }
 
-      if (data?.accessToken) {
-        headers['Authorization'] = `Bearer ${data?.accessToken}`
-      }
-
-      const res = await fetch(`/v1beta/chat/completions`, {
+      const res = await fetcher(`/v1beta/chat/completions`, {
         ...options,
         method: 'POST',
-        headers
+        headers,
+        customFetch: fetch,
+        responseFormatter(response) {
+          const stream = OpenAIStream(response, undefined)
+          return new StreamingTextResponse(stream)
+        }
       })
 
-      const stream = OpenAIStream(res, undefined)
-      return new StreamingTextResponse(stream)
+      return res
+    }
+
+    return () => {
+      if (window?._originFetch) {
+        window.fetch = window._originFetch
+        window._originFetch = undefined
+      }
     }
-  }, [data?.accessToken])
+  }, [])
 }

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

@@ -8,23 +8,36 @@ import {
 } from './auth'
 import { client } from './gql'
 
+interface FetcherOptions extends RequestInit {
+  responseFormat?: 'json' | 'blob'
+  responseFormatter?: (response: Response) => any
+  customFetch?: (
+    input: RequestInfo | URL,
+    init?: RequestInit | undefined
+  ) => Promise<Response>
+}
 interface PendingRequest {
   url: string
-  init?: RequestInit & { responseFormat?: 'json' | 'blob' }
+  options?: FetcherOptions
   resolve: Function
 }
 let refreshing = false
 const queue: PendingRequest[] = []
 
-export default async function tokenFetcher(
+export default async function authEnhancedFetch(
   url: string,
-  init?: PendingRequest['init']
+  options?: FetcherOptions
 ): Promise<any> {
-  const response: Response = await fetch(url, addAuthToRequest(init))
+  const currentFetcher = options?.customFetch ?? window.fetch
+  const response: Response = await currentFetcher(
+    url,
+    addAuthToRequest(options)
+  )
+
   if (response.status === 401) {
     if (refreshing) {
       return new Promise(resolve => {
-        queue.push({ url, init, resolve })
+        queue.push({ url, options, resolve })
       })
     }
 
@@ -47,31 +60,29 @@ export default async function tokenFetcher(
       refreshing = false
       while (queue.length) {
         const q = queue.shift()
-        q?.resolve(requestWithAuth(q.url, q.init))
+        q?.resolve(requestWithAuth(q.url, q.options))
       }
 
-      return requestWithAuth(url, init)
+      return requestWithAuth(url, options)
     } else {
       refreshing = false
       queue.length = 0
       clearAuthToken()
     }
   } else {
-    return init?.responseFormat === 'blob' ? response.blob() : response.json()
+    return formatResponse(response, options)
   }
 }
 
-function addAuthToRequest(
-  init?: PendingRequest['init']
-): PendingRequest['init'] {
-  const headers = new Headers(init?.headers)
+function addAuthToRequest(options?: FetcherOptions): FetcherOptions {
+  const headers = new Headers(options?.headers)
 
   if (typeof window !== 'undefined') {
     headers.append('authorization', `Bearer ${getAuthToken()?.accessToken}`)
   }
 
   return {
-    ...(init || {}),
+    ...(options || {}),
     headers
   }
 }
@@ -85,8 +96,24 @@ async function refreshAuth(refreshToken: string) {
   return client.executeMutation(refreshAuth)
 }
 
-function requestWithAuth(url: string, init?: PendingRequest['init']) {
-  return fetch(url, addAuthToRequest(init)).then(response => {
-    return init?.responseFormat === 'blob' ? response.blob() : response.json()
+function requestWithAuth(url: string, options?: FetcherOptions) {
+  const currentFetcher = options?.customFetch ?? window.fetch
+  return currentFetcher(url, addAuthToRequest(options)).then(x => {
+    return formatResponse(x, options)
   })
 }
+
+function formatResponse(
+  response: Response,
+  options?: Pick<FetcherOptions, 'responseFormat' | 'responseFormatter'>
+) {
+  if (options?.responseFormatter) {
+    return options.responseFormatter(response)
+  }
+
+  if (options?.responseFormat === 'blob') {
+    return response.blob()
+  }
+
+  return response.json()
+}

ee/tabby-ui/lib/types/index.ts

@@ -1,3 +1,9 @@
 export * from './common'
 export * from './chat'
 export * from './repositories'
+
+declare global {
+  interface Window {
+    _originFetch?: Window['fetch'] | undefined
+  }
+}