Files
sub2api/backend/internal/service/gateway_bedrock.go
T
李建琦 6d655c9903
Release / update-version (push) Has been cancelled
Release / build-frontend (push) Has been cancelled
Release / release (push) Has been cancelled
Release / sync-version-file (push) Has been cancelled
CI / shell (push) Canceled after 0s
CI / test (push) Canceled after 0s
CI / frontend (push) Canceled after 0s
CI / golangci-lint (push) Canceled after 0s
Security Scan / backend-security (push) Canceled after 0s
Security Scan / frontend-security (push) Canceled after 0s
Sub2API v1.0 - AI API 网关(二开初始版本,基于上游 Wei-Shaw/sub2api)
2026-08-21 18:30:13 +08:00

413 lines
13 KiB
Go
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
package service
// 本文件由 gateway_service.go 纯移动拆分而来:Bedrock 上游转发(CC 兼容转换、
// 请求构建、错误处理与非流式响应)。仅做代码搬迁,无任何行为变更。
import (
"bytes"
"context"
"errors"
"fmt"
"io"
"net/http"
"strings"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
"github.com/gin-gonic/gin"
)
// ApplyBedrockCCCompat 应用 Bedrock CC 兼容转换(渠道级模型映射后调用)
// 清理 body 中 Anthropic API 专有字段、修复 thinking/tool_use ID、过滤 beta token
// 同时过滤 HTTP header 中的 anthropic-beta(防止 Passthrough 路径透传不支持的 token)。
func (s *GatewayService) ApplyBedrockCCCompat(c *gin.Context, body []byte, model string, account *Account, groupID *int64) []byte {
if !s.isBedrockCCCompatEnabled(c.Request.Context(), account, groupID) {
return body
}
body = sanitizeBedrockCCFields(body)
body = sanitizeBedrockThinking(body, model)
body = sanitizeBedrockToolUseIDs(body)
body = sanitizeBedrockCCBetaTokens(body, model)
// 过滤 HTTP header 中的 anthropic-beta,只保留 Bedrock 支持的 token
if betaHeader := c.GetHeader("anthropic-beta"); betaHeader != "" {
if filtered := ResolveBedrockBetaTokens(betaHeader, body, model); len(filtered) > 0 {
c.Request.Header.Set("anthropic-beta", strings.Join(filtered, ", "))
} else {
c.Request.Header.Del("anthropic-beta")
}
}
return body
}
// isBedrockCCCompatEnabled 检查渠道是否启用了 Bedrock CC 兼容模式
func (s *GatewayService) isBedrockCCCompatEnabled(ctx context.Context, account *Account, groupID *int64) bool {
if groupID == nil || s.channelService == nil {
return false
}
ch, err := s.channelService.GetChannelForGroup(ctx, *groupID)
if err != nil || ch == nil {
return false
}
return ch.IsBedrockCCCompatEnabled(account.Platform)
}
// forwardBedrock 转发请求到 AWS Bedrock
func (s *GatewayService) forwardBedrock(
ctx context.Context,
c *gin.Context,
account *Account,
parsed *ParsedRequest,
startTime time.Time,
) (*ForwardResult, error) {
reqModel := parsed.Model
reqStream := parsed.Stream
body := parsed.Body.Bytes()
region := bedrockRuntimeRegion(account)
mappedModel, ok := ResolveBedrockModelID(account, reqModel)
if !ok {
return nil, fmt.Errorf("unsupported bedrock model: %s", reqModel)
}
if mappedModel != reqModel {
logger.LegacyPrintf("service.gateway", "[Bedrock] Model mapping: %s -> %s (account: %s)", reqModel, mappedModel, account.Name)
}
betaHeader := ""
if c != nil && c.Request != nil {
betaHeader = c.GetHeader("anthropic-beta")
}
// 准备请求体(注入 anthropic_version/anthropic_beta,移除 Bedrock 不支持的字段,清理 cache_control
betaTokens, err := s.resolveBedrockBetaTokensForRequest(ctx, account, betaHeader, body, mappedModel)
if err != nil {
return nil, err
}
bedrockBody, err := PrepareBedrockRequestBodyWithTokens(body, mappedModel, betaTokens, false)
if err != nil {
return nil, fmt.Errorf("prepare bedrock request body: %w", err)
}
proxyURL := ""
if account.ProxyID != nil && account.Proxy != nil {
proxyURL = account.Proxy.URL()
}
logger.LegacyPrintf("service.gateway", "[Bedrock] 命中 Bedrock 分支: account=%d name=%s model=%s->%s stream=%v",
account.ID, account.Name, reqModel, mappedModel, reqStream)
// 根据账号类型选择认证方式
var signer *BedrockSigner
var bedrockAPIKey string
if account.IsBedrockAPIKey() {
bedrockAPIKey = account.GetCredential("api_key")
if bedrockAPIKey == "" {
return nil, fmt.Errorf("api_key not found in bedrock credentials")
}
} else {
signer, err = NewBedrockSignerFromAccount(account)
if err != nil {
return nil, fmt.Errorf("create bedrock signer: %w", err)
}
}
// 执行上游请求(含重试)
resp, err := s.executeBedrockUpstream(ctx, c, account, bedrockBody, mappedModel, region, reqStream, signer, bedrockAPIKey, proxyURL)
if err != nil {
return nil, err
}
defer func() { _ = resp.Body.Close() }()
// 将 Bedrock 的 x-amzn-requestid 映射到 x-request-id
// 使通用错误处理函数(handleErrorResponse、handleRetryExhaustedError)能正确提取 AWS request ID。
if awsReqID := resp.Header.Get("x-amzn-requestid"); awsReqID != "" && resp.Header.Get("x-request-id") == "" {
resp.Header.Set("x-request-id", awsReqID)
}
// 错误/failover 处理
if resp.StatusCode >= 400 {
return s.handleBedrockUpstreamErrors(ctx, resp, c, account)
}
// Bedrock 分支绕过通用 Forward 成功路径,这里保持上游接受回调语义一致。
if parsed.OnUpstreamAccepted != nil {
parsed.OnUpstreamAccepted()
}
// 响应处理
var usage *ClaudeUsage
var firstTokenMs *int
var clientDisconnect bool
if reqStream {
streamResult, err := s.handleBedrockStreamingResponse(ctx, resp, c, account, startTime, reqModel)
if err != nil {
return nil, err
}
usage = streamResult.usage
firstTokenMs = streamResult.firstTokenMs
clientDisconnect = streamResult.clientDisconnect
} else {
usage, err = s.handleBedrockNonStreamingResponse(ctx, resp, c, account)
if err != nil {
return nil, err
}
}
if usage == nil {
usage = &ClaudeUsage{}
}
return &ForwardResult{
RequestID: resp.Header.Get("x-amzn-requestid"),
Usage: *usage,
Model: reqModel,
UpstreamModel: mappedModel,
Stream: reqStream,
Duration: time.Since(startTime),
FirstTokenMs: firstTokenMs,
ClientDisconnect: clientDisconnect,
}, nil
}
// executeBedrockUpstream 执行 Bedrock 上游请求(含重试逻辑)
func (s *GatewayService) executeBedrockUpstream(
ctx context.Context,
c *gin.Context,
account *Account,
body []byte,
modelID string,
region string,
stream bool,
signer *BedrockSigner,
apiKey string,
proxyURL string,
) (*http.Response, error) {
var resp *http.Response
var err error
retryStart := time.Now()
for attempt := 1; attempt <= maxRetryAttempts; attempt++ {
var upstreamReq *http.Request
if account.IsBedrockAPIKey() {
upstreamReq, err = s.buildUpstreamRequestBedrockAPIKey(ctx, body, modelID, region, stream, apiKey)
} else {
upstreamReq, err = s.buildUpstreamRequestBedrock(ctx, body, modelID, region, stream, signer)
}
if err != nil {
return nil, err
}
resp, err = s.httpUpstream.DoWithTLS(upstreamReq, proxyURL, account.ID, account.Concurrency, nil)
if err != nil {
if resp != nil && resp.Body != nil {
_ = resp.Body.Close()
}
safeErr := sanitizeUpstreamErrorMessage(err.Error())
setOpsUpstreamError(c, 0, safeErr, "")
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
Platform: account.Platform,
AccountID: account.ID,
AccountName: account.Name,
UpstreamStatusCode: 0,
UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()),
Kind: "request_error",
Message: safeErr,
})
c.JSON(http.StatusBadGateway, gin.H{
"type": "error",
"error": gin.H{
"type": "upstream_error",
"message": "Upstream request failed",
},
})
return nil, fmt.Errorf("upstream request failed: %s", safeErr)
}
if resp.StatusCode >= 400 && resp.StatusCode != 400 && s.shouldRetryUpstreamError(account, resp.StatusCode) {
if attempt < maxRetryAttempts {
elapsed := time.Since(retryStart)
if elapsed >= maxRetryElapsed {
break
}
delay := retryBackoffDelay(attempt)
remaining := maxRetryElapsed - elapsed
if delay > remaining {
delay = remaining
}
if delay <= 0 {
break
}
respBody, _ := s.readUpstreamErrorBody(resp)
_ = resp.Body.Close()
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
Platform: account.Platform,
AccountID: account.ID,
AccountName: account.Name,
UpstreamStatusCode: resp.StatusCode,
UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()),
Kind: "retry",
Message: extractUpstreamErrorMessage(respBody),
Detail: func() string {
if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody {
return truncateString(string(respBody), s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes)
}
return ""
}(),
})
logger.LegacyPrintf("service.gateway", "[Bedrock] account %d: upstream error %d, retry %d/%d after %v",
account.ID, resp.StatusCode, attempt, maxRetryAttempts, delay)
if err := sleepWithContext(ctx, delay); err != nil {
return nil, err
}
continue
}
break
}
break
}
if resp == nil || resp.Body == nil {
return nil, errors.New("upstream request failed: empty response")
}
return resp, nil
}
// handleBedrockUpstreamErrors 处理 Bedrock 上游 4xx/5xx 错误(failover + 错误响应)
func (s *GatewayService) handleBedrockUpstreamErrors(
ctx context.Context,
resp *http.Response,
c *gin.Context,
account *Account,
) (*ForwardResult, error) {
// retry exhausted + failover
if s.shouldRetryUpstreamError(account, resp.StatusCode) {
if s.shouldFailoverUpstreamError(resp.StatusCode) {
respBody, _ := s.readUpstreamErrorBody(resp)
_ = resp.Body.Close()
resp.Body = io.NopCloser(bytes.NewReader(respBody))
logger.LegacyPrintf("service.gateway", "[Bedrock] Upstream error (retry exhausted, failover): Account=%d(%s) Status=%d Body=%s",
account.ID, account.Name, resp.StatusCode, truncateString(string(respBody), 1000))
s.handleRetryExhaustedSideEffects(ctx, resp, account)
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
Platform: account.Platform,
AccountID: account.ID,
AccountName: account.Name,
UpstreamStatusCode: resp.StatusCode,
Kind: "retry_exhausted_failover",
Message: extractUpstreamErrorMessage(respBody),
})
return nil, &UpstreamFailoverError{
StatusCode: resp.StatusCode,
ResponseBody: respBody,
RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode),
}
}
return s.handleRetryExhaustedError(ctx, resp, c, account)
}
// non-retryable failover
if s.shouldFailoverUpstreamError(resp.StatusCode) {
respBody, _ := s.readUpstreamErrorBody(resp)
_ = resp.Body.Close()
resp.Body = io.NopCloser(bytes.NewReader(respBody))
s.handleFailoverSideEffects(ctx, resp, account)
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
Platform: account.Platform,
AccountID: account.ID,
AccountName: account.Name,
UpstreamStatusCode: resp.StatusCode,
Kind: "failover",
Message: extractUpstreamErrorMessage(respBody),
})
return nil, &UpstreamFailoverError{
StatusCode: resp.StatusCode,
ResponseBody: respBody,
RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode),
}
}
// other errors
return s.handleErrorResponse(ctx, resp, c, account)
}
// buildUpstreamRequestBedrock 构建 Bedrock 上游请求
func (s *GatewayService) buildUpstreamRequestBedrock(
ctx context.Context,
body []byte,
modelID string,
region string,
stream bool,
signer *BedrockSigner,
) (*http.Request, error) {
targetURL := BuildBedrockURL(region, modelID, stream)
req, err := http.NewRequestWithContext(ctx, http.MethodPost, targetURL, bytes.NewReader(body))
if err != nil {
return nil, err
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Accept", "application/json")
// SigV4 签名
if err := signer.SignRequest(ctx, req, body); err != nil {
return nil, fmt.Errorf("sign bedrock request: %w", err)
}
return req, nil
}
// buildUpstreamRequestBedrockAPIKey 构建 Bedrock API Key (Bearer Token) 上游请求
func (s *GatewayService) buildUpstreamRequestBedrockAPIKey(
ctx context.Context,
body []byte,
modelID string,
region string,
stream bool,
apiKey string,
) (*http.Request, error) {
targetURL := BuildBedrockURL(region, modelID, stream)
req, err := http.NewRequestWithContext(ctx, http.MethodPost, targetURL, bytes.NewReader(body))
if err != nil {
return nil, err
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Accept", "application/json")
req.Header.Set("Authorization", "Bearer "+apiKey)
return req, nil
}
// handleBedrockNonStreamingResponse 处理 Bedrock 非流式响应
// Bedrock InvokeModel 非流式响应的 body 格式与 Claude API 兼容
func (s *GatewayService) handleBedrockNonStreamingResponse(
ctx context.Context,
resp *http.Response,
c *gin.Context,
account *Account,
) (*ClaudeUsage, error) {
body, err := ReadUpstreamResponseBody(resp.Body, s.cfg, c, anthropicTooLargeError)
if err != nil {
return nil, err
}
// 转换 Bedrock 特有的 amazon-bedrock-invocationMetrics 为标准 Anthropic usage 格式
// 并移除该字段避免透传给客户端
body = transformBedrockInvocationMetrics(body)
usage := parseClaudeUsageFromResponseBody(body)
c.Header("Content-Type", "application/json")
if v := resp.Header.Get("x-amzn-requestid"); v != "" {
c.Header("x-request-id", v)
}
c.Data(resp.StatusCode, "application/json", body)
return usage, nil
}