From eeb504d447dcd32c552c98e972a48e8f9ff8c3b6 Mon Sep 17 00:00:00 2001
From: MyPrototypeWhat <43230886+MyPrototypeWhat@users.noreply.github.com>
Date: Tue, 17 Jun 2025 21:20:52 +0800
Subject: [PATCH] fix: refactor provider middleware (#7164)

---
 src/renderer/src/aiCore/index.ts              |  2 +
 .../common/ErrorHandlerMiddleware.ts          | 42 +++++++++----------
 .../feat/ImageGenerationMiddleware.ts         |  1 -
 3 files changed, 21 insertions(+), 24 deletions(-)

diff --git a/src/renderer/src/aiCore/index.ts b/src/renderer/src/aiCore/index.ts
index c7bffa17de..f5e33ad34c 100644
--- a/src/renderer/src/aiCore/index.ts
+++ b/src/renderer/src/aiCore/index.ts
@@ -11,6 +11,7 @@ import { AnthropicAPIClient } from './clients/anthropic/AnthropicAPIClient'
 import { OpenAIResponseAPIClient } from './clients/openai/OpenAIResponseAPIClient'
 import { CompletionsMiddlewareBuilder } from './middleware/builder'
 import { MIDDLEWARE_NAME as AbortHandlerMiddlewareName } from './middleware/common/AbortHandlerMiddleware'
+import { MIDDLEWARE_NAME as ErrorHandlerMiddlewareName } from './middleware/common/ErrorHandlerMiddleware'
 import { MIDDLEWARE_NAME as FinalChunkConsumerMiddlewareName } from './middleware/common/FinalChunkConsumerMiddleware'
 import { applyCompletionsMiddlewares } from './middleware/composer'
 import { MIDDLEWARE_NAME as McpToolChunkMiddlewareName } from './middleware/core/McpToolChunkMiddleware'
@@ -62,6 +63,7 @@ export default class AiProvider {
       builder.clear()
       builder
         .add(MiddlewareRegistry[FinalChunkConsumerMiddlewareName])
+        .add(MiddlewareRegistry[ErrorHandlerMiddlewareName])
         .add(MiddlewareRegistry[AbortHandlerMiddlewareName])
         .add(MiddlewareRegistry[ImageGenerationMiddlewareName])
     } else {
diff --git a/src/renderer/src/aiCore/middleware/common/ErrorHandlerMiddleware.ts b/src/renderer/src/aiCore/middleware/common/ErrorHandlerMiddleware.ts
index 2dd5aa9833..8875a0b627 100644
--- a/src/renderer/src/aiCore/middleware/common/ErrorHandlerMiddleware.ts
+++ b/src/renderer/src/aiCore/middleware/common/ErrorHandlerMiddleware.ts
@@ -1,5 +1,4 @@
 import { Chunk } from '@renderer/types/chunk'
-import { isAbortError } from '@renderer/utils/error'
 
 import { CompletionsResult } from '../schemas'
 import { CompletionsContext } from '../types'
@@ -26,30 +25,27 @@ export const ErrorHandlerMiddleware =
       // 尝试执行下一个中间件
       return await next(ctx, params)
     } catch (error: any) {
-      let errorStream: ReadableStream<Chunk> | undefined
-      // 有些sdk的abort error 是直接抛出的
-      if (!isAbortError(error)) {
-        // 1. 使用通用的工具函数将错误解析为标准格式
-        const errorChunk = createErrorChunk(error)
-        // 2. 调用从外部传入的 onError 回调
-        if (params.onError) {
-          params.onError(error)
-        }
-
-        // 3. 根据配置决定是重新抛出错误，还是将其作为流的一部分向下传递
-        if (shouldThrow) {
-          throw error
-        }
-
-        // 如果不抛出，则创建一个只包含该错误块的流并向下传递
-        errorStream = new ReadableStream<Chunk>({
-          start(controller) {
-            controller.enqueue(errorChunk)
-            controller.close()
-          }
-        })
+      console.log('ErrorHandlerMiddleware_error', error)
+      // 1. 使用通用的工具函数将错误解析为标准格式
+      const errorChunk = createErrorChunk(error)
+      // 2. 调用从外部传入的 onError 回调
+      if (params.onError) {
+        params.onError(error)
       }
 
+      // 3. 根据配置决定是重新抛出错误，还是将其作为流的一部分向下传递
+      if (shouldThrow) {
+        throw error
+      }
+
+      // 如果不抛出，则创建一个只包含该错误块的流并向下传递
+      const errorStream = new ReadableStream<Chunk>({
+        start(controller) {
+          controller.enqueue(errorChunk)
+          controller.close()
+        }
+      })
+
       return {
         rawOutput: undefined,
         stream: errorStream, // 将包含错误的流传递下去
diff --git a/src/renderer/src/aiCore/middleware/feat/ImageGenerationMiddleware.ts b/src/renderer/src/aiCore/middleware/feat/ImageGenerationMiddleware.ts
index 324382b918..d0a4dc4903 100644
--- a/src/renderer/src/aiCore/middleware/feat/ImageGenerationMiddleware.ts
+++ b/src/renderer/src/aiCore/middleware/feat/ImageGenerationMiddleware.ts
@@ -17,7 +17,6 @@ export const ImageGenerationMiddleware: CompletionsMiddleware =
     const { assistant, messages } = params
     const client = context.apiClientInstance as BaseApiClient<OpenAI>
     const signal = context._internal?.flowControl?.abortSignal
-
     if (!assistant.model || !isDedicatedImageGenerationModel(assistant.model) || typeof messages === 'string') {
       return next(context, params)
     }