feat: add mirrors select on confirming installation

hawa130 · hawa130 · commit a8aea4b1abd5 · 2025-08-07T15:55:02.000+08:00
diff --git a/app/install-env/ssh/atoms.ts b/app/install-env/ssh/atoms.ts
@@ -0,0 +1,20 @@
+/**
+ * Copyright (c) 2025 EM-GeekLab
+ * LLMOne is licensed under Mulan PSL v2.
+ * You can use this software according to the terms and conditions of the Mulan PSL v2.
+ * You may obtain a copy of Mulan PSL v2 at:
+ *          http://license.coscl.org.cn/MulanPSL2
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
+ * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
+ * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
+ * See the Mulan PSL v2 for more details.
+ */
+
+import { enableMapSet } from 'immer'
+import { atom } from 'jotai'
+
+import type { InstallOptions } from '@/lib/ssh/ssh-deployer'
+
+enableMapSet()
+
+export const installConfigAtom = atom<Map<string, InstallOptions>>(new Map())
diff --git a/app/install-env/ssh/hooks.tsx b/app/install-env/ssh/hooks.tsx
@@ -10,8 +10,11 @@
  * See the Mulan PSL v2 for more details.
  */
 
+import { useCallback } from 'react'
 import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query'
+import { useAtomValue } from 'jotai/react'
 
+import { installConfigAtom } from '@/app/install-env/ssh/atoms'
 import { useGlobalStoreApi } from '@/stores'
 import { useTRPC, useTRPCClient } from '@/trpc/client'
 
@@ -20,13 +23,14 @@ export function useInstallStartTrigger(): { start: () => void } {
   const trpcClient = useTRPCClient()
   const trpc = useTRPC()
   const queryClient = useQueryClient()
+  const installConfigs = useAtomValue(installConfigAtom)
 
   const { mutate: start } = useMutation({
     mutationFn: async () => {
       const { finalSshHosts } = storeApi.getState()
       await Promise.all(
         finalSshHosts.map(async ({ ip: host }) => {
-          await trpcClient.sshDeploy.install.triggerOnce.mutate(host)
+          await trpcClient.sshDeploy.install.triggerOnce.mutate({ host, options: installConfigs.get(host) })
         }),
       )
     },
@@ -47,10 +51,11 @@ export function useInstallStartTrigger(): { start: () => void } {
 export function useInstallRetryTrigger(): { retry: (host: string) => void } {
   const trpc = useTRPC()
   const queryClient = useQueryClient()
+  const installConfigs = useAtomValue(installConfigAtom)
 
-  const { mutate: retry } = useMutation(
+  const { mutate } = useMutation(
     trpc.sshDeploy.install.triggerOnce.mutationOptions({
-      onSuccess: (_, host) => {
+      onSuccess: (_, { host }) => {
         return Promise.all([
           queryClient.resetQueries({ queryKey: trpc.sshDeploy.install.status.queryKey(host) }),
           queryClient.resetQueries({ queryKey: trpc.sshDeploy.install.statusAll.queryKey() }),
@@ -59,6 +64,11 @@ export function useInstallRetryTrigger(): { retry: (host: string) => void } {
     }),
   )
 
+  const retry = useCallback(
+    (host: string) => mutate({ host, options: installConfigs.get(host) }),
+    [installConfigs, mutate],
+  )
+
   return { retry }
 }
 
diff --git a/app/install-env/ssh/jotai-provider.tsx b/app/install-env/ssh/jotai-provider.tsx
@@ -0,0 +1,7 @@
+'use client'
+
+import { Provider } from 'jotai/react'
+
+export function JotaiProvider({ children }: { children: React.ReactNode }) {
+  return <Provider>{children}</Provider>
+}
diff --git a/app/install-env/ssh/ssh-hosts-confirm.tsx b/app/install-env/ssh/ssh-hosts-confirm.tsx
@@ -13,17 +13,20 @@
 'use client'
 
 import * as React from 'react'
-import { ReactNode } from 'react'
+import { ComponentProps, ReactNode, useId } from 'react'
 import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query'
+import { useAtom } from 'jotai/react'
 import { AlertTriangleIcon, ArrowRightIcon, CheckIcon } from 'lucide-react'
 import { match } from 'ts-pattern'
 
-import type { InstallFlag, SshDeployerInfo } from '@/lib/ssh/ssh-deployer'
+import type { InstallFlag, InstallOptions, SshDeployerInfo } from '@/lib/ssh/ssh-deployer'
+import { cn } from '@/lib/utils'
 import { AppCardContent, AppCardFooter, AppCardSection } from '@/components/app/app-card'
 import { Callout } from '@/components/base/callout'
 import { NavButton } from '@/components/base/nav-button'
 import { Badge } from '@/components/ui/badge'
 import { Button } from '@/components/ui/button'
+import { installConfigAtom } from '@/app/install-env/ssh/atoms'
 import { useGlobalStore } from '@/stores'
 import { useTRPC, useTRPCClient } from '@/trpc/client'
 
@@ -112,13 +115,17 @@ function HostConfirmList({ host }: { host: SshDeployerInfo }) {
         flag={flags.updateSources}
         plannedMessage="即将更新软件源镜像"
         completedMessage="已更新软件源镜像"
-      />
+      >
+        <PackageMirrorSelect host={host.host} />
+      </HostConfirmItem>
       <HostConfirmItem
         flag={flags.installDependencies}
         plannedMessage="即将安装基础依赖"
         completedMessage="已安装基础依赖"
       />
-      <HostConfirmItem flag={flags.installDocker} plannedMessage="即将安装 Docker" completedMessage="已安装 Docker" />
+      <HostConfirmItem flag={flags.installDocker} plannedMessage="即将安装 Docker" completedMessage="已安装 Docker">
+        <DockerPackageMirrorSelect host={host.host} />
+      </HostConfirmItem>
       {flags.installNvidiaGpu && (
         <HostConfirmItem
           flag={flags.installNvidiaGpu}
@@ -145,6 +152,68 @@ function HostConfirmList({ host }: { host: SshDeployerInfo }) {
   )
 }
 
+function PackageMirrorSelect({ host }: { host: string }) {
+  const id = useId()
+  const [options, setOptions] = useAtom(installConfigAtom)
+  const value = options.get(host)?.packageMirror || 'none'
+  const setValue = (e: React.ChangeEvent<HTMLSelectElement>) => {
+    const newMirror = e.target.value as InstallOptions['packageMirror']
+    setOptions((prev) => {
+      const newOptions = new Map(prev)
+      const currentConfig = newOptions.get(host) || {}
+      newOptions.set(host, { ...currentConfig, packageMirror: newMirror })
+      return newOptions
+    })
+  }
+
+  return (
+    <NativeSelect id={id} name="package-mirror" value={value} onChange={setValue}>
+      <option value="none">不修改源</option>
+      <option value="mirrors.aliyun.com">阿里源</option>
+      <option value="mirrors.tencent.com">腾讯源</option>
+      <option value="mirrors.tuna.tsinghua.edu.cn">清华源</option>
+      <option value="mirrors.ustc.edu.cn">中科大源</option>
+    </NativeSelect>
+  )
+}
+
+function DockerPackageMirrorSelect({ host }: { host: string }) {
+  const id = useId()
+  const [options, setOptions] = useAtom(installConfigAtom)
+  const value = options.get(host)?.dockerPackageMirror || 'mirrors.aliyun.com/docker-ce'
+  const setValue = (e: React.ChangeEvent<HTMLSelectElement>) => {
+    const newMirror = e.target.value as InstallOptions['dockerPackageMirror']
+    setOptions((prev) => {
+      const newOptions = new Map(prev)
+      const currentConfig = newOptions.get(host) || {}
+      newOptions.set(host, { ...currentConfig, dockerPackageMirror: newMirror })
+      return newOptions
+    })
+  }
+
+  return (
+    <NativeSelect id={id} name="docker-package-mirror" value={value} onChange={setValue}>
+      <option value="mirrors.aliyun.com/docker-ce">阿里源</option>
+      <option value="mirrors.tencent.com/docker-ce">腾讯源</option>
+      <option value="mirrors.tuna.tsinghua.edu.cn/docker-ce">清华源</option>
+      <option value="mirrors.ustc.edu.cn/docker-ce">中科大源</option>
+      <option value="download.docker.com">官方源</option>
+    </NativeSelect>
+  )
+}
+
+function NativeSelect({ className, ...props }: ComponentProps<'select'>) {
+  return (
+    <select
+      className={cn(
+        'shrink-0 rounded-sm border bg-background py-px text-xs shadow-xs outline-none group-data-completed:hidden hover:border-accent-foreground/25 hover:bg-accent hover:text-accent-foreground focus-visible:border-ring focus-visible:ring-[2px] focus-visible:ring-ring/50 disabled:hover:border-border disabled:hover:bg-background disabled:hover:text-foreground dark:border-input dark:bg-input/30 dark:hover:bg-input/50 dark:disabled:hover:bg-input/30',
+        className,
+      )}
+      {...props}
+    />
+  )
+}
+
 function HostConfirmItem({
   children,
   customIcon,
diff --git a/app/install-env/ssh/ssh-page.tsx b/app/install-env/ssh/ssh-page.tsx
@@ -14,14 +14,15 @@
 
 import * as React from 'react'
 
+import { JotaiProvider } from './jotai-provider'
 import { SshHostsConfirm } from './ssh-hosts-confirm'
 import { SshInstallPage } from './ssh-install-page'
 
 export function SshPage() {
   return (
-    <>
+    <JotaiProvider>
       <SshHostsConfirm />
       <SshInstallPage />
-    </>
+    </JotaiProvider>
   )
 }
diff --git a/lib/ssh/ssh-deployer.ts b/lib/ssh/ssh-deployer.ts
@@ -22,7 +22,7 @@ import { logger } from '@/lib/logger'
 import { messageError, wrapError } from '@/lib/utils/error'
 import { SshFinalConnectionInfo } from '@/app/connect-info/schemas'
 
-import { PackageManager } from './pm'
+import { DockerPackageMirrors, PackageManager, PackageMirrors } from './pm'
 import { MxaCtl, SystemInfo } from './ssh-controller'
 
 const log = logger.child({ module: 'ssh.deployer' })
@@ -53,6 +53,11 @@ export type InstallStepFlags = {
   installNvidiaCtk?: InstallFlag
 }
 
+export type InstallOptions = {
+  packageMirror?: PackageMirrors | 'none'
+  dockerPackageMirror?: DockerPackageMirrors
+}
+
 export type SshDeployerStatus = 'idle' | 'failed' | 'installing' | 'completed'
 
 export type SshDeployerInfo = {
@@ -242,14 +247,14 @@ echo "{\
     this.ee.on('install:completed', listener)
   }
 
-  async install({ onProgress }: { onProgress?: LogListener } = {}) {
+  async install({ onProgress, ...options }: { onProgress?: LogListener } & InstallOptions = {}) {
     if (onProgress) {
       this.onLog(onProgress)
     }
     this.beforeInstall()
-    await this.updateSources()
+    await this.updateSources(options.packageMirror)
     await this.installDependencies()
-    await this.installDocker()
+    await this.installDocker(options.dockerPackageMirror)
     if (this.installFlags.installNvidiaGpu) {
       await this.installNvidiaGpu()
     }
@@ -313,20 +318,20 @@ echo "{\
     }
   }
 
-  private async updatePmIndex() {
-    await this.execInstallScript({
-      script: this.pm.updateIndex(),
-      flag: this.installFlags.updateSources,
-      initLog: '更新软件包索引',
-      successLog: '软件包索引更新完成',
-      errorLog: '软件包索引更新失败',
-      errorMessage: '软件包索引更新失败，请检查网络连接或手动更新软件包索引',
-    })
-  }
-
-  private async updateSources() {
+  private async updateSources(mirror?: InstallOptions['packageMirror']) {
+    if (mirror === 'none') {
+      await this.execInstallScript({
+        script: this.pm.updateIndex(),
+        flag: this.installFlags.updateSources,
+        initLog: '更新软件包索引',
+        successLog: '软件包索引更新完成',
+        errorLog: '软件包索引更新失败',
+        errorMessage: '软件包索引更新失败，请检查网络连接或手动更新软件包索引',
+      })
+      return
+    }
     await this.execInstallScript({
-      script: this.pm.updateSources(),
+      script: this.pm.updateSources(mirror),
       flag: this.installFlags.updateSources,
       initLog: '更新软件源',
       successLog: '软件源更新完成',
@@ -353,9 +358,9 @@ echo "{\
     })
   }
 
-  private async installDocker() {
+  private async installDocker(mirror?: InstallOptions['dockerPackageMirror']) {
     await this.execInstallScript({
-      script: this.pm.installDocker(),
+      script: this.pm.installDocker(mirror),
       flag: this.installFlags.installDocker,
       initLog: '安装 Docker',
       successLog: 'Docker 安装完成',
@@ -508,17 +513,17 @@ export class SshDeployerManager {
     return deployer
   }
 
-  async installTrigger(host: string) {
+  async installTrigger(host: string, options: InstallOptions = {}) {
     const deployer = this.getDeployer(host)
-    await deployer.install()
+    await deployer.install(options)
   }
 
-  installStream(host: string): EventIterator<string> {
+  installStream(host: string, options: InstallOptions = {}): EventIterator<string> {
     const deployer = this.getDeployer(host)
     return new EventIterator<string>(({ push, stop, fail }) => {
       deployer.getLogs().forEach((log) => push(log))
       deployer
-        .install({ onProgress: (data) => push(data) })
+        .install({ onProgress: (data) => push(data), ...options })
         .then(() => stop())
         .catch((err) => fail(err))
     })
diff --git a/trpc/router/ssh-deploy.ts b/trpc/router/ssh-deploy.ts
@@ -13,6 +13,7 @@
 import { TRPCError } from '@trpc/server'
 
 import { SshDeployerManager } from '@/lib/ssh/ssh-deployer'
+import { z } from '@/lib/zod'
 import { sshHostsListSchema } from '@/app/connect-info/schemas'
 import { baseProcedure, createRouter } from '@/trpc/init'
 
@@ -29,6 +30,35 @@ function getSshDeployerManager() {
 
 const installStatusCache = new Map<string, Promise<void>>()
 
+const triggerInputSchema = z.object({
+  host: z.string(),
+  options: z
+    .object({
+      packageMirror: z
+        .enum([
+          'none',
+          'mirrors.aliyun.com',
+          'mirrors.tencent.com',
+          'mirrors.tuna.tsinghua.edu.cn',
+          'mirrors.ustc.edu.cn',
+        ])
+        .optional(),
+      dockerPackageMirror: z
+        .enum([
+          'mirrors.aliyun.com/docker-ce',
+          'mirrors.tencent.com/docker-ce',
+          'mirrors.tuna.tsinghua.edu.cn/docker-ce',
+          'mirrors.ustc.edu.cn/docker-ce',
+          'download.docker.com',
+        ])
+        .optional(),
+    })
+    .default({
+      packageMirror: 'none',
+      dockerPackageMirror: 'mirrors.aliyun.com/docker-ce',
+    }),
+})
+
 export const sshDeployRouter = createRouter({
   deployer: {
     init: baseProcedure.input(sshHostsListSchema).mutation(async ({ input }) => {
@@ -53,13 +83,13 @@ export const sshDeployRouter = createRouter({
     }),
   },
   install: {
-    trigger: baseProcedure.input(inputType<string>).mutation(async ({ input: host }) => {
-      const promise = getSshDeployerManager().installTrigger(host)
+    trigger: baseProcedure.input(triggerInputSchema).mutation(async ({ input: { host, options } }) => {
+      const promise = getSshDeployerManager().installTrigger(host, options)
       installStatusCache.set(host, promise)
       return await promise
     }),
-    triggerOnce: baseProcedure.input(inputType<string>).mutation(async ({ input: host }) => {
-      const promise = getSshDeployerManager().installTrigger(host)
+    triggerOnce: baseProcedure.input(triggerInputSchema).mutation(async ({ input: { host, options } }) => {
+      const promise = getSshDeployerManager().installTrigger(host, options)
       installStatusCache.set(host, promise)
     }),
     status: baseProcedure.input(inputType<string>).query(async ({ input: host }) => {
