feat: display remote workers

ee/tabby-ui/app/(dashboard)/components/worker-card.tsx

@@ -1,33 +1,26 @@
-import {
+import { CardTitle, CardHeader, CardContent, Card } from '@/components/ui/card'
-  CardTitle,
+import { Worker, WorkerKind } from '@/lib/gql/generates/graphql'
-  CardHeader,
-  CardContent,
-  Card,
-  CardDescription
-} from '@/components/ui/card'
-import { HealthInfo } from '@/lib/hooks/use-health'
 
-type RunnerType = 'completion' | 'chat' | 'index'
+type RunnerType = WorkerKind | 'INDEX'
 
-interface RunnerCardProps {
+interface RunnerCardProps extends Partial<Omit<Worker, '__typename' | 'kind'>> {
-  source: string
+  kind: RunnerType
-  name: string
-  type: RunnerType
-  health: HealthInfo
 }
 
 export default function RunnerCard({
-  source,
+  addr,
   name,
-  type,
+  kind,
-  health
+  device,
+  cudaDevices,
+  cpuCount,
+  cpuInfo
 }: RunnerCardProps) {
-  const { device, cuda_devices } = health
   return (
     <Card className="rounded-xl p-2 shadow-md">
       <CardHeader className="p-0 px-4 pb-2 pt-4">
         <CardTitle className="text-md flex items-center font-normal">
-          <ModelIcon type={type} />
+          <ModelIcon type={kind} />
           <p className="ml-2">{name}</p>
         </CardTitle>
       </CardHeader>
@@ -51,7 +44,7 @@ export default function RunnerCard({
             <path d="M5 16v-3a1 1 0 0 1 1-1h12a1 1 0 0 1 1 1v3" />
             <path d="M12 12V8" />
           </svg>
-          <p className="ml-2">{source}</p>
+          <p className="ml-2">{addr}</p>
         </Info>
         <Info>
           <svg
@@ -78,11 +71,12 @@ export default function RunnerCard({
             <path d="M9 20v2" />
           </svg>
           <p className="ml-2">
-            {health.cpu_info} ({health.cpu_count} cores)
+            {cpuInfo} ({cpuCount} cores)
           </p>
         </Info>
         {device == 'cuda' &&
-          cuda_devices.map((x, i) => (
+          cudaDevices?.length &&
+          cudaDevices.map((x, i) => (
             <Info key={i}>
               <svg
                 className=" h-5 w-5 text-gray-400"
@@ -128,7 +122,7 @@ function Info({ children }: InfoProps) {
 
 function ModelIcon({ type }: { type: RunnerType }) {
   const className = 'h-5 w-5'
-  if (type == 'completion') {
+  if (type == WorkerKind.Completion) {
     return (
       <svg
         xmlns="http://www.w3.org/2000/svg"
@@ -148,7 +142,7 @@ function ModelIcon({ type }: { type: RunnerType }) {
         <path d="m14 17 2-2-2-2" />
       </svg>
     )
-  } else if (type == 'chat') {
+  } else if (type == WorkerKind.Chat) {
     return (
       <svg
         xmlns="http://www.w3.org/2000/svg"
@@ -166,7 +160,7 @@ function ModelIcon({ type }: { type: RunnerType }) {
         <path d="M18 9h2a2 2 0 0 1 2 2v11l-4-4h-6a2 2 0 0 1-2-2v-1" />
       </svg>
     )
-  } else if (type == 'index') {
+  } else if (type == 'INDEX') {
     return (
       <svg
         xmlns="http://www.w3.org/2000/svg"

ee/tabby-ui/app/(dashboard)/page.tsx

@@ -14,6 +14,8 @@ import { Separator } from '@/components/ui/separator'
 import { useHealth } from '@/lib/hooks/use-health'
 import { PropsWithChildren, useEffect, useState } from 'react'
 import WorkerCard from './components/worker-card'
+import { useMergedWorkers } from '@/lib/hooks/use-remote-worker'
+import { WorkerKind } from '@/lib/gql/generates/graphql'
 
 const COMMUNITY_DIALOG_SHOWN_KEY = 'community-dialog-shown'
 
@@ -72,6 +74,7 @@ function toBadgeString(str: string) {
 
 function MainPanel() {
   const { data: healthInfo } = useHealth()
+  const workers = useMergedWorkers(healthInfo)
 
   if (!healthInfo) return
 
@@ -98,26 +101,29 @@ function MainPanel() {
       <div className="mt-4 rounded-lg bg-zinc-100 p-4 dark:bg-zinc-800">
         <span className="font-bold">Workers</span>
         <div className="mt-4 flex flex-col gap-4 lg:flex-row lg:flex-wrap">
-          {healthInfo.model &&
+          {!!workers?.[WorkerKind.Completion] && (
-            <WorkerCard
+            <>
-              source="localhost"
+              {workers[WorkerKind.Completion].map((worker, i) => {
-              name={healthInfo.model}
+                return <WorkerCard key={i} {...worker} />
-              type="completion"
+              })}
-              health={healthInfo}
+            </>
-            />}
+          )}
-          {healthInfo.chat_model && (
+          {!!workers?.[WorkerKind.Chat] && (
-            <WorkerCard
+            <>
-              source="localhost"
+              {workers[WorkerKind.Chat].map((worker, i) => {
-              name={healthInfo.chat_model}
+                return <WorkerCard key={i} {...worker} />
-              type="chat"
+              })}
-              health={healthInfo}
+            </>
-            />
           )}
           <WorkerCard
-            source="localhost"
+            addr="localhost"
             name="Code Search Index"
-            type="index"
+            kind="INDEX"
-            health={healthInfo}
+            arch=""
+            device={healthInfo.device}
+            cudaDevices={healthInfo.cuda_devices}
+            cpuCount={healthInfo.cpu_count}
+            cpuInfo={healthInfo.cpu_info}
           />
         </div>
       </div>

ee/tabby-ui/codegen.ts

@@ -0,0 +1,19 @@
+
+import type { CodegenConfig } from '@graphql-codegen/cli';
+
+const config: CodegenConfig = {
+  overwrite: true,
+  schema: "../tabby-webserver/graphql/schema.graphql",
+  documents: "./**/*.tsx",
+  // documents: "./**/*.graphql",
+  ignoreNoDocuments: true,
+  generates: {
+    "lib/gql/generates/": {
+      preset: "client",
+      plugins: []
+    }
+  },
+  hooks: { afterAllFileWrite: ['prettier --write'] }
+};
+
+export default config;

ee/tabby-ui/components/header.tsx

@@ -1,7 +1,6 @@
 'use client'
 
 import * as React from 'react'
-
 import { cn } from '@/lib/utils'
 import { buttonVariants } from '@/components/ui/button'
 import { IconGitHub, IconNotice } from '@/components/ui/icons'
@@ -10,6 +9,9 @@ import Link from 'next/link'
 import { useHealth } from '@/lib/hooks/use-health'
 import { ReleaseInfo, useLatestRelease } from '@/lib/hooks/use-latest-release'
 import { compare } from 'compare-versions'
+import { useMergedWorkers } from '@/lib/hooks/use-remote-worker'
+import { WorkerKind } from '@/lib/gql/generates/graphql'
+import { has } from 'lodash-es'
 
 const ThemeToggle = dynamic(
   () => import('@/components/theme-toggle').then(x => x.ThemeToggle),
@@ -18,7 +20,8 @@ const ThemeToggle = dynamic(
 
 export function Header() {
   const { data } = useHealth()
-  const isChatEnabled = !!data?.chat_model
+  const workers = useMergedWorkers(data)
+  const isChatEnabled = has(workers, WorkerKind.Chat)
   const version = data?.version?.git_describe
   const { data: latestRelease } = useLatestRelease()
   const newVersionAvailable = isNewVersionAvailable(version, latestRelease)

ee/tabby-ui/lib/gql/generates/fragment-masking.ts

@@ -0,0 +1,85 @@
+import {
+  ResultOf,
+  DocumentTypeDecoration,
+  TypedDocumentNode
+} from '@graphql-typed-document-node/core'
+import { FragmentDefinitionNode } from 'graphql'
+import { Incremental } from './graphql'
+
+export type FragmentType<
+  TDocumentType extends DocumentTypeDecoration<any, any>
+> = TDocumentType extends DocumentTypeDecoration<infer TType, any>
+  ? [TType] extends [{ ' $fragmentName'?: infer TKey }]
+    ? TKey extends string
+      ? { ' $fragmentRefs'?: { [key in TKey]: TType } }
+      : never
+    : never
+  : never
+
+// return non-nullable if `fragmentType` is non-nullable
+export function useFragment<TType>(
+  _documentNode: DocumentTypeDecoration<TType, any>,
+  fragmentType: FragmentType<DocumentTypeDecoration<TType, any>>
+): TType
+// return nullable if `fragmentType` is nullable
+export function useFragment<TType>(
+  _documentNode: DocumentTypeDecoration<TType, any>,
+  fragmentType:
+    | FragmentType<DocumentTypeDecoration<TType, any>>
+    | null
+    | undefined
+): TType | null | undefined
+// return array of non-nullable if `fragmentType` is array of non-nullable
+export function useFragment<TType>(
+  _documentNode: DocumentTypeDecoration<TType, any>,
+  fragmentType: ReadonlyArray<FragmentType<DocumentTypeDecoration<TType, any>>>
+): ReadonlyArray<TType>
+// return array of nullable if `fragmentType` is array of nullable
+export function useFragment<TType>(
+  _documentNode: DocumentTypeDecoration<TType, any>,
+  fragmentType:
+    | ReadonlyArray<FragmentType<DocumentTypeDecoration<TType, any>>>
+    | null
+    | undefined
+): ReadonlyArray<TType> | null | undefined
+export function useFragment<TType>(
+  _documentNode: DocumentTypeDecoration<TType, any>,
+  fragmentType:
+    | FragmentType<DocumentTypeDecoration<TType, any>>
+    | ReadonlyArray<FragmentType<DocumentTypeDecoration<TType, any>>>
+    | null
+    | undefined
+): TType | ReadonlyArray<TType> | null | undefined {
+  return fragmentType as any
+}
+
+export function makeFragmentData<
+  F extends DocumentTypeDecoration<any, any>,
+  FT extends ResultOf<F>
+>(data: FT, _fragment: F): FragmentType<F> {
+  return data as FragmentType<F>
+}
+export function isFragmentReady<TQuery, TFrag>(
+  queryNode: DocumentTypeDecoration<TQuery, any>,
+  fragmentNode: TypedDocumentNode<TFrag>,
+  data:
+    | FragmentType<TypedDocumentNode<Incremental<TFrag>, any>>
+    | null
+    | undefined
+): data is FragmentType<typeof fragmentNode> {
+  const deferredFields = (
+    queryNode as {
+      __meta__?: { deferredFields: Record<string, (keyof TFrag)[]> }
+    }
+  ).__meta__?.deferredFields
+
+  if (!deferredFields) return true
+
+  const fragDef = fragmentNode.definitions[0] as
+    | FragmentDefinitionNode
+    | undefined
+  const fragName = fragDef?.name?.value
+
+  const fields = (fragName && deferredFields[fragName]) || []
+  return fields.length > 0 && fields.every(field => data && field in data)
+}

ee/tabby-ui/lib/gql/generates/gql.ts

@@ -0,0 +1,46 @@
+/* eslint-disable */
+import * as types from './graphql'
+import { TypedDocumentNode as DocumentNode } from '@graphql-typed-document-node/core'
+
+/**
+ * Map of all GraphQL operations in the project.
+ *
+ * This map has several performance disadvantages:
+ * 1. It is not tree-shakeable, so it will include all operations in the project.
+ * 2. It is not minifiable, so the string of a GraphQL query will be multiple times inside the bundle.
+ * 3. It does not support dead code elimination, so it will add unused operations.
+ *
+ * Therefore it is highly recommended to use the babel or swc plugin for production.
+ */
+const 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':
+    types.GetWorkersDocument
+}
+
+/**
+ * The graphql function is used to parse GraphQL queries into a document that can be used by GraphQL clients.
+ *
+ *
+ * @example
+ * ```ts
+ * const query = graphql(`query GetUser($id: ID!) { user(id: $id) { name } }`);
+ * ```
+ *
+ * The query argument is unknown!
+ * Please regenerate the types.
+ */
+export function graphql(source: string): unknown
+
+/**
+ * The graphql function is used to parse GraphQL queries into a document that can be used by GraphQL clients.
+ */
+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']
+
+export function graphql(source: string) {
+  return (documents as any)[source] ?? {}
+}
+
+export type DocumentType<TDocumentNode extends DocumentNode<any, any>> =
+  TDocumentNode extends DocumentNode<infer TType, any> ? TType : never

ee/tabby-ui/lib/gql/generates/graphql.ts

@@ -0,0 +1,108 @@
+/* eslint-disable */
+import { TypedDocumentNode as DocumentNode } from '@graphql-typed-document-node/core'
+export type Maybe<T> = T | null
+export type InputMaybe<T> = Maybe<T>
+export type Exact<T extends { [key: string]: unknown }> = {
+  [K in keyof T]: T[K]
+}
+export type MakeOptional<T, K extends keyof T> = Omit<T, K> & {
+  [SubKey in K]?: Maybe<T[SubKey]>
+}
+export type MakeMaybe<T, K extends keyof T> = Omit<T, K> & {
+  [SubKey in K]: Maybe<T[SubKey]>
+}
+export type MakeEmpty<
+  T extends { [key: string]: unknown },
+  K extends keyof T
+> = { [_ in K]?: never }
+export type Incremental<T> =
+  | T
+  | {
+      [P in keyof T]?: P extends ' $fragmentName' | '__typename' ? T[P] : never
+    }
+/** All built-in and custom scalars, mapped to their actual values */
+export type Scalars = {
+  ID: { input: string; output: string }
+  String: { input: string; output: string }
+  Boolean: { input: boolean; output: boolean }
+  Int: { input: number; output: number }
+  Float: { input: number; output: number }
+}
+
+export type Mutation = {
+  __typename?: 'Mutation'
+  resetRegistrationToken: Scalars['String']['output']
+}
+
+export type Query = {
+  __typename?: 'Query'
+  registrationToken: Scalars['String']['output']
+  workers: Array<Worker>
+}
+
+export type Worker = {
+  __typename?: 'Worker'
+  addr: Scalars['String']['output']
+  arch: Scalars['String']['output']
+  cpuCount: Scalars['Int']['output']
+  cpuInfo: Scalars['String']['output']
+  cudaDevices: Array<Scalars['String']['output']>
+  device: Scalars['String']['output']
+  kind: WorkerKind
+  name: Scalars['String']['output']
+}
+
+export enum WorkerKind {
+  Chat = 'CHAT',
+  Completion = 'COMPLETION'
+}
+
+export type GetWorkersQueryVariables = Exact<{ [key: string]: never }>
+
+export type GetWorkersQuery = {
+  __typename?: 'Query'
+  workers: Array<{
+    __typename?: 'Worker'
+    kind: WorkerKind
+    name: string
+    addr: string
+    device: string
+    arch: string
+    cpuInfo: string
+    cpuCount: number
+    cudaDevices: Array<string>
+  }>
+}
+
+export const GetWorkersDocument = {
+  kind: 'Document',
+  definitions: [
+    {
+      kind: 'OperationDefinition',
+      operation: 'query',
+      name: { kind: 'Name', value: 'GetWorkers' },
+      selectionSet: {
+        kind: 'SelectionSet',
+        selections: [
+          {
+            kind: 'Field',
+            name: { kind: 'Name', value: 'workers' },
+            selectionSet: {
+              kind: 'SelectionSet',
+              selections: [
+                { kind: 'Field', name: { kind: 'Name', value: 'kind' } },
+                { kind: 'Field', name: { kind: 'Name', value: 'name' } },
+                { kind: 'Field', name: { kind: 'Name', value: 'addr' } },
+                { kind: 'Field', name: { kind: 'Name', value: 'device' } },
+                { kind: 'Field', name: { kind: 'Name', value: 'arch' } },
+                { kind: 'Field', name: { kind: 'Name', value: 'cpuInfo' } },
+                { kind: 'Field', name: { kind: 'Name', value: 'cpuCount' } },
+                { kind: 'Field', name: { kind: 'Name', value: 'cudaDevices' } }
+              ]
+            }
+          }
+        ]
+      }
+    }
+  ]
+} as unknown as DocumentNode<GetWorkersQuery, GetWorkersQueryVariables>

ee/tabby-ui/lib/gql/generates/index.ts

@@ -0,0 +1,2 @@
+export * from './fragment-masking'
+export * from './gql'

ee/tabby-ui/lib/gql/request-documents.tsx

@@ -0,0 +1,16 @@
+import { graphql } from './generates'
+
+export const GetAllWorkers = graphql(/* GraphQL */ `
+  query GetWorkers {
+    workers {
+      kind
+      name
+      addr
+      device
+      arch
+      cpuInfo
+      cpuCount
+      cudaDevices
+    }
+  }
+`)

ee/tabby-ui/lib/hooks/use-graphql.ts

@@ -0,0 +1,28 @@
+import useSWR, { SWRResponse, SWRConfiguration } from 'swr'
+import { request } from '@/lib/tabby-gql-client'
+import { Variables } from 'graphql-request'
+import { TypedDocumentNode } from '@graphql-typed-document-node/core'
+import { ASTNode, Kind, OperationDefinitionNode } from 'graphql'
+
+const isOperationDefinition = (def: ASTNode): def is OperationDefinitionNode =>
+  def.kind === Kind.OPERATION_DEFINITION
+
+function useGraphQL<TResult, TVariables extends Variables | undefined>(
+  document: TypedDocumentNode<TResult, TVariables>,
+  variables?: TVariables,
+  options?: SWRConfiguration<TResult>
+): SWRResponse<TResult> {
+  return useSWR(
+    [
+      document.definitions.find(isOperationDefinition)?.name?.value,
+      document,
+      variables
+    ],
+    ([_key, document, variables]) => {
+      return request({ document, variables })
+    },
+    options
+  )
+}
+
+export { useGraphQL }

ee/tabby-ui/lib/hooks/use-remote-worker.ts

@@ -0,0 +1,52 @@
+import { groupBy, findIndex } from 'lodash-es'
+import {
+  GetWorkersDocument,
+  Worker,
+  WorkerKind
+} from '@/lib/gql/generates/graphql'
+import { useGraphQL } from './use-graphql'
+import type { HealthInfo } from './use-health'
+
+function useRemoteWorkers() {
+  return useGraphQL(GetWorkersDocument)
+}
+
+const modelNameMap: Record<WorkerKind, 'chat_model' | 'model'> = {
+  [WorkerKind.Chat]: 'chat_model',
+  [WorkerKind.Completion]: 'model'
+}
+function transformHealthInfoToWorker(
+  healthInfo: HealthInfo,
+  kind: WorkerKind
+): Worker {
+  return {
+    kind,
+    device: healthInfo.device,
+    addr: 'localhost',
+    arch: '',
+    cpuInfo: healthInfo.cpu_info,
+    name: healthInfo?.[modelNameMap[kind]] ?? '',
+    cpuCount: healthInfo.cpu_count,
+    cudaDevices: healthInfo.cuda_devices
+  }
+}
+
+function useMergedWorkers(healthInfo: HealthInfo | undefined) {
+  const { data } = useRemoteWorkers()
+  let workers = data?.workers || []
+
+  const haveRemoteCompletionWorkers =
+    findIndex(workers, { kind: WorkerKind.Completion }) > -1
+  const haveRemoteChatWorkers =
+    findIndex(workers, { kind: WorkerKind.Chat }) > -1
+
+  if (!haveRemoteCompletionWorkers && healthInfo?.model) {
+    workers.push(transformHealthInfoToWorker(healthInfo, WorkerKind.Completion))
+  }
+  if (!haveRemoteChatWorkers && healthInfo?.chat_model) {
+    workers.push(transformHealthInfoToWorker(healthInfo, WorkerKind.Chat))
+  }
+  return groupBy(workers, worker => worker.kind)
+}
+
+export { useRemoteWorkers, useMergedWorkers }

ee/tabby-ui/lib/hooks/use-store.ts

@@ -1,11 +1,16 @@
 import React from 'react'
+import { useHydrated } from './use-hydration'
 
-export const useStore = <T, F>(
+export const useStore = <T extends { _hasHydrated?: boolean }, F>(
   store: (callback: (state: T) => unknown) => unknown,
   callback: (state: T) => F
 ) => {
+  const hydrated = useHydrated()
+  const _hasZustandHydrated = store((state: T) => state?._hasHydrated)
   const result = store(callback) as F
-  const [data, setData] = React.useState<F>()
+  const [data, setData] = React.useState<F>(
+    hydrated && _hasZustandHydrated ? result : (undefined as F)
+  )
 
   React.useEffect(() => {
     setData(result)

ee/tabby-ui/lib/tabby-gql-client.ts

@@ -0,0 +1,15 @@
+import { GraphQLClient, Variables, RequestOptions } from 'graphql-request'
+
+export const graphQLClient = new GraphQLClient(
+  `${process.env.NEXT_PUBLIC_TABBY_SERVER_URL ?? ''}/graphql`,
+  {
+    credentials: 'include',
+    mode: 'cors'
+  }
+)
+
+export function request<T, V extends Variables = Variables>(
+  options: RequestOptions<V, T>
+) {
+  return graphQLClient.request(options)
+}

ee/tabby-ui/package.json

@@ -11,7 +11,9 @@
     "preview": "next build && next start",
     "type-check": "tsc --noEmit",
     "format:write": "prettier --write \"{app,lib,components}/**/*.{ts,tsx,mdx}\" --cache",
-    "format:check": "prettier --check \"{app,lib,components}**/*.{ts,tsx,mdx}\" --cache"
+    "format:check": "prettier --check \"{app,lib,components}**/*.{ts,tsx,mdx}\" --cache",
+    "codegen": "graphql-codegen --config codegen.ts",
+    "codegen:watch": "graphql-codegen --config codegen.ts --watch"
   },
   "dependencies": {
     "@radix-ui/react-alert-dialog": "1.0.4",
@@ -33,6 +35,8 @@
     "compare-versions": "^6.1.0",
     "downshift": "^8.2.2",
     "focus-trap-react": "^10.1.1",
+    "graphql": "^16.8.1",
+    "graphql-request": "^6.1.0",
     "lodash-es": "^4.17.21",
     "nanoid": "^4.0.2",
     "next": "^13.4.7",
@@ -52,6 +56,8 @@
     "zustand": "^4.4.6"
   },
   "devDependencies": {
+    "@graphql-codegen/cli": "5.0.0",
+    "@graphql-codegen/client-preset": "4.1.0",
     "@tailwindcss/typography": "^0.5.9",
     "@types/lodash-es": "^4.17.10",
     "@types/node": "^17.0.12",