diff --git a/backend/internal/handler/gemini_v1beta_handler.go b/backend/internal/handler/gemini_v1beta_handler.go index 8627e0d139..95c550302e 100644 --- a/backend/internal/handler/gemini_v1beta_handler.go +++ b/backend/internal/handler/gemini_v1beta_handler.go @@ -39,9 +39,11 @@ func (h *GatewayHandler) GeminiV1BetaListModels(c *gin.Context) { googleError(c, http.StatusUnauthorized, "Invalid API key") return } - // 检查平台：优先使用强制平台（/antigravity 路由），否则要求 gemini 分组 + // 检查平台：优先使用强制平台（/antigravity 路由），否则要求 gemini 分组（composite 分组允许路由到 antigravity） forcePlatform, hasForcePlatform := middleware.GetForcePlatformFromContext(c) - if !hasForcePlatform && effectiveAPIKeyPlatform(c, apiKey) != service.PlatformGemini { + effectivePlatform := effectiveAPIKeyPlatform(c, apiKey) + if !hasForcePlatform && effectivePlatform != service.PlatformGemini && + !servesCompositeAntigravityGemini(apiKey, effectivePlatform) { googleError(c, http.StatusBadRequest, "API key group platform is not gemini") return } @@ -223,9 +225,11 @@ func (h *GatewayHandler) GeminiV1BetaGetModel(c *gin.Context) { googleError(c, http.StatusUnauthorized, "Invalid API key") return } - // 检查平台：优先使用强制平台（/antigravity 路由），否则要求 gemini 分组 + // 检查平台：优先使用强制平台（/antigravity 路由），否则要求 gemini 分组（composite 分组允许路由到 antigravity） forcePlatform, hasForcePlatform := middleware.GetForcePlatformFromContext(c) - if !hasForcePlatform && effectiveAPIKeyPlatform(c, apiKey) != service.PlatformGemini { + effectivePlatform := effectiveAPIKeyPlatform(c, apiKey) + if !hasForcePlatform && effectivePlatform != service.PlatformGemini && + !servesCompositeAntigravityGemini(apiKey, effectivePlatform) { googleError(c, http.StatusBadRequest, "API key group platform is not gemini") return } @@ -244,9 +248,13 @@ func (h *GatewayHandler) GeminiV1BetaGetModel(c *gin.Context) { if resolvedModel, ok := service.ResolvedUpstreamModelFromContext(c.Request.Context()); ok && strings.TrimSpace(resolvedModel) != "" { modelName = strings.TrimSpace(resolvedModel) } + if !hasForcePlatform && servesCompositeAntigravityGemini(apiKey, effectivePlatform) && !isAntigravityGeminiModel(modelName) { + googleError(c, http.StatusBadRequest, "API key group platform is not gemini") + return + } - // 强制 antigravity 模式：返回 antigravity 模型信息 - if forcePlatform == service.PlatformAntigravity { + // 强制 antigravity 模式或 composite 路由到 antigravity：返回 antigravity 模型信息 + if forcePlatform == service.PlatformAntigravity || effectivePlatform == service.PlatformAntigravity { c.JSON(http.StatusOK, antigravity.FallbackGeminiModel(modelName)) return } @@ -299,9 +307,11 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) { zap.Any("group_id", apiKey.GroupID), ) - // 检查平台：优先使用强制平台（/antigravity 路由，中间件已设置 request.Context），否则要求 gemini 分组 + // 检查平台：优先使用强制平台（/antigravity 路由，中间件已设置 request.Context），否则要求 gemini 分组（composite 分组允许路由到 antigravity） + effectivePlatform := effectiveAPIKeyPlatform(c, apiKey) if !middleware.HasForcePlatform(c) { - if effectiveAPIKeyPlatform(c, apiKey) != service.PlatformGemini { + if effectivePlatform != service.PlatformGemini && + !servesCompositeAntigravityGemini(apiKey, effectivePlatform) { googleError(c, http.StatusBadRequest, "API key group platform is not gemini") return } @@ -321,6 +331,10 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) { if resolvedModel, ok := service.ResolvedUpstreamModelFromContext(c.Request.Context()); ok && strings.TrimSpace(resolvedModel) != "" { modelName = strings.TrimSpace(resolvedModel) } + if !middleware.HasForcePlatform(c) && servesCompositeAntigravityGemini(apiKey, effectivePlatform) && !isAntigravityGeminiModel(modelName) { + googleError(c, http.StatusBadRequest, "API key group platform is not gemini") + return + } stream := action == "streamGenerateContent" reqLog = reqLog.With(zap.String("model", modelName), zap.String("action", action), zap.Bool("stream", stream)) @@ -567,7 +581,11 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) { return } if len(fs.FailedAccountIDs) == 0 { - cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, modelName, modelName, service.PlatformGemini) + diagPlatform := service.PlatformGemini + if p := effectiveAPIKeyPlatform(c, apiKey); p == service.PlatformAntigravity { + diagPlatform = service.PlatformAntigravity + } + cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, modelName, modelName, diagPlatform) if !cls.ModelNotFound { markOpsRoutingCapacityLimitedIfNoAvailable(c, err) } @@ -1053,3 +1071,16 @@ func derefGroupID(groupID *int64) int64 { } return *groupID } + +func servesCompositeAntigravityGemini(apiKey *service.APIKey, effectivePlatform string) bool { + return apiKey.Group != nil && apiKey.Group.Platform == service.PlatformComposite && effectivePlatform == service.PlatformAntigravity +} + +func isAntigravityGeminiModel(model string) bool { + trimmed := strings.ToLower(strings.TrimSpace(strings.TrimPrefix(model, "models/"))) + if trimmed == "" || strings.HasPrefix(trimmed, "claude-") { + return false + } + detected, ok := service.DetectModelPlatform(trimmed) + return ok && detected == service.PlatformGemini +} diff --git a/backend/internal/handler/gemini_v1beta_handler_test.go b/backend/internal/handler/gemini_v1beta_handler_test.go index 44d66d5c97..4f84879599 100644 --- a/backend/internal/handler/gemini_v1beta_handler_test.go +++ b/backend/internal/handler/gemini_v1beta_handler_test.go @@ -216,3 +216,75 @@ func TestShouldFallbackGeminiModel_DelegatesScopeFallback(t *testing.T) { } require.True(t, shouldFallbackGeminiModel("gemini-future-model", res)) } + +func TestGeminiV1BetaGetModel_CompositeAntigravityRoute(t *testing.T) { + gin.SetMode(gin.TestMode) + + t.Run("composite route to antigravity gemini model succeeds", func(t *testing.T) { + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + req := httptest.NewRequest(http.MethodGet, "/v1beta/models/gemini-4-argon", nil) + req = req.WithContext(service.WithCompositeRouteDecision(req.Context(), service.CompositeRouteDecision{ + Matched: true, + TargetPlatform: service.PlatformAntigravity, + UpstreamModel: "gemini-4-argon-high", + })) + c.Request = req + c.Params = gin.Params{{Key: "model", Value: "gemini-4-argon"}} + groupID := int64(4) + c.Set(string(middleware.ContextKeyAPIKey), &service.APIKey{ + ID: 1, + GroupID: &groupID, + Group: &service.Group{ID: groupID, Platform: service.PlatformComposite}, + }) + + (&GatewayHandler{}).GeminiV1BetaGetModel(c) + require.Equal(t, http.StatusOK, rec.Code, rec.Body.String()) + require.Contains(t, rec.Body.String(), "models/gemini-4-argon-high") + }) + + t.Run("composite route to antigravity claude model is rejected on v1beta", func(t *testing.T) { + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + req := httptest.NewRequest(http.MethodGet, "/v1beta/models/claude-sonnet-5-5-high", nil) + req = req.WithContext(service.WithCompositeRouteDecision(req.Context(), service.CompositeRouteDecision{ + Matched: true, + TargetPlatform: service.PlatformAntigravity, + UpstreamModel: "claude-sonnet-5-5-high", + })) + c.Request = req + c.Params = gin.Params{{Key: "model", Value: "claude-sonnet-5-5-high"}} + groupID := int64(4) + c.Set(string(middleware.ContextKeyAPIKey), &service.APIKey{ + ID: 1, + GroupID: &groupID, + Group: &service.Group{ID: groupID, Platform: service.PlatformComposite}, + }) + + (&GatewayHandler{}).GeminiV1BetaGetModel(c) + require.Equal(t, http.StatusBadRequest, rec.Code, rec.Body.String()) + }) + + t.Run("composite route to antigravity claude model is rejected on generateContent", func(t *testing.T) { + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + req := httptest.NewRequest(http.MethodPost, "/v1beta/models/claude-sonnet-5-5-high:generateContent", nil) + req = req.WithContext(service.WithCompositeRouteDecision(req.Context(), service.CompositeRouteDecision{ + Matched: true, + TargetPlatform: service.PlatformAntigravity, + UpstreamModel: "claude-sonnet-5-5-high", + })) + c.Request = req + c.Params = gin.Params{{Key: "modelAction", Value: "/claude-sonnet-5-5-high:generateContent"}} + groupID := int64(4) + c.Set(string(middleware.ContextKeyAPIKey), &service.APIKey{ + ID: 1, + GroupID: &groupID, + Group: &service.Group{ID: groupID, Platform: service.PlatformComposite}, + }) + c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: 1, Concurrency: 1}) + + (&GatewayHandler{}).GeminiV1BetaModels(c) + require.Equal(t, http.StatusBadRequest, rec.Code, rec.Body.String()) + }) +} diff --git a/backend/internal/pkg/antigravity/request_transformer.go b/backend/internal/pkg/antigravity/request_transformer.go index 9349e54ee8..a13748926e 100644 --- a/backend/internal/pkg/antigravity/request_transformer.go +++ b/backend/internal/pkg/antigravity/request_transformer.go @@ -588,6 +588,7 @@ func buildParts(content json.RawMessage, toolIDToName map[string]string, allowDu ID: block.ToolUseID, }, }) + parts = append(parts, extractToolResultInlineParts(block.Content)...) } } @@ -638,6 +639,31 @@ func parseToolResultContent(content json.RawMessage, isError bool) string { return string(content) } +func extractToolResultInlineParts(content json.RawMessage) []GeminiPart { + if len(content) == 0 { + return nil + } + var innerBlocks []ContentBlock + if err := json.Unmarshal(content, &innerBlocks); err != nil { + return nil + } + var inlineParts []GeminiPart + for _, item := range innerBlocks { + if (item.Type == "image" || item.Type == "document") && + item.Source != nil && + item.Source.Type == "base64" && + strings.TrimSpace(item.Source.Data) != "" { + inlineParts = append(inlineParts, GeminiPart{ + InlineData: &GeminiInlineData{ + MimeType: item.Source.MediaType, + Data: item.Source.Data, + }, + }) + } + } + return inlineParts +} + // buildGenerationConfig 构建 generationConfig const ( defaultMaxOutputTokens = 64000 diff --git a/backend/internal/pkg/antigravity/request_transformer_test.go b/backend/internal/pkg/antigravity/request_transformer_test.go index 1aff9f5775..7bdb40914b 100644 --- a/backend/internal/pkg/antigravity/request_transformer_test.go +++ b/backend/internal/pkg/antigravity/request_transformer_test.go @@ -694,3 +694,64 @@ func TestToolConfigAlwaysPresent(t *testing.T) { }) } } + +func TestTransformClaudeToGemini_ToolResultImagesReachUpstream(t *testing.T) { + claudeReq := &ClaudeRequest{ + Model: "claude-sonnet-4-5", + Messages: []ClaudeMessage{ + { + Role: "assistant", + Content: json.RawMessage(`[ + {"type":"tool_use","id":"toolu_img_1","name":"read_image","input":{"path":"colors.png"}} + ]`), + }, + { + Role: "user", + Content: json.RawMessage(`[ + { + "type":"tool_result", + "tool_use_id":"toolu_img_1", + "content":[ + {"type":"text","text":"loaded colors.png"}, + {"type":"image","source":{"type":"base64","media_type":"image/png","data":"aGVsbG8="}} + ] + }, + { + "type":"tool_result", + "tool_use_id":"toolu_img_1", + "content":[ + {"type":"image","source":{"type":"base64","media_type":"image/jpeg","data":"d29ybGQ="}} + ] + } + ]`), + }, + }, + } + + body, err := TransformClaudeToGeminiWithOptions(claudeReq, "project-1", "claude-sonnet-4-5", DefaultTransformOptions()) + require.NoError(t, err) + + var req V1InternalRequest + require.NoError(t, json.Unmarshal(body, &req)) + require.Len(t, req.Request.Contents, 2) + + userParts := req.Request.Contents[1].Parts + require.Len(t, userParts, 4) + + require.NotNil(t, userParts[0].FunctionResponse) + require.Equal(t, "read_image", userParts[0].FunctionResponse.Name) + require.Equal(t, "toolu_img_1", userParts[0].FunctionResponse.ID) + require.Equal(t, "loaded colors.png", userParts[0].FunctionResponse.Response["result"]) + + require.NotNil(t, userParts[1].InlineData) + require.Equal(t, "image/png", userParts[1].InlineData.MimeType) + require.Equal(t, "aGVsbG8=", userParts[1].InlineData.Data) + + require.NotNil(t, userParts[2].FunctionResponse) + require.Equal(t, "read_image", userParts[2].FunctionResponse.Name) + require.Equal(t, "Command executed successfully.", userParts[2].FunctionResponse.Response["result"]) + + require.NotNil(t, userParts[3].InlineData) + require.Equal(t, "image/jpeg", userParts[3].InlineData.MimeType) + require.Equal(t, "d29ybGQ=", userParts[3].InlineData.Data) +} diff --git a/backend/internal/service/antigravity_gateway_gemini.go b/backend/internal/service/antigravity_gateway_gemini.go index 4a8c70dcfb..b6396cb08a 100644 --- a/backend/internal/service/antigravity_gateway_gemini.go +++ b/backend/internal/service/antigravity_gateway_gemini.go @@ -141,6 +141,8 @@ func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Co injectedBody = reconciled } + injectedBody = ensureGeminiFunctionCallThoughtSignatures(injectedBody) + // 包装请求 wrappedBody, err := s.wrapV1InternalRequest(projectID, mappedModel, injectedBody) if err != nil { diff --git a/backend/internal/service/antigravity_gateway_service_test.go b/backend/internal/service/antigravity_gateway_service_test.go index 79ce35f596..92735948aa 100644 --- a/backend/internal/service/antigravity_gateway_service_test.go +++ b/backend/internal/service/antigravity_gateway_service_test.go @@ -2012,3 +2012,175 @@ func generateLargeUnwrapJSON(minSize int) []byte { b, _ := json.Marshal(outer) return b } + +func TestAntigravityGatewayService_ForwardGemini_NonStreamPreservesFunctionCallAndThoughtSignature(t *testing.T) { + gin.SetMode(gin.TestMode) + body := []byte(`{"contents":[{"role":"user","parts":[{"text":"Call get_weather for Tokyo."}]}],"tools":[{"functionDeclarations":[{"name":"get_weather","description":"Get weather","parameters":{"type":"object","properties":{"city":{"type":"string"}},"required":["city"]}}]}]}`) + + writer := httptest.NewRecorder() + c, _ := gin.CreateTestContext(writer) + c.Request = httptest.NewRequest(http.MethodPost, "/v1beta/models/gemini-4-argon:generateContent", bytes.NewReader(body)) + + sseStream := strings.Join([]string{ + `data: {"response":{"candidates":[{"content":{"role":"model","parts":[{"thoughtSignature":"sig_fc_1","functionCall":{"name":"get_weather","args":{"city":"Tokyo"},"id":"call_1"}}]}}],"usageMetadata":{"promptTokenCount":300,"candidatesTokenCount":16,"totalTokenCount":328,"thoughtsTokenCount":12},"modelVersion":"gemini-4-argon","responseId":"resp_1"}}`, + "", + `data: {"response":{"candidates":[{"content":{"role":"model","parts":[{"text":""}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":300,"candidatesTokenCount":22,"totalTokenCount":334,"thoughtsTokenCount":12},"modelVersion":"gemini-4-argon","responseId":"resp_1"}}`, + "", + }, "\n") + + upstream := &queuedHTTPUpstreamStub{ + responses: []*http.Response{{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(strings.NewReader(sseStream)), + }}, + } + svc := &AntigravityGatewayService{ + settingService: NewSettingService(&antigravitySettingRepoStub{}, &config.Config{Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize}}), + tokenProvider: &AntigravityTokenProvider{}, + httpUpstream: upstream, + } + account := &Account{ + ID: 109, + Name: "antigravity-fc-nonstream", + Platform: PlatformAntigravity, + Type: AccountTypeOAuth, + Status: StatusActive, + Concurrency: 1, + Credentials: map[string]any{ + "access_token": "test-access-token", + "project_id": "test-project", + "model_mapping": map[string]any{ + "gemini-4-argon": "gemini-4-argon", + }, + }, + } + + result, err := svc.ForwardGemini(context.Background(), c, account, "gemini-4-argon", "generateContent", false, body, false) + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, http.StatusOK, writer.Code) + + var respPayload map[string]any + require.NoError(t, json.Unmarshal(writer.Body.Bytes(), &respPayload)) + candidates, ok := respPayload["candidates"].([]any) + require.True(t, ok) + require.Len(t, candidates, 1) + candidate, ok := candidates[0].(map[string]any) + require.True(t, ok) + require.Equal(t, "STOP", candidate["finishReason"]) + + content, ok := candidate["content"].(map[string]any) + require.True(t, ok) + parts, ok := content["parts"].([]any) + require.True(t, ok) + require.Len(t, parts, 1, "trailing empty text part should not overwrite or pollute functionCall part") + + part0, ok := parts[0].(map[string]any) + require.True(t, ok) + require.Equal(t, "sig_fc_1", part0["thoughtSignature"]) + fc, ok := part0["functionCall"].(map[string]any) + require.True(t, ok) + require.Equal(t, "get_weather", fc["name"]) + require.Equal(t, "call_1", fc["id"]) + require.Equal(t, map[string]any{"city": "Tokyo"}, fc["args"]) +} + +func TestAntigravityGatewayService_ForwardGemini_StreamAndFunctionResponseFollowUp(t *testing.T) { + gin.SetMode(gin.TestMode) + + // 1) Stream request with functionDeclarations returns unwrapped functionCall and thoughtSignature + streamReqBody := []byte(`{"contents":[{"role":"user","parts":[{"text":"Call get_weather for Tokyo."}]}],"tools":[{"functionDeclarations":[{"name":"get_weather","parameters":{"type":"object","properties":{"city":{"type":"string"}}}}]}]}`) + streamWriter := httptest.NewRecorder() + streamCtx, _ := gin.CreateTestContext(streamWriter) + streamCtx.Request = httptest.NewRequest(http.MethodPost, "/v1beta/models/gemini-4-argon:streamGenerateContent?alt=sse", bytes.NewReader(streamReqBody)) + + streamSSE := strings.Join([]string{ + `data: {"response":{"candidates":[{"content":{"role":"model","parts":[{"thoughtSignature":"sig_stream_1","functionCall":{"name":"get_weather","args":{"city":"Tokyo"},"id":"call_s1"}}]}}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":5}}}`, + "", + `data: {"response":{"candidates":[{"content":{"role":"model","parts":[{"text":""}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":6}}}`, + "", + }, "\n") + + // 2) Non-stream follow-up with functionCall (missing thoughtSignature) + functionResponse (with inlineData image) + followUpReqBody := []byte(`{"contents":[{"role":"user","parts":[{"text":"Call get_weather for Tokyo."}]},{"role":"model","parts":[{"functionCall":{"name":"get_weather","args":{"city":"Tokyo"},"id":"call_s1"}}]},{"role":"user","parts":[{"functionResponse":{"name":"get_weather","id":"call_s1","response":{"weather":"Sunny"},"parts":[{"inlineData":{"mimeType":"image/png","data":"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg=="}}]}}]}],"tools":[{"functionDeclarations":[{"name":"get_weather","parameters":{"type":"object","properties":{"city":{"type":"string"}}}}]}]}`) + followUpWriter := httptest.NewRecorder() + followUpCtx, _ := gin.CreateTestContext(followUpWriter) + followUpCtx.Request = httptest.NewRequest(http.MethodPost, "/v1beta/models/gemini-4-argon:generateContent", bytes.NewReader(followUpReqBody)) + + followUpSSE := strings.Join([]string{ + `data: {"response":{"candidates":[{"content":{"role":"model","parts":[{"text":"It is sunny in "}]}}],"usageMetadata":{"promptTokenCount":20,"candidatesTokenCount":4}}}`, + "", + `data: {"response":{"candidates":[{"content":{"role":"model","parts":[{"text":"Tokyo."}]}}],"usageMetadata":{"promptTokenCount":20,"candidatesTokenCount":6}}}`, + "", + `data: {"response":{"candidates":[{"content":{"role":"model","parts":[{"thoughtSignature":"sig_final","text":""}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":20,"candidatesTokenCount":8}}}`, + "", + }, "\n") + + upstream := &queuedHTTPUpstreamStub{ + responses: []*http.Response{ + { + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(strings.NewReader(streamSSE)), + }, + { + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(strings.NewReader(followUpSSE)), + }, + }, + } + + svc := &AntigravityGatewayService{ + settingService: NewSettingService(&antigravitySettingRepoStub{}, &config.Config{Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize}}), + tokenProvider: &AntigravityTokenProvider{}, + httpUpstream: upstream, + } + account := &Account{ + ID: 110, + Name: "antigravity-fc-stream", + Platform: PlatformAntigravity, + Type: AccountTypeOAuth, + Status: StatusActive, + Concurrency: 1, + Credentials: map[string]any{ + "access_token": "test-access-token", + "project_id": "test-project", + "model_mapping": map[string]any{ + "gemini-4-argon": "gemini-4-argon", + }, + }, + } + + res1, err := svc.ForwardGemini(context.Background(), streamCtx, account, "gemini-4-argon", "streamGenerateContent", true, streamReqBody, false) + require.NoError(t, err) + require.NotNil(t, res1) + require.Contains(t, streamWriter.Body.String(), `"functionCall":{"name":"get_weather","args":{"city":"Tokyo"},"id":"call_s1"}`) + require.Contains(t, streamWriter.Body.String(), `"thoughtSignature":"sig_stream_1"`) + + res2, err := svc.ForwardGemini(context.Background(), followUpCtx, account, "gemini-4-argon", "generateContent", false, followUpReqBody, false) + require.NoError(t, err) + require.NotNil(t, res2) + require.Len(t, upstream.requestBodies, 2) + + // Verify follow-up request had dummy thoughtSignature injected on functionCall and preserved functionResponse image part + secondSent := string(upstream.requestBodies[1]) + require.Contains(t, secondSent, `"thoughtSignature":"skip_thought_signature_validator"`) + require.Contains(t, secondSent, `"mimeType":"image/png"`) + + var followUpResp map[string]any + require.NoError(t, json.Unmarshal(followUpWriter.Body.Bytes(), &followUpResp)) + cands, ok := followUpResp["candidates"].([]any) + require.True(t, ok) + cand0, ok := cands[0].(map[string]any) + require.True(t, ok) + content, ok := cand0["content"].(map[string]any) + require.True(t, ok) + parts, ok := content["parts"].([]any) + require.True(t, ok) + require.NotEmpty(t, parts) + firstPart, ok := parts[0].(map[string]any) + require.True(t, ok) + require.Equal(t, "It is sunny in Tokyo.", firstPart["text"]) +} diff --git a/backend/internal/service/antigravity_gateway_streaming.go b/backend/internal/service/antigravity_gateway_streaming.go index 675f0d7f87..58d5b9ce4f 100644 --- a/backend/internal/service/antigravity_gateway_streaming.go +++ b/backend/internal/service/antigravity_gateway_streaming.go @@ -350,8 +350,7 @@ func (s *AntigravityGatewayService) handleGeminiStreamToNonStreaming(c *gin.Cont var firstTokenMs *int var last map[string]any var lastWithParts map[string]any - var collectedImageParts []map[string]any // 收集所有包含图片的 parts - var collectedTextParts []string // 收集所有文本片段 + var collectedParts []map[string]any // 收集所有 parts（text、thinking、functionCall、inlineData 等） type scanEvent struct { line string @@ -467,19 +466,10 @@ func (s *AntigravityGatewayService) handleGeminiStreamToNonStreaming(c *gin.Cont } } - // 保留最后一个有 parts 的响应 + // 保留最后一个有 parts 的响应，并收集所有 parts if parts := extractGeminiParts(parsed); len(parts) > 0 { lastWithParts = parsed - // 收集包含图片和文本的 parts - for _, part := range parts { - if inlineData, ok := part["inlineData"].(map[string]any); ok { - collectedImageParts = append(collectedImageParts, part) - _ = inlineData // 避免 unused 警告 - } - if text, ok := part["text"].(string); ok && text != "" { - collectedTextParts = append(collectedTextParts, text) - } - } + collectedParts = append(collectedParts, parts...) } case <-intervalCh: @@ -506,14 +496,9 @@ returnResponse: } } - // 如果收集到了图片 parts，需要合并到最终响应中 - if len(collectedImageParts) > 0 { - finalResponse = mergeImagePartsToResponse(finalResponse, collectedImageParts) - } - - // 如果收集到了文本，需要合并到最终响应中 - if len(collectedTextParts) > 0 { - finalResponse = mergeTextPartsToResponse(finalResponse, collectedTextParts) + // 将收集的所有 parts 合并到最终响应中 + if len(collectedParts) > 0 { + finalResponse = mergeCollectedPartsToResponse(finalResponse, collectedParts) } respBody, err := json.Marshal(finalResponse) @@ -576,21 +561,27 @@ func mergeCollectedPartsToResponse(response map[string]any, collectedParts []map return response } - result, _, setParts := getOrCreateGeminiParts(response) + result, existingParts, setParts := getOrCreateGeminiParts(response) // 合并策略： // 1. 保持原始顺序 - // 2. 连续的普通 text parts 合并为一个 + // 2. 连续的普通 text parts 合并为一个，并保留关联的 thoughtSignature // 3. thinking、functionCall、inlineData 等保持原样 var mergedParts []any var textBuffer strings.Builder + var textSignature string flushTextBuffer := func() { - if textBuffer.Len() > 0 { - mergedParts = append(mergedParts, map[string]any{ + if textBuffer.Len() > 0 || textSignature != "" { + part := map[string]any{ "text": textBuffer.String(), - }) + } + if textSignature != "" { + part["thoughtSignature"] = textSignature + } + mergedParts = append(mergedParts, part) textBuffer.Reset() + textSignature = "" } } @@ -605,6 +596,9 @@ func mergeCollectedPartsToResponse(response map[string]any, collectedParts []map } else { // 普通 text，累积到 buffer _, _ = textBuffer.WriteString(text) + if sig, ok := part["thoughtSignature"].(string); ok && sig != "" { + textSignature = sig + } } } else { // 非 text part（functionCall、inlineData 等），先刷新 text buffer，然后保留原样 @@ -616,76 +610,14 @@ func mergeCollectedPartsToResponse(response map[string]any, collectedParts []map // 刷新剩余的 text flushTextBuffer() + if len(mergedParts) == 0 { + mergedParts = existingParts + } + setParts(mergedParts) return result } -// mergeImagePartsToResponse 将收集到的图片 parts 合并到 Gemini 响应中 -func mergeImagePartsToResponse(response map[string]any, imageParts []map[string]any) map[string]any { - if len(imageParts) == 0 { - return response - } - - result, existingParts, setParts := getOrCreateGeminiParts(response) - - // 检查现有 parts 中是否已经有图片 - for _, p := range existingParts { - if pm, ok := p.(map[string]any); ok { - if _, hasInline := pm["inlineData"]; hasInline { - return result // 已有图片，不重复添加 - } - } - } - - // 添加收集到的图片 parts - for _, imgPart := range imageParts { - existingParts = append(existingParts, imgPart) - } - setParts(existingParts) - return result -} - -// mergeTextPartsToResponse 将收集到的文本合并到 Gemini 响应中 -func mergeTextPartsToResponse(response map[string]any, textParts []string) map[string]any { - if len(textParts) == 0 { - return response - } - - mergedText := strings.Join(textParts, "") - result, existingParts, setParts := getOrCreateGeminiParts(response) - - // 查找并更新第一个 text part，或创建新的 - newParts := make([]any, 0, len(existingParts)+1) - textUpdated := false - - for _, p := range existingParts { - pm, ok := p.(map[string]any) - if !ok { - newParts = append(newParts, p) - continue - } - if _, hasText := pm["text"]; hasText && !textUpdated { - // 用累积的文本替换 - newPart := make(map[string]any) - for k, v := range pm { - newPart[k] = v - } - newPart["text"] = mergedText - newParts = append(newParts, newPart) - textUpdated = true - } else { - newParts = append(newParts, pm) - } - } - - if !textUpdated { - newParts = append([]any{map[string]any{"text": mergedText}}, newParts...) - } - - setParts(newParts) - return result -} - // antigravityStatusClientClosed 是客户端在上游响应前断开时回写的状态码（499 client closed request）。 const antigravityStatusClientClosed = 499