fix: some proxy errors in main process (#2294)

* formatter proxy config type

* fix: some proxy errors in main process

* chore: Remove debug logging in ProxyManager

package.json

@@ -77,6 +77,7 @@
     "markdown-it": "^14.1.0",
     "officeparser": "^4.1.1",
     "tokenx": "^0.4.1",
+    "undici": "^7.3.0",
     "webdav": "4.11.4"
   },
   "devDependencies": {

src/main/ipc.ts

@@ -2,7 +2,7 @@ import fs from 'node:fs'
 import path from 'node:path'
 
 import { Shortcut, ThemeMode } from '@types'
-import { BrowserWindow, ipcMain, ProxyConfig, session, shell } from 'electron'
+import { BrowserWindow, ipcMain, session, shell } from 'electron'
 import log from 'electron-log'
 
 import { titleBarOverlayDark, titleBarOverlayLight } from './config'
@@ -14,12 +14,12 @@ import FileService from './services/FileService'
 import FileStorage from './services/FileStorage'
 import { GeminiService } from './services/GeminiService'
 import KnowledgeService from './services/KnowledgeService'
+import { ProxyConfig, proxyManager } from './services/ProxyManager'
 import { registerShortcuts, unregisterAllShortcuts } from './services/ShortcutService'
 import { TrayService } from './services/TrayService'
 import { windowService } from './services/WindowService'
 import { getResourcePath } from './utils'
-import { decrypt } from './utils/aes'
-import { encrypt } from './utils/aes'
+import { decrypt, encrypt } from './utils/aes'
 import { compress, decompress } from './utils/zip'
 
 const fileManager = new FileStorage()
@@ -40,9 +40,9 @@ export function registerIpc(mainWindow: BrowserWindow, app: Electron.App) {
   }))
 
   ipcMain.handle('app:proxy', async (_, proxy: string) => {
-    const sessions = [session.defaultSession, session.fromPartition('persist:webview')]
-    const proxyConfig: ProxyConfig = proxy === 'system' ? { mode: 'system' } : proxy ? { proxyRules: proxy } : {}
-    await Promise.all(sessions.map((session) => session.setProxy(proxyConfig)))
+    const proxyConfig: ProxyConfig =
+      proxy === 'system' ? { mode: 'system' } : proxy ? { mode: 'custom', url: proxy } : { mode: 'none' }
+    await proxyManager.configureProxy(proxyConfig)
   })
 
   ipcMain.handle('app:reload', () => mainWindow.reload())

src/main/services/GeminiService.ts

@@ -1,14 +1,16 @@
 import { FileMetadataResponse, FileState, GoogleAIFileManager } from '@google/generative-ai/server'
 import { FileType } from '@types'
 import fs from 'fs'
+import { ProxyAgent, setGlobalDispatcher } from 'undici'
 
 import { CacheService } from './CacheService'
-
+import { proxyManager } from './ProxyManager'
 export class GeminiService {
   private static readonly FILE_LIST_CACHE_KEY = 'gemini_file_list'
   private static readonly CACHE_DURATION = 3000
 
   static async uploadFile(_: Electron.IpcMainInvokeEvent, file: FileType, apiKey: string) {
+    setGlobalDispatcher(new ProxyAgent(proxyManager.getProxyUrl() || ''))
     const fileManager = new GoogleAIFileManager(apiKey)
     const uploadResult = await fileManager.uploadFile(file.path, {
       mimeType: 'application/pdf',
@@ -29,6 +31,7 @@ export class GeminiService {
     file: FileType,
     apiKey: string
   ): Promise<FileMetadataResponse | undefined> {
+    setGlobalDispatcher(new ProxyAgent(proxyManager.getProxyUrl() || ''))
     const fileManager = new GoogleAIFileManager(apiKey)
 
     const cachedResponse = CacheService.get<any>(GeminiService.FILE_LIST_CACHE_KEY)
@@ -52,11 +55,13 @@ export class GeminiService {
   }
 
   static async listFiles(_: Electron.IpcMainInvokeEvent, apiKey: string) {
+    setGlobalDispatcher(new ProxyAgent(proxyManager.getProxyUrl() || ''))
     const fileManager = new GoogleAIFileManager(apiKey)
     return await fileManager.listFiles()
   }
 
   static async deleteFile(_: Electron.IpcMainInvokeEvent, apiKey: string, fileId: string) {
+    setGlobalDispatcher(new ProxyAgent(proxyManager.getProxyUrl() || ''))
     const fileManager = new GoogleAIFileManager(apiKey)
     await fileManager.deleteFile(fileId)
   }

src/main/services/KnowledgeService.ts

@@ -15,6 +15,7 @@ import { FileType, KnowledgeBaseParams, KnowledgeItem } from '@types'
 import { app } from 'electron'
 import { v4 as uuidv4 } from 'uuid'
 
+import { proxyManager } from './ProxyManager'
 import { windowService } from './WindowService'
 
 class KnowledgeService {
@@ -48,13 +49,14 @@ class KnowledgeService {
               azureOpenAIApiVersion: apiVersion,
               azureOpenAIApiDeploymentName: model,
               azureOpenAIApiInstanceName: getInstanceName(baseURL),
+              configuration: { httpAgent: proxyManager.getProxyAgent() },
               dimensions,
               batchSize
             })
           : new OpenAiEmbeddings({
               model,
               apiKey,
-              configuration: { baseURL },
+              configuration: { baseURL, httpAgent: proxyManager.getProxyAgent() },
               dimensions,
               batchSize
             })

src/main/services/ProxyManager.ts

@@ -0,0 +1,115 @@
+import { ProxyConfig as _ProxyConfig, session } from 'electron'
+import { HttpsProxyAgent } from 'https-proxy-agent'
+
+type ProxyMode = 'system' | 'custom' | 'none'
+
+export interface ProxyConfig {
+  mode: ProxyMode
+  url?: string | null
+}
+
+export class ProxyManager {
+  private config: ProxyConfig
+  private proxyAgent: HttpsProxyAgent | null = null
+  private proxyUrl: string | null = null
+
+  constructor() {
+    this.config = {
+      mode: 'system',
+      url: ''
+    }
+  }
+
+  private async setSessionsProxy(config: _ProxyConfig): Promise<void> {
+    const sessions = [session.defaultSession, session.fromPartition('persist:webview')]
+    await Promise.all(sessions.map((session) => session.setProxy(config)))
+  }
+
+  async configureProxy(config: ProxyConfig): Promise<void> {
+    try {
+      this.config = config
+      if (this.config.mode === 'system') {
+        await this.setSystemProxy()
+      } else if (this.config.mode == 'custom') {
+        await this.setCustomProxy()
+      } else {
+        await this.clearProxy()
+      }
+    } catch (error) {
+      console.error('Failed to config proxy:', error)
+      throw error
+    }
+  }
+
+  private setEnvironment(url: string): void {
+    process.env.grpc_proxy = url
+    process.env.HTTP_PROXY = url
+    process.env.HTTPS_PROXY = url
+    process.env.http_proxy = url
+    process.env.https_proxy = url
+  }
+
+  private async setSystemProxy(): Promise<void> {
+    try {
+      this.proxyUrl = await this.resolveSystemProxy()
+      if (this.proxyUrl) {
+        this.proxyAgent = new HttpsProxyAgent(this.proxyUrl)
+        this.setEnvironment(this.proxyUrl)
+        await this.setSessionsProxy({ mode: 'system' })
+      }
+    } catch (error) {
+      console.error('Failed to set system proxy:', error)
+      throw error
+    }
+  }
+
+  private async setCustomProxy(): Promise<void> {
+    try {
+      if (this.config.url) {
+        this.proxyAgent = new HttpsProxyAgent(this.config.url)
+        this.setEnvironment(this.config.url)
+        await this.setSessionsProxy({ proxyRules: this.config.url })
+      }
+    } catch (error) {
+      console.error('Failed to set custom proxy:', error)
+      throw error
+    }
+  }
+
+  private async clearProxy(): Promise<void> {
+    delete process.env.HTTP_PROXY
+    delete process.env.HTTPS_PROXY
+    await this.setSessionsProxy({})
+    this.config = { mode: 'none' }
+  }
+
+  private async resolveSystemProxy(): Promise<string | null> {
+    try {
+      return await this.resolveElectronProxy()
+    } catch (error) {
+      console.error('Failed to resolve system proxy:', error)
+      return null
+    }
+  }
+
+  private async resolveElectronProxy(): Promise<string | null> {
+    try {
+      const proxyString = await session.defaultSession.resolveProxy('https://dummy.com')
+      const [protocol, address] = proxyString.split(';')[0].split(' ')
+      return protocol === 'PROXY' ? `http://${address}` : null
+    } catch (error) {
+      console.error('Failed to resolve electron proxy:', error)
+      return null
+    }
+  }
+
+  getProxyAgent(): HttpsProxyAgent | null {
+    return this.proxyAgent
+  }
+
+  getProxyUrl(): string | null {
+    return this.proxyUrl
+  }
+}
+
+export const proxyManager = new ProxyManager()

src/main/services/WebDav.ts

@@ -1,8 +1,12 @@
 import { WebDavConfig } from '@types'
 import Logger from 'electron-log'
+import { HttpProxyAgent } from 'http-proxy-agent'
+import { HttpsProxyAgent } from 'https-proxy-agent'
 import Stream from 'stream'
 import { BufferLike, createClient, GetFileContentsOptions, PutFileContentsOptions, WebDAVClient } from 'webdav'
 
+import { proxyManager } from './ProxyManager'
+
 export default class WebDav {
   public instance: WebDAVClient | undefined
   private webdavPath: string
@@ -10,11 +14,16 @@ export default class WebDav {
   constructor(params: WebDavConfig) {
     this.webdavPath = params.webdavPath
 
+    const httpAgent = new HttpProxyAgent(proxyManager.getProxyUrl() || '')
+    const httpsAgent = new HttpsProxyAgent(proxyManager.getProxyUrl() || '')
+
     this.instance = createClient(params.webdavHost, {
       username: params.webdavUser,
       password: params.webdavPass,
       maxBodyLength: Infinity,
-      maxContentLength: Infinity
+      maxContentLength: Infinity,
+      httpAgent: httpAgent,
+      httpsAgent: httpsAgent
     })
 
     this.putFileContents = this.putFileContents.bind(this)

yarn.lock

@@ -3094,6 +3094,7 @@ __metadata:
     tinycolor2: "npm:^1.6.0"
     tokenx: "npm:^0.4.1"
     typescript: "npm:^5.6.2"
+    undici: "npm:^7.3.0"
     uuid: "npm:^10.0.0"
     vite: "npm:^5.0.12"
     webdav: "npm:4.11.4"
@@ -13943,6 +13944,13 @@ __metadata:
   languageName: node
   linkType: hard
 
+"undici@npm:^7.3.0":
+  version: 7.3.0
+  resolution: "undici@npm:7.3.0"
+  checksum: 10c0/62c5e335725cadb02e19950932c7823fc330cbfd80106e6836daa6db1379aa727510b77de0a4e6f912087b288ded93f7daf4b8c154ad36fd5c9c4b96b26888b8
+  languageName: node
+  linkType: hard
+
 "unified@npm:^11.0.0":
   version: 11.0.5
   resolution: "unified@npm:11.0.5"