antigravity_gateway_service.go 23 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655
  1. package service
  2. import (
  3. "bytes"
  4. "context"
  5. "encoding/json"
  6. "errors"
  7. "fmt"
  8. "io"
  9. "log/slog"
  10. "net/http"
  11. "strings"
  12. "sync"
  13. "time"
  14. "github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
  15. "github.com/Wei-Shaw/sub2api/internal/pkg/logger"
  16. "github.com/google/uuid"
  17. "github.com/tidwall/gjson"
  18. )
  19. const (
  20. antigravityStickySessionTTL = time.Hour
  21. antigravityMaxRetries = 3
  22. antigravityRetryBaseDelay = 1 * time.Second
  23. antigravityRetryMaxDelay = 16 * time.Second
  24. // 限流相关常量
  25. // antigravityRateLimitThreshold 限流等待/切换阈值
  26. // - 智能重试:retryDelay < 此阈值时等待后重试,>= 此阈值时直接限流模型
  27. // - 预检查:剩余限流时间 < 此阈值时等待,>= 此阈值时切换账号
  28. antigravityRateLimitThreshold = 7 * time.Second
  29. antigravitySmartRetryMinWait = 1 * time.Second // 智能重试最小等待时间
  30. antigravitySmartRetryMaxAttempts = 1 // 智能重试最大次数(仅重试 1 次,防止重复限流/长期等待)
  31. antigravityDefaultRateLimitDuration = 30 * time.Second // 默认限流时间(无 retryDelay 时使用)
  32. // MODEL_CAPACITY_EXHAUSTED 专用重试参数
  33. // 模型容量不足时,所有账号共享同一容量池,切换账号无意义
  34. // 使用固定 1s 间隔重试,最多重试 60 次
  35. antigravityModelCapacityRetryMaxAttempts = 60
  36. antigravityModelCapacityRetryWait = 1 * time.Second
  37. // Google RPC 状态和类型常量
  38. googleRPCStatusResourceExhausted = "RESOURCE_EXHAUSTED"
  39. googleRPCStatusUnavailable = "UNAVAILABLE"
  40. googleRPCTypeRetryInfo = "type.googleapis.com/google.rpc.RetryInfo"
  41. googleRPCTypeErrorInfo = "type.googleapis.com/google.rpc.ErrorInfo"
  42. googleRPCReasonModelCapacityExhausted = "MODEL_CAPACITY_EXHAUSTED"
  43. googleRPCReasonRateLimitExceeded = "RATE_LIMIT_EXCEEDED"
  44. // 单账号 503 退避重试:Service 层原地重试的最大次数
  45. // 在 handleSmartRetry 中,对于 shouldRateLimitModel(长延迟 ≥ 7s)的情况,
  46. // 多账号模式下会设限流+切换账号;但单账号模式下改为原地等待+重试。
  47. antigravitySingleAccountSmartRetryMaxAttempts = 3
  48. // 单账号 503 退避重试:原地重试时单次最大等待时间
  49. // 防止上游返回过长的 retryDelay 导致请求卡住太久
  50. antigravitySingleAccountSmartRetryMaxWait = 15 * time.Second
  51. // 单账号 503 退避重试:原地重试的总累计等待时间上限
  52. // 超过此上限将不再重试,直接返回 503
  53. antigravitySingleAccountSmartRetryTotalMaxWait = 30 * time.Second
  54. // MODEL_CAPACITY_EXHAUSTED 全局去重:重试全部失败后的 cooldown 时间
  55. antigravityModelCapacityCooldown = 10 * time.Second
  56. )
  57. // antigravityPassthroughErrorMessages 透传给客户端的错误消息白名单(小写)
  58. // 匹配时使用 strings.Contains,无需完全匹配
  59. var antigravityPassthroughErrorMessages = []string{
  60. "prompt is too long",
  61. }
  62. // MODEL_CAPACITY_EXHAUSTED 全局去重:避免多个并发请求同时对同一模型进行容量耗尽重试
  63. var (
  64. modelCapacityExhaustedMu sync.RWMutex
  65. modelCapacityExhaustedUntil = make(map[string]time.Time) // modelName -> cooldown until
  66. )
  67. const (
  68. antigravityForwardBaseURLEnv = "GATEWAY_ANTIGRAVITY_FORWARD_BASE_URL"
  69. antigravityFallbackSecondsEnv = "GATEWAY_ANTIGRAVITY_FALLBACK_COOLDOWN_SECONDS"
  70. )
  71. const antigravityProjectIDFallbackCredentialKey = "antigravity_project_id"
  72. var errAntigravityProjectIDRequired = errors.New("该 standard-tier Antigravity 账号需配置 project_id")
  73. // AntigravityAccountSwitchError 账号切换信号
  74. // 当账号限流时间超过阈值时,通知上层切换账号
  75. type AntigravityAccountSwitchError struct {
  76. OriginalAccountID int64
  77. RateLimitedModel string
  78. IsStickySession bool // 是否为粘性会话切换(决定是否缓存计费)
  79. }
  80. func (e *AntigravityAccountSwitchError) Error() string {
  81. return fmt.Sprintf("account %d model %s rate limited, need switch",
  82. e.OriginalAccountID, e.RateLimitedModel)
  83. }
  84. // IsAntigravityAccountSwitchError 检查错误是否为账号切换信号
  85. func IsAntigravityAccountSwitchError(err error) (*AntigravityAccountSwitchError, bool) {
  86. var switchErr *AntigravityAccountSwitchError
  87. if errors.As(err, &switchErr) {
  88. return switchErr, true
  89. }
  90. return nil, false
  91. }
  92. // PromptTooLongError 表示上游明确返回 prompt too long
  93. type PromptTooLongError struct {
  94. StatusCode int
  95. RequestID string
  96. Body []byte
  97. }
  98. func (e *PromptTooLongError) Error() string {
  99. return fmt.Sprintf("prompt too long: status=%d", e.StatusCode)
  100. }
  101. // AntigravityGatewayService 处理 Antigravity 平台的 API 转发
  102. type AntigravityGatewayService struct {
  103. accountRepo AccountRepository
  104. tokenProvider *AntigravityTokenProvider
  105. rateLimitService *RateLimitService
  106. httpUpstream HTTPUpstream
  107. settingService *SettingService
  108. cache GatewayCache // 用于模型级限流时清除粘性会话绑定
  109. schedulerSnapshot *SchedulerSnapshotService
  110. internal500Cache Internal500CounterCache // INTERNAL 500 渐进惩罚计数器
  111. }
  112. func (s *AntigravityGatewayService) upstreamErrorBodyReadLimit() int64 {
  113. limit := gatewayUpstreamErrorBodyReadLimit
  114. if s != nil && s.settingService != nil && s.settingService.cfg != nil && s.settingService.cfg.Gateway.LogUpstreamErrorBody && s.settingService.cfg.Gateway.LogUpstreamErrorBodyMaxBytes > int(limit) {
  115. limit = int64(s.settingService.cfg.Gateway.LogUpstreamErrorBodyMaxBytes)
  116. }
  117. return limit
  118. }
  119. func (s *AntigravityGatewayService) readUpstreamErrorBody(resp *http.Response) []byte {
  120. if resp == nil || resp.Body == nil {
  121. return nil
  122. }
  123. body, _ := io.ReadAll(io.LimitReader(resp.Body, s.upstreamErrorBodyReadLimit()))
  124. return body
  125. }
  126. func NewAntigravityGatewayService(
  127. accountRepo AccountRepository,
  128. cache GatewayCache,
  129. schedulerSnapshot *SchedulerSnapshotService,
  130. tokenProvider *AntigravityTokenProvider,
  131. rateLimitService *RateLimitService,
  132. httpUpstream HTTPUpstream,
  133. settingService *SettingService,
  134. internal500Cache Internal500CounterCache,
  135. ) *AntigravityGatewayService {
  136. return &AntigravityGatewayService{
  137. accountRepo: accountRepo,
  138. tokenProvider: tokenProvider,
  139. rateLimitService: rateLimitService,
  140. httpUpstream: httpUpstream,
  141. settingService: settingService,
  142. cache: cache,
  143. schedulerSnapshot: schedulerSnapshot,
  144. internal500Cache: internal500Cache,
  145. }
  146. }
  147. // GetTokenProvider 返回 token provider
  148. func (s *AntigravityGatewayService) GetTokenProvider() *AntigravityTokenProvider {
  149. return s.tokenProvider
  150. }
  151. // getLogConfig 获取上游错误日志配置
  152. // 返回是否记录日志体和最大字节数
  153. func (s *AntigravityGatewayService) getLogConfig() (logBody bool, maxBytes int) {
  154. maxBytes = 2048 // 默认值
  155. if s.settingService == nil || s.settingService.cfg == nil {
  156. return false, maxBytes
  157. }
  158. cfg := s.settingService.cfg.Gateway
  159. if cfg.LogUpstreamErrorBodyMaxBytes > 0 {
  160. maxBytes = cfg.LogUpstreamErrorBodyMaxBytes
  161. }
  162. return cfg.LogUpstreamErrorBody, maxBytes
  163. }
  164. // getUpstreamErrorDetail 获取上游错误详情(用于日志记录)
  165. func (s *AntigravityGatewayService) getUpstreamErrorDetail(body []byte) string {
  166. logBody, maxBytes := s.getLogConfig()
  167. if !logBody {
  168. return ""
  169. }
  170. return truncateString(string(body), maxBytes)
  171. }
  172. // checkErrorPolicy nil 安全的包装
  173. func (s *AntigravityGatewayService) checkErrorPolicy(ctx context.Context, account *Account, statusCode int, body []byte, requestedModel ...string) ErrorPolicyResult {
  174. if s.rateLimitService == nil {
  175. return ErrorPolicyNone
  176. }
  177. return s.rateLimitService.CheckErrorPolicy(ctx, account, statusCode, body, firstRequestedModel(requestedModel))
  178. }
  179. // applyErrorPolicy 应用错误策略结果,返回是否应终止当前循环及应返回的状态码。
  180. // ErrorPolicySkipped 时 outStatus 为 500(前端约定:未命中的错误返回 500)。
  181. func (s *AntigravityGatewayService) applyErrorPolicy(p antigravityRetryLoopParams, statusCode int, headers http.Header, respBody []byte) (handled bool, outStatus int, retErr error) {
  182. modelKey := resolveFinalAntigravityModelKey(p.ctx, p.account, p.requestedModel)
  183. switch s.checkErrorPolicy(p.ctx, p.account, statusCode, respBody, modelKey) {
  184. case ErrorPolicySkipped:
  185. if s.handleAntigravityModelRateLimitBeforePolicy(p, statusCode, headers, respBody) {
  186. return true, statusCode, nil
  187. }
  188. return true, http.StatusInternalServerError, nil
  189. case ErrorPolicyMatched:
  190. if s.handleAntigravityModelRateLimitBeforePolicy(p, statusCode, headers, respBody) {
  191. return true, statusCode, nil
  192. }
  193. _ = p.handleError(p.ctx, p.prefix, p.account, statusCode, headers, respBody,
  194. p.requestedModel, p.groupID, p.sessionHash, p.isStickySession)
  195. return true, statusCode, nil
  196. case ErrorPolicyTempUnscheduled:
  197. slog.Info("temp_unschedulable_matched",
  198. "prefix", p.prefix, "status_code", statusCode, "account_id", p.account.ID)
  199. return true, statusCode, &AntigravityAccountSwitchError{OriginalAccountID: p.account.ID, RateLimitedModel: p.requestedModel, IsStickySession: p.isStickySession}
  200. }
  201. return false, statusCode, nil
  202. }
  203. func (s *AntigravityGatewayService) handleAntigravityModelRateLimitBeforePolicy(p antigravityRetryLoopParams, statusCode int, headers http.Header, respBody []byte) bool {
  204. if statusCode != http.StatusTooManyRequests && statusCode != http.StatusServiceUnavailable {
  205. return false
  206. }
  207. if p.account == nil || p.account.Platform != PlatformAntigravity {
  208. return false
  209. }
  210. _, shouldRateLimitModel, waitDuration, modelName, isModelCapacityExhausted := shouldTriggerAntigravitySmartRetry(p.account, respBody)
  211. if isModelCapacityExhausted || !shouldRateLimitModel || strings.TrimSpace(modelName) == "" {
  212. return false
  213. }
  214. rateLimitDuration := waitDuration
  215. if rateLimitDuration <= 0 {
  216. rateLimitDuration = antigravityDefaultRateLimitDuration
  217. }
  218. resetAt := time.Now().Add(rateLimitDuration)
  219. if !s.setAntigravityModelRateLimits(p.ctx, p.accountRepo, p.account, modelName, p.prefix, statusCode, resetAt, false) {
  220. return false
  221. }
  222. s.clearStickySession(p.ctx, p.groupID, p.sessionHash)
  223. logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d model_rate_limited_before_error_policy model=%s account=%d reset_in=%v",
  224. p.prefix, statusCode, modelName, p.account.ID, rateLimitDuration)
  225. return true
  226. }
  227. // mapAntigravityModel 获取映射后的模型名
  228. // 完全依赖映射配置:账户映射(通配符)→ 默认映射兜底(DefaultAntigravityModelMapping)
  229. // 注意:返回空字符串表示模型不被支持,调度时会过滤掉该账号
  230. func mapAntigravityModel(account *Account, requestedModel string) string {
  231. if account == nil {
  232. return ""
  233. }
  234. requestedModel = strings.TrimPrefix(requestedModel, "models/")
  235. // 获取映射表(未配置时自动使用 DefaultAntigravityModelMapping)
  236. mapping := account.GetModelMapping()
  237. if len(mapping) == 0 {
  238. return "" // 无映射配置(非 Antigravity 平台)
  239. }
  240. // 通过映射表查询(支持精确匹配 + 通配符)
  241. mapped := account.GetMappedModel(requestedModel)
  242. // 判断是否映射成功(mapped != requestedModel 说明找到了映射规则)
  243. if mapped != requestedModel {
  244. return mapped
  245. }
  246. // 如果 mapped == requestedModel,检查是否在映射表中配置(精确或通配符)
  247. // 这区分两种情况:
  248. // 1. 映射表中有 "model-a": "model-a"(显式透传)→ 返回 model-a
  249. // 2. 通配符匹配 "claude-*": "claude-sonnet-4-5" 恰好目标等于请求名 → 返回 model-a
  250. // 3. 映射表中没有 model-a 的配置 → 返回空(不支持)
  251. if account.IsModelSupported(requestedModel) {
  252. return requestedModel
  253. }
  254. // 未在映射表中配置的模型,返回空字符串(不支持)
  255. return ""
  256. }
  257. // getMappedModel 获取映射后的模型名
  258. // 完全依赖映射配置:账户映射(通配符)→ 默认映射兜底
  259. func (s *AntigravityGatewayService) getMappedModel(account *Account, requestedModel string) string {
  260. return mapAntigravityModel(account, requestedModel)
  261. }
  262. func resolveAntigravityProjectID(account *Account) (string, error) {
  263. if account == nil {
  264. return "", errAntigravityProjectIDRequired
  265. }
  266. if projectID := strings.TrimSpace(account.GetCredential("project_id")); projectID != "" {
  267. return projectID, nil
  268. }
  269. if projectID := strings.TrimSpace(account.GetCredential(antigravityProjectIDFallbackCredentialKey)); projectID != "" {
  270. return projectID, nil
  271. }
  272. if projectID := strings.TrimSpace(account.GetExtraString(antigravityProjectIDFallbackCredentialKey)); projectID != "" {
  273. return projectID, nil
  274. }
  275. return "", errAntigravityProjectIDRequired
  276. }
  277. // applyThinkingModelSuffix 根据 thinking 配置调整模型名
  278. // 当映射结果是 claude-sonnet-4-5 且请求开启了 thinking 时,改为 claude-sonnet-4-5-thinking
  279. func applyThinkingModelSuffix(mappedModel string, thinkingEnabled bool) string {
  280. if !thinkingEnabled {
  281. return mappedModel
  282. }
  283. if mappedModel == "claude-sonnet-4-5" {
  284. return "claude-sonnet-4-5-thinking"
  285. }
  286. return mappedModel
  287. }
  288. // IsModelSupported 检查模型是否被支持
  289. // 所有 claude- 和 gemini- 前缀的模型都能通过映射或透传支持
  290. func (s *AntigravityGatewayService) IsModelSupported(requestedModel string) bool {
  291. return strings.HasPrefix(requestedModel, "claude-") ||
  292. strings.HasPrefix(requestedModel, "gemini-")
  293. }
  294. // TestConnectionResult 测试连接结果
  295. type TestConnectionResult struct {
  296. Text string // 响应文本
  297. MappedModel string // 实际使用的模型
  298. }
  299. // TestConnection 测试 Antigravity 账号连接。
  300. // 复用 antigravityRetryLoop 的完整重试 / credits overages / 智能重试逻辑,
  301. // 与真实调度行为一致。差异:不做账号切换(测试指定账号)、不记录 ops 错误。
  302. func (s *AntigravityGatewayService) TestConnection(ctx context.Context, account *Account, modelID string) (*TestConnectionResult, error) {
  303. // 获取 token
  304. if s.tokenProvider == nil {
  305. return nil, errors.New("antigravity token provider not configured")
  306. }
  307. accessToken, err := s.tokenProvider.GetAccessToken(ctx, account)
  308. if err != nil {
  309. return nil, fmt.Errorf("获取 access_token 失败: %w", err)
  310. }
  311. projectID, err := resolveAntigravityProjectID(account)
  312. if err != nil {
  313. return nil, err
  314. }
  315. // 模型映射
  316. mappedModel := s.getMappedModel(account, modelID)
  317. if mappedModel == "" {
  318. return nil, fmt.Errorf("model %s not in whitelist", modelID)
  319. }
  320. // 构建请求体
  321. var requestBody []byte
  322. if strings.HasPrefix(modelID, "gemini-") {
  323. requestBody, err = s.buildGeminiTestRequest(projectID, mappedModel)
  324. } else {
  325. requestBody, err = s.buildClaudeTestRequest(projectID, mappedModel)
  326. }
  327. if err != nil {
  328. return nil, fmt.Errorf("构建请求失败: %w", err)
  329. }
  330. if intelligentPrompt(ctx) != "" {
  331. requestBody, err = intelligentAntigravityBody(ctx, requestBody)
  332. if err != nil {
  333. return nil, err
  334. }
  335. }
  336. // 代理 URL
  337. proxyURL := ""
  338. if account.ProxyID != nil && account.Proxy != nil {
  339. proxyURL = account.Proxy.URL()
  340. }
  341. // 复用 antigravityRetryLoop:完整的重试 / credits overages / 智能重试
  342. prefix := fmt.Sprintf("[antigravity-Test] account=%d(%s)", account.ID, account.Name)
  343. p := antigravityRetryLoopParams{
  344. ctx: ctx,
  345. prefix: prefix,
  346. account: account,
  347. proxyURL: proxyURL,
  348. accessToken: accessToken,
  349. action: "streamGenerateContent",
  350. body: requestBody,
  351. c: nil, // 无 gin.Context → 跳过 ops 追踪
  352. httpUpstream: s.httpUpstream,
  353. settingService: s.settingService,
  354. accountRepo: s.accountRepo,
  355. requestedModel: modelID,
  356. handleError: testConnectionHandleError,
  357. }
  358. var result *antigravityRetryLoopResult
  359. if run := intelligentContext(ctx); run != nil {
  360. // Capability observations make one request. The production retry loop
  361. // may apply account error policies, change cache scheduling and enable
  362. // credit fallback; those side effects do not belong to a test result.
  363. var resp *http.Response
  364. resp, err = s.runIntelligentAntigravity(ctx, account, accessToken, requestBody)
  365. result = &antigravityRetryLoopResult{resp: resp}
  366. } else {
  367. result, err = s.antigravityRetryLoop(p)
  368. }
  369. if err != nil {
  370. // AccountSwitchError → 测试时不切换账号,返回友好提示
  371. var switchErr *AntigravityAccountSwitchError
  372. if errors.As(err, &switchErr) {
  373. return nil, fmt.Errorf("该账号模型 %s 当前限流中,请稍后重试", switchErr.RateLimitedModel)
  374. }
  375. return nil, err
  376. }
  377. if result == nil || result.resp == nil {
  378. return nil, errors.New("upstream returned empty response")
  379. }
  380. defer func() { _ = result.resp.Body.Close() }()
  381. respBody, err := io.ReadAll(io.LimitReader(result.resp.Body, s.upstreamErrorBodyReadLimit()))
  382. if err != nil {
  383. return nil, fmt.Errorf("读取响应失败: %w", err)
  384. }
  385. if result.resp.StatusCode >= 400 {
  386. return nil, fmt.Errorf("API 返回 %d: %s", result.resp.StatusCode, string(respBody))
  387. }
  388. text := extractTextFromSSEResponse(respBody)
  389. return &TestConnectionResult{Text: text, MappedModel: mappedModel}, nil
  390. }
  391. // testConnectionHandleError 是 TestConnection 使用的轻量 handleError 回调。
  392. // 仅记录日志,不做 ops 错误追踪或粘性会话清除。
  393. func testConnectionHandleError(
  394. _ context.Context, prefix string, account *Account,
  395. statusCode int, _ http.Header, body []byte,
  396. requestedModel string, _ int64, _ string, _ bool,
  397. ) *handleModelRateLimitResult {
  398. logger.LegacyPrintf("service.antigravity_gateway",
  399. "%s test_handle_error status=%d model=%s account=%d body=%s",
  400. prefix, statusCode, requestedModel, account.ID, truncateForLog(body, 200))
  401. return nil
  402. }
  403. // buildGeminiTestRequest 构建 Gemini 格式测试请求
  404. // 使用最小 token 消耗:输入 "." + maxOutputTokens: 1
  405. func (s *AntigravityGatewayService) buildGeminiTestRequest(projectID, model string) ([]byte, error) {
  406. payload := map[string]any{
  407. "contents": []map[string]any{
  408. {
  409. "role": "user",
  410. "parts": []map[string]any{
  411. {"text": "."},
  412. },
  413. },
  414. },
  415. // Antigravity 上游要求必须包含身份提示词
  416. "systemInstruction": map[string]any{
  417. "parts": []map[string]any{
  418. {"text": antigravity.GetDefaultIdentityPatch()},
  419. },
  420. },
  421. "generationConfig": map[string]any{
  422. "maxOutputTokens": 1,
  423. },
  424. }
  425. payloadBytes, _ := json.Marshal(payload)
  426. return s.wrapV1InternalRequest(projectID, model, payloadBytes)
  427. }
  428. // buildClaudeTestRequest 构建 Claude 格式测试请求并转换为 Gemini 格式
  429. // 使用最小 token 消耗:输入 "." + MaxTokens: 1
  430. func (s *AntigravityGatewayService) buildClaudeTestRequest(projectID, mappedModel string) ([]byte, error) {
  431. claudeReq := &antigravity.ClaudeRequest{
  432. Model: mappedModel,
  433. Messages: []antigravity.ClaudeMessage{
  434. {
  435. Role: "user",
  436. Content: json.RawMessage(`"."`),
  437. },
  438. },
  439. MaxTokens: 1,
  440. Stream: false,
  441. }
  442. return antigravity.TransformClaudeToGemini(claudeReq, projectID, mappedModel)
  443. }
  444. func (s *AntigravityGatewayService) getClaudeTransformOptions(ctx context.Context) antigravity.TransformOptions {
  445. opts := antigravity.DefaultTransformOptions()
  446. if s.settingService == nil {
  447. return opts
  448. }
  449. opts.EnableIdentityPatch = s.settingService.IsIdentityPatchEnabled(ctx)
  450. opts.IdentityPatch = s.settingService.GetIdentityPatchPrompt(ctx)
  451. return opts
  452. }
  453. // extractTextFromSSEResponse 从 SSE 流式响应中提取文本
  454. func extractTextFromSSEResponse(respBody []byte) string {
  455. var texts []string
  456. lines := bytes.Split(respBody, []byte("\n"))
  457. for _, line := range lines {
  458. line = bytes.TrimSpace(line)
  459. if len(line) == 0 {
  460. continue
  461. }
  462. // 跳过 SSE 前缀
  463. if bytes.HasPrefix(line, []byte("data:")) {
  464. line = bytes.TrimPrefix(line, []byte("data:"))
  465. line = bytes.TrimSpace(line)
  466. }
  467. // 跳过非 JSON 行
  468. if len(line) == 0 || line[0] != '{' {
  469. continue
  470. }
  471. // 解析 JSON
  472. var data map[string]any
  473. if err := json.Unmarshal(line, &data); err != nil {
  474. continue
  475. }
  476. // 尝试从 response.candidates[0].content.parts[].text 提取
  477. response, ok := data["response"].(map[string]any)
  478. if !ok {
  479. // 尝试直接从 candidates 提取(某些响应格式)
  480. response = data
  481. }
  482. candidates, ok := response["candidates"].([]any)
  483. if !ok || len(candidates) == 0 {
  484. continue
  485. }
  486. candidate, ok := candidates[0].(map[string]any)
  487. if !ok {
  488. continue
  489. }
  490. content, ok := candidate["content"].(map[string]any)
  491. if !ok {
  492. continue
  493. }
  494. parts, ok := content["parts"].([]any)
  495. if !ok {
  496. continue
  497. }
  498. for _, part := range parts {
  499. if partMap, ok := part.(map[string]any); ok {
  500. if text, ok := partMap["text"].(string); ok && text != "" {
  501. texts = append(texts, text)
  502. }
  503. }
  504. }
  505. }
  506. return strings.Join(texts, "")
  507. }
  508. // injectIdentityPatchToGeminiRequest 为 Gemini 格式请求注入身份提示词
  509. // 如果请求中已包含 "You are Antigravity" 则不重复注入
  510. func injectIdentityPatchToGeminiRequest(body []byte) ([]byte, error) {
  511. var request map[string]any
  512. if err := json.Unmarshal(body, &request); err != nil {
  513. return nil, fmt.Errorf("解析 Gemini 请求失败: %w", err)
  514. }
  515. // 检查现有 systemInstruction 是否已包含身份提示词
  516. if sysInst, ok := request["systemInstruction"].(map[string]any); ok {
  517. if parts, ok := sysInst["parts"].([]any); ok {
  518. for _, part := range parts {
  519. if partMap, ok := part.(map[string]any); ok {
  520. if text, ok := partMap["text"].(string); ok {
  521. if strings.Contains(text, "You are Antigravity") {
  522. // 已包含身份提示词,直接返回原始请求
  523. return body, nil
  524. }
  525. }
  526. }
  527. }
  528. }
  529. }
  530. // 获取默认身份提示词
  531. identityPatch := antigravity.GetDefaultIdentityPatch()
  532. // 构建新的 systemInstruction
  533. newPart := map[string]any{"text": identityPatch}
  534. if existing, ok := request["systemInstruction"].(map[string]any); ok {
  535. // 已有 systemInstruction,在开头插入身份提示词
  536. if parts, ok := existing["parts"].([]any); ok {
  537. existing["parts"] = append([]any{newPart}, parts...)
  538. } else {
  539. existing["parts"] = []any{newPart}
  540. }
  541. } else {
  542. // 没有 systemInstruction,创建新的
  543. request["systemInstruction"] = map[string]any{
  544. "parts": []any{newPart},
  545. }
  546. }
  547. return json.Marshal(request)
  548. }
  549. // wrapV1InternalRequest 包装请求为 v1internal 格式
  550. func (s *AntigravityGatewayService) wrapV1InternalRequest(projectID, model string, originalBody []byte) ([]byte, error) {
  551. var request any
  552. if err := json.Unmarshal(originalBody, &request); err != nil {
  553. return nil, fmt.Errorf("解析请求体失败: %w", err)
  554. }
  555. projectID = strings.TrimSpace(projectID)
  556. if projectID == "" {
  557. return nil, errAntigravityProjectIDRequired
  558. }
  559. wrapped := map[string]any{
  560. "project": projectID,
  561. "requestId": "agent-" + uuid.New().String(),
  562. "userAgent": "antigravity", // 固定值,与官方客户端一致
  563. "requestType": "agent",
  564. "model": model,
  565. "request": request,
  566. }
  567. return json.Marshal(wrapped)
  568. }
  569. // unwrapV1InternalResponse 解包 v1internal 响应
  570. // 使用 gjson 零拷贝提取 response 字段,避免 Unmarshal+Marshal 双重开销
  571. func (s *AntigravityGatewayService) unwrapV1InternalResponse(body []byte) ([]byte, error) {
  572. result := gjson.GetBytes(body, "response")
  573. if result.Exists() {
  574. return []byte(result.Raw), nil
  575. }
  576. return body, nil
  577. }