antigravity_gateway_streaming.go 40 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229
  1. package service
  2. import (
  3. "bufio"
  4. "context"
  5. "encoding/json"
  6. "errors"
  7. "fmt"
  8. "net/http"
  9. "strings"
  10. "sync/atomic"
  11. "time"
  12. "github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
  13. "github.com/Wei-Shaw/sub2api/internal/pkg/logger"
  14. "github.com/gin-gonic/gin"
  15. )
  16. type antigravityStreamResult struct {
  17. usage *ClaudeUsage
  18. firstTokenMs *int
  19. clientDisconnect bool // 客户端是否在流式传输过程中断开
  20. }
  21. func (s *AntigravityGatewayService) observeAntigravityGeminiSSELine(c *gin.Context, line string) {
  22. observer := upstreamResponseModelObserverFromContext(c)
  23. if observer == nil {
  24. observer = beginUpstreamResponseModelObservation(c)
  25. }
  26. trimmed := strings.TrimSpace(line)
  27. if !strings.HasPrefix(trimmed, "data:") {
  28. return
  29. }
  30. payload := strings.TrimSpace(strings.TrimPrefix(trimmed, "data:"))
  31. if payload == "" || payload == "[DONE]" {
  32. return
  33. }
  34. // Observe the original payload: ObserveGemini supports both the v1internal
  35. // wrapper and direct Gemini response shapes. The main stream handler will
  36. // unwrap the same line for business processing, so unwrapping here would be
  37. // duplicate work on every SSE event.
  38. observer.ObserveGemini([]byte(payload))
  39. }
  40. // antigravityClientWriter 封装流式响应的客户端写入,自动检测断开并标记。
  41. // 断开后所有写入操作变为 no-op,调用方通过 Disconnected() 判断是否继续 drain 上游。
  42. type antigravityClientWriter struct {
  43. w gin.ResponseWriter
  44. flusher http.Flusher
  45. disconnected bool
  46. prefix string // 日志前缀,标识来源方法
  47. beforeFirstWrite func()
  48. }
  49. func newAntigravityClientWriter(w gin.ResponseWriter, flusher http.Flusher, prefix string) *antigravityClientWriter {
  50. return &antigravityClientWriter{w: w, flusher: flusher, prefix: prefix}
  51. }
  52. // Write 写入数据到客户端,写入失败时标记断开并返回 false
  53. func (cw *antigravityClientWriter) Write(p []byte) bool {
  54. if cw.disconnected {
  55. return false
  56. }
  57. cw.prepareFirstWrite()
  58. if _, err := cw.w.Write(p); err != nil {
  59. cw.markDisconnected()
  60. return false
  61. }
  62. cw.flusher.Flush()
  63. return true
  64. }
  65. // Fprintf 格式化写入数据到客户端,写入失败时标记断开并返回 false
  66. func (cw *antigravityClientWriter) Fprintf(format string, args ...any) bool {
  67. if cw.disconnected {
  68. return false
  69. }
  70. cw.prepareFirstWrite()
  71. if _, err := fmt.Fprintf(cw.w, format, args...); err != nil {
  72. cw.markDisconnected()
  73. return false
  74. }
  75. cw.flusher.Flush()
  76. return true
  77. }
  78. func (cw *antigravityClientWriter) Disconnected() bool { return cw.disconnected }
  79. func (cw *antigravityClientWriter) prepareFirstWrite() {
  80. if cw.beforeFirstWrite == nil {
  81. return
  82. }
  83. prepare := cw.beforeFirstWrite
  84. cw.beforeFirstWrite = nil
  85. prepare()
  86. }
  87. func (cw *antigravityClientWriter) markDisconnected() {
  88. cw.disconnected = true
  89. logger.LegacyPrintf("service.antigravity_gateway", "Client disconnected during streaming (%s), continuing to drain upstream for billing", cw.prefix)
  90. }
  91. // handleStreamReadError 处理上游读取错误的通用逻辑。
  92. // 返回 (clientDisconnect, handled):handled=true 表示错误已处理,调用方应返回已收集的 usage。
  93. func handleStreamReadError(err error, clientDisconnected bool, prefix string) (disconnect bool, handled bool) {
  94. if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
  95. logger.LegacyPrintf("service.antigravity_gateway", "Context canceled during streaming (%s), returning collected usage", prefix)
  96. return true, true
  97. }
  98. if clientDisconnected {
  99. logger.LegacyPrintf("service.antigravity_gateway", "Upstream read error after client disconnect (%s): %v, returning collected usage", prefix, err)
  100. return true, true
  101. }
  102. return false, false
  103. }
  104. func (s *AntigravityGatewayService) handleGeminiStreamingResponse(c *gin.Context, resp *http.Response, startTime time.Time) (*antigravityStreamResult, error) {
  105. if upstreamResponseModelObserverFromContext(c) == nil {
  106. beginUpstreamResponseModelObservation(c)
  107. }
  108. c.Status(resp.StatusCode)
  109. c.Header("Cache-Control", "no-cache")
  110. c.Header("Connection", "keep-alive")
  111. c.Header("X-Accel-Buffering", "no")
  112. contentType := resp.Header.Get("Content-Type")
  113. if contentType == "" {
  114. contentType = "text/event-stream; charset=utf-8"
  115. }
  116. c.Header("Content-Type", contentType)
  117. flusher, ok := c.Writer.(http.Flusher)
  118. if !ok {
  119. return nil, errors.New("streaming not supported")
  120. }
  121. // 使用 Scanner 并限制单行大小,避免 ReadString 无上限导致 OOM
  122. scanner := bufio.NewScanner(resp.Body)
  123. maxLineSize := defaultMaxLineSize
  124. if s.settingService.cfg != nil && s.settingService.cfg.Gateway.MaxLineSize > 0 {
  125. maxLineSize = s.settingService.cfg.Gateway.MaxLineSize
  126. }
  127. scanBuf := getSSEScannerBuf64K()
  128. scanner.Buffer(scanBuf[:0], maxLineSize)
  129. usage := &ClaudeUsage{}
  130. var firstTokenMs *int
  131. type scanEvent struct {
  132. line string
  133. err error
  134. }
  135. // 独立 goroutine 读取上游,避免读取阻塞影响超时处理
  136. events := make(chan scanEvent, 16)
  137. done := make(chan struct{})
  138. sendEvent := func(ev scanEvent) bool {
  139. select {
  140. case events <- ev:
  141. return true
  142. case <-done:
  143. return false
  144. }
  145. }
  146. var lastReadAt int64
  147. atomic.StoreInt64(&lastReadAt, time.Now().UnixNano())
  148. go func(scanBuf *sseScannerBuf64K) {
  149. defer putSSEScannerBuf64K(scanBuf)
  150. defer close(events)
  151. for scanner.Scan() {
  152. atomic.StoreInt64(&lastReadAt, time.Now().UnixNano())
  153. if !sendEvent(scanEvent{line: scanner.Text()}) {
  154. return
  155. }
  156. }
  157. if err := scanner.Err(); err != nil {
  158. _ = sendEvent(scanEvent{err: err})
  159. }
  160. }(scanBuf)
  161. defer close(done)
  162. // 上游数据间隔超时保护(防止上游挂起长期占用连接)
  163. streamInterval := time.Duration(0)
  164. if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamDataIntervalTimeout > 0 {
  165. streamInterval = time.Duration(s.settingService.cfg.Gateway.StreamDataIntervalTimeout) * time.Second
  166. }
  167. var intervalTicker *time.Ticker
  168. if streamInterval > 0 {
  169. intervalTicker = time.NewTicker(streamInterval)
  170. defer intervalTicker.Stop()
  171. }
  172. var intervalCh <-chan time.Time
  173. if intervalTicker != nil {
  174. intervalCh = intervalTicker.C
  175. }
  176. // 下游 keepalive:防止代理/Cloudflare Tunnel 因连接空闲而断开
  177. keepaliveInterval := time.Duration(0)
  178. if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamKeepaliveInterval > 0 {
  179. keepaliveInterval = time.Duration(s.settingService.cfg.Gateway.StreamKeepaliveInterval) * time.Second
  180. }
  181. var keepaliveTicker *time.Ticker
  182. if keepaliveInterval > 0 {
  183. keepaliveTicker = time.NewTicker(keepaliveInterval)
  184. defer keepaliveTicker.Stop()
  185. }
  186. var keepaliveCh <-chan time.Time
  187. if keepaliveTicker != nil {
  188. keepaliveCh = keepaliveTicker.C
  189. }
  190. lastDataAt := time.Now()
  191. cw := newAntigravityClientWriter(c.Writer, flusher, "antigravity gemini")
  192. // 仅发送一次错误事件,避免多次写入导致协议混乱
  193. errorEventSent := false
  194. sendErrorEvent := func(reason string) {
  195. if errorEventSent || cw.Disconnected() {
  196. return
  197. }
  198. errorEventSent = true
  199. _, _ = fmt.Fprintf(c.Writer, "event: error\ndata: {\"error\":\"%s\"}\n\n", reason)
  200. flusher.Flush()
  201. }
  202. for {
  203. select {
  204. case ev, ok := <-events:
  205. if !ok {
  206. return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: cw.Disconnected()}, nil
  207. }
  208. if ev.err != nil {
  209. if disconnect, handled := handleStreamReadError(ev.err, cw.Disconnected(), "antigravity gemini"); handled {
  210. return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: disconnect}, nil
  211. }
  212. if errors.Is(ev.err, bufio.ErrTooLong) {
  213. logger.LegacyPrintf("service.antigravity_gateway", "SSE line too long (antigravity): max_size=%d error=%v", maxLineSize, ev.err)
  214. sendErrorEvent("response_too_large")
  215. return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs}, ev.err
  216. }
  217. sendErrorEvent("stream_read_error")
  218. return nil, ev.err
  219. }
  220. lastDataAt = time.Now()
  221. line := ev.line
  222. s.observeAntigravityGeminiSSELine(c, line)
  223. trimmed := strings.TrimRight(line, "\r\n")
  224. if strings.HasPrefix(trimmed, "data:") {
  225. payload := strings.TrimSpace(strings.TrimPrefix(trimmed, "data:"))
  226. if payload == "" || payload == "[DONE]" {
  227. cw.Fprintf("%s\n", line)
  228. continue
  229. }
  230. // 解包 v1internal 响应
  231. inner, parseErr := s.unwrapV1InternalResponse([]byte(payload))
  232. if parseErr == nil && inner != nil {
  233. payload = string(inner)
  234. }
  235. // 解析 usage
  236. if u := extractGeminiUsage(inner); u != nil {
  237. usage = u
  238. }
  239. var parsed map[string]any
  240. if json.Unmarshal(inner, &parsed) == nil {
  241. // Check for MALFORMED_FUNCTION_CALL
  242. if candidates, ok := parsed["candidates"].([]any); ok && len(candidates) > 0 {
  243. if cand, ok := candidates[0].(map[string]any); ok {
  244. if fr, ok := cand["finishReason"].(string); ok && fr == "MALFORMED_FUNCTION_CALL" {
  245. logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity] MALFORMED_FUNCTION_CALL detected in forward stream")
  246. if content, ok := cand["content"]; ok {
  247. if b, err := json.Marshal(content); err == nil {
  248. logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity] Malformed content: %s", string(b))
  249. }
  250. }
  251. }
  252. }
  253. }
  254. }
  255. if firstTokenMs == nil {
  256. ms := int(time.Since(startTime).Milliseconds())
  257. firstTokenMs = &ms
  258. }
  259. cw.Fprintf("data: %s\n\n", payload)
  260. continue
  261. }
  262. cw.Fprintf("%s\n", line)
  263. case <-intervalCh:
  264. lastRead := time.Unix(0, atomic.LoadInt64(&lastReadAt))
  265. if time.Since(lastRead) < streamInterval {
  266. continue
  267. }
  268. if cw.Disconnected() {
  269. logger.LegacyPrintf("service.antigravity_gateway", "Upstream timeout after client disconnect (antigravity gemini), returning collected usage")
  270. return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: true}, nil
  271. }
  272. logger.LegacyPrintf("service.antigravity_gateway", "Stream data interval timeout (antigravity)")
  273. sendErrorEvent("stream_timeout")
  274. return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs}, fmt.Errorf("stream data interval timeout")
  275. case <-keepaliveCh:
  276. if cw.Disconnected() {
  277. continue
  278. }
  279. if time.Since(lastDataAt) < keepaliveInterval {
  280. continue
  281. }
  282. // SSE ping/keepalive:保持连接活跃防止 Cloudflare Tunnel 等代理断开
  283. if !cw.Fprintf(":\n\n") {
  284. logger.LegacyPrintf("service.antigravity_gateway", "Client disconnected during keepalive ping (antigravity gemini), continuing to drain upstream for billing")
  285. continue
  286. }
  287. }
  288. }
  289. }
  290. // handleGeminiStreamToNonStreaming 读取上游流式响应,合并为非流式响应返回给客户端
  291. // Gemini 流式响应是增量的,需要累积所有 chunk 的内容
  292. func (s *AntigravityGatewayService) handleGeminiStreamToNonStreaming(c *gin.Context, resp *http.Response, startTime time.Time) (*antigravityStreamResult, error) {
  293. if upstreamResponseModelObserverFromContext(c) == nil {
  294. beginUpstreamResponseModelObservation(c)
  295. }
  296. scanner := bufio.NewScanner(resp.Body)
  297. maxLineSize := defaultMaxLineSize
  298. if s.settingService.cfg != nil && s.settingService.cfg.Gateway.MaxLineSize > 0 {
  299. maxLineSize = s.settingService.cfg.Gateway.MaxLineSize
  300. }
  301. scanBuf := getSSEScannerBuf64K()
  302. scanner.Buffer(scanBuf[:0], maxLineSize)
  303. usage := &ClaudeUsage{}
  304. var firstTokenMs *int
  305. var last map[string]any
  306. var lastWithParts map[string]any
  307. var collectedImageParts []map[string]any // 收集所有包含图片的 parts
  308. var collectedTextParts []string // 收集所有文本片段
  309. type scanEvent struct {
  310. line string
  311. err error
  312. }
  313. // 独立 goroutine 读取上游,避免读取阻塞影响超时处理
  314. events := make(chan scanEvent, 16)
  315. done := make(chan struct{})
  316. sendEvent := func(ev scanEvent) bool {
  317. select {
  318. case events <- ev:
  319. return true
  320. case <-done:
  321. return false
  322. }
  323. }
  324. var lastReadAt int64
  325. atomic.StoreInt64(&lastReadAt, time.Now().UnixNano())
  326. go func(scanBuf *sseScannerBuf64K) {
  327. defer putSSEScannerBuf64K(scanBuf)
  328. defer close(events)
  329. for scanner.Scan() {
  330. atomic.StoreInt64(&lastReadAt, time.Now().UnixNano())
  331. if !sendEvent(scanEvent{line: scanner.Text()}) {
  332. return
  333. }
  334. }
  335. if err := scanner.Err(); err != nil {
  336. _ = sendEvent(scanEvent{err: err})
  337. }
  338. }(scanBuf)
  339. defer close(done)
  340. // 上游数据间隔超时保护(防止上游挂起长期占用连接)
  341. streamInterval := time.Duration(0)
  342. if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamDataIntervalTimeout > 0 {
  343. streamInterval = time.Duration(s.settingService.cfg.Gateway.StreamDataIntervalTimeout) * time.Second
  344. }
  345. var intervalTicker *time.Ticker
  346. if streamInterval > 0 {
  347. intervalTicker = time.NewTicker(streamInterval)
  348. defer intervalTicker.Stop()
  349. }
  350. var intervalCh <-chan time.Time
  351. if intervalTicker != nil {
  352. intervalCh = intervalTicker.C
  353. }
  354. for {
  355. select {
  356. case ev, ok := <-events:
  357. if !ok {
  358. // 流结束,返回收集的响应
  359. goto returnResponse
  360. }
  361. if ev.err != nil {
  362. if errors.Is(ev.err, bufio.ErrTooLong) {
  363. logger.LegacyPrintf("service.antigravity_gateway", "SSE line too long (antigravity non-stream): max_size=%d error=%v", maxLineSize, ev.err)
  364. }
  365. return nil, ev.err
  366. }
  367. line := ev.line
  368. s.observeAntigravityGeminiSSELine(c, line)
  369. trimmed := strings.TrimRight(line, "\r\n")
  370. if !strings.HasPrefix(trimmed, "data:") {
  371. continue
  372. }
  373. payload := strings.TrimSpace(strings.TrimPrefix(trimmed, "data:"))
  374. if payload == "" || payload == "[DONE]" {
  375. continue
  376. }
  377. // 解包 v1internal 响应
  378. inner, parseErr := s.unwrapV1InternalResponse([]byte(payload))
  379. if parseErr != nil {
  380. continue
  381. }
  382. var parsed map[string]any
  383. if err := json.Unmarshal(inner, &parsed); err != nil {
  384. continue
  385. }
  386. // 记录首 token 时间
  387. if firstTokenMs == nil {
  388. ms := int(time.Since(startTime).Milliseconds())
  389. firstTokenMs = &ms
  390. }
  391. last = parsed
  392. // 提取 usage
  393. if u := extractGeminiUsage(inner); u != nil {
  394. usage = u
  395. }
  396. // Check for MALFORMED_FUNCTION_CALL
  397. if candidates, ok := parsed["candidates"].([]any); ok && len(candidates) > 0 {
  398. if cand, ok := candidates[0].(map[string]any); ok {
  399. if fr, ok := cand["finishReason"].(string); ok && fr == "MALFORMED_FUNCTION_CALL" {
  400. logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity] MALFORMED_FUNCTION_CALL detected in forward non-stream collect")
  401. if content, ok := cand["content"]; ok {
  402. if b, err := json.Marshal(content); err == nil {
  403. logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity] Malformed content: %s", string(b))
  404. }
  405. }
  406. }
  407. }
  408. }
  409. // 保留最后一个有 parts 的响应
  410. if parts := extractGeminiParts(parsed); len(parts) > 0 {
  411. lastWithParts = parsed
  412. // 收集包含图片和文本的 parts
  413. for _, part := range parts {
  414. if inlineData, ok := part["inlineData"].(map[string]any); ok {
  415. collectedImageParts = append(collectedImageParts, part)
  416. _ = inlineData // 避免 unused 警告
  417. }
  418. if text, ok := part["text"].(string); ok && text != "" {
  419. collectedTextParts = append(collectedTextParts, text)
  420. }
  421. }
  422. }
  423. case <-intervalCh:
  424. lastRead := time.Unix(0, atomic.LoadInt64(&lastReadAt))
  425. if time.Since(lastRead) < streamInterval {
  426. continue
  427. }
  428. logger.LegacyPrintf("service.antigravity_gateway", "Stream data interval timeout (antigravity non-stream)")
  429. return nil, fmt.Errorf("stream data interval timeout")
  430. }
  431. }
  432. returnResponse:
  433. // 选择最后一个有效响应
  434. finalResponse := pickGeminiCollectResult(last, lastWithParts)
  435. // 处理空响应情况 — 触发同账号重试 + failover 切换账号
  436. if last == nil && lastWithParts == nil {
  437. logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Forward] warning: empty stream response (gemini non-stream), triggering failover")
  438. return nil, &UpstreamFailoverError{
  439. StatusCode: http.StatusBadGateway,
  440. ResponseBody: []byte(`{"error":"empty stream response from upstream"}`),
  441. RetryableOnSameAccount: true,
  442. }
  443. }
  444. // 如果收集到了图片 parts,需要合并到最终响应中
  445. if len(collectedImageParts) > 0 {
  446. finalResponse = mergeImagePartsToResponse(finalResponse, collectedImageParts)
  447. }
  448. // 如果收集到了文本,需要合并到最终响应中
  449. if len(collectedTextParts) > 0 {
  450. finalResponse = mergeTextPartsToResponse(finalResponse, collectedTextParts)
  451. }
  452. respBody, err := json.Marshal(finalResponse)
  453. if err != nil {
  454. return nil, fmt.Errorf("failed to marshal response: %w", err)
  455. }
  456. c.Data(http.StatusOK, "application/json", respBody)
  457. return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs}, nil
  458. }
  459. // getOrCreateGeminiParts 获取 Gemini 响应的 parts 结构,返回深拷贝和更新回调
  460. func getOrCreateGeminiParts(response map[string]any) (result map[string]any, existingParts []any, setParts func([]any)) {
  461. // 深拷贝 response
  462. result = make(map[string]any)
  463. for k, v := range response {
  464. result[k] = v
  465. }
  466. // 获取或创建 candidates
  467. candidates, ok := result["candidates"].([]any)
  468. if !ok || len(candidates) == 0 {
  469. candidates = []any{map[string]any{}}
  470. }
  471. // 获取第一个 candidate
  472. candidate, ok := candidates[0].(map[string]any)
  473. if !ok {
  474. candidate = make(map[string]any)
  475. candidates[0] = candidate
  476. }
  477. // 获取或创建 content
  478. content, ok := candidate["content"].(map[string]any)
  479. if !ok {
  480. content = map[string]any{"role": "model"}
  481. candidate["content"] = content
  482. }
  483. // 获取现有 parts
  484. existingParts, ok = content["parts"].([]any)
  485. if !ok {
  486. existingParts = []any{}
  487. }
  488. // 返回更新回调
  489. setParts = func(newParts []any) {
  490. content["parts"] = newParts
  491. result["candidates"] = candidates
  492. }
  493. return result, existingParts, setParts
  494. }
  495. // mergeCollectedPartsToResponse 将收集的所有 parts 合并到 Gemini 响应中
  496. // 这个函数会合并所有类型的 parts:text、thinking、functionCall、inlineData 等
  497. // 保持原始顺序,只合并连续的普通 text parts
  498. func mergeCollectedPartsToResponse(response map[string]any, collectedParts []map[string]any) map[string]any {
  499. if len(collectedParts) == 0 {
  500. return response
  501. }
  502. result, _, setParts := getOrCreateGeminiParts(response)
  503. // 合并策略:
  504. // 1. 保持原始顺序
  505. // 2. 连续的普通 text parts 合并为一个
  506. // 3. thinking、functionCall、inlineData 等保持原样
  507. var mergedParts []any
  508. var textBuffer strings.Builder
  509. flushTextBuffer := func() {
  510. if textBuffer.Len() > 0 {
  511. mergedParts = append(mergedParts, map[string]any{
  512. "text": textBuffer.String(),
  513. })
  514. textBuffer.Reset()
  515. }
  516. }
  517. for _, part := range collectedParts {
  518. // 检查是否是普通 text part
  519. if text, ok := part["text"].(string); ok {
  520. // 检查是否有 thought 标记
  521. if thought, _ := part["thought"].(bool); thought {
  522. // thinking part,先刷新 text buffer,然后保留原样
  523. flushTextBuffer()
  524. mergedParts = append(mergedParts, part)
  525. } else {
  526. // 普通 text,累积到 buffer
  527. _, _ = textBuffer.WriteString(text)
  528. }
  529. } else {
  530. // 非 text part(functionCall、inlineData 等),先刷新 text buffer,然后保留原样
  531. flushTextBuffer()
  532. mergedParts = append(mergedParts, part)
  533. }
  534. }
  535. // 刷新剩余的 text
  536. flushTextBuffer()
  537. setParts(mergedParts)
  538. return result
  539. }
  540. // mergeImagePartsToResponse 将收集到的图片 parts 合并到 Gemini 响应中
  541. func mergeImagePartsToResponse(response map[string]any, imageParts []map[string]any) map[string]any {
  542. if len(imageParts) == 0 {
  543. return response
  544. }
  545. result, existingParts, setParts := getOrCreateGeminiParts(response)
  546. // 检查现有 parts 中是否已经有图片
  547. for _, p := range existingParts {
  548. if pm, ok := p.(map[string]any); ok {
  549. if _, hasInline := pm["inlineData"]; hasInline {
  550. return result // 已有图片,不重复添加
  551. }
  552. }
  553. }
  554. // 添加收集到的图片 parts
  555. for _, imgPart := range imageParts {
  556. existingParts = append(existingParts, imgPart)
  557. }
  558. setParts(existingParts)
  559. return result
  560. }
  561. // mergeTextPartsToResponse 将收集到的文本合并到 Gemini 响应中
  562. func mergeTextPartsToResponse(response map[string]any, textParts []string) map[string]any {
  563. if len(textParts) == 0 {
  564. return response
  565. }
  566. mergedText := strings.Join(textParts, "")
  567. result, existingParts, setParts := getOrCreateGeminiParts(response)
  568. // 查找并更新第一个 text part,或创建新的
  569. newParts := make([]any, 0, len(existingParts)+1)
  570. textUpdated := false
  571. for _, p := range existingParts {
  572. pm, ok := p.(map[string]any)
  573. if !ok {
  574. newParts = append(newParts, p)
  575. continue
  576. }
  577. if _, hasText := pm["text"]; hasText && !textUpdated {
  578. // 用累积的文本替换
  579. newPart := make(map[string]any)
  580. for k, v := range pm {
  581. newPart[k] = v
  582. }
  583. newPart["text"] = mergedText
  584. newParts = append(newParts, newPart)
  585. textUpdated = true
  586. } else {
  587. newParts = append(newParts, pm)
  588. }
  589. }
  590. if !textUpdated {
  591. newParts = append([]any{map[string]any{"text": mergedText}}, newParts...)
  592. }
  593. setParts(newParts)
  594. return result
  595. }
  596. func (s *AntigravityGatewayService) writeClaudeError(c *gin.Context, status int, errType, message string) error {
  597. MarkResponseCommitted(c)
  598. c.JSON(status, gin.H{
  599. "type": "error",
  600. "error": gin.H{"type": errType, "message": message},
  601. })
  602. return fmt.Errorf("%s", message)
  603. }
  604. // WriteMappedClaudeError 导出版本,供 handler 层使用(如 fallback 错误处理)
  605. func (s *AntigravityGatewayService) WriteMappedClaudeError(c *gin.Context, account *Account, upstreamStatus int, upstreamRequestID string, body []byte) error {
  606. return s.writeMappedClaudeError(c, account, upstreamStatus, upstreamRequestID, body)
  607. }
  608. func (s *AntigravityGatewayService) writeMappedClaudeError(c *gin.Context, account *Account, upstreamStatus int, upstreamRequestID string, body []byte) error {
  609. MarkResponseCommitted(c)
  610. upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(body))
  611. upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
  612. logBody, maxBytes := s.getLogConfig()
  613. upstreamDetail := s.getUpstreamErrorDetail(body)
  614. setOpsUpstreamError(c, upstreamStatus, upstreamMsg, upstreamDetail)
  615. appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
  616. ProxyID: opsUpstreamProxyID(account),
  617. ProxyName: opsUpstreamProxyName(account),
  618. Platform: account.Platform,
  619. AccountID: account.ID,
  620. AccountName: account.Name,
  621. UpstreamStatusCode: upstreamStatus,
  622. UpstreamRequestID: upstreamRequestID,
  623. Kind: "http_error",
  624. Message: upstreamMsg,
  625. Detail: upstreamDetail,
  626. })
  627. // 记录上游错误详情便于排障(可选:由配置控制;不回显到客户端)
  628. if logBody {
  629. logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Forward] upstream_error status=%d body=%s", upstreamStatus, truncateForLog(body, maxBytes))
  630. }
  631. // 检查错误透传规则
  632. if ptStatus, ptErrType, ptErrMsg, matched := applyErrorPassthroughRule(
  633. c, account.Platform, upstreamStatus, body,
  634. 0, "", "",
  635. ); matched {
  636. c.JSON(ptStatus, gin.H{
  637. "type": "error",
  638. "error": gin.H{"type": ptErrType, "message": ptErrMsg},
  639. })
  640. if upstreamMsg == "" {
  641. return fmt.Errorf("upstream error: %d", upstreamStatus)
  642. }
  643. return fmt.Errorf("upstream error: %d message=%s", upstreamStatus, upstreamMsg)
  644. }
  645. var statusCode int
  646. var errType, errMsg string
  647. switch upstreamStatus {
  648. case 400:
  649. statusCode = http.StatusBadRequest
  650. errType = "invalid_request_error"
  651. errMsg = getPassthroughOrDefault(upstreamMsg, "Invalid request")
  652. case 401:
  653. statusCode = http.StatusBadGateway
  654. errType = "authentication_error"
  655. errMsg = "Upstream authentication failed"
  656. case 403:
  657. statusCode = http.StatusBadGateway
  658. errType = "permission_error"
  659. errMsg = "Upstream access forbidden"
  660. case 429:
  661. statusCode = http.StatusTooManyRequests
  662. errType = "rate_limit_error"
  663. errMsg = "Upstream rate limit exceeded"
  664. case 529:
  665. statusCode = http.StatusServiceUnavailable
  666. errType = "overloaded_error"
  667. errMsg = "Upstream service overloaded"
  668. default:
  669. statusCode = http.StatusBadGateway
  670. errType = "upstream_error"
  671. errMsg = "Upstream request failed"
  672. }
  673. c.JSON(statusCode, gin.H{
  674. "type": "error",
  675. "error": gin.H{"type": errType, "message": errMsg},
  676. })
  677. if upstreamMsg == "" {
  678. return fmt.Errorf("upstream error: %d", upstreamStatus)
  679. }
  680. return fmt.Errorf("upstream error: %d message=%s", upstreamStatus, upstreamMsg)
  681. }
  682. func (s *AntigravityGatewayService) writeGoogleError(c *gin.Context, status int, message string) error {
  683. MarkResponseCommitted(c)
  684. statusStr := "UNKNOWN"
  685. switch status {
  686. case 400:
  687. statusStr = "INVALID_ARGUMENT"
  688. case 404:
  689. statusStr = "NOT_FOUND"
  690. case 429:
  691. statusStr = "RESOURCE_EXHAUSTED"
  692. case 500:
  693. statusStr = "INTERNAL"
  694. case 502, 503:
  695. statusStr = "UNAVAILABLE"
  696. }
  697. c.JSON(status, gin.H{
  698. "error": gin.H{
  699. "code": status,
  700. "message": message,
  701. "status": statusStr,
  702. },
  703. })
  704. return fmt.Errorf("%s", message)
  705. }
  706. // collectClaudeStreamResponse 收集上游流式响应,转换为 Claude 非流式格式返回
  707. // 用于处理客户端非流式请求但上游只支持流式的情况
  708. func (s *AntigravityGatewayService) collectClaudeStreamResponse(c *gin.Context, resp *http.Response, startTime time.Time, originalModel string) ([]byte, *antigravityStreamResult, error) {
  709. if upstreamResponseModelObserverFromContext(c) == nil {
  710. beginUpstreamResponseModelObservation(c)
  711. }
  712. scanner := bufio.NewScanner(resp.Body)
  713. maxLineSize := defaultMaxLineSize
  714. if s.settingService.cfg != nil && s.settingService.cfg.Gateway.MaxLineSize > 0 {
  715. maxLineSize = s.settingService.cfg.Gateway.MaxLineSize
  716. }
  717. scanBuf := getSSEScannerBuf64K()
  718. scanner.Buffer(scanBuf[:0], maxLineSize)
  719. var firstTokenMs *int
  720. var last map[string]any
  721. var lastWithParts map[string]any
  722. var collectedParts []map[string]any // 收集所有 parts(包括 text、thinking、functionCall、inlineData 等)
  723. var meaningfulResponse bool
  724. type scanEvent struct {
  725. line string
  726. err error
  727. }
  728. // 独立 goroutine 读取上游,避免读取阻塞影响超时处理
  729. events := make(chan scanEvent, 16)
  730. done := make(chan struct{})
  731. sendEvent := func(ev scanEvent) bool {
  732. select {
  733. case events <- ev:
  734. return true
  735. case <-done:
  736. return false
  737. }
  738. }
  739. var lastReadAt int64
  740. atomic.StoreInt64(&lastReadAt, time.Now().UnixNano())
  741. go func(scanBuf *sseScannerBuf64K) {
  742. defer putSSEScannerBuf64K(scanBuf)
  743. defer close(events)
  744. for scanner.Scan() {
  745. atomic.StoreInt64(&lastReadAt, time.Now().UnixNano())
  746. if !sendEvent(scanEvent{line: scanner.Text()}) {
  747. return
  748. }
  749. }
  750. if err := scanner.Err(); err != nil {
  751. _ = sendEvent(scanEvent{err: err})
  752. }
  753. }(scanBuf)
  754. defer close(done)
  755. // 上游数据间隔超时保护(防止上游挂起长期占用连接)
  756. streamInterval := time.Duration(0)
  757. if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamDataIntervalTimeout > 0 {
  758. streamInterval = time.Duration(s.settingService.cfg.Gateway.StreamDataIntervalTimeout) * time.Second
  759. }
  760. var intervalTicker *time.Ticker
  761. if streamInterval > 0 {
  762. intervalTicker = time.NewTicker(streamInterval)
  763. defer intervalTicker.Stop()
  764. }
  765. var intervalCh <-chan time.Time
  766. if intervalTicker != nil {
  767. intervalCh = intervalTicker.C
  768. }
  769. for {
  770. select {
  771. case ev, ok := <-events:
  772. if !ok {
  773. // 流结束,转换并返回响应
  774. goto returnResponse
  775. }
  776. if ev.err != nil {
  777. if errors.Is(ev.err, bufio.ErrTooLong) {
  778. logger.LegacyPrintf("service.antigravity_gateway", "SSE line too long (antigravity claude non-stream): max_size=%d error=%v", maxLineSize, ev.err)
  779. }
  780. return nil, nil, ev.err
  781. }
  782. line := ev.line
  783. trimmed := strings.TrimRight(line, "\r\n")
  784. if !strings.HasPrefix(trimmed, "data:") {
  785. continue
  786. }
  787. payload := strings.TrimSpace(strings.TrimPrefix(trimmed, "data:"))
  788. if payload == "" || payload == "[DONE]" {
  789. continue
  790. }
  791. // 解包 v1internal 响应
  792. inner, parseErr := s.unwrapV1InternalResponse([]byte(payload))
  793. if parseErr != nil {
  794. continue
  795. }
  796. upstreamResponseModelObserverFromContext(c).ObserveGemini(inner)
  797. var parsed map[string]any
  798. if err := json.Unmarshal(inner, &parsed); err != nil {
  799. continue
  800. }
  801. last = parsed
  802. // 保留最后一个有 parts 的响应,并收集所有 parts
  803. parts := extractGeminiParts(parsed)
  804. if len(parts) > 0 {
  805. lastWithParts = parsed
  806. // 收集所有 parts(text、thinking、functionCall、inlineData 等)
  807. collectedParts = append(collectedParts, parts...)
  808. }
  809. if len(parts) > 0 || strings.TrimSpace(extractGeminiFinishReason(parsed)) != "" {
  810. meaningfulResponse = true
  811. if firstTokenMs == nil {
  812. ms := int(time.Since(startTime).Milliseconds())
  813. firstTokenMs = &ms
  814. }
  815. }
  816. case <-intervalCh:
  817. lastRead := time.Unix(0, atomic.LoadInt64(&lastReadAt))
  818. if time.Since(lastRead) < streamInterval {
  819. continue
  820. }
  821. logger.LegacyPrintf("service.antigravity_gateway", "Stream data interval timeout (antigravity claude non-stream)")
  822. return nil, nil, fmt.Errorf("stream data interval timeout")
  823. }
  824. }
  825. returnResponse:
  826. // 处理空响应情况 — 触发同账号重试 + failover 切换账号
  827. if !meaningfulResponse {
  828. logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Forward] warning: empty stream response (claude non-stream), triggering failover")
  829. return nil, nil, &UpstreamFailoverError{
  830. StatusCode: http.StatusBadGateway,
  831. ResponseBody: []byte(`{"error":"empty stream response from upstream"}`),
  832. RetryableOnSameAccount: true,
  833. }
  834. }
  835. // 选择最后一个有效响应
  836. finalResponse := pickGeminiCollectResult(last, lastWithParts)
  837. // 将收集的所有 parts 合并到最终响应中
  838. if len(collectedParts) > 0 {
  839. finalResponse = mergeCollectedPartsToResponse(finalResponse, collectedParts)
  840. }
  841. // 序列化为 JSON(Gemini 格式)
  842. geminiBody, err := json.Marshal(finalResponse)
  843. if err != nil {
  844. return nil, nil, fmt.Errorf("failed to marshal gemini response: %w", err)
  845. }
  846. // 转换 Gemini 响应为 Claude 格式
  847. claudeResp, agUsage, err := antigravity.TransformGeminiToClaude(geminiBody, originalModel)
  848. if err != nil {
  849. logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Forward] transform_error error=%v body=%s", err, string(geminiBody))
  850. return nil, nil, fmt.Errorf("failed to parse upstream response: %w", err)
  851. }
  852. // 转换为 service.ClaudeUsage
  853. usage := &ClaudeUsage{
  854. InputTokens: agUsage.InputTokens,
  855. OutputTokens: agUsage.OutputTokens,
  856. CacheCreationInputTokens: agUsage.CacheCreationInputTokens,
  857. CacheReadInputTokens: agUsage.CacheReadInputTokens,
  858. ImageOutputTokens: agUsage.ImageOutputTokens,
  859. }
  860. return claudeResp, &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs}, nil
  861. }
  862. // handleClaudeStreamToNonStreaming 收集上游流式响应,转换为 Claude 非流式格式返回
  863. // 用于处理客户端非流式请求但上游只支持流式的情况
  864. func (s *AntigravityGatewayService) handleClaudeStreamToNonStreaming(c *gin.Context, resp *http.Response, startTime time.Time, originalModel string) (*antigravityStreamResult, error) {
  865. claudeResp, streamRes, err := s.collectClaudeStreamResponse(c, resp, startTime, originalModel)
  866. if err != nil {
  867. var failoverErr *UpstreamFailoverError
  868. if errors.As(err, &failoverErr) {
  869. return nil, err
  870. }
  871. if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
  872. return nil, err
  873. }
  874. errMsg := "Failed to parse upstream response"
  875. errType := "upstream_error"
  876. if strings.Contains(err.Error(), "stream data interval timeout") {
  877. errMsg = "Upstream stream data interval timeout"
  878. errType = "upstream_timeout"
  879. } else if errors.Is(err, bufio.ErrTooLong) {
  880. errMsg = "Upstream response line too long"
  881. errType = "response_too_large"
  882. }
  883. return nil, s.writeClaudeError(c, http.StatusBadGateway, errType, errMsg)
  884. }
  885. c.Data(http.StatusOK, "application/json", claudeResp)
  886. return streamRes, nil
  887. }
  888. // handleClaudeStreamingResponse 处理 Claude 流式响应(Gemini SSE → Claude SSE 转换)
  889. func (s *AntigravityGatewayService) handleClaudeStreamingResponse(c *gin.Context, resp *http.Response, startTime time.Time, originalModel string) (*antigravityStreamResult, error) {
  890. c.Header("Content-Type", "text/event-stream")
  891. c.Header("Cache-Control", "no-cache")
  892. c.Header("Connection", "keep-alive")
  893. c.Header("X-Accel-Buffering", "no")
  894. c.Status(http.StatusOK)
  895. flusher, ok := c.Writer.(http.Flusher)
  896. if !ok {
  897. return nil, errors.New("streaming not supported")
  898. }
  899. processor := antigravity.NewStreamingProcessor(originalModel)
  900. var firstTokenMs *int
  901. // 使用 Scanner 并限制单行大小,避免 ReadString 无上限导致 OOM
  902. scanner := bufio.NewScanner(resp.Body)
  903. maxLineSize := defaultMaxLineSize
  904. if s.settingService.cfg != nil && s.settingService.cfg.Gateway.MaxLineSize > 0 {
  905. maxLineSize = s.settingService.cfg.Gateway.MaxLineSize
  906. }
  907. scanBuf := getSSEScannerBuf64K()
  908. scanner.Buffer(scanBuf[:0], maxLineSize)
  909. // 辅助函数:转换 antigravity.ClaudeUsage 到 service.ClaudeUsage
  910. convertUsage := func(agUsage *antigravity.ClaudeUsage) *ClaudeUsage {
  911. if agUsage == nil {
  912. return &ClaudeUsage{}
  913. }
  914. return &ClaudeUsage{
  915. InputTokens: agUsage.InputTokens,
  916. OutputTokens: agUsage.OutputTokens,
  917. CacheCreationInputTokens: agUsage.CacheCreationInputTokens,
  918. CacheReadInputTokens: agUsage.CacheReadInputTokens,
  919. }
  920. }
  921. type scanEvent struct {
  922. line string
  923. err error
  924. }
  925. // 独立 goroutine 读取上游,避免读取阻塞影响超时处理
  926. events := make(chan scanEvent, 16)
  927. done := make(chan struct{})
  928. sendEvent := func(ev scanEvent) bool {
  929. select {
  930. case events <- ev:
  931. return true
  932. case <-done:
  933. return false
  934. }
  935. }
  936. var lastReadAt int64
  937. atomic.StoreInt64(&lastReadAt, time.Now().UnixNano())
  938. go func(scanBuf *sseScannerBuf64K) {
  939. defer putSSEScannerBuf64K(scanBuf)
  940. defer close(events)
  941. for scanner.Scan() {
  942. atomic.StoreInt64(&lastReadAt, time.Now().UnixNano())
  943. if !sendEvent(scanEvent{line: scanner.Text()}) {
  944. return
  945. }
  946. }
  947. if err := scanner.Err(); err != nil {
  948. _ = sendEvent(scanEvent{err: err})
  949. }
  950. }(scanBuf)
  951. defer close(done)
  952. streamInterval := time.Duration(0)
  953. if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamDataIntervalTimeout > 0 {
  954. streamInterval = time.Duration(s.settingService.cfg.Gateway.StreamDataIntervalTimeout) * time.Second
  955. }
  956. var intervalTicker *time.Ticker
  957. if streamInterval > 0 {
  958. intervalTicker = time.NewTicker(streamInterval)
  959. defer intervalTicker.Stop()
  960. }
  961. var intervalCh <-chan time.Time
  962. if intervalTicker != nil {
  963. intervalCh = intervalTicker.C
  964. }
  965. // 下游 keepalive:防止代理/Cloudflare Tunnel 因连接空闲而断开
  966. keepaliveInterval := time.Duration(0)
  967. if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamKeepaliveInterval > 0 {
  968. keepaliveInterval = time.Duration(s.settingService.cfg.Gateway.StreamKeepaliveInterval) * time.Second
  969. }
  970. var keepaliveTicker *time.Ticker
  971. if keepaliveInterval > 0 {
  972. keepaliveTicker = time.NewTicker(keepaliveInterval)
  973. defer keepaliveTicker.Stop()
  974. }
  975. var keepaliveCh <-chan time.Time
  976. if keepaliveTicker != nil {
  977. keepaliveCh = keepaliveTicker.C
  978. }
  979. lastDataAt := time.Now()
  980. cw := newAntigravityClientWriter(c.Writer, flusher, "antigravity claude")
  981. // 仅发送一次错误事件,避免多次写入导致协议混乱
  982. errorEventSent := false
  983. sendErrorEvent := func(reason string) {
  984. if errorEventSent || cw.Disconnected() {
  985. return
  986. }
  987. errorEventSent = true
  988. _, _ = fmt.Fprintf(c.Writer, "event: error\ndata: {\"error\":\"%s\"}\n\n", reason)
  989. flusher.Flush()
  990. }
  991. // finishUsage 是获取 processor 最终 usage 的辅助函数
  992. finishUsage := func() *ClaudeUsage {
  993. _, agUsage := processor.Finish()
  994. return convertUsage(agUsage)
  995. }
  996. for {
  997. select {
  998. case ev, ok := <-events:
  999. if !ok {
  1000. // 上游完成,发送结束事件
  1001. finalEvents, agUsage := processor.Finish()
  1002. if len(finalEvents) > 0 {
  1003. cw.Write(finalEvents)
  1004. } else if !processor.MessageStartSent() && !cw.Disconnected() {
  1005. // 整个流未收到任何可解析的上游数据(全部 SSE 行均无法被 JSON 解析),
  1006. // 触发 failover 在同账号重试,避免向客户端发出缺少 message_start 的残缺流
  1007. logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Claude-Stream] empty stream response (no valid events parsed), triggering failover")
  1008. return nil, &UpstreamFailoverError{
  1009. StatusCode: http.StatusBadGateway,
  1010. ResponseBody: []byte(`{"error":"empty stream response from upstream"}`),
  1011. RetryableOnSameAccount: true,
  1012. }
  1013. }
  1014. return &antigravityStreamResult{usage: convertUsage(agUsage), firstTokenMs: firstTokenMs, clientDisconnect: cw.Disconnected()}, nil
  1015. }
  1016. if ev.err != nil {
  1017. if disconnect, handled := handleStreamReadError(ev.err, cw.Disconnected(), "antigravity claude"); handled {
  1018. return &antigravityStreamResult{usage: finishUsage(), firstTokenMs: firstTokenMs, clientDisconnect: disconnect}, nil
  1019. }
  1020. if errors.Is(ev.err, bufio.ErrTooLong) {
  1021. logger.LegacyPrintf("service.antigravity_gateway", "SSE line too long (antigravity): max_size=%d error=%v", maxLineSize, ev.err)
  1022. sendErrorEvent("response_too_large")
  1023. return &antigravityStreamResult{usage: convertUsage(nil), firstTokenMs: firstTokenMs}, ev.err
  1024. }
  1025. sendErrorEvent("stream_read_error")
  1026. return nil, fmt.Errorf("stream read error: %w", ev.err)
  1027. }
  1028. lastDataAt = time.Now()
  1029. s.observeAntigravityGeminiSSELine(c, ev.line)
  1030. // 处理 SSE 行,转换为 Claude 格式
  1031. claudeEvents := processor.ProcessLine(strings.TrimRight(ev.line, "\r\n"))
  1032. if len(claudeEvents) > 0 {
  1033. if firstTokenMs == nil {
  1034. ms := int(time.Since(startTime).Milliseconds())
  1035. firstTokenMs = &ms
  1036. }
  1037. cw.Write(claudeEvents)
  1038. }
  1039. case <-intervalCh:
  1040. lastRead := time.Unix(0, atomic.LoadInt64(&lastReadAt))
  1041. if time.Since(lastRead) < streamInterval {
  1042. continue
  1043. }
  1044. if cw.Disconnected() {
  1045. logger.LegacyPrintf("service.antigravity_gateway", "Upstream timeout after client disconnect (antigravity claude), returning collected usage")
  1046. return &antigravityStreamResult{usage: finishUsage(), firstTokenMs: firstTokenMs, clientDisconnect: true}, nil
  1047. }
  1048. logger.LegacyPrintf("service.antigravity_gateway", "Stream data interval timeout (antigravity)")
  1049. sendErrorEvent("stream_timeout")
  1050. return &antigravityStreamResult{usage: convertUsage(nil), firstTokenMs: firstTokenMs}, fmt.Errorf("stream data interval timeout")
  1051. case <-keepaliveCh:
  1052. if cw.Disconnected() {
  1053. continue
  1054. }
  1055. if time.Since(lastDataAt) < keepaliveInterval {
  1056. continue
  1057. }
  1058. // SSE ping 事件:Anthropic 原生格式,客户端会正确处理,
  1059. // 同时保持连接活跃防止 Cloudflare Tunnel 等代理断开
  1060. if !cw.Fprintf("event: ping\ndata: {\"type\": \"ping\"}\n\n") {
  1061. logger.LegacyPrintf("service.antigravity_gateway", "Client disconnected during keepalive ping (antigravity claude), continuing to drain upstream for billing")
  1062. continue
  1063. }
  1064. }
  1065. }
  1066. }
  1067. func (s *AntigravityGatewayService) extractImageInputSize(body []byte) string {
  1068. var req antigravity.GeminiRequest
  1069. if err := json.Unmarshal(body, &req); err != nil {
  1070. return ""
  1071. }
  1072. if req.GenerationConfig != nil && req.GenerationConfig.ImageConfig != nil {
  1073. return strings.TrimSpace(req.GenerationConfig.ImageConfig.ImageSize)
  1074. }
  1075. return ""
  1076. }
  1077. // isImageGenerationModel 判断模型是否为图片生成模型
  1078. // 支持的模型:gemini-3.1-flash-image, gemini-3-pro-image, gemini-2.5-flash-image 等
  1079. func isImageGenerationModel(model string) bool {
  1080. modelLower := strings.ToLower(model)
  1081. // 移除 models/ 前缀
  1082. modelLower = strings.TrimPrefix(modelLower, "models/")
  1083. // 精确匹配或前缀匹配
  1084. return modelLower == "gemini-3.1-flash-image" ||
  1085. modelLower == "gemini-3.1-flash-image-preview" ||
  1086. strings.HasPrefix(modelLower, "gemini-3.1-flash-image-") ||
  1087. modelLower == "gemini-3-pro-image" ||
  1088. modelLower == "gemini-3-pro-image-preview" ||
  1089. strings.HasPrefix(modelLower, "gemini-3-pro-image-") ||
  1090. modelLower == "gemini-2.5-flash-image" ||
  1091. modelLower == "gemini-2.5-flash-image-preview" ||
  1092. strings.HasPrefix(modelLower, "gemini-2.5-flash-image-")
  1093. }