refactor(GeminiProvider): streamline abort signal handling and improve stream processing (#5276)

package.json

@@ -120,7 +120,7 @@
     "@emotion/is-prop-valid": "^1.3.1",
     "@eslint-react/eslint-plugin": "^1.36.1",
     "@eslint/js": "^9.22.0",
-    "@google/genai": "patch:@google/genai@npm%3A0.8.0#~/.yarn/patches/@google-genai-npm-0.8.0-450d0d9a7d.patch",
+    "@google/genai": "^0.10.0",
     "@hello-pangea/dnd": "^16.6.0",
     "@kangfenmao/keyv-storage": "^0.1.0",
     "@modelcontextprotocol/sdk": "^1.9.0",

src/renderer/src/providers/AiProvider/GeminiProvider.ts

@@ -320,32 +320,12 @@ export default class GeminiProvider extends BaseProvider {
     const start_time_millsec = new Date().getTime()
 
     const { cleanup, abortController } = this.createAbortController(userLastMessage?.id, true)
-    const signalProxy = {
-      _originalSignal: abortController.signal,
-
-      addEventListener: (eventName: string, listener: () => void) => {
-        if (eventName === 'abort') {
-          abortController.signal.addEventListener('abort', listener)
-        }
-      },
-      removeEventListener: (eventName: string, listener: () => void) => {
-        if (eventName === 'abort') {
-          abortController.signal.removeEventListener('abort', listener)
-        }
-      },
-      get aborted() {
-        return abortController.signal.aborted
-      }
-    }
-
     if (!streamOutput) {
       const response = await chat.sendMessage({
         message: messageContents as PartUnion,
         config: {
           ...generateContentConfig,
-          httpOptions: {
+          abortSignal: abortController.signal
-            signal: signalProxy as any
-          }
         }
       })
       const time_completion_millsec = new Date().getTime() - start_time_millsec
@@ -371,9 +351,7 @@ export default class GeminiProvider extends BaseProvider {
       message: messageContents as PartUnion,
       config: {
         ...generateContentConfig,
-        httpOptions: {
+        abortSignal: abortController.signal
-          signal: signalProxy as any
-        }
       }
     })
     let time_first_token_millsec = 0
@@ -399,9 +377,7 @@ export default class GeminiProvider extends BaseProvider {
           message: flatten(toolResults.map((ts) => (ts as Content).parts)) as PartUnion,
           config: {
             ...generateContentConfig,
-            httpOptions: {
+            abortSignal: abortController.signal
-              signal: signalProxy as any
-            }
           }
         })
         await processStream(newStream, idx + 1)
@@ -410,38 +386,43 @@ export default class GeminiProvider extends BaseProvider {
 
     const processStream = async (stream: AsyncGenerator<GenerateContentResponse>, idx: number) => {
       let content = ''
-      for await (const chunk of stream) {
+      try {
-        if (window.keyv.get(EVENT_NAMES.CHAT_COMPLETION_PAUSED)) break
+        for await (const chunk of stream) {
+          if (window.keyv.get(EVENT_NAMES.CHAT_COMPLETION_PAUSED)) break
 
-        if (time_first_token_millsec == 0) {
+          if (time_first_token_millsec == 0) {
-          time_first_token_millsec = new Date().getTime() - start_time_millsec
+            time_first_token_millsec = new Date().getTime() - start_time_millsec
+          }
+
+          const time_completion_millsec = new Date().getTime() - start_time_millsec
+
+          if (chunk.text !== undefined) {
+            content += chunk.text
+          }
+          await processToolUses(content, idx)
+          const generateImage = this.processGeminiImageResponse(chunk)
+
+          onChunk({
+            text: chunk.text !== undefined ? chunk.text : '',
+            usage: {
+              prompt_tokens: chunk.usageMetadata?.promptTokenCount || 0,
+              completion_tokens: chunk.usageMetadata?.candidatesTokenCount || 0,
+              thoughts_tokens: chunk.usageMetadata?.thoughtsTokenCount || 0,
+              total_tokens: chunk.usageMetadata?.totalTokenCount || 0
+            },
+            metrics: {
+              completion_tokens: chunk.usageMetadata?.candidatesTokenCount,
+              time_completion_millsec,
+              time_first_token_millsec
+            },
+            search: chunk.candidates?.[0]?.groundingMetadata,
+            mcpToolResponse: toolResponses,
+            generateImage: generateImage
+          })
         }
-
+      } catch (error) {
-        const time_completion_millsec = new Date().getTime() - start_time_millsec
+        console.error('Error processing stream chunk:', error)
-
+        throw error
-        if (chunk.text !== undefined) {
-          content += chunk.text
-        }
-        await processToolUses(content, idx)
-        const generateImage = this.processGeminiImageResponse(chunk)
-
-        onChunk({
-          text: chunk.text !== undefined ? chunk.text : '',
-          usage: {
-            prompt_tokens: chunk.usageMetadata?.promptTokenCount || 0,
-            completion_tokens: chunk.usageMetadata?.candidatesTokenCount || 0,
-            thoughts_tokens: chunk.usageMetadata?.thoughtsTokenCount || 0,
-            total_tokens: chunk.usageMetadata?.totalTokenCount || 0
-          },
-          metrics: {
-            completion_tokens: chunk.usageMetadata?.candidatesTokenCount,
-            time_completion_millsec,
-            time_first_token_millsec
-          },
-          search: chunk.candidates?.[0]?.groundingMetadata,
-          mcpToolResponse: toolResponses,
-          generateImage: generateImage
-        })
       }
     }
 

yarn.lock

@@ -1489,23 +1489,15 @@ __metadata:
   languageName: node
   linkType: hard
 
-"@google/genai@npm:0.8.0":
+"@google/genai@npm:^0.10.0":
-  version: 0.8.0
+  version: 0.10.0
-  resolution: "@google/genai@npm:0.8.0"
+  resolution: "@google/genai@npm:0.10.0"
   dependencies:
     google-auth-library: "npm:^9.14.2"
     ws: "npm:^8.18.0"
-  checksum: 10c0/8a26a7dd1ab26aeeef5b5610612965ab271142460912c31b12f201cf6e00f5a4965910b195033992bdee1a7ee2b88c55f55d3a2727e09e4cd8d30ecbd0d655d0
+    zod: "npm:^3.22.4"
-  languageName: node
+    zod-to-json-schema: "npm:^3.22.4"
-  linkType: hard
+  checksum: 10c0/89a80eebdd053f2643fb977b939c3de1ece6889549f7d030b81a21dae0e0e0750233e453dafe665f8d337c642b4ba7a2f19b06bc270f3e2de8aeb803af1851d1
-
-"@google/genai@patch:@google/genai@npm%3A0.8.0#~/.yarn/patches/@google-genai-npm-0.8.0-450d0d9a7d.patch":
-  version: 0.8.0
-  resolution: "@google/genai@patch:@google/genai@npm%3A0.8.0#~/.yarn/patches/@google-genai-npm-0.8.0-450d0d9a7d.patch::version=0.8.0&hash=e4f82c"
-  dependencies:
-    google-auth-library: "npm:^9.14.2"
-    ws: "npm:^8.18.0"
-  checksum: 10c0/b2082599cbdb9ffe18020461a2290ac60b0912f21106254bb618a1a1faa6e400caf16f92479fdb877928fd9670617ec0131f5787b4e3db2e5918daf2c7b87def
   languageName: node
   linkType: hard
 
@@ -4315,7 +4307,7 @@ __metadata:
     "@emotion/is-prop-valid": "npm:^1.3.1"
     "@eslint-react/eslint-plugin": "npm:^1.36.1"
     "@eslint/js": "npm:^9.22.0"
-    "@google/genai": "patch:@google/genai@npm%3A0.8.0#~/.yarn/patches/@google-genai-npm-0.8.0-450d0d9a7d.patch"
+    "@google/genai": "npm:^0.10.0"
     "@hello-pangea/dnd": "npm:^16.6.0"
     "@kangfenmao/keyv-storage": "npm:^0.1.0"
     "@langchain/community": "npm:^0.3.36"
@@ -18032,7 +18024,7 @@ __metadata:
   languageName: node
   linkType: hard
 
-"zod-to-json-schema@npm:^3.22.3, zod-to-json-schema@npm:^3.22.5, zod-to-json-schema@npm:^3.24.1":
+"zod-to-json-schema@npm:^3.22.3, zod-to-json-schema@npm:^3.22.4, zod-to-json-schema@npm:^3.22.5, zod-to-json-schema@npm:^3.24.1":
   version: 3.24.5
   resolution: "zod-to-json-schema@npm:3.24.5"
   peerDependencies: