Merge remote-tracking branch 'upstream/main'

2026-01-08 21:35:34 +08:00
parent fd1b14fd1d b015a3bd8a
commit b46b3c5c3c
8 changed files with 645 additions and 156 deletions
--- a/backend/internal/service/antigravity_gateway_service.go
+++ b/backend/internal/service/antigravity_gateway_service.go
@@ -181,12 +181,15 @@ func (s *AntigravityGatewayService) TestConnection(ctx context.Context, account
 		return nil, fmt.Errorf("构建请求失败: %w", err)
 	}

-	// 构建 HTTP 请求（非流式）
-	req, err := antigravity.NewAPIRequest(ctx, "generateContent", accessToken, requestBody)
+	// 构建 HTTP 请求（总是使用流式 endpoint，与官方客户端一致）
+	req, err := antigravity.NewAPIRequest(ctx, "streamGenerateContent", accessToken, requestBody)
 	if err != nil {
 		return nil, err
 	}

+	// 调试日志：Test 请求信息
+	log.Printf("[antigravity-Test] account=%s request_size=%d url=%s", account.Name, len(requestBody), req.URL.String())
+
 	// 代理 URL
 	proxyURL := ""
 	if account.ProxyID != nil && account.Proxy != nil {
@@ -210,14 +213,8 @@ func (s *AntigravityGatewayService) TestConnection(ctx context.Context, account
 		return nil, fmt.Errorf("API 返回 %d: %s", resp.StatusCode, string(respBody))
 	}

-	// 解包 v1internal 响应
-	unwrapped, err := s.unwrapV1InternalResponse(respBody)
-	if err != nil {
-		return nil, fmt.Errorf("解包响应失败: %w", err)
-	}
-
-	// 提取响应文本
-	text := extractGeminiResponseText(unwrapped)
+	// 解析流式响应，提取文本
+	text := extractTextFromSSEResponse(respBody)

 	return &TestConnectionResult{
 		Text:        text,
@@ -236,6 +233,12 @@ func (s *AntigravityGatewayService) buildGeminiTestRequest(projectID, model stri
 				},
 			},
 		},
+		// Antigravity 上游要求必须包含身份提示词
+		"systemInstruction": map[string]any{
+			"parts": []map[string]any{
+				{"text": antigravity.GetDefaultIdentityPatch()},
+			},
+		},
 	}
 	payloadBytes, _ := json.Marshal(payload)
 	return s.wrapV1InternalRequest(projectID, model, payloadBytes)
@@ -267,38 +270,66 @@ func (s *AntigravityGatewayService) getClaudeTransformOptions(ctx context.Contex
 	return opts
 }

-// extractGeminiResponseText 从 Gemini 响应中提取文本
-func extractGeminiResponseText(respBody []byte) string {
-	var resp map[string]any
-	if err := json.Unmarshal(respBody, &resp); err != nil {
-		return ""
-	}
-
-	candidates, ok := resp["candidates"].([]any)
-	if !ok || len(candidates) == 0 {
-		return ""
-	}
-
-	candidate, ok := candidates[0].(map[string]any)
-	if !ok {
-		return ""
-	}
-
-	content, ok := candidate["content"].(map[string]any)
-	if !ok {
-		return ""
-	}
-
-	parts, ok := content["parts"].([]any)
-	if !ok {
-		return ""
-	}
-
+// extractTextFromSSEResponse 从 SSE 流式响应中提取文本
+func extractTextFromSSEResponse(respBody []byte) string {
 	var texts []string
-	for _, part := range parts {
-		if partMap, ok := part.(map[string]any); ok {
-			if text, ok := partMap["text"].(string); ok && text != "" {
-				texts = append(texts, text)
+	lines := bytes.Split(respBody, []byte("\n"))
+
+	for _, line := range lines {
+		line = bytes.TrimSpace(line)
+		if len(line) == 0 {
+			continue
+		}
+
+		// 跳过 SSE 前缀
+		if bytes.HasPrefix(line, []byte("data:")) {
+			line = bytes.TrimPrefix(line, []byte("data:"))
+			line = bytes.TrimSpace(line)
+		}
+
+		// 跳过非 JSON 行
+		if len(line) == 0 || line[0] != '{' {
+			continue
+		}
+
+		// 解析 JSON
+		var data map[string]any
+		if err := json.Unmarshal(line, &data); err != nil {
+			continue
+		}
+
+		// 尝试从 response.candidates[0].content.parts[].text 提取
+		response, ok := data["response"].(map[string]any)
+		if !ok {
+			// 尝试直接从 candidates 提取（某些响应格式）
+			response = data
+		}
+
+		candidates, ok := response["candidates"].([]any)
+		if !ok || len(candidates) == 0 {
+			continue
+		}
+
+		candidate, ok := candidates[0].(map[string]any)
+		if !ok {
+			continue
+		}
+
+		content, ok := candidate["content"].(map[string]any)
+		if !ok {
+			continue
+		}
+
+		parts, ok := content["parts"].([]any)
+		if !ok {
+			continue
+		}
+
+		for _, part := range parts {
+			if partMap, ok := part.(map[string]any); ok {
+				if text, ok := partMap["text"].(string); ok && text != "" {
+					texts = append(texts, text)
+				}
 			}
 		}
 	}
@@ -306,6 +337,53 @@ func extractGeminiResponseText(respBody []byte) string {
 	return strings.Join(texts, "")
 }

+// injectIdentityPatchToGeminiRequest 为 Gemini 格式请求注入身份提示词
+// 如果请求中已包含 "You are Antigravity" 则不重复注入
+func injectIdentityPatchToGeminiRequest(body []byte) ([]byte, error) {
+	var request map[string]any
+	if err := json.Unmarshal(body, &request); err != nil {
+		return nil, fmt.Errorf("解析 Gemini 请求失败: %w", err)
+	}
+
+	// 检查现有 systemInstruction 是否已包含身份提示词
+	if sysInst, ok := request["systemInstruction"].(map[string]any); ok {
+		if parts, ok := sysInst["parts"].([]any); ok {
+			for _, part := range parts {
+				if partMap, ok := part.(map[string]any); ok {
+					if text, ok := partMap["text"].(string); ok {
+						if strings.Contains(text, "You are Antigravity") {
+							// 已包含身份提示词，直接返回原始请求
+							return body, nil
+						}
+					}
+				}
+			}
+		}
+	}
+
+	// 获取默认身份提示词
+	identityPatch := antigravity.GetDefaultIdentityPatch()
+
+	// 构建新的 systemInstruction
+	newPart := map[string]any{"text": identityPatch}
+
+	if existing, ok := request["systemInstruction"].(map[string]any); ok {
+		// 已有 systemInstruction，在开头插入身份提示词
+		if parts, ok := existing["parts"].([]any); ok {
+			existing["parts"] = append([]any{newPart}, parts...)
+		} else {
+			existing["parts"] = []any{newPart}
+		}
+	} else {
+		// 没有 systemInstruction，创建新的
+		request["systemInstruction"] = map[string]any{
+			"parts": []any{newPart},
+		}
+	}
+
+	return json.Marshal(request)
+}
+
 // wrapV1InternalRequest 包装请求为 v1internal 格式
 func (s *AntigravityGatewayService) wrapV1InternalRequest(projectID, model string, originalBody []byte) ([]byte, error) {
 	var request any
@@ -316,7 +394,7 @@ func (s *AntigravityGatewayService) wrapV1InternalRequest(projectID, model strin
 	wrapped := map[string]any{
 		"project":     projectID,
 		"requestId":   "agent-" + uuid.New().String(),
-		"userAgent":   "sub2api",
+		"userAgent":   "antigravity", // 固定值，与官方客户端一致
 		"requestType": "agent",
 		"model":       model,
 		"request":     request,
@@ -391,17 +469,20 @@ func (s *AntigravityGatewayService) Forward(ctx context.Context, c *gin.Context,
 		proxyURL = account.Proxy.URL()
 	}

+	// 获取转换选项
+	// Antigravity 上游要求必须包含身份提示词，否则会返回 429
+	transformOpts := s.getClaudeTransformOptions(ctx)
+	transformOpts.EnableIdentityPatch = true // 强制启用，Antigravity 上游必需
+
 	// 转换 Claude 请求为 Gemini 格式
-	geminiBody, err := antigravity.TransformClaudeToGeminiWithOptions(&claudeReq, projectID, mappedModel, s.getClaudeTransformOptions(ctx))
+	geminiBody, err := antigravity.TransformClaudeToGeminiWithOptions(&claudeReq, projectID, mappedModel, transformOpts)
 	if err != nil {
 		return nil, fmt.Errorf("transform request: %w", err)
 	}

-	// 构建上游 action
-	action := "generateContent"
-	if claudeReq.Stream {
-		action = "streamGenerateContent?alt=sse"
-	}
+	// Antigravity 上游只支持流式请求，统一使用 streamGenerateContent
+	// 如果客户端请求非流式，在响应处理阶段会收集完整流式响应后转换返回
+	action := "streamGenerateContent"

 	// 重试循环
 	var resp *http.Response
@@ -438,7 +519,7 @@ func (s *AntigravityGatewayService) Forward(ctx context.Context, c *gin.Context,
 			_ = resp.Body.Close()

 			if attempt < antigravityMaxRetries {
-				log.Printf("%s status=%d retry=%d/%d", prefix, resp.StatusCode, attempt, antigravityMaxRetries)
+				log.Printf("%s status=%d retry=%d/%d body=%s", prefix, resp.StatusCode, attempt, antigravityMaxRetries, truncateForLog(respBody, 500))
 				if !sleepAntigravityBackoffWithContext(ctx, attempt) {
 					log.Printf("%s status=context_canceled_during_backoff", prefix)
 					return nil, ctx.Err()
@@ -557,6 +638,7 @@ func (s *AntigravityGatewayService) Forward(ctx context.Context, c *gin.Context,
 	var usage *ClaudeUsage
 	var firstTokenMs *int
 	if claudeReq.Stream {
+		// 客户端要求流式，直接透传转换
 		streamRes, err := s.handleClaudeStreamingResponse(c, resp, startTime, originalModel)
 		if err != nil {
 			log.Printf("%s status=stream_error error=%v", prefix, err)
@@ -565,10 +647,14 @@ func (s *AntigravityGatewayService) Forward(ctx context.Context, c *gin.Context,
 		usage = streamRes.usage
 		firstTokenMs = streamRes.firstTokenMs
 	} else {
-		usage, err = s.handleClaudeNonStreamingResponse(c, resp, originalModel)
+		// 客户端要求非流式，收集流式响应后转换返回
+		streamRes, err := s.handleClaudeStreamToNonStreaming(c, resp, startTime, originalModel)
 		if err != nil {
+			log.Printf("%s status=stream_collect_error error=%v", prefix, err)
 			return nil, err
 		}
+		usage = streamRes.usage
+		firstTokenMs = streamRes.firstTokenMs
 	}

 	return &ForwardResult{
@@ -901,21 +987,22 @@ func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Co
 		proxyURL = account.Proxy.URL()
 	}

-	// 包装请求
-	wrappedBody, err := s.wrapV1InternalRequest(projectID, mappedModel, body)
+	// Antigravity 上游要求必须包含身份提示词，注入到请求中
+	injectedBody, err := injectIdentityPatchToGeminiRequest(body)
 	if err != nil {
 		return nil, err
 	}

-	// 构建上游 action
-	upstreamAction := action
-	if action == "generateContent" && stream {
-		upstreamAction = "streamGenerateContent"
-	}
-	if stream || upstreamAction == "streamGenerateContent" {
-		upstreamAction += "?alt=sse"
+	// 包装请求
+	wrappedBody, err := s.wrapV1InternalRequest(projectID, mappedModel, injectedBody)
+	if err != nil {
+		return nil, err
 	}

+	// Antigravity 上游只支持流式请求，统一使用 streamGenerateContent
+	// 如果客户端请求非流式，在响应处理阶段会收集完整流式响应后返回
+	upstreamAction := "streamGenerateContent"
+
 	// 重试循环
 	var resp *http.Response
 	for attempt := 1; attempt <= antigravityMaxRetries; attempt++ {
@@ -992,7 +1079,7 @@ func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Co
 			if fallbackModel != "" && fallbackModel != mappedModel {
 				log.Printf("[Antigravity] Model not found (%s), retrying with fallback model %s (account: %s)", mappedModel, fallbackModel, account.Name)

-				fallbackWrapped, err := s.wrapV1InternalRequest(projectID, fallbackModel, body)
+				fallbackWrapped, err := s.wrapV1InternalRequest(projectID, fallbackModel, injectedBody)
 				if err == nil {
 					fallbackReq, err := antigravity.NewAPIRequest(ctx, upstreamAction, accessToken, fallbackWrapped)
 					if err == nil {
@@ -1042,7 +1129,8 @@ handleSuccess:
 	var usage *ClaudeUsage
 	var firstTokenMs *int

-	if stream || upstreamAction == "streamGenerateContent" {
+	if stream {
+		// 客户端要求流式，直接透传
 		streamRes, err := s.handleGeminiStreamingResponse(c, resp, startTime)
 		if err != nil {
 			log.Printf("%s status=stream_error error=%v", prefix, err)
@@ -1051,11 +1139,14 @@ handleSuccess:
 		usage = streamRes.usage
 		firstTokenMs = streamRes.firstTokenMs
 	} else {
-		usageResp, err := s.handleGeminiNonStreamingResponse(c, resp)
+		// 客户端要求非流式，收集流式响应后返回
+		streamRes, err := s.handleGeminiStreamToNonStreaming(c, resp, startTime)
 		if err != nil {
+			log.Printf("%s status=stream_collect_error error=%v", prefix, err)
 			return nil, err
 		}
-		usage = usageResp
+		usage = streamRes.usage
+		firstTokenMs = streamRes.firstTokenMs
 	}

 	if usage == nil {
@@ -1102,9 +1193,9 @@ func (s *AntigravityGatewayService) shouldFailoverUpstreamError(statusCode int)
 // sleepAntigravityBackoffWithContext 带 context 取消检查的退避等待
 // 返回 true 表示正常完成等待，false 表示 context 已取消
 func sleepAntigravityBackoffWithContext(ctx context.Context, attempt int) bool {
-	delay := geminiRetryBaseDelay * time.Duration(1<<uint(attempt-1))
-	if delay > geminiRetryMaxDelay {
-		delay = geminiRetryMaxDelay
+	delay := antigravityRetryBaseDelay * time.Duration(1<<uint(attempt-1))
+	if delay > antigravityRetryMaxDelay {
+		delay = antigravityRetryMaxDelay
 	}

 	// +/- 20% jitter
@@ -1316,25 +1407,150 @@ func (s *AntigravityGatewayService) handleGeminiStreamingResponse(c *gin.Context
 	}
 }

-func (s *AntigravityGatewayService) handleGeminiNonStreamingResponse(c *gin.Context, resp *http.Response) (*ClaudeUsage, error) {
-	respBody, err := io.ReadAll(resp.Body)
-	if err != nil {
-		return nil, err
+// handleGeminiStreamToNonStreaming 读取上游流式响应，合并为非流式响应返回给客户端
+// Gemini 流式响应中每个 chunk 都包含累积的完整文本，只需保留最后一个有效响应
+func (s *AntigravityGatewayService) handleGeminiStreamToNonStreaming(c *gin.Context, resp *http.Response, startTime time.Time) (*antigravityStreamResult, error) {
+	scanner := bufio.NewScanner(resp.Body)
+	maxLineSize := defaultMaxLineSize
+	if s.settingService.cfg != nil && s.settingService.cfg.Gateway.MaxLineSize > 0 {
+		maxLineSize = s.settingService.cfg.Gateway.MaxLineSize
+	}
+	scanner.Buffer(make([]byte, 64*1024), maxLineSize)
+
+	usage := &ClaudeUsage{}
+	var firstTokenMs *int
+	var last map[string]any
+	var lastWithParts map[string]any
+
+	type scanEvent struct {
+		line string
+		err  error
 	}

-	// 解包 v1internal 响应
-	unwrapped, _ := s.unwrapV1InternalResponse(respBody)
-
-	var parsed map[string]any
-	if json.Unmarshal(unwrapped, &parsed) == nil {
-		if u := extractGeminiUsage(parsed); u != nil {
-			c.Data(resp.StatusCode, "application/json", unwrapped)
-			return u, nil
+	// 独立 goroutine 读取上游，避免读取阻塞影响超时处理
+	events := make(chan scanEvent, 16)
+	done := make(chan struct{})
+	sendEvent := func(ev scanEvent) bool {
+		select {
+		case events <- ev:
+			return true
+		case <-done:
+			return false
 		}
 	}

-	c.Data(resp.StatusCode, "application/json", unwrapped)
-	return &ClaudeUsage{}, nil
+	var lastReadAt int64
+	atomic.StoreInt64(&lastReadAt, time.Now().UnixNano())
+	go func() {
+		defer close(events)
+		for scanner.Scan() {
+			atomic.StoreInt64(&lastReadAt, time.Now().UnixNano())
+			if !sendEvent(scanEvent{line: scanner.Text()}) {
+				return
+			}
+		}
+		if err := scanner.Err(); err != nil {
+			_ = sendEvent(scanEvent{err: err})
+		}
+	}()
+	defer close(done)
+
+	// 上游数据间隔超时保护（防止上游挂起长期占用连接）
+	streamInterval := time.Duration(0)
+	if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamDataIntervalTimeout > 0 {
+		streamInterval = time.Duration(s.settingService.cfg.Gateway.StreamDataIntervalTimeout) * time.Second
+	}
+	var intervalTicker *time.Ticker
+	if streamInterval > 0 {
+		intervalTicker = time.NewTicker(streamInterval)
+		defer intervalTicker.Stop()
+	}
+	var intervalCh <-chan time.Time
+	if intervalTicker != nil {
+		intervalCh = intervalTicker.C
+	}
+
+	for {
+		select {
+		case ev, ok := <-events:
+			if !ok {
+				// 流结束，返回收集的响应
+				goto returnResponse
+			}
+			if ev.err != nil {
+				if errors.Is(ev.err, bufio.ErrTooLong) {
+					log.Printf("SSE line too long (antigravity non-stream): max_size=%d error=%v", maxLineSize, ev.err)
+				}
+				return nil, ev.err
+			}
+
+			line := ev.line
+			trimmed := strings.TrimRight(line, "\r\n")
+
+			if !strings.HasPrefix(trimmed, "data:") {
+				continue
+			}
+
+			payload := strings.TrimSpace(strings.TrimPrefix(trimmed, "data:"))
+			if payload == "" || payload == "[DONE]" {
+				continue
+			}
+
+			// 解包 v1internal 响应
+			inner, parseErr := s.unwrapV1InternalResponse([]byte(payload))
+			if parseErr != nil {
+				continue
+			}
+
+			var parsed map[string]any
+			if err := json.Unmarshal(inner, &parsed); err != nil {
+				continue
+			}
+
+			// 记录首 token 时间
+			if firstTokenMs == nil {
+				ms := int(time.Since(startTime).Milliseconds())
+				firstTokenMs = &ms
+			}
+
+			last = parsed
+
+			// 提取 usage
+			if u := extractGeminiUsage(parsed); u != nil {
+				usage = u
+			}
+
+			// 保留最后一个有 parts 的响应
+			if parts := extractGeminiParts(parsed); len(parts) > 0 {
+				lastWithParts = parsed
+			}
+
+		case <-intervalCh:
+			lastRead := time.Unix(0, atomic.LoadInt64(&lastReadAt))
+			if time.Since(lastRead) < streamInterval {
+				continue
+			}
+			log.Printf("Stream data interval timeout (antigravity non-stream)")
+			return nil, fmt.Errorf("stream data interval timeout")
+		}
+	}
+
+returnResponse:
+	// 选择最后一个有效响应
+	finalResponse := pickGeminiCollectResult(last, lastWithParts)
+
+	// 处理空响应情况
+	if last == nil && lastWithParts == nil {
+		log.Printf("[antigravity-Forward] warning: empty stream response, no valid chunks received")
+	}
+
+	respBody, err := json.Marshal(finalResponse)
+	if err != nil {
+		return nil, fmt.Errorf("failed to marshal response: %w", err)
+	}
+	c.Data(http.StatusOK, "application/json", respBody)
+
+	return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs}, nil
 }

 func (s *AntigravityGatewayService) writeClaudeError(c *gin.Context, status int, errType, message string) error {
@@ -1411,17 +1627,148 @@ func (s *AntigravityGatewayService) writeGoogleError(c *gin.Context, status int,
 	return fmt.Errorf("%s", message)
 }

-// handleClaudeNonStreamingResponse 处理 Claude 非流式响应（Gemini → Claude 转换）
-func (s *AntigravityGatewayService) handleClaudeNonStreamingResponse(c *gin.Context, resp *http.Response, originalModel string) (*ClaudeUsage, error) {
-	body, err := io.ReadAll(io.LimitReader(resp.Body, 8<<20))
+// handleClaudeStreamToNonStreaming 收集上游流式响应，转换为 Claude 非流式格式返回
+// 用于处理客户端非流式请求但上游只支持流式的情况
+func (s *AntigravityGatewayService) handleClaudeStreamToNonStreaming(c *gin.Context, resp *http.Response, startTime time.Time, originalModel string) (*antigravityStreamResult, error) {
+	scanner := bufio.NewScanner(resp.Body)
+	maxLineSize := defaultMaxLineSize
+	if s.settingService.cfg != nil && s.settingService.cfg.Gateway.MaxLineSize > 0 {
+		maxLineSize = s.settingService.cfg.Gateway.MaxLineSize
+	}
+	scanner.Buffer(make([]byte, 64*1024), maxLineSize)
+
+	var firstTokenMs *int
+	var last map[string]any
+	var lastWithParts map[string]any
+
+	type scanEvent struct {
+		line string
+		err  error
+	}
+
+	// 独立 goroutine 读取上游，避免读取阻塞影响超时处理
+	events := make(chan scanEvent, 16)
+	done := make(chan struct{})
+	sendEvent := func(ev scanEvent) bool {
+		select {
+		case events <- ev:
+			return true
+		case <-done:
+			return false
+		}
+	}
+
+	var lastReadAt int64
+	atomic.StoreInt64(&lastReadAt, time.Now().UnixNano())
+	go func() {
+		defer close(events)
+		for scanner.Scan() {
+			atomic.StoreInt64(&lastReadAt, time.Now().UnixNano())
+			if !sendEvent(scanEvent{line: scanner.Text()}) {
+				return
+			}
+		}
+		if err := scanner.Err(); err != nil {
+			_ = sendEvent(scanEvent{err: err})
+		}
+	}()
+	defer close(done)
+
+	// 上游数据间隔超时保护（防止上游挂起长期占用连接）
+	streamInterval := time.Duration(0)
+	if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamDataIntervalTimeout > 0 {
+		streamInterval = time.Duration(s.settingService.cfg.Gateway.StreamDataIntervalTimeout) * time.Second
+	}
+	var intervalTicker *time.Ticker
+	if streamInterval > 0 {
+		intervalTicker = time.NewTicker(streamInterval)
+		defer intervalTicker.Stop()
+	}
+	var intervalCh <-chan time.Time
+	if intervalTicker != nil {
+		intervalCh = intervalTicker.C
+	}
+
+	for {
+		select {
+		case ev, ok := <-events:
+			if !ok {
+				// 流结束，转换并返回响应
+				goto returnResponse
+			}
+			if ev.err != nil {
+				if errors.Is(ev.err, bufio.ErrTooLong) {
+					log.Printf("SSE line too long (antigravity claude non-stream): max_size=%d error=%v", maxLineSize, ev.err)
+				}
+				return nil, ev.err
+			}
+
+			line := ev.line
+			trimmed := strings.TrimRight(line, "\r\n")
+
+			if !strings.HasPrefix(trimmed, "data:") {
+				continue
+			}
+
+			payload := strings.TrimSpace(strings.TrimPrefix(trimmed, "data:"))
+			if payload == "" || payload == "[DONE]" {
+				continue
+			}
+
+			// 解包 v1internal 响应
+			inner, parseErr := s.unwrapV1InternalResponse([]byte(payload))
+			if parseErr != nil {
+				continue
+			}
+
+			var parsed map[string]any
+			if err := json.Unmarshal(inner, &parsed); err != nil {
+				continue
+			}
+
+			// 记录首 token 时间
+			if firstTokenMs == nil {
+				ms := int(time.Since(startTime).Milliseconds())
+				firstTokenMs = &ms
+			}
+
+			last = parsed
+
+			// 保留最后一个有 parts 的响应
+			if parts := extractGeminiParts(parsed); len(parts) > 0 {
+				lastWithParts = parsed
+			}
+
+		case <-intervalCh:
+			lastRead := time.Unix(0, atomic.LoadInt64(&lastReadAt))
+			if time.Since(lastRead) < streamInterval {
+				continue
+			}
+			log.Printf("Stream data interval timeout (antigravity claude non-stream)")
+			return nil, fmt.Errorf("stream data interval timeout")
+		}
+	}
+
+returnResponse:
+	// 选择最后一个有效响应
+	finalResponse := pickGeminiCollectResult(last, lastWithParts)
+
+	// 处理空响应情况
+	if last == nil && lastWithParts == nil {
+		log.Printf("[antigravity-Forward] warning: empty stream response, no valid chunks received")
+		return nil, s.writeClaudeError(c, http.StatusBadGateway, "upstream_error", "Empty response from upstream")
+	}
+
+	// 序列化为 JSON（Gemini 格式）
+	geminiBody, err := json.Marshal(finalResponse)
 	if err != nil {
-		return nil, s.writeClaudeError(c, http.StatusBadGateway, "upstream_error", "Failed to read upstream response")
+		return nil, fmt.Errorf("failed to marshal gemini response: %w", err)
 	}

 	// 转换 Gemini 响应为 Claude 格式
-	claudeResp, agUsage, err := antigravity.TransformGeminiToClaude(body, originalModel)
+	claudeResp, agUsage, err := antigravity.TransformGeminiToClaude(geminiBody, originalModel)
 	if err != nil {
-		log.Printf("[antigravity-Forward] transform_error error=%v body=%s", err, string(body))
+		log.Printf("[antigravity-Forward] transform_error error=%v body=%s", err, string(geminiBody))
 		return nil, s.writeClaudeError(c, http.StatusBadGateway, "upstream_error", "Failed to parse upstream response")
 	}

@@ -1434,7 +1781,8 @@ func (s *AntigravityGatewayService) handleClaudeNonStreamingResponse(c *gin.Cont
 		CacheCreationInputTokens: agUsage.CacheCreationInputTokens,
 		CacheReadInputTokens:     agUsage.CacheReadInputTokens,
 	}
-	return usage, nil
+
+	return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs}, nil
 }

 // handleClaudeStreamingResponse 处理 Claude 流式响应（Gemini SSE → Claude SSE 转换）
--- a/backend/internal/service/gateway_service.go
+++ b/backend/internal/service/gateway_service.go
@@ -109,12 +109,13 @@ type ClaudeUsage struct {

 // ForwardResult 转发结果
 type ForwardResult struct {
-	RequestID    string
-	Usage        ClaudeUsage
-	Model        string
-	Stream       bool
-	Duration     time.Duration
-	FirstTokenMs *int // 首字时间（流式请求）
+	RequestID        string
+	Usage            ClaudeUsage
+	Model            string
+	Stream           bool
+	Duration         time.Duration
+	FirstTokenMs     *int // 首字时间（流式请求）
+	ClientDisconnect bool // 客户端是否在流式传输过程中断开

 	// 图片生成计费字段（仅 gemini-3-pro-image 使用）
 	ImageCount int    // 生成的图片数量
@@ -1465,6 +1466,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A
 	// 处理正常响应
 	var usage *ClaudeUsage
 	var firstTokenMs *int
+	var clientDisconnect bool
 	if reqStream {
 		streamResult, err := s.handleStreamingResponse(ctx, resp, c, account, startTime, originalModel, reqModel)
 		if err != nil {
@@ -1477,6 +1479,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A
 		}
 		usage = streamResult.usage
 		firstTokenMs = streamResult.firstTokenMs
+		clientDisconnect = streamResult.clientDisconnect
 	} else {
 		usage, err = s.handleNonStreamingResponse(ctx, resp, c, account, originalModel, reqModel)
 		if err != nil {
@@ -1485,12 +1488,13 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A
 	}

 	return &ForwardResult{
-		RequestID:    resp.Header.Get("x-request-id"),
-		Usage:        *usage,
-		Model:        originalModel, // 使用原始模型用于计费和日志
-		Stream:       reqStream,
-		Duration:     time.Since(startTime),
-		FirstTokenMs: firstTokenMs,
+		RequestID:        resp.Header.Get("x-request-id"),
+		Usage:            *usage,
+		Model:            originalModel, // 使用原始模型用于计费和日志
+		Stream:           reqStream,
+		Duration:         time.Since(startTime),
+		FirstTokenMs:     firstTokenMs,
+		ClientDisconnect: clientDisconnect,
 	}, nil
 }

@@ -1845,8 +1849,9 @@ func (s *GatewayService) handleRetryExhaustedError(ctx context.Context, resp *ht

 // streamingResult 流式响应结果
 type streamingResult struct {
-	usage        *ClaudeUsage
-	firstTokenMs *int
+	usage            *ClaudeUsage
+	firstTokenMs     *int
+	clientDisconnect bool // 客户端是否在流式传输过程中断开
 }

 func (s *GatewayService) handleStreamingResponse(ctx context.Context, resp *http.Response, c *gin.Context, account *Account, startTime time.Time, originalModel, mappedModel string) (*streamingResult, error) {
@@ -1942,14 +1947,27 @@ func (s *GatewayService) handleStreamingResponse(ctx context.Context, resp *http
 	}

 	needModelReplace := originalModel != mappedModel
+	clientDisconnected := false // 客户端断开标志，断开后继续读取上游以获取完整usage

 	for {
 		select {
 		case ev, ok := <-events:
 			if !ok {
-				return &streamingResult{usage: usage, firstTokenMs: firstTokenMs}, nil
+				// 上游完成，返回结果
+				return &streamingResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: clientDisconnected}, nil
 			}
 			if ev.err != nil {
+				// 检测 context 取消（客户端断开会导致 context 取消，进而影响上游读取）
+				if errors.Is(ev.err, context.Canceled) || errors.Is(ev.err, context.DeadlineExceeded) {
+					log.Printf("Context canceled during streaming, returning collected usage")
+					return &streamingResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: true}, nil
+				}
+				// 客户端已通过写入失败检测到断开，上游也出错了，返回已收集的 usage
+				if clientDisconnected {
+					log.Printf("Upstream read error after client disconnect: %v, returning collected usage", ev.err)
+					return &streamingResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: true}, nil
+				}
+				// 客户端未断开，正常的错误处理
 				if errors.Is(ev.err, bufio.ErrTooLong) {
 					log.Printf("SSE line too long: account=%d max_size=%d error=%v", account.ID, maxLineSize, ev.err)
 					sendErrorEvent("response_too_large")
@@ -1960,38 +1978,40 @@ func (s *GatewayService) handleStreamingResponse(ctx context.Context, resp *http
 			}
 			line := ev.line
 			if line == "event: error" {
+				// 上游返回错误事件，如果客户端已断开仍返回已收集的 usage
+				if clientDisconnected {
+					return &streamingResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: true}, nil
+				}
 				return nil, errors.New("have error in stream")
 			}

 			// Extract data from SSE line (supports both "data: " and "data:" formats)
+			var data string
 			if sseDataRe.MatchString(line) {
-				data := sseDataRe.ReplaceAllString(line, "")
-
+				data = sseDataRe.ReplaceAllString(line, "")
 				// 如果有模型映射，替换响应中的model字段
 				if needModelReplace {
 					line = s.replaceModelInSSELine(line, mappedModel, originalModel)
 				}
+			}

-				// 转发行
+			// 写入客户端（统一处理 data 行和非 data 行）
+			if !clientDisconnected {
 				if _, err := fmt.Fprintf(w, "%s\n", line); err != nil {
-					sendErrorEvent("write_failed")
-					return &streamingResult{usage: usage, firstTokenMs: firstTokenMs}, err
+					clientDisconnected = true
+					log.Printf("Client disconnected during streaming, continuing to drain upstream for billing")
+				} else {
+					flusher.Flush()
 				}
-				flusher.Flush()
+			}

-				// 记录首字时间：第一个有效的 content_block_delta 或 message_start
-				if firstTokenMs == nil && data != "" && data != "[DONE]" {
+			// 无论客户端是否断开，都解析 usage（仅对 data 行）
+			if data != "" {
+				if firstTokenMs == nil && data != "[DONE]" {
 					ms := int(time.Since(startTime).Milliseconds())
 					firstTokenMs = &ms
 				}
 				s.parseSSEUsage(data, usage)
-			} else {
-				// 非 data 行直接转发
-				if _, err := fmt.Fprintf(w, "%s\n", line); err != nil {
-					sendErrorEvent("write_failed")
-					return &streamingResult{usage: usage, firstTokenMs: firstTokenMs}, err
-				}
-				flusher.Flush()
 			}

 		case <-intervalCh:
@@ -1999,6 +2019,11 @@ func (s *GatewayService) handleStreamingResponse(ctx context.Context, resp *http
 			if time.Since(lastRead) < streamInterval {
 				continue
 			}
+			if clientDisconnected {
+				// 客户端已断开，上游也超时了，返回已收集的 usage
+				log.Printf("Upstream timeout after client disconnect, returning collected usage")
+				return &streamingResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: true}, nil
+			}
 			log.Printf("Stream data interval timeout: account=%d model=%s interval=%s", account.ID, originalModel, streamInterval)
 			sendErrorEvent("stream_timeout")
 			return &streamingResult{usage: usage, firstTokenMs: firstTokenMs}, fmt.Errorf("stream data interval timeout")