Sub2API v1.0 - AI API 网关（二开初始版本，基于上游 Wei-Shaw/sub2api）

2026-08-21 18:30:13 +08:00
commit 6d655c9903
3584 changed files with 1270640 additions and 0 deletions
@@ -0,0 +1,565 @@
+package service
+
+import (
+	"bufio"
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"net/http"
+	"strings"
+	"time"
+
+	"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
+	"github.com/Wei-Shaw/sub2api/internal/pkg/apicompat"
+	"github.com/gin-gonic/gin"
+)
+
+type antigravityCompatProtocol uint8
+
+const (
+	antigravityCompatChatCompletions antigravityCompatProtocol = iota
+	antigravityCompatResponses
+)
+
+const (
+	// AntigravityCredentialRejectedClientMessage 是可安全返回给客户端的认证修复提示。
+	AntigravityCredentialRejectedClientMessage = "Antigravity rejected the OAuth credential after refresh; reauthorize the account and verify project_id"
+	// AntigravityCredentialRejectedReason 标识上游拒绝已刷新 OAuth 凭据。
+	AntigravityCredentialRejectedReason GatewayFailureReason = "antigravity_oauth_credential_rejected"
+)
+
+type antigravityCompatRequest struct {
+	protocol        antigravityCompatProtocol
+	originalBody    []byte
+	claudeBody      []byte
+	originalModel   string
+	clientStream    bool
+	includeUsage    bool
+	startTime       time.Time
+	reasoningEffort *string
+}
+
+type antigravityCompatUpstreamCall struct {
+	request      antigravityCompatRequest
+	billingModel string
+	prefix       string
+	proxyURL     string
+	accessToken  string
+	geminiBody   []byte
+}
+
+// ForwardAsChatCompletions 使用 Antigravity 原生 OAuth 账号转发 Chat Completions 请求。
+func (s *AntigravityGatewayService) ForwardAsChatCompletions(
+	ctx context.Context,
+	c *gin.Context,
+	account *Account,
+	body []byte,
+	_ *ParsedRequest,
+) (*ForwardResult, error) {
+	if err := s.validateAntigravityCompatAccount(c, account); err != nil {
+		return nil, err
+	}
+
+	var request apicompat.ChatCompletionsRequest
+	if json.Unmarshal(body, &request) != nil {
+		return nil, s.writeAntigravityCompatError(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body")
+	}
+	if strings.TrimSpace(request.Model) == "" {
+		return nil, s.writeAntigravityCompatError(c, http.StatusBadRequest, "invalid_request_error", "model is required")
+	}
+
+	responsesRequest, err := apicompat.ChatCompletionsToResponses(&request)
+	if err != nil {
+		return nil, s.writeAntigravityCompatError(c, http.StatusBadRequest, "invalid_request_error", err.Error())
+	}
+	claudeRequest, err := apicompat.ResponsesToAnthropicRequest(responsesRequest)
+	if err != nil {
+		return nil, s.writeAntigravityCompatError(c, http.StatusBadRequest, "invalid_request_error", err.Error())
+	}
+	preserveChatCompletionTokenLimit(&request, claudeRequest)
+	claudeRequest.Stream = request.Stream
+	claudeBody, err := json.Marshal(claudeRequest)
+	if err != nil {
+		return nil, fmt.Errorf("marshal anthropic request: %w", err)
+	}
+
+	return s.forwardAntigravityCompat(ctx, c, account, antigravityCompatRequest{
+		protocol:        antigravityCompatChatCompletions,
+		originalBody:    body,
+		claudeBody:      claudeBody,
+		originalModel:   request.Model,
+		clientStream:    request.Stream,
+		includeUsage:    request.StreamOptions != nil && request.StreamOptions.IncludeUsage,
+		startTime:       time.Now(),
+		reasoningEffort: extractCCReasoningEffortFromBody(body),
+	})
+}
+
+// ForwardAsResponses 使用 Antigravity 原生 OAuth 账号转发 Responses 请求。
+func (s *AntigravityGatewayService) ForwardAsResponses(
+	ctx context.Context,
+	c *gin.Context,
+	account *Account,
+	body []byte,
+	_ *ParsedRequest,
+) (*ForwardResult, error) {
+	if err := s.validateAntigravityCompatAccount(c, account); err != nil {
+		return nil, err
+	}
+
+	var request apicompat.ResponsesRequest
+	if json.Unmarshal(body, &request) != nil {
+		return nil, s.writeAntigravityCompatError(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body")
+	}
+	if strings.TrimSpace(request.Model) == "" {
+		return nil, s.writeAntigravityCompatError(c, http.StatusBadRequest, "invalid_request_error", "model is required")
+	}
+
+	claudeRequest, err := apicompat.ResponsesToAnthropicRequest(&request)
+	if err != nil {
+		return nil, s.writeAntigravityCompatError(c, http.StatusBadRequest, "invalid_request_error", err.Error())
+	}
+	claudeRequest.Stream = request.Stream
+	claudeBody, err := json.Marshal(claudeRequest)
+	if err != nil {
+		return nil, fmt.Errorf("marshal anthropic request: %w", err)
+	}
+
+	return s.forwardAntigravityCompat(ctx, c, account, antigravityCompatRequest{
+		protocol:        antigravityCompatResponses,
+		originalBody:    body,
+		claudeBody:      claudeBody,
+		originalModel:   request.Model,
+		clientStream:    request.Stream,
+		startTime:       time.Now(),
+		reasoningEffort: ExtractResponsesReasoningEffortFromBody(body),
+	})
+}
+
+func (s *AntigravityGatewayService) validateAntigravityCompatAccount(c *gin.Context, account *Account) error {
+	if account != nil && account.Platform == PlatformAntigravity && account.Type == AccountTypeOAuth {
+		return nil
+	}
+	return s.writeAntigravityCompatError(
+		c,
+		http.StatusBadRequest,
+		"invalid_request_error",
+		"native OAuth account required for antigravity compatibility mode",
+	)
+}
+
+func preserveChatCompletionTokenLimit(request *apicompat.ChatCompletionsRequest, claudeRequest *apicompat.AnthropicRequest) {
+	if request == nil || claudeRequest == nil {
+		return
+	}
+	limit := request.MaxTokens
+	if request.MaxCompletionTokens != nil {
+		limit = request.MaxCompletionTokens
+	}
+	if limit != nil && *limit > 0 {
+		claudeRequest.MaxTokens = *limit
+	}
+}
+
+func (s *AntigravityGatewayService) forwardAntigravityCompat(
+	ctx context.Context,
+	c *gin.Context,
+	account *Account,
+	request antigravityCompatRequest,
+) (*ForwardResult, error) {
+	beginUpstreamResponseModelObservation(c)
+	call, err := s.prepareAntigravityCompatCall(ctx, c, account, request)
+	if err != nil {
+		return nil, err
+	}
+
+	result, err := s.antigravityRetryLoop(antigravityRetryLoopParams{
+		ctx:             ctx,
+		prefix:          call.prefix,
+		account:         account,
+		proxyURL:        call.proxyURL,
+		accessToken:     call.accessToken,
+		action:          "streamGenerateContent",
+		body:            call.geminiBody,
+		c:               c,
+		httpUpstream:    s.httpUpstream,
+		settingService:  s.settingService,
+		accountRepo:     s.accountRepo,
+		handleError:     s.handleUpstreamError,
+		requestedModel:  request.originalModel,
+		isStickySession: false,
+		groupID:         0,
+		sessionHash:     "",
+	})
+	if err != nil {
+		return nil, s.handleAntigravityCompatTransportError(c, err)
+	}
+
+	return s.consumeAntigravityCompatResponse(ctx, c, account, call, result.resp)
+}
+
+func (s *AntigravityGatewayService) prepareAntigravityCompatCall(
+	ctx context.Context,
+	c *gin.Context,
+	account *Account,
+	request antigravityCompatRequest,
+) (*antigravityCompatUpstreamCall, error) {
+	var claudeRequest antigravity.ClaudeRequest
+	if json.Unmarshal(request.claudeBody, &claudeRequest) != nil {
+		return nil, s.writeAntigravityCompatError(c, http.StatusBadRequest, "invalid_request_error", "Invalid request body")
+	}
+
+	mappedModel := s.getMappedModel(account, request.originalModel)
+	if mappedModel == "" {
+		MarkOpsClientBusinessLimited(c, OpsClientBusinessLimitedReasonLocalFeatureGate)
+		message := fmt.Sprintf("model %s not in whitelist", request.originalModel)
+		return nil, s.writeAntigravityCompatError(c, http.StatusForbidden, "permission_error", message)
+	}
+	thinkingEnabled := claudeRequest.Thinking != nil &&
+		(claudeRequest.Thinking.Type == "enabled" || claudeRequest.Thinking.Type == "adaptive")
+	mappedModel = applyThinkingModelSuffix(mappedModel, thinkingEnabled)
+
+	if s.tokenProvider == nil {
+		return nil, s.writeAntigravityCompatError(c, http.StatusBadGateway, "api_error", "Antigravity token provider not configured")
+	}
+	accessToken, err := s.tokenProvider.GetAccessToken(ctx, account)
+	if err != nil {
+		return nil, &UpstreamFailoverError{
+			StatusCode:   http.StatusBadGateway,
+			ResponseBody: []byte(`{"error":{"type":"authentication_error","message":"Failed to get upstream access token"},"type":"error"}`),
+		}
+	}
+
+	projectID, err := resolveAntigravityProjectID(account)
+	if err != nil {
+		_ = s.writeAntigravityCompatError(c, http.StatusBadRequest, "invalid_request_error", err.Error())
+		return nil, err
+	}
+	geminiBody, err := s.buildAntigravityCompatGeminiBody(ctx, request.claudeBody, &claudeRequest, projectID, mappedModel)
+	if err != nil {
+		return nil, s.writeAntigravityCompatError(c, http.StatusBadRequest, "invalid_request_error", "Invalid request")
+	}
+
+	request.reasoningEffort = ApplyThinkingEnabledFallback(request.reasoningEffort, request.originalBody, mappedModel)
+	return &antigravityCompatUpstreamCall{
+		request:      request,
+		billingModel: mappedModel,
+		prefix:       logPrefix(getSessionID(c), account.Name),
+		proxyURL:     antigravityCompatProxyURL(account),
+		accessToken:  accessToken,
+		geminiBody:   geminiBody,
+	}, nil
+}
+
+func (s *AntigravityGatewayService) buildAntigravityCompatGeminiBody(
+	ctx context.Context,
+	claudeBody []byte,
+	claudeRequest *antigravity.ClaudeRequest,
+	projectID string,
+	mappedModel string,
+) ([]byte, error) {
+	if strings.HasPrefix(strings.ToLower(mappedModel), "gemini-") {
+		body, err := convertClaudeMessagesToGeminiGenerateContent(claudeBody)
+		if err != nil {
+			return nil, err
+		}
+		body, err = enableMixedGeminiToolInvocations(body)
+		if err != nil {
+			return nil, err
+		}
+		body = ensureGeminiFunctionCallThoughtSignatures(body)
+		body, err = injectIdentityPatchToGeminiRequest(body)
+		if err != nil {
+			return nil, err
+		}
+		if cleaned, cleanErr := cleanGeminiRequest(body); cleanErr == nil {
+			body = cleaned
+		}
+		return s.wrapV1InternalRequest(projectID, mappedModel, body)
+	}
+
+	options := s.getClaudeTransformOptions(ctx)
+	options.EnableIdentityPatch = true
+	return antigravity.TransformClaudeToGeminiWithOptions(claudeRequest, projectID, mappedModel, options)
+}
+
+func enableMixedGeminiToolInvocations(body []byte) ([]byte, error) {
+	var request map[string]any
+	if err := json.Unmarshal(body, &request); err != nil {
+		return nil, err
+	}
+
+	var hasGoogleSearch, hasFunctionDeclarations bool
+	if tools, ok := request["tools"].([]any); ok {
+		for _, rawTool := range tools {
+			tool, ok := rawTool.(map[string]any)
+			if !ok {
+				continue
+			}
+			_, hasSearch := tool["googleSearch"]
+			declarations, hasFunctions := tool["functionDeclarations"].([]any)
+			hasGoogleSearch = hasGoogleSearch || hasSearch
+			hasFunctionDeclarations = hasFunctionDeclarations || hasFunctions && len(declarations) > 0
+		}
+	}
+	if !hasGoogleSearch || !hasFunctionDeclarations {
+		return body, nil
+	}
+
+	toolConfig, _ := request["toolConfig"].(map[string]any)
+	if toolConfig == nil {
+		toolConfig = make(map[string]any)
+		request["toolConfig"] = toolConfig
+	}
+	toolConfig["includeServerSideToolInvocations"] = true
+	return json.Marshal(request)
+}
+
+func antigravityCompatProxyURL(account *Account) string {
+	if account.ProxyID == nil || account.Proxy == nil {
+		return ""
+	}
+	return account.Proxy.URL()
+}
+
+func (s *AntigravityGatewayService) handleAntigravityCompatTransportError(c *gin.Context, err error) error {
+	if switchErr, ok := IsAntigravityAccountSwitchError(err); ok {
+		return &UpstreamFailoverError{
+			StatusCode:        http.StatusServiceUnavailable,
+			ForceCacheBilling: switchErr.IsStickySession,
+		}
+	}
+	if c.Request.Context().Err() != nil {
+		return s.writeAntigravityCompatError(c, http.StatusBadGateway, "client_disconnected", "Client disconnected before upstream response")
+	}
+	return s.writeAntigravityCompatError(c, http.StatusBadGateway, "upstream_error", "Upstream request failed after retries")
+}
+
+func (s *AntigravityGatewayService) consumeAntigravityCompatResponse(
+	ctx context.Context,
+	c *gin.Context,
+	account *Account,
+	call *antigravityCompatUpstreamCall,
+	resp *http.Response,
+) (*ForwardResult, error) {
+	defer func() { _ = resp.Body.Close() }()
+	if resp.StatusCode >= http.StatusBadRequest {
+		return nil, s.handleAntigravityCompatHTTPError(ctx, c, account, call, resp)
+	}
+
+	requestID := resp.Header.Get("x-request-id")
+	if requestID != "" {
+		c.Header("x-request-id", requestID)
+	}
+	streamResult, err := s.consumeAntigravityCompatSuccess(c, call, resp)
+	if err != nil {
+		return nil, err
+	}
+	if streamResult.usage == nil {
+		streamResult.usage = &ClaudeUsage{}
+	}
+
+	return &ForwardResult{
+		RequestID:                     requestID,
+		Usage:                         *streamResult.usage,
+		Model:                         call.request.originalModel,
+		UpstreamModel:                 call.billingModel,
+		UpstreamResponseModel:         observedUpstreamResponseModel(c),
+		UpstreamResponseModelConflict: observedUpstreamResponseModelConflict(c),
+		Stream:                        call.request.clientStream,
+		Duration:                      time.Since(call.request.startTime),
+		FirstTokenMs:                  streamResult.firstTokenMs,
+		ReasoningEffort:               call.request.reasoningEffort,
+		ClientDisconnect:              streamResult.clientDisconnect,
+	}, nil
+}
+
+func (s *AntigravityGatewayService) consumeAntigravityCompatSuccess(
+	c *gin.Context,
+	call *antigravityCompatUpstreamCall,
+	resp *http.Response,
+) (*antigravityStreamResult, error) {
+	if call.request.clientStream {
+		if call.request.protocol == antigravityCompatChatCompletions {
+			return s.handleChatCompletionsStreamingFromAntigravity(
+				c,
+				resp,
+				call.request.startTime,
+				call.request.originalModel,
+				call.request.includeUsage,
+			)
+		}
+		return s.handleResponsesStreamingFromAntigravity(c, resp, call.request.startTime, call.request.originalModel)
+	}
+
+	if call.request.protocol == antigravityCompatChatCompletions {
+		return s.handleChatCompletionsNonStreamingFromAntigravity(c, resp, call.request.startTime, call.request.originalModel)
+	}
+	return s.handleResponsesNonStreamingFromAntigravity(c, resp, call.request.startTime, call.request.originalModel)
+}
+
+func (s *AntigravityGatewayService) handleAntigravityCompatHTTPError(
+	ctx context.Context,
+	c *gin.Context,
+	account *Account,
+	call *antigravityCompatUpstreamCall,
+	resp *http.Response,
+) error {
+	body := s.readUpstreamErrorBody(resp)
+	s.handleUpstreamError(
+		ctx,
+		call.prefix,
+		account,
+		resp.StatusCode,
+		resp.Header,
+		body,
+		call.request.originalModel,
+		0,
+		"",
+		false,
+	)
+	if s.shouldFailoverUpstreamError(resp.StatusCode) {
+		message := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractAntigravityErrorMessage(body)))
+		event := OpsUpstreamErrorEvent{
+			Platform:           account.Platform,
+			AccountID:          account.ID,
+			AccountName:        account.Name,
+			UpstreamStatusCode: resp.StatusCode,
+			UpstreamRequestID:  resp.Header.Get("x-request-id"),
+			Kind:               "failover",
+			Message:            message,
+			Detail:             s.getUpstreamErrorDetail(body),
+		}
+		if resp.StatusCode == http.StatusUnauthorized {
+			event.Stage = string(GatewayFailureStageAccountAuth)
+			event.Scope = string(GatewayFailureScopeAccount)
+			event.Reason = string(AntigravityCredentialRejectedReason)
+			appendOpsUpstreamError(c, event)
+			return antigravityCredentialRejectedError(resp, body)
+		}
+		appendOpsUpstreamError(c, event)
+		return &UpstreamFailoverError{
+			StatusCode:      resp.StatusCode,
+			ResponseBody:    body,
+			ResponseHeaders: resp.Header.Clone(),
+		}
+	}
+	return s.writeMappedAntigravityCompatError(c, account, resp.StatusCode, resp.Header.Get("x-request-id"), body)
+}
+
+func antigravityCredentialRejectedError(resp *http.Response, body []byte) *UpstreamFailoverError {
+	return &UpstreamFailoverError{
+		StatusCode:        resp.StatusCode,
+		ResponseBody:      body,
+		ResponseHeaders:   resp.Header.Clone(),
+		Stage:             GatewayFailureStageAccountAuth,
+		Scope:             GatewayFailureScopeAccount,
+		Reason:            AntigravityCredentialRejectedReason,
+		NextAccountAction: NextAccountRetry,
+		ClientStatusCode:  http.StatusBadGateway,
+		ClientMessage:     AntigravityCredentialRejectedClientMessage,
+	}
+}
+
+func (s *AntigravityGatewayService) writeAntigravityCompatError(
+	c *gin.Context,
+	status int,
+	errType string,
+	message string,
+) error {
+	MarkResponseCommitted(c)
+	c.JSON(status, gin.H{
+		"error": gin.H{
+			"message": message,
+			"type":    errType,
+			"param":   nil,
+			"code":    nil,
+		},
+	})
+	return errors.New(message)
+}
+
+func (s *AntigravityGatewayService) writeMappedAntigravityCompatError(
+	c *gin.Context,
+	account *Account,
+	upstreamStatus int,
+	upstreamRequestID string,
+	body []byte,
+) error {
+	MarkResponseCommitted(c)
+	message := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractAntigravityErrorMessage(body)))
+	setOpsUpstreamError(c, upstreamStatus, message, s.getUpstreamErrorDetail(body))
+	appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
+		Platform:           account.Platform,
+		AccountID:          account.ID,
+		AccountName:        account.Name,
+		UpstreamStatusCode: upstreamStatus,
+		UpstreamRequestID:  upstreamRequestID,
+		Kind:               "http_error",
+		Message:            message,
+	})
+	c.JSON(mapUpstreamStatusCode(upstreamStatus), gin.H{
+		"error": gin.H{
+			"message": getPassthroughOrDefault(message, "Upstream request failed"),
+			"type":    "upstream_error",
+			"param":   nil,
+			"code":    nil,
+		},
+	})
+	return fmt.Errorf("upstream error: %d %s", upstreamStatus, message)
+}
+
+func (s *AntigravityGatewayService) handleChatCompletionsNonStreamingFromAntigravity(
+	c *gin.Context,
+	resp *http.Response,
+	startTime time.Time,
+	originalModel string,
+) (*antigravityStreamResult, error) {
+	claudeResponse, result, err := s.collectClaudeStreamResponse(c, resp, startTime, originalModel)
+	if err != nil {
+		return nil, s.mapAntigravityCompatCollectionError(c, err)
+	}
+	var anthropicResponse apicompat.AnthropicResponse
+	if json.Unmarshal(claudeResponse, &anthropicResponse) != nil {
+		return nil, s.writeAntigravityCompatError(c, http.StatusBadGateway, "upstream_error", "Failed to parse upstream response")
+	}
+	responsesResponse := apicompat.AnthropicToResponsesResponse(&anthropicResponse)
+	c.JSON(http.StatusOK, apicompat.ResponsesToChatCompletions(responsesResponse, originalModel))
+	return result, nil
+}
+
+func (s *AntigravityGatewayService) handleResponsesNonStreamingFromAntigravity(
+	c *gin.Context,
+	resp *http.Response,
+	startTime time.Time,
+	originalModel string,
+) (*antigravityStreamResult, error) {
+	claudeResponse, result, err := s.collectClaudeStreamResponse(c, resp, startTime, originalModel)
+	if err != nil {
+		return nil, s.mapAntigravityCompatCollectionError(c, err)
+	}
+	var anthropicResponse apicompat.AnthropicResponse
+	if json.Unmarshal(claudeResponse, &anthropicResponse) != nil {
+		return nil, s.writeAntigravityCompatError(c, http.StatusBadGateway, "upstream_error", "Failed to parse upstream response")
+	}
+	c.JSON(http.StatusOK, apicompat.AnthropicToResponsesResponse(&anthropicResponse))
+	return result, nil
+}
+
+func (s *AntigravityGatewayService) mapAntigravityCompatCollectionError(c *gin.Context, err error) error {
+	var failoverError *UpstreamFailoverError
+	if errors.As(err, &failoverError) {
+		return err
+	}
+	if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
+		return err
+	}
+	if strings.Contains(err.Error(), "stream data interval timeout") {
+		return s.writeAntigravityCompatError(c, http.StatusBadGateway, "upstream_timeout", "Upstream stream data interval timeout")
+	}
+	if errors.Is(err, bufio.ErrTooLong) {
+		return s.writeAntigravityCompatError(c, http.StatusBadGateway, "response_too_large", "Upstream response line too long")
+	}
+	return s.writeAntigravityCompatError(c, http.StatusBadGateway, "upstream_error", "Failed to parse upstream response")
+}