Merge pull request #678 from zengqinglei/fix/gemini-sse-transparent-relay

修复 Gemini 流式请求通过 HTTP 代理时的 120 秒超时问题
2026-01-22 16:43:35 +00:00 · 2025-11-20 00:53:29 -05:00
parent fb7a8f841a 9eccc7da49
commit 1b7016c3dc
5 changed files with 270 additions and 129 deletions
--- a/.github/workflows/auto-release-pipeline.yml
+++ b/.github/workflows/auto-release-pipeline.yml
@@ -4,6 +4,7 @@ on:
  push:
    branches:
      - main
+  workflow_dispatch:  # 支持手动触发

 permissions:
  contents: write
--- a/src/routes/geminiRoutes.js
+++ b/src/routes/geminiRoutes.js
@@ -924,86 +924,104 @@ async function handleStreamGenerateContent(req, res) {
    res.setHeader('X-Accel-Buffering', 'no')

    // 处理流式响应并捕获usage数据
-    let streamBuffer = '' // 统一的流处理缓冲区
+    // 方案 A++：透明转发 + 异步 usage 提取 + SSE 心跳机制
+    let streamBuffer = '' // 缓冲区用于处理不完整的行
    let totalUsage = {
      promptTokenCount: 0,
      candidatesTokenCount: 0,
      totalTokenCount: 0
    }
-    const usageReported = false
+    let usageReported = false // 修复：改为 let 以便后续修改
+
+    // SSE 心跳机制：防止 Clash 等代理 120 秒超时
+    let heartbeatTimer = null
+    let lastDataTime = Date.now()
+    const HEARTBEAT_INTERVAL = 15000 // 15 秒
+
+    const sendHeartbeat = () => {
+      const timeSinceLastData = Date.now() - lastDataTime
+      if (timeSinceLastData >= HEARTBEAT_INTERVAL && !res.destroyed) {
+        res.write('\n') // 发送空行保持连接活跃
+        logger.info(`💓 Sent SSE keepalive (gap: ${(timeSinceLastData / 1000).toFixed(1)}s)`)
+      }
+    }
+
+    heartbeatTimer = setInterval(sendHeartbeat, HEARTBEAT_INTERVAL)

    streamResponse.on('data', (chunk) => {
      try {
-        const chunkStr = chunk.toString()
+        // 更新最后数据时间
+        lastDataTime = Date.now()

-        if (!chunkStr.trim()) {
-          return
+        // 1️⃣ 立即转发原始数据（零延迟，最高优先级）
+        // 对所有版本（v1beta 和 v1internal）都采用透明转发
+        if (!res.destroyed) {
+          res.write(chunk) // 直接转发 Buffer，无需转换和序列化
        }

-        // 使用统一缓冲区处理不完整的行
-        streamBuffer += chunkStr
-        const lines = streamBuffer.split('\n')
-        streamBuffer = lines.pop() || '' // 保留最后一个不完整的行
+        // 2️⃣ 异步提取 usage 数据（不阻塞转发）
+        // 使用 setImmediate 将解析放到下一个事件循环
+        setImmediate(() => {
+          try {
+            const chunkStr = chunk.toString()
+            if (!chunkStr.trim()) {
+              return
+            }

-        const processedLines = []
+            // 快速检查是否包含 usage 数据（避免不必要的解析）
+            if (!chunkStr.includes('usageMetadata')) {
+              return
+            }

-        for (const line of lines) {
-          if (!line.trim()) {
-            continue // 跳过空行，不添加到处理队列
-          }
+            // 处理不完整的行
+            streamBuffer += chunkStr
+            const lines = streamBuffer.split('\n')
+            streamBuffer = lines.pop() || ''

-          // 解析 SSE 行
-          const parsed = parseSSELine(line)
-
-          // 提取 usage 数据（适用于所有版本）
-          if (parsed.type === 'data' && parsed.data.response?.usageMetadata) {
-            totalUsage = parsed.data.response.usageMetadata
-            logger.debug('📊 Captured Gemini usage data:', totalUsage)
-          }
-
-          // 根据版本处理输出
-          if (version === 'v1beta') {
-            if (parsed.type === 'data') {
-              if (parsed.data.response) {
-                // 有 response 字段，只返回 response 的内容
-                processedLines.push(`data: ${JSON.stringify(parsed.data.response)}`)
-              } else {
-                // 没有 response 字段，返回整个数据对象
-                processedLines.push(`data: ${JSON.stringify(parsed.data)}`)
+            // 仅解析包含 usage 的行
+            for (const line of lines) {
+              if (!line.trim() || !line.includes('usageMetadata')) {
+                continue
              }
-            } else if (parsed.type === 'control') {
-              // 控制消息（如 [DONE]）保持原样
-              processedLines.push(line)
-            }
-            // 跳过其他类型的行（'other', 'invalid'）
-          }
-        }

-        // 发送数据到客户端
-        if (version === 'v1beta') {
-          for (const line of processedLines) {
-            if (!res.destroyed) {
-              res.write(`${line}\n\n`)
+              try {
+                const parsed = parseSSELine(line)
+                if (parsed.type === 'data' && parsed.data.response?.usageMetadata) {
+                  totalUsage = parsed.data.response.usageMetadata
+                  logger.debug('📊 Captured Gemini usage data:', totalUsage)
+                }
+              } catch (parseError) {
+                // 解析失败但不影响转发
+                logger.warn('⚠️ Failed to parse usage line:', parseError.message)
+              }
            }
+          } catch (error) {
+            // 提取失败但不影响转发
+            logger.warn('⚠️ Error extracting usage data:', error.message)
          }
-        } else {
-          // v1internal 直接转发原始数据
-          if (!res.destroyed) {
-            res.write(chunkStr)
-          }
-        }
+        })
      } catch (error) {
        logger.error('Error processing stream chunk:', error)
+        // 不中断流，继续处理后续数据
      }
    })

-    streamResponse.on('end', async () => {
+    streamResponse.on('end', () => {
      logger.info('Stream completed successfully')

-      // 记录使用统计
+      // 清理心跳定时器
+      if (heartbeatTimer) {
+        clearInterval(heartbeatTimer)
+        heartbeatTimer = null
+      }
+
+      // 立即结束响应，不阻塞
+      res.end()
+
+      // 异步记录使用统计（不阻塞响应）
      if (!usageReported && totalUsage.totalTokenCount > 0) {
-        try {
-          await apiKeyService.recordUsage(
+        Promise.all([
+          apiKeyService.recordUsage(
            req.apiKey.id,
            totalUsage.promptTokenCount || 0,
            totalUsage.candidatesTokenCount || 0,
@@ -1011,12 +1029,8 @@ async function handleStreamGenerateContent(req, res) {
            0, // cacheReadTokens
            model,
            account.id
-          )
-          logger.info(
-            `📊 Recorded Gemini stream usage - Input: ${totalUsage.promptTokenCount}, Output: ${totalUsage.candidatesTokenCount}, Total: ${totalUsage.totalTokenCount}`
-          )
-
-          await applyRateLimitTracking(
+          ),
+          applyRateLimitTracking(
            req,
            {
              inputTokens: totalUsage.promptTokenCount || 0,
@@ -1027,17 +1041,30 @@ async function handleStreamGenerateContent(req, res) {
            model,
            'gemini-stream'
          )
-        } catch (error) {
-          logger.error('Failed to record Gemini usage:', error)
-        }
+        ])
+          .then(() => {
+            logger.info(
+              `📊 Recorded Gemini stream usage - Input: ${totalUsage.promptTokenCount}, Output: ${totalUsage.candidatesTokenCount}, Total: ${totalUsage.totalTokenCount}`
+            )
+            usageReported = true
+          })
+          .catch((error) => {
+            logger.error('Failed to record Gemini usage:', error)
+          })
      }
-
-      res.end()
    })

    streamResponse.on('error', (error) => {
      logger.error('Stream error:', error)
+
+      // 清理心跳定时器
+      if (heartbeatTimer) {
+        clearInterval(heartbeatTimer)
+        heartbeatTimer = null
+      }
+
      if (!res.headersSent) {
+        // 如果还没发送响应头，可以返回正常的错误响应
        res.status(500).json({
          error: {
            message: error.message || 'Stream error',
@@ -1045,6 +1072,27 @@ async function handleStreamGenerateContent(req, res) {
          }
        })
      } else {
+        // 如果已经开始流式传输，发送 SSE 格式的错误事件和结束标记
+        // 这样客户端可以正确识别流的结束，避免 "Premature close" 错误
+        if (!res.destroyed) {
+          try {
+            // 发送错误事件（SSE 格式）
+            res.write(
+              `data: ${JSON.stringify({
+                error: {
+                  message: error.message || 'Stream error',
+                  type: 'stream_error',
+                  code: error.code
+                }
+              })}\n\n`
+            )
+
+            // 发送 SSE 结束标记
+            res.write('data: [DONE]\n\n')
+          } catch (writeError) {
+            logger.error('Error sending error event:', writeError)
+          }
+        }
        res.end()
      }
    })
--- a/src/routes/openaiGeminiRoutes.js
+++ b/src/routes/openaiGeminiRoutes.js
@@ -386,7 +386,7 @@ router.post('/v1/chat/completions', authenticateApiKey, async (req, res) => {
        candidatesTokenCount: 0,
        totalTokenCount: 0
      }
-      const usageReported = false
+      let usageReported = false // 修复：改为 let 以便后续修改

      streamResponse.on('data', (chunk) => {
        try {
@@ -512,6 +512,9 @@ router.post('/v1/chat/completions', authenticateApiKey, async (req, res) => {
            logger.info(
              `📊 Recorded Gemini stream usage - Input: ${totalUsage.promptTokenCount}, Output: ${totalUsage.candidatesTokenCount}, Total: ${totalUsage.totalTokenCount}`
            )
+
+            // 修复：标记 usage 已上报，避免重复上报
+            usageReported = true
          } catch (error) {
            logger.error('Failed to record Gemini usage:', error)
          }
@@ -534,8 +537,23 @@ router.post('/v1/chat/completions', authenticateApiKey, async (req, res) => {
          })
        } else {
          // 如果已经开始发送流数据，发送错误事件
-          res.write(`data: {"error": {"message": "${error.message || 'Stream error'}"}}\n\n`)
-          res.write('data: [DONE]\n\n')
+          // 修复：使用 JSON.stringify 避免字符串插值导致的格式错误
+          if (!res.destroyed) {
+            try {
+              res.write(
+                `data: ${JSON.stringify({
+                  error: {
+                    message: error.message || 'Stream error',
+                    type: 'stream_error',
+                    code: error.code
+                  }
+                })}\n\n`
+              )
+              res.write('data: [DONE]\n\n')
+            } catch (writeError) {
+              logger.error('Error sending error event:', writeError)
+            }
+          }
          res.end()
        }
      })
--- a/src/routes/standardGeminiRoutes.js
+++ b/src/routes/standardGeminiRoutes.js
@@ -510,76 +510,102 @@ async function handleStandardStreamGenerateContent(req, res) {
    res.setHeader('X-Accel-Buffering', 'no')

    // 处理流式响应并捕获usage数据
-    let streamBuffer = '' // 统一的流处理缓冲区
+    // 方案 A++：透明转发 + 异步 usage 提取 + SSE 心跳机制
+    let streamBuffer = '' // 缓冲区用于处理不完整的行
    let totalUsage = {
      promptTokenCount: 0,
      candidatesTokenCount: 0,
      totalTokenCount: 0
    }

+    // SSE 心跳机制：防止 Clash 等代理 120 秒超时
+    let heartbeatTimer = null
+    let lastDataTime = Date.now()
+    const HEARTBEAT_INTERVAL = 15000 // 15 秒
+
+    const sendHeartbeat = () => {
+      const timeSinceLastData = Date.now() - lastDataTime
+      if (timeSinceLastData >= HEARTBEAT_INTERVAL && !res.destroyed) {
+        res.write('\n') // 发送空行保持连接活跃
+        logger.info(`💓 Sent SSE keepalive (gap: ${(timeSinceLastData / 1000).toFixed(1)}s)`)
+      }
+    }
+
+    heartbeatTimer = setInterval(sendHeartbeat, HEARTBEAT_INTERVAL)
+
    streamResponse.on('data', (chunk) => {
      try {
-        const chunkStr = chunk.toString()
+        // 更新最后数据时间
+        lastDataTime = Date.now()

-        if (!chunkStr.trim()) {
-          return
+        // 1️⃣ 立即转发原始数据（零延迟，最高优先级）
+        if (!res.destroyed) {
+          res.write(chunk) // 直接转发 Buffer，无需转换和序列化
        }

-        // 使用统一缓冲区处理不完整的行
-        streamBuffer += chunkStr
-        const lines = streamBuffer.split('\n')
-        streamBuffer = lines.pop() || '' // 保留最后一个不完整的行
-
-        for (const line of lines) {
-          if (!line.trim()) {
-            continue // 跳过空行
-          }
-
-          // 解析 SSE 行
-          const parsed = parseSSELine(line)
-
-          // 记录无效的解析（用于调试）
-          if (parsed.type === 'invalid') {
-            logger.warn('Failed to parse SSE line:', {
-              line: parsed.line.substring(0, 100),
-              error: parsed.error.message
-            })
-            continue
-          }
-
-          // 捕获 usage 数据
-          if (parsed.type === 'data' && parsed.data.response?.usageMetadata) {
-            totalUsage = parsed.data.response.usageMetadata
-            logger.debug('📊 Captured Gemini usage data:', totalUsage)
-          }
-
-          // 转换格式并发送
-          if (!res.destroyed) {
-            if (parsed.type === 'data') {
-              // 转换格式：移除 response 包装，直接返回标准 Gemini API 格式
-              if (parsed.data.response) {
-                res.write(`data: ${JSON.stringify(parsed.data.response)}\n\n`)
-              } else {
-                res.write(`data: ${JSON.stringify(parsed.data)}\n\n`)
-              }
-            } else if (parsed.type === 'control') {
-              // 保持控制消息（如 [DONE]）原样
-              res.write(`${parsed.line}\n\n`)
+        // 2️⃣ 异步提取 usage 数据（不阻塞转发）
+        // 使用 setImmediate 将解析放到下一个事件循环
+        setImmediate(() => {
+          try {
+            const chunkStr = chunk.toString()
+            if (!chunkStr.trim()) {
+              return
            }
+
+            // 快速检查是否包含 usage 数据（避免不必要的解析）
+            if (!chunkStr.includes('usageMetadata')) {
+              return
+            }
+
+            // 处理不完整的行
+            streamBuffer += chunkStr
+            const lines = streamBuffer.split('\n')
+            streamBuffer = lines.pop() || ''
+
+            // 仅解析包含 usage 的行
+            for (const line of lines) {
+              if (!line.trim() || !line.includes('usageMetadata')) {
+                continue
+              }
+
+              try {
+                const parsed = parseSSELine(line)
+                if (parsed.type === 'data' && parsed.data.response?.usageMetadata) {
+                  totalUsage = parsed.data.response.usageMetadata
+                  logger.debug('📊 Captured Gemini usage data:', totalUsage)
+                }
+              } catch (parseError) {
+                // 解析失败但不影响转发
+                logger.warn('⚠️ Failed to parse usage line:', parseError.message)
+              }
+            }
+          } catch (error) {
+            // 提取失败但不影响转发
+            logger.warn('⚠️ Error extracting usage data:', error.message)
          }
-        }
+        })
      } catch (error) {
        logger.error('Error processing stream chunk:', error)
+        // 不中断流，继续处理后续数据
      }
    })

-    streamResponse.on('end', async () => {
+    streamResponse.on('end', () => {
      logger.info('Stream completed successfully')

-      // 记录使用统计
+      // 清理心跳定时器
+      if (heartbeatTimer) {
+        clearInterval(heartbeatTimer)
+        heartbeatTimer = null
+      }
+
+      // 立即结束响应，不阻塞
+      res.end()
+
+      // 异步记录使用统计（不阻塞响应）
      if (totalUsage.totalTokenCount > 0) {
-        try {
-          await apiKeyService.recordUsage(
+        apiKeyService
+          .recordUsage(
            req.apiKey.id,
            totalUsage.promptTokenCount || 0,
            totalUsage.candidatesTokenCount || 0,
@@ -588,24 +614,32 @@ async function handleStandardStreamGenerateContent(req, res) {
            model,
            account.id
          )
-          logger.info(
-            `📊 Recorded Gemini stream usage - Input: ${totalUsage.promptTokenCount}, Output: ${totalUsage.candidatesTokenCount}, Total: ${totalUsage.totalTokenCount}`
-          )
-        } catch (error) {
-          logger.error('Failed to record Gemini usage:', error)
-        }
+          .then(() => {
+            logger.info(
+              `📊 Recorded Gemini stream usage - Input: ${totalUsage.promptTokenCount}, Output: ${totalUsage.candidatesTokenCount}, Total: ${totalUsage.totalTokenCount}`
+            )
+          })
+          .catch((error) => {
+            logger.error('Failed to record Gemini usage:', error)
+          })
      } else {
        logger.warn(
          `⚠️ Stream completed without usage data - totalTokenCount: ${totalUsage.totalTokenCount}`
        )
      }
-
-      res.end()
    })

    streamResponse.on('error', (error) => {
      logger.error('Stream error:', error)
+
+      // 清理心跳定时器
+      if (heartbeatTimer) {
+        clearInterval(heartbeatTimer)
+        heartbeatTimer = null
+      }
+
      if (!res.headersSent) {
+        // 如果还没发送响应头，可以返回正常的错误响应
        res.status(500).json({
          error: {
            message: error.message || 'Stream error',
@@ -613,6 +647,27 @@ async function handleStandardStreamGenerateContent(req, res) {
          }
        })
      } else {
+        // 如果已经开始流式传输，发送 SSE 格式的错误事件和结束标记
+        // 这样客户端可以正确识别流的结束，避免 "Premature close" 错误
+        if (!res.destroyed) {
+          try {
+            // 发送错误事件（SSE 格式）
+            res.write(
+              `data: ${JSON.stringify({
+                error: {
+                  message: error.message || 'Stream error',
+                  type: 'stream_error',
+                  code: error.code
+                }
+              })}\n\n`
+            )
+
+            // 发送 SSE 结束标记
+            res.write('data: [DONE]\n\n')
+          } catch (writeError) {
+            logger.error('Error sending error event:', writeError)
+          }
+        }
        res.end()
      }
    })
--- a/src/services/geminiAccountService.js
+++ b/src/services/geminiAccountService.js
@@ -1,6 +1,7 @@
 const redisClient = require('../models/redis')
 const { v4: uuidv4 } = require('uuid')
 const crypto = require('crypto')
+const https = require('https')
 const config = require('../../config/config')
 const logger = require('../utils/logger')
 const { OAuth2Client } = require('google-auth-library')
@@ -21,6 +22,18 @@ const OAUTH_CLIENT_ID = '681255809395-oo8ft2oprdrnp9e3aqf6av3hmdib135j.apps.goog
 const OAUTH_CLIENT_SECRET = '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'
 const OAUTH_SCOPES = ['https://www.googleapis.com/auth/cloud-platform']

+// 🌐 TCP Keep-Alive Agent 配置
+// 解决长时间流式请求中 NAT/防火墙空闲超时导致的连接中断问题
+const keepAliveAgent = new https.Agent({
+  keepAlive: true,
+  keepAliveMsecs: 30000, // 每30秒发送一次 keep-alive 探测
+  timeout: 120000, // 120秒连接超时
+  maxSockets: 100, // 最大并发连接数
+  maxFreeSockets: 10 // 保持的空闲连接数
+})
+
+logger.info('🌐 Gemini HTTPS Agent initialized with TCP Keep-Alive support')
+
 // 加密相关常量
 const ALGORITHM = 'aes-256-cbc'
 const ENCRYPTION_SALT = 'gemini-account-salt'
@@ -1472,7 +1485,7 @@ async function generateContent(
      'Content-Type': 'application/json'
    },
    data: request,
-    timeout: 60000 // 生成内容可能需要更长时间
+    timeout: 600000 // 生成内容可能需要更长时间
  }

  // 添加代理配置
@@ -1485,7 +1498,10 @@ async function generateContent(
      `🌐 Using proxy for Gemini generateContent: ${ProxyHelper.getProxyDescription(proxyConfig)}`
    )
  } else {
-    logger.debug('🌐 No proxy configured for Gemini generateContent')
+    // 没有代理时，使用 keepAlive agent 防止长时间请求被中断
+    axiosConfig.httpsAgent = keepAliveAgent
+    axiosConfig.httpAgent = keepAliveAgent
+    logger.debug('🌐 Using keepAlive agent for Gemini generateContent')
  }

  const response = await axios(axiosConfig)
@@ -1548,7 +1564,7 @@ async function generateContentStream(
    },
    data: request,
    responseType: 'stream',
-    timeout: 60000
+    timeout: 0 // 流式请求不设置超时限制，由 keepAlive 和 AbortSignal 控制
  }

  // 添加代理配置
@@ -1561,7 +1577,10 @@ async function generateContentStream(
      `🌐 Using proxy for Gemini streamGenerateContent: ${ProxyHelper.getProxyDescription(proxyConfig)}`
    )
  } else {
-    logger.debug('🌐 No proxy configured for Gemini streamGenerateContent')
+    // 没有代理时，使用 keepAlive agent 防止长时间流式请求被中断
+    axiosConfig.httpsAgent = keepAliveAgent
+    axiosConfig.httpAgent = keepAliveAgent
+    logger.debug('🌐 Using keepAlive agent for Gemini streamGenerateContent')
  }

  // 如果提供了中止信号，添加到配置中