gateway_bedrock.go 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403
  1. package service
  2. // 本文件由 gateway_service.go 纯移动拆分而来:Bedrock 上游转发(CC 兼容转换、
  3. // 请求构建、错误处理与非流式响应)。仅做代码搬迁,无任何行为变更。
  4. import (
  5. "bytes"
  6. "context"
  7. "errors"
  8. "fmt"
  9. "io"
  10. "net/http"
  11. "strings"
  12. "time"
  13. "github.com/Wei-Shaw/sub2api/internal/pkg/logger"
  14. "github.com/gin-gonic/gin"
  15. )
  16. // ApplyBedrockCCCompat 应用 Bedrock CC 兼容转换(渠道级模型映射后调用)
  17. // 清理 body 中 Anthropic API 专有字段、修复 thinking/tool_use ID、过滤 beta token,
  18. // 同时过滤 HTTP header 中的 anthropic-beta(防止 Passthrough 路径透传不支持的 token)。
  19. func (s *GatewayService) ApplyBedrockCCCompat(c *gin.Context, body []byte, model string, account *Account, groupID *int64) []byte {
  20. if !s.isBedrockCCCompatEnabled(c.Request.Context(), account, groupID) {
  21. return body
  22. }
  23. body = sanitizeBedrockCCFields(body)
  24. body = sanitizeBedrockThinking(body, model)
  25. body = sanitizeBedrockToolUseIDs(body)
  26. body = sanitizeBedrockCCBetaTokens(body, model)
  27. // 过滤 HTTP header 中的 anthropic-beta,只保留 Bedrock 支持的 token
  28. if betaHeader := c.GetHeader("anthropic-beta"); betaHeader != "" {
  29. if filtered := ResolveBedrockBetaTokens(betaHeader, body, model); len(filtered) > 0 {
  30. c.Request.Header.Set("anthropic-beta", strings.Join(filtered, ", "))
  31. } else {
  32. c.Request.Header.Del("anthropic-beta")
  33. }
  34. }
  35. return body
  36. }
  37. // isBedrockCCCompatEnabled 检查渠道是否启用了 Bedrock CC 兼容模式
  38. func (s *GatewayService) isBedrockCCCompatEnabled(ctx context.Context, account *Account, groupID *int64) bool {
  39. if groupID == nil || s.channelService == nil {
  40. return false
  41. }
  42. ch, err := s.channelService.GetChannelForGroup(ctx, *groupID)
  43. if err != nil || ch == nil {
  44. return false
  45. }
  46. return ch.IsBedrockCCCompatEnabled(account.Platform)
  47. }
  48. // forwardBedrock 转发请求到 AWS Bedrock
  49. func (s *GatewayService) forwardBedrock(
  50. ctx context.Context,
  51. c *gin.Context,
  52. account *Account,
  53. parsed *ParsedRequest,
  54. startTime time.Time,
  55. ) (*ForwardResult, error) {
  56. reqModel := parsed.Model
  57. reqStream := parsed.Stream
  58. body := parsed.Body.Bytes()
  59. region := bedrockRuntimeRegion(account)
  60. mappedModel, ok := ResolveBedrockModelID(account, reqModel)
  61. if !ok {
  62. return nil, fmt.Errorf("unsupported bedrock model: %s", reqModel)
  63. }
  64. if mappedModel != reqModel {
  65. logger.LegacyPrintf("service.gateway", "[Bedrock] Model mapping: %s -> %s (account: %s)", reqModel, mappedModel, account.Name)
  66. }
  67. betaHeader := ""
  68. if c != nil && c.Request != nil {
  69. betaHeader = c.GetHeader("anthropic-beta")
  70. }
  71. // 准备请求体(注入 anthropic_version/anthropic_beta,移除 Bedrock 不支持的字段,清理 cache_control)
  72. betaTokens, err := s.resolveBedrockBetaTokensForRequest(ctx, account, betaHeader, body, mappedModel)
  73. if err != nil {
  74. return nil, err
  75. }
  76. bedrockBody, err := PrepareBedrockRequestBodyWithTokens(body, mappedModel, betaTokens, false)
  77. if err != nil {
  78. return nil, fmt.Errorf("prepare bedrock request body: %w", err)
  79. }
  80. proxyURL := ""
  81. if account.ProxyID != nil && account.Proxy != nil {
  82. proxyURL = account.Proxy.URL()
  83. }
  84. logger.LegacyPrintf("service.gateway", "[Bedrock] 命中 Bedrock 分支: account=%d name=%s model=%s->%s stream=%v",
  85. account.ID, account.Name, reqModel, mappedModel, reqStream)
  86. // 根据账号类型选择认证方式
  87. var signer *BedrockSigner
  88. var bedrockAPIKey string
  89. if account.IsBedrockAPIKey() {
  90. bedrockAPIKey = account.GetCredential("api_key")
  91. if bedrockAPIKey == "" {
  92. return nil, fmt.Errorf("api_key not found in bedrock credentials")
  93. }
  94. } else {
  95. signer, err = NewBedrockSignerFromAccount(account)
  96. if err != nil {
  97. return nil, fmt.Errorf("create bedrock signer: %w", err)
  98. }
  99. }
  100. // 执行上游请求(含重试)
  101. resp, err := s.executeBedrockUpstream(ctx, c, account, bedrockBody, mappedModel, region, reqStream, signer, bedrockAPIKey, proxyURL)
  102. if err != nil {
  103. return nil, err
  104. }
  105. defer func() { _ = resp.Body.Close() }()
  106. // 将 Bedrock 的 x-amzn-requestid 映射到 x-request-id,
  107. // 使通用错误处理函数(handleErrorResponse、handleRetryExhaustedError)能正确提取 AWS request ID。
  108. if awsReqID := resp.Header.Get("x-amzn-requestid"); awsReqID != "" && resp.Header.Get("x-request-id") == "" {
  109. resp.Header.Set("x-request-id", awsReqID)
  110. }
  111. // 错误/failover 处理
  112. if resp.StatusCode >= 400 {
  113. return s.handleBedrockUpstreamErrors(ctx, resp, c, account)
  114. }
  115. // Bedrock 分支绕过通用 Forward 成功路径,这里保持上游接受回调语义一致。
  116. if parsed.OnUpstreamAccepted != nil {
  117. parsed.OnUpstreamAccepted()
  118. }
  119. // 响应处理
  120. var usage *ClaudeUsage
  121. var firstTokenMs *int
  122. var clientDisconnect bool
  123. if reqStream {
  124. streamResult, err := s.handleBedrockStreamingResponse(ctx, resp, c, account, startTime, reqModel)
  125. if err != nil {
  126. return nil, err
  127. }
  128. usage = streamResult.usage
  129. firstTokenMs = streamResult.firstTokenMs
  130. clientDisconnect = streamResult.clientDisconnect
  131. } else {
  132. usage, err = s.handleBedrockNonStreamingResponse(ctx, resp, c, account)
  133. if err != nil {
  134. return nil, err
  135. }
  136. }
  137. if usage == nil {
  138. usage = &ClaudeUsage{}
  139. }
  140. return &ForwardResult{
  141. RequestID: resp.Header.Get("x-amzn-requestid"),
  142. UpstreamHeaders: resp.Header,
  143. Usage: *usage,
  144. Model: reqModel,
  145. UpstreamModel: mappedModel,
  146. Stream: reqStream,
  147. Duration: time.Since(startTime),
  148. FirstTokenMs: firstTokenMs,
  149. ClientDisconnect: clientDisconnect,
  150. }, nil
  151. }
  152. // executeBedrockUpstream 执行 Bedrock 上游请求(含重试逻辑)
  153. func (s *GatewayService) executeBedrockUpstream(
  154. ctx context.Context,
  155. c *gin.Context,
  156. account *Account,
  157. body []byte,
  158. modelID string,
  159. region string,
  160. stream bool,
  161. signer *BedrockSigner,
  162. apiKey string,
  163. proxyURL string,
  164. ) (*http.Response, error) {
  165. var resp *http.Response
  166. var err error
  167. retryStart := time.Now()
  168. for attempt := 1; attempt <= maxRetryAttempts; attempt++ {
  169. var upstreamReq *http.Request
  170. if account.IsBedrockAPIKey() {
  171. upstreamReq, err = s.buildUpstreamRequestBedrockAPIKey(ctx, body, modelID, region, stream, apiKey)
  172. } else {
  173. upstreamReq, err = s.buildUpstreamRequestBedrock(ctx, body, modelID, region, stream, signer)
  174. }
  175. if err != nil {
  176. return nil, err
  177. }
  178. resp, err = s.httpUpstream.DoWithTLS(WithAccountTrafficRequest(upstreamReq, account), proxyURL, account.ID, account.Concurrency, nil)
  179. if err != nil {
  180. if resp != nil && resp.Body != nil {
  181. _ = resp.Body.Close()
  182. }
  183. return nil, s.handleUpstreamTransportError(ctx, c, account, err, OpsUpstreamErrorEvent{
  184. UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()),
  185. })
  186. }
  187. if resp.StatusCode >= 400 && resp.StatusCode != 400 && s.shouldRetryUpstreamError(account, resp.StatusCode) {
  188. if attempt < maxRetryAttempts {
  189. elapsed := time.Since(retryStart)
  190. if elapsed >= maxRetryElapsed {
  191. break
  192. }
  193. delay := retryBackoffDelay(attempt)
  194. remaining := maxRetryElapsed - elapsed
  195. if delay > remaining {
  196. delay = remaining
  197. }
  198. if delay <= 0 {
  199. break
  200. }
  201. respBody, _ := s.readUpstreamErrorBody(resp)
  202. _ = resp.Body.Close()
  203. appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
  204. ProxyID: opsUpstreamProxyID(account),
  205. ProxyName: opsUpstreamProxyName(account),
  206. Platform: account.Platform,
  207. AccountID: account.ID,
  208. AccountName: account.Name,
  209. UpstreamStatusCode: resp.StatusCode,
  210. UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()),
  211. Kind: "retry",
  212. Message: extractUpstreamErrorMessage(respBody),
  213. Detail: func() string {
  214. if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody {
  215. return truncateString(string(respBody), s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes)
  216. }
  217. return ""
  218. }(),
  219. })
  220. logger.LegacyPrintf("service.gateway", "[Bedrock] account %d: upstream error %d, retry %d/%d after %v",
  221. account.ID, resp.StatusCode, attempt, maxRetryAttempts, delay)
  222. if err := sleepWithContext(ctx, delay); err != nil {
  223. return nil, err
  224. }
  225. continue
  226. }
  227. break
  228. }
  229. break
  230. }
  231. if resp == nil || resp.Body == nil {
  232. return nil, errors.New("upstream request failed: empty response")
  233. }
  234. return resp, nil
  235. }
  236. // handleBedrockUpstreamErrors 处理 Bedrock 上游 4xx/5xx 错误(failover + 错误响应)
  237. func (s *GatewayService) handleBedrockUpstreamErrors(
  238. ctx context.Context,
  239. resp *http.Response,
  240. c *gin.Context,
  241. account *Account,
  242. ) (*ForwardResult, error) {
  243. // retry exhausted + failover
  244. if s.shouldRetryUpstreamError(account, resp.StatusCode) {
  245. if s.shouldFailoverUpstreamError(resp.StatusCode) {
  246. respBody, _ := s.readUpstreamErrorBody(resp)
  247. _ = resp.Body.Close()
  248. resp.Body = io.NopCloser(bytes.NewReader(respBody))
  249. logger.LegacyPrintf("service.gateway", "[Bedrock] Upstream error (retry exhausted, failover): Account=%d(%s) Status=%d Body=%s",
  250. account.ID, account.Name, resp.StatusCode, truncateString(string(respBody), 1000))
  251. s.handleRetryExhaustedSideEffects(ctx, resp, account)
  252. appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
  253. ProxyID: opsUpstreamProxyID(account),
  254. ProxyName: opsUpstreamProxyName(account),
  255. Platform: account.Platform,
  256. AccountID: account.ID,
  257. AccountName: account.Name,
  258. UpstreamStatusCode: resp.StatusCode,
  259. Kind: "retry_exhausted_failover",
  260. Message: extractUpstreamErrorMessage(respBody),
  261. })
  262. return nil, &UpstreamFailoverError{
  263. StatusCode: resp.StatusCode,
  264. ResponseBody: respBody,
  265. RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode),
  266. }
  267. }
  268. return s.handleRetryExhaustedError(ctx, resp, c, account)
  269. }
  270. // non-retryable failover
  271. if s.shouldFailoverUpstreamError(resp.StatusCode) {
  272. respBody, _ := s.readUpstreamErrorBody(resp)
  273. _ = resp.Body.Close()
  274. resp.Body = io.NopCloser(bytes.NewReader(respBody))
  275. s.handleFailoverSideEffects(ctx, resp, account)
  276. appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
  277. ProxyID: opsUpstreamProxyID(account),
  278. ProxyName: opsUpstreamProxyName(account),
  279. Platform: account.Platform,
  280. AccountID: account.ID,
  281. AccountName: account.Name,
  282. UpstreamStatusCode: resp.StatusCode,
  283. Kind: "failover",
  284. Message: extractUpstreamErrorMessage(respBody),
  285. })
  286. return nil, &UpstreamFailoverError{
  287. StatusCode: resp.StatusCode,
  288. ResponseBody: respBody,
  289. RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode),
  290. }
  291. }
  292. // other errors
  293. return s.handleErrorResponse(ctx, resp, c, account)
  294. }
  295. // buildUpstreamRequestBedrock 构建 Bedrock 上游请求
  296. func (s *GatewayService) buildUpstreamRequestBedrock(
  297. ctx context.Context,
  298. body []byte,
  299. modelID string,
  300. region string,
  301. stream bool,
  302. signer *BedrockSigner,
  303. ) (*http.Request, error) {
  304. targetURL := BuildBedrockURL(region, modelID, stream)
  305. req, err := http.NewRequestWithContext(ctx, http.MethodPost, targetURL, bytes.NewReader(body))
  306. if err != nil {
  307. return nil, err
  308. }
  309. req.Header.Set("Content-Type", "application/json")
  310. req.Header.Set("Accept", "application/json")
  311. // SigV4 签名
  312. if err := signer.SignRequest(ctx, req, body); err != nil {
  313. return nil, fmt.Errorf("sign bedrock request: %w", err)
  314. }
  315. return req, nil
  316. }
  317. // buildUpstreamRequestBedrockAPIKey 构建 Bedrock API Key (Bearer Token) 上游请求
  318. func (s *GatewayService) buildUpstreamRequestBedrockAPIKey(
  319. ctx context.Context,
  320. body []byte,
  321. modelID string,
  322. region string,
  323. stream bool,
  324. apiKey string,
  325. ) (*http.Request, error) {
  326. targetURL := BuildBedrockURL(region, modelID, stream)
  327. req, err := http.NewRequestWithContext(ctx, http.MethodPost, targetURL, bytes.NewReader(body))
  328. if err != nil {
  329. return nil, err
  330. }
  331. req.Header.Set("Content-Type", "application/json")
  332. req.Header.Set("Accept", "application/json")
  333. req.Header.Set("Authorization", "Bearer "+apiKey)
  334. return req, nil
  335. }
  336. // handleBedrockNonStreamingResponse 处理 Bedrock 非流式响应
  337. // Bedrock InvokeModel 非流式响应的 body 格式与 Claude API 兼容
  338. func (s *GatewayService) handleBedrockNonStreamingResponse(
  339. ctx context.Context,
  340. resp *http.Response,
  341. c *gin.Context,
  342. account *Account,
  343. ) (*ClaudeUsage, error) {
  344. body, err := ReadUpstreamResponseBody(resp.Body, s.cfg, c, anthropicTooLargeError)
  345. if err != nil {
  346. return nil, err
  347. }
  348. // 转换 Bedrock 特有的 amazon-bedrock-invocationMetrics 为标准 Anthropic usage 格式
  349. // 并移除该字段避免透传给客户端
  350. body = transformBedrockInvocationMetrics(body)
  351. usage := parseClaudeUsageFromResponseBody(body)
  352. c.Header("Content-Type", "application/json")
  353. if v := resp.Header.Get("x-amzn-requestid"); v != "" {
  354. c.Header("x-request-id", v)
  355. }
  356. c.Data(resp.StatusCode, "application/json", body)
  357. return usage, nil
  358. }