| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650165116521653165416551656165716581659166016611662166316641665166616671668166916701671167216731674167516761677167816791680168116821683168416851686168716881689169016911692169316941695169616971698169917001701170217031704170517061707170817091710171117121713171417151716171717181719172017211722172317241725172617271728172917301731173217331734173517361737173817391740174117421743174417451746174717481749175017511752175317541755175617571758175917601761176217631764176517661767176817691770177117721773177417751776177717781779178017811782178317841785178617871788178917901791179217931794179517961797179817991800180118021803180418051806180718081809181018111812181318141815181618171818181918201821182218231824182518261827182818291830183118321833183418351836183718381839184018411842184318441845184618471848184918501851185218531854185518561857185818591860186118621863186418651866186718681869187018711872187318741875187618771878187918801881188218831884188518861887188818891890189118921893189418951896189718981899190019011902190319041905190619071908190919101911191219131914191519161917191819191920192119221923192419251926192719281929193019311932193319341935193619371938193919401941194219431944194519461947194819491950195119521953195419551956195719581959196019611962196319641965196619671968196919701971197219731974197519761977197819791980198119821983198419851986198719881989199019911992199319941995199619971998199920002001200220032004200520062007200820092010201120122013201420152016201720182019202020212022202320242025202620272028202920302031203220332034203520362037203820392040204120422043204420452046204720482049205020512052205320542055205620572058205920602061206220632064206520662067206820692070207120722073207420752076207720782079208020812082208320842085208620872088208920902091209220932094209520962097209820992100210121022103210421052106210721082109211021112112211321142115211621172118211921202121212221232124212521262127212821292130213121322133213421352136213721382139214021412142214321442145214621472148214921502151215221532154215521562157215821592160216121622163216421652166216721682169217021712172217321742175217621772178217921802181218221832184218521862187218821892190219121922193219421952196219721982199220022012202220322042205220622072208220922102211221222132214221522162217221822192220222122222223222422252226222722282229223022312232223322342235223622372238223922402241224222432244224522462247224822492250225122522253225422552256225722582259226022612262226322642265226622672268226922702271227222732274227522762277227822792280228122822283228422852286228722882289229022912292229322942295229622972298229923002301230223032304230523062307230823092310231123122313231423152316231723182319232023212322232323242325232623272328232923302331233223332334233523362337233823392340234123422343234423452346234723482349235023512352235323542355235623572358235923602361236223632364236523662367236823692370237123722373237423752376237723782379238023812382238323842385238623872388238923902391239223932394239523962397239823992400240124022403240424052406240724082409241024112412241324142415241624172418241924202421242224232424242524262427242824292430243124322433243424352436243724382439244024412442244324442445244624472448244924502451245224532454245524562457245824592460246124622463246424652466246724682469247024712472247324742475247624772478247924802481248224832484248524862487248824892490249124922493249424952496249724982499250025012502250325042505250625072508250925102511251225132514251525162517251825192520252125222523252425252526252725282529253025312532253325342535253625372538253925402541254225432544254525462547254825492550255125522553255425552556255725582559256025612562256325642565256625672568256925702571257225732574257525762577257825792580258125822583258425852586258725882589259025912592259325942595259625972598259926002601260226032604260526062607260826092610261126122613261426152616261726182619262026212622262326242625262626272628262926302631263226332634263526362637263826392640264126422643264426452646264726482649265026512652265326542655265626572658265926602661266226632664266526662667266826692670267126722673267426752676267726782679268026812682268326842685268626872688268926902691269226932694269526962697269826992700270127022703270427052706270727082709271027112712271327142715271627172718271927202721272227232724272527262727272827292730273127322733273427352736273727382739274027412742274327442745274627472748274927502751275227532754275527562757275827592760276127622763276427652766276727682769277027712772277327742775277627772778277927802781278227832784278527862787278827892790279127922793279427952796279727982799280028012802280328042805280628072808280928102811281228132814281528162817281828192820282128222823282428252826282728282829283028312832283328342835283628372838283928402841284228432844284528462847284828492850285128522853285428552856285728582859286028612862286328642865286628672868286928702871287228732874287528762877287828792880288128822883288428852886288728882889289028912892289328942895289628972898289929002901290229032904290529062907290829092910291129122913291429152916291729182919292029212922292329242925292629272928292929302931293229332934293529362937293829392940294129422943294429452946294729482949295029512952295329542955295629572958295929602961296229632964296529662967296829692970297129722973297429752976297729782979298029812982298329842985298629872988298929902991299229932994299529962997299829993000300130023003300430053006300730083009301030113012301330143015301630173018301930203021302230233024302530263027302830293030303130323033303430353036303730383039304030413042304330443045304630473048304930503051305230533054305530563057305830593060306130623063306430653066306730683069307030713072307330743075307630773078307930803081308230833084308530863087308830893090309130923093309430953096309730983099310031013102310331043105310631073108310931103111311231133114311531163117311831193120312131223123312431253126312731283129313031313132313331343135313631373138313931403141314231433144314531463147314831493150315131523153315431553156315731583159316031613162316331643165316631673168316931703171317231733174317531763177317831793180318131823183318431853186318731883189319031913192319331943195319631973198319932003201320232033204320532063207320832093210321132123213321432153216321732183219322032213222322332243225322632273228322932303231323232333234323532363237323832393240324132423243324432453246324732483249325032513252325332543255325632573258325932603261326232633264326532663267326832693270327132723273327432753276327732783279328032813282328332843285328632873288328932903291329232933294329532963297329832993300330133023303330433053306330733083309331033113312331333143315331633173318331933203321332233233324332533263327332833293330333133323333333433353336333733383339334033413342334333443345334633473348334933503351335233533354335533563357335833593360336133623363336433653366336733683369337033713372337333743375337633773378337933803381338233833384338533863387338833893390339133923393339433953396339733983399340034013402340334043405340634073408340934103411341234133414341534163417341834193420342134223423342434253426342734283429343034313432343334343435343634373438343934403441344234433444344534463447344834493450345134523453345434553456345734583459346034613462346334643465346634673468346934703471347234733474347534763477347834793480348134823483348434853486348734883489349034913492349334943495349634973498349935003501350235033504350535063507350835093510351135123513351435153516351735183519352035213522352335243525352635273528352935303531353235333534353535363537353835393540354135423543354435453546354735483549355035513552355335543555355635573558355935603561356235633564356535663567356835693570357135723573357435753576357735783579358035813582358335843585358635873588358935903591359235933594359535963597359835993600360136023603360436053606360736083609361036113612361336143615361636173618361936203621362236233624362536263627362836293630363136323633363436353636363736383639364036413642364336443645364636473648364936503651365236533654365536563657365836593660366136623663366436653666366736683669367036713672367336743675367636773678367936803681368236833684368536863687368836893690369136923693369436953696369736983699370037013702370337043705370637073708370937103711371237133714371537163717371837193720372137223723372437253726372737283729373037313732373337343735373637373738373937403741374237433744374537463747374837493750375137523753375437553756375737583759376037613762376337643765376637673768376937703771377237733774377537763777377837793780378137823783378437853786378737883789379037913792379337943795379637973798 |
- package service
- import (
- "bufio"
- "bytes"
- "context"
- "crypto/rand"
- "encoding/base64"
- "encoding/hex"
- "encoding/json"
- "errors"
- "fmt"
- "io"
- "log"
- "math"
- mathrand "math/rand"
- "net/http"
- "regexp"
- "strings"
- "time"
- "github.com/Wei-Shaw/sub2api/internal/config"
- "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
- "github.com/Wei-Shaw/sub2api/internal/pkg/geminicli"
- "github.com/Wei-Shaw/sub2api/internal/pkg/googleapi"
- "github.com/Wei-Shaw/sub2api/internal/pkg/logger"
- "github.com/Wei-Shaw/sub2api/internal/util/responseheaders"
- "github.com/Wei-Shaw/sub2api/internal/util/urlvalidator"
- "github.com/gin-gonic/gin"
- "github.com/tidwall/gjson"
- )
- const geminiStickySessionTTL = time.Hour
- const (
- geminiMaxRetries = 5
- geminiRetryBaseDelay = 1 * time.Second
- geminiRetryMaxDelay = 16 * time.Second
- )
- // Gemini tool calling now requires `thoughtSignature` in parts that include `functionCall`.
- // Many clients don't send it; we inject a known dummy signature to satisfy the validator.
- // Ref: https://ai.google.dev/gemini-api/docs/thought-signatures
- const geminiDummyThoughtSignature = "skip_thought_signature_validator"
- type GeminiMessagesCompatService struct {
- accountRepo AccountRepository
- groupRepo GroupRepository
- cache GatewayCache
- schedulerSnapshot *SchedulerSnapshotService
- tokenProvider *GeminiTokenProvider
- rateLimitService *RateLimitService
- httpUpstream HTTPUpstream
- antigravityGatewayService *AntigravityGatewayService
- cfg *config.Config
- responseHeaderFilter *responseheaders.CompiledHeaderFilter
- }
- func (s *GeminiMessagesCompatService) readUpstreamErrorBody(resp *http.Response) []byte {
- if resp == nil || resp.Body == nil {
- return nil
- }
- limit := gatewayUpstreamErrorBodyReadLimit
- if s != nil && s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody && s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes > int(limit) {
- limit = int64(s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes)
- }
- body, _ := io.ReadAll(io.LimitReader(resp.Body, limit))
- return body
- }
- func NewGeminiMessagesCompatService(
- accountRepo AccountRepository,
- groupRepo GroupRepository,
- cache GatewayCache,
- schedulerSnapshot *SchedulerSnapshotService,
- tokenProvider *GeminiTokenProvider,
- rateLimitService *RateLimitService,
- httpUpstream HTTPUpstream,
- antigravityGatewayService *AntigravityGatewayService,
- cfg *config.Config,
- ) *GeminiMessagesCompatService {
- return &GeminiMessagesCompatService{
- accountRepo: accountRepo,
- groupRepo: groupRepo,
- cache: cache,
- schedulerSnapshot: schedulerSnapshot,
- tokenProvider: tokenProvider,
- rateLimitService: rateLimitService,
- httpUpstream: httpUpstream,
- antigravityGatewayService: antigravityGatewayService,
- cfg: cfg,
- responseHeaderFilter: compileResponseHeaderFilter(cfg),
- }
- }
- // GetTokenProvider returns the token provider for OAuth accounts
- func (s *GeminiMessagesCompatService) GetTokenProvider() *GeminiTokenProvider {
- return s.tokenProvider
- }
- func (s *GeminiMessagesCompatService) SelectAccountForModel(ctx context.Context, groupID *int64, sessionHash string, requestedModel string) (*Account, error) {
- return s.SelectAccountForModelWithExclusions(ctx, groupID, sessionHash, requestedModel, nil)
- }
- func (s *GeminiMessagesCompatService) SelectAccountForModelWithExclusions(ctx context.Context, groupID *int64, sessionHash string, requestedModel string, excludedIDs map[int64]struct{}) (*Account, error) {
- // 1. 确定目标平台和调度模式
- // Determine target platform and scheduling mode
- platform, useMixedScheduling, hasForcePlatform, err := s.resolvePlatformAndSchedulingMode(ctx, groupID)
- if err != nil {
- return nil, err
- }
- cacheKey := "gemini:" + sessionHash
- // 2. 尝试粘性会话命中
- // Try sticky session hit
- if account := s.tryStickySessionHit(ctx, groupID, sessionHash, cacheKey, requestedModel, excludedIDs, platform, useMixedScheduling); account != nil {
- return s.hydrateSelectedAccount(ctx, account)
- }
- // 3. 查询可调度账户(强制平台模式:优先按分组查找,找不到再查全部)
- // Query schedulable accounts (force platform mode: try group first, fallback to all)
- accounts, err := s.listSchedulableAccountsOnce(ctx, groupID, platform, hasForcePlatform)
- if err != nil {
- return nil, fmt.Errorf("query accounts failed: %w", err)
- }
- // 强制平台模式下,分组中找不到账户时回退查询全部
- if len(accounts) == 0 && groupID != nil && hasForcePlatform {
- accounts, err = s.listSchedulableAccountsOnce(ctx, nil, platform, hasForcePlatform)
- if err != nil {
- return nil, fmt.Errorf("query accounts failed: %w", err)
- }
- }
- // 4. 按优先级 + LRU 选择最佳账号
- // Select best account by priority + LRU
- selected := s.selectBestGeminiAccount(ctx, accounts, requestedModel, excludedIDs, platform, useMixedScheduling)
- if selected == nil {
- if requestedModel != "" {
- return nil, fmt.Errorf("no available Gemini accounts supporting model: %s", requestedModel)
- }
- return nil, errors.New("no available Gemini accounts")
- }
- // 5. 设置粘性会话绑定
- // Set sticky session binding
- if sessionHash != "" {
- _ = s.cache.SetSessionAccountID(ctx, derefGroupID(groupID), cacheKey, selected.ID, geminiStickySessionTTL)
- }
- return s.hydrateSelectedAccount(ctx, selected)
- }
- // resolvePlatformAndSchedulingMode 解析目标平台和调度模式。
- // 返回:平台名称、是否使用混合调度、是否强制平台、错误。
- //
- // resolvePlatformAndSchedulingMode resolves target platform and scheduling mode.
- // Returns: platform name, whether to use mixed scheduling, whether force platform, error.
- func (s *GeminiMessagesCompatService) resolvePlatformAndSchedulingMode(ctx context.Context, groupID *int64) (platform string, useMixedScheduling bool, hasForcePlatform bool, err error) {
- // 优先检查 context 中的强制平台(/antigravity 路由)
- forcePlatform, hasForcePlatform := ctx.Value(ctxkey.ForcePlatform).(string)
- if hasForcePlatform && forcePlatform != "" {
- return forcePlatform, false, true, nil
- }
- if groupID != nil {
- // 根据分组 platform 决定查询哪种账号
- var group *Group
- if ctxGroup, ok := ctx.Value(ctxkey.Group).(*Group); ok && IsGroupContextValid(ctxGroup) && ctxGroup.ID == *groupID {
- group = ctxGroup
- } else {
- group, err = s.groupRepo.GetByIDLite(ctx, *groupID)
- if err != nil {
- return "", false, false, fmt.Errorf("get group failed: %w", err)
- }
- }
- // gemini 分组支持混合调度(包含启用了 mixed_scheduling 的 antigravity 账户)
- return group.Platform, group.Platform == PlatformGemini, false, nil
- }
- // 无分组时只使用原生 gemini 平台
- return PlatformGemini, true, false, nil
- }
- // tryStickySessionHit 尝试从粘性会话获取账号。
- // 如果命中且账号可用则返回账号;如果账号不可用则清理会话并返回 nil。
- //
- // tryStickySessionHit attempts to get account from sticky session.
- // Returns account if hit and usable; clears session and returns nil if account unavailable.
- func (s *GeminiMessagesCompatService) tryStickySessionHit(
- ctx context.Context,
- groupID *int64,
- sessionHash, cacheKey, requestedModel string,
- excludedIDs map[int64]struct{},
- platform string,
- useMixedScheduling bool,
- ) *Account {
- if sessionHash == "" {
- return nil
- }
- accountID, err := s.cache.GetSessionAccountID(ctx, derefGroupID(groupID), cacheKey)
- if err != nil || accountID <= 0 {
- return nil
- }
- if _, excluded := excludedIDs[accountID]; excluded {
- return nil
- }
- account, err := s.getSchedulableAccount(ctx, accountID)
- if err != nil {
- return nil
- }
- // 检查账号是否需要清理粘性会话
- // Check if sticky session should be cleared
- if shouldClearStickySession(account, requestedModel) {
- _ = s.cache.DeleteSessionAccountID(ctx, derefGroupID(groupID), cacheKey)
- return nil
- }
- // 验证账号是否可用于当前请求
- // Verify account is usable for current request
- if !s.isAccountUsableForRequest(ctx, account, requestedModel, platform, useMixedScheduling) {
- return nil
- }
- // 刷新会话 TTL 并返回账号
- // Refresh session TTL and return account
- _ = s.cache.RefreshSessionTTL(ctx, derefGroupID(groupID), cacheKey, geminiStickySessionTTL)
- return account
- }
- // isAccountUsableForRequest 检查账号是否可用于当前请求。
- // 验证:模型调度、模型支持、平台匹配、速率限制预检。
- //
- // isAccountUsableForRequest checks if account is usable for current request.
- // Validates: model scheduling, model support, platform matching, rate limit precheck.
- func (s *GeminiMessagesCompatService) isAccountUsableForRequest(
- ctx context.Context,
- account *Account,
- requestedModel, platform string,
- useMixedScheduling bool,
- ) bool {
- return s.isAccountUsableForRequestWithPrecheck(ctx, account, requestedModel, platform, useMixedScheduling, nil)
- }
- func (s *GeminiMessagesCompatService) isAccountUsableForRequestWithPrecheck(
- ctx context.Context,
- account *Account,
- requestedModel, platform string,
- useMixedScheduling bool,
- precheckResult map[int64]bool,
- ) bool {
- // 检查模型调度能力
- // Check model scheduling capability
- if !account.IsSchedulableForModelWithContext(ctx, requestedModel) {
- return false
- }
- // 检查模型支持
- // Check model support
- if requestedModel != "" && !s.isModelSupportedByAccount(account, requestedModel) {
- return false
- }
- // 检查平台匹配
- // Check platform matching
- if !s.isAccountValidForPlatform(account, platform, useMixedScheduling) {
- return false
- }
- // 速率限制预检
- // Rate limit precheck
- if !s.passesRateLimitPreCheckWithCache(ctx, account, requestedModel, precheckResult) {
- return false
- }
- return true
- }
- // isAccountValidForPlatform 检查账号是否匹配目标平台。
- // 原生平台直接匹配;混合调度模式下 antigravity 需要启用 mixed_scheduling。
- //
- // isAccountValidForPlatform checks if account matches target platform.
- // Native platform matches directly; mixed scheduling mode requires antigravity to enable mixed_scheduling.
- func (s *GeminiMessagesCompatService) isAccountValidForPlatform(account *Account, platform string, useMixedScheduling bool) bool {
- if account.Platform == platform {
- return true
- }
- if useMixedScheduling && account.Platform == PlatformAntigravity && account.IsMixedSchedulingEnabled() {
- return true
- }
- return false
- }
- func (s *GeminiMessagesCompatService) passesRateLimitPreCheckWithCache(ctx context.Context, account *Account, requestedModel string, precheckResult map[int64]bool) bool {
- if s.rateLimitService == nil || requestedModel == "" {
- return true
- }
- if precheckResult != nil {
- if ok, exists := precheckResult[account.ID]; exists {
- return ok
- }
- }
- ok, err := s.rateLimitService.PreCheckUsage(ctx, account, requestedModel)
- if err != nil {
- logger.LegacyPrintf("service.gemini_messages_compat", "[Gemini PreCheck] Account %d precheck error: %v", account.ID, err)
- }
- return ok
- }
- // selectBestGeminiAccount 从候选账号中选择最佳账号(优先级 + LRU + OAuth 优先)。
- // 返回 nil 表示无可用账号。
- //
- // selectBestGeminiAccount selects best account from candidates (priority + LRU + OAuth preferred).
- // Returns nil if no available account.
- func (s *GeminiMessagesCompatService) selectBestGeminiAccount(
- ctx context.Context,
- accounts []Account,
- requestedModel string,
- excludedIDs map[int64]struct{},
- platform string,
- useMixedScheduling bool,
- ) *Account {
- var selected *Account
- precheckResult := s.buildPreCheckUsageResultMap(ctx, accounts, requestedModel)
- for i := range accounts {
- acc := &accounts[i]
- // 跳过被排除的账号
- if _, excluded := excludedIDs[acc.ID]; excluded {
- continue
- }
- // 检查账号是否可用于当前请求
- if !s.isAccountUsableForRequestWithPrecheck(ctx, acc, requestedModel, platform, useMixedScheduling, precheckResult) {
- continue
- }
- // 选择最佳账号
- if selected == nil {
- selected = acc
- continue
- }
- if s.isBetterGeminiAccount(acc, selected) {
- selected = acc
- }
- }
- return selected
- }
- func (s *GeminiMessagesCompatService) buildPreCheckUsageResultMap(ctx context.Context, accounts []Account, requestedModel string) map[int64]bool {
- if s.rateLimitService == nil || requestedModel == "" || len(accounts) == 0 {
- return nil
- }
- candidates := make([]*Account, 0, len(accounts))
- for i := range accounts {
- candidates = append(candidates, &accounts[i])
- }
- result, err := s.rateLimitService.PreCheckUsageBatch(ctx, candidates, requestedModel)
- if err != nil {
- logger.LegacyPrintf("service.gemini_messages_compat", "[Gemini PreCheckBatch] failed: %v", err)
- }
- return result
- }
- // isBetterGeminiAccount 判断 candidate 是否比 current 更优。
- // 规则:优先级更高(数值更小)优先;同优先级时,未使用过的优先(OAuth > 非 OAuth),其次是最久未使用的。
- //
- // isBetterGeminiAccount checks if candidate is better than current.
- // Rules: higher priority (lower value) wins; same priority: never used (OAuth > non-OAuth) > least recently used.
- func (s *GeminiMessagesCompatService) isBetterGeminiAccount(candidate, current *Account) bool {
- // 优先级更高(数值更小)
- if candidate.Priority < current.Priority {
- return true
- }
- if candidate.Priority > current.Priority {
- return false
- }
- // 同优先级,比较最后使用时间
- switch {
- case candidate.LastUsedAt == nil && current.LastUsedAt != nil:
- // candidate 从未使用,优先
- return true
- case candidate.LastUsedAt != nil && current.LastUsedAt == nil:
- // current 从未使用,保持
- return false
- case candidate.LastUsedAt == nil && current.LastUsedAt == nil:
- // 都未使用,优先选择 OAuth 账号(更兼容 Code Assist 流程)
- return candidate.Type == AccountTypeOAuth && current.Type != AccountTypeOAuth
- default:
- // 都使用过,选择最久未使用的
- return candidate.LastUsedAt.Before(*current.LastUsedAt)
- }
- }
- // isModelSupportedByAccount 根据账户平台检查模型支持
- func (s *GeminiMessagesCompatService) isModelSupportedByAccount(account *Account, requestedModel string) bool {
- if account.Platform == PlatformAntigravity {
- if strings.TrimSpace(requestedModel) == "" {
- return true
- }
- return mapAntigravityModel(account, requestedModel) != ""
- }
- return account.IsModelSupported(requestedModel)
- }
- // GetAntigravityGatewayService 返回 AntigravityGatewayService
- func (s *GeminiMessagesCompatService) GetAntigravityGatewayService() *AntigravityGatewayService {
- return s.antigravityGatewayService
- }
- func (s *GeminiMessagesCompatService) getSchedulableAccount(ctx context.Context, accountID int64) (*Account, error) {
- if s.schedulerSnapshot != nil {
- return s.schedulerSnapshot.GetAccount(ctx, accountID)
- }
- return s.accountRepo.GetByID(ctx, accountID)
- }
- func (s *GeminiMessagesCompatService) hydrateSelectedAccount(ctx context.Context, account *Account) (*Account, error) {
- if account == nil {
- return nil, nil
- }
- hydrated := account
- if s.schedulerSnapshot != nil {
- var err error
- hydrated, err = s.schedulerSnapshot.GetAccount(ctx, account.ID)
- if err != nil {
- return nil, err
- }
- if hydrated == nil {
- return nil, fmt.Errorf("selected gemini account %d not found during hydration", account.ID)
- }
- }
- if err := ResolveRandomProxyFromSource(ctx, hydrated, s.accountRepo); err != nil {
- return nil, err
- }
- return hydrated, nil
- }
- func (s *GeminiMessagesCompatService) listSchedulableAccountsOnce(ctx context.Context, groupID *int64, platform string, hasForcePlatform bool) ([]Account, error) {
- var accounts []Account
- var err error
- if s.schedulerSnapshot != nil {
- accounts, _, err = s.schedulerSnapshot.ListSchedulableAccounts(ctx, groupID, platform, hasForcePlatform)
- } else {
- useMixedScheduling := platform == PlatformGemini && !hasForcePlatform
- queryPlatforms := []string{platform}
- if useMixedScheduling {
- queryPlatforms = []string{platform, PlatformAntigravity}
- }
- if groupID != nil {
- accounts, err = s.accountRepo.ListSchedulableByGroupIDAndPlatforms(ctx, *groupID, queryPlatforms)
- } else if s.cfg != nil && s.cfg.RunMode == config.RunModeSimple {
- accounts, err = s.accountRepo.ListSchedulableByPlatforms(ctx, queryPlatforms)
- } else {
- accounts, err = s.accountRepo.ListSchedulableUngroupedByPlatforms(ctx, queryPlatforms)
- }
- }
- if err != nil {
- return accounts, err
- }
- return FilterTierPoolAccounts(accounts, TierPoolFromContext(ctx)), nil
- }
- func (s *GeminiMessagesCompatService) validateUpstreamBaseURL(raw string) (string, error) {
- if s.cfg != nil && !s.cfg.Security.URLAllowlist.Enabled {
- normalized, err := urlvalidator.ValidateURLFormat(raw, s.cfg.Security.URLAllowlist.AllowInsecureHTTP)
- if err != nil {
- return "", fmt.Errorf("invalid base_url: %w", err)
- }
- return normalized, nil
- }
- normalized, err := urlvalidator.ValidateHTTPSURL(raw, urlvalidator.ValidationOptions{
- AllowedHosts: s.cfg.Security.URLAllowlist.UpstreamHosts,
- RequireAllowlist: true,
- AllowPrivate: s.cfg.Security.URLAllowlist.AllowPrivateHosts,
- })
- if err != nil {
- return "", fmt.Errorf("invalid base_url: %w", err)
- }
- return normalized, nil
- }
- // HasAntigravityAccounts 检查是否有可用的 antigravity 账户
- func (s *GeminiMessagesCompatService) HasAntigravityAccounts(ctx context.Context, groupID *int64) (bool, error) {
- accounts, err := s.listSchedulableAccountsOnce(ctx, groupID, PlatformAntigravity, false)
- if err != nil {
- return false, err
- }
- return len(accounts) > 0, nil
- }
- // SelectAccountForAIStudioEndpoints selects an account that is likely to succeed against
- // generativelanguage.googleapis.com (e.g. GET /v1beta/models).
- //
- // Preference order:
- // 1) API key accounts (AI Studio)
- // 2) OAuth accounts without project_id (AI Studio OAuth)
- // 3) OAuth accounts explicitly marked as ai_studio
- // 4) Any remaining Gemini accounts (fallback)
- func (s *GeminiMessagesCompatService) SelectAccountForAIStudioEndpoints(ctx context.Context, groupID *int64) (*Account, error) {
- accounts, err := s.listSchedulableAccountsOnce(ctx, groupID, PlatformGemini, true)
- if err != nil {
- return nil, fmt.Errorf("query accounts failed: %w", err)
- }
- if len(accounts) == 0 {
- return nil, errors.New("no available Gemini accounts")
- }
- rank := func(a *Account) int {
- if a == nil {
- return 999
- }
- switch a.Type {
- case AccountTypeAPIKey:
- if strings.TrimSpace(a.GetCredential("api_key")) != "" {
- return 0
- }
- return 9
- case AccountTypeOAuth:
- if strings.TrimSpace(a.GetCredential("project_id")) == "" {
- return 1
- }
- if strings.TrimSpace(a.GetCredential("oauth_type")) == "ai_studio" {
- return 2
- }
- // Code Assist OAuth tokens often lack AI Studio scopes for models listing.
- return 3
- case AccountTypeServiceAccount:
- // Vertex service accounts use aiplatform.googleapis.com, not the AI Studio
- // endpoint (generativelanguage.googleapis.com), so they cannot serve these requests.
- return 999
- default:
- return 10
- }
- }
- var selected *Account
- for i := range accounts {
- acc := &accounts[i]
- if selected == nil {
- selected = acc
- continue
- }
- r1, r2 := rank(acc), rank(selected)
- if r1 < r2 {
- selected = acc
- continue
- }
- if r1 > r2 {
- continue
- }
- if acc.Priority < selected.Priority {
- selected = acc
- } else if acc.Priority == selected.Priority {
- switch {
- case acc.LastUsedAt == nil && selected.LastUsedAt != nil:
- selected = acc
- case acc.LastUsedAt != nil && selected.LastUsedAt == nil:
- // keep selected
- case acc.LastUsedAt == nil && selected.LastUsedAt == nil:
- if acc.Type == AccountTypeOAuth && selected.Type != AccountTypeOAuth {
- selected = acc
- }
- default:
- if acc.LastUsedAt.Before(*selected.LastUsedAt) {
- selected = acc
- }
- }
- }
- }
- if selected == nil {
- return nil, errors.New("no available Gemini accounts")
- }
- return s.hydrateSelectedAccount(ctx, selected)
- }
- func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Context, account *Account, body []byte) (*ForwardResult, error) {
- beginUpstreamResponseModelObservation(c)
- beginGeminiImageOutputObservation(c)
- startTime := time.Now()
- var req struct {
- Model string `json:"model"`
- Stream bool `json:"stream"`
- }
- if err := json.Unmarshal(body, &req); err != nil {
- return nil, fmt.Errorf("parse request: %w", err)
- }
- if strings.TrimSpace(req.Model) == "" {
- return nil, fmt.Errorf("missing model")
- }
- originalModel := req.Model
- mappedModel := req.Model
- if account.Type == AccountTypeAPIKey || account.Type == AccountTypeServiceAccount {
- mappedModel = account.GetMappedModel(req.Model)
- }
- geminiReq, err := convertClaudeMessagesToGeminiGenerateContent(body)
- if err != nil {
- return nil, s.writeClaudeError(c, http.StatusBadRequest, "invalid_request_error", err.Error())
- }
- geminiReq = ensureGeminiFunctionCallThoughtSignatures(geminiReq)
- originalClaudeBody := body
- proxyURL := ""
- if account.ProxyID != nil && account.Proxy != nil {
- proxyURL = account.Proxy.URL()
- }
- var requestIDHeader string
- var buildReq func(ctx context.Context) (*http.Request, string, error)
- useUpstreamStream := req.Stream
- if account.Type == AccountTypeOAuth && !req.Stream && strings.TrimSpace(account.GetCredential("project_id")) != "" {
- // Code Assist's non-streaming generateContent may return no content; use streaming upstream and aggregate.
- useUpstreamStream = true
- }
- switch account.Type {
- case AccountTypeAPIKey:
- buildReq = func(ctx context.Context) (*http.Request, string, error) {
- apiKey := account.GetCredential("api_key")
- if strings.TrimSpace(apiKey) == "" {
- return nil, "", errors.New("gemini api_key not configured")
- }
- baseURL := account.GetGeminiBaseURL(geminicli.AIStudioBaseURL)
- normalizedBaseURL, err := s.validateUpstreamBaseURL(baseURL)
- if err != nil {
- return nil, "", err
- }
- action := "generateContent"
- if req.Stream {
- action = "streamGenerateContent"
- }
- fullURL, err := buildGeminiAIStudioModelActionURL(normalizedBaseURL, mappedModel, action, req.Stream)
- if err != nil {
- return nil, "", err
- }
- restGeminiReq := normalizeGeminiRequestForAIStudio(geminiReq)
- upstreamReq, err := http.NewRequestWithContext(ctx, http.MethodPost, fullURL, bytes.NewReader(restGeminiReq))
- if err != nil {
- return nil, "", err
- }
- upstreamReq.Header.Set("Content-Type", "application/json")
- upstreamReq.Header.Set("x-goog-api-key", apiKey)
- return upstreamReq, "x-request-id", nil
- }
- requestIDHeader = "x-request-id"
- case AccountTypeOAuth:
- buildReq = func(ctx context.Context) (*http.Request, string, error) {
- if s.tokenProvider == nil {
- return nil, "", errors.New("gemini token provider not configured")
- }
- accessToken, err := s.tokenProvider.GetAccessToken(ctx, account)
- if err != nil {
- return nil, "", err
- }
- projectID := strings.TrimSpace(account.GetCredential("project_id"))
- action := "generateContent"
- if useUpstreamStream {
- action = "streamGenerateContent"
- }
- // Two modes for OAuth:
- // 1. With project_id -> Code Assist API (wrapped request)
- // 2. Without project_id -> AI Studio API (direct OAuth, like API key but with Bearer token)
- if projectID != "" {
- // Mode 1: Code Assist API
- baseURL, err := s.validateUpstreamBaseURL(geminicli.GeminiCliBaseURL)
- if err != nil {
- return nil, "", err
- }
- fullURL := fmt.Sprintf("%s/v1internal:%s", strings.TrimRight(baseURL, "/"), action)
- if useUpstreamStream {
- fullURL += "?alt=sse"
- }
- wrapped := map[string]any{
- "model": mappedModel,
- "project": projectID,
- }
- var inner any
- if err := json.Unmarshal(geminiReq, &inner); err != nil {
- return nil, "", fmt.Errorf("failed to parse gemini request: %w", err)
- }
- wrapped["request"] = inner
- wrappedBytes, _ := json.Marshal(wrapped)
- upstreamReq, err := http.NewRequestWithContext(ctx, http.MethodPost, fullURL, bytes.NewReader(wrappedBytes))
- if err != nil {
- return nil, "", err
- }
- upstreamReq.Header.Set("Content-Type", "application/json")
- upstreamReq.Header.Set("Authorization", "Bearer "+accessToken)
- upstreamReq.Header.Set("User-Agent", geminicli.GeminiCLIUserAgent)
- return upstreamReq, "x-request-id", nil
- } else {
- // Mode 2: AI Studio API with OAuth (like API key mode, but using Bearer token)
- baseURL := account.GetGeminiBaseURL(geminicli.AIStudioBaseURL)
- normalizedBaseURL, err := s.validateUpstreamBaseURL(baseURL)
- if err != nil {
- return nil, "", err
- }
- fullURL, err := buildGeminiAIStudioModelActionURL(normalizedBaseURL, mappedModel, action, useUpstreamStream)
- if err != nil {
- return nil, "", err
- }
- restGeminiReq := normalizeGeminiRequestForAIStudio(geminiReq)
- upstreamReq, err := http.NewRequestWithContext(ctx, http.MethodPost, fullURL, bytes.NewReader(restGeminiReq))
- if err != nil {
- return nil, "", err
- }
- upstreamReq.Header.Set("Content-Type", "application/json")
- upstreamReq.Header.Set("Authorization", "Bearer "+accessToken)
- return upstreamReq, "x-request-id", nil
- }
- }
- requestIDHeader = "x-request-id"
- case AccountTypeServiceAccount:
- buildReq = func(ctx context.Context) (*http.Request, string, error) {
- if s.tokenProvider == nil {
- return nil, "", errors.New("gemini token provider not configured")
- }
- accessToken, err := s.tokenProvider.GetAccessToken(ctx, account)
- if err != nil {
- return nil, "", err
- }
- action := "generateContent"
- if req.Stream {
- action = "streamGenerateContent"
- }
- fullURL, err := buildVertexGeminiURL(account.VertexProjectID(), account.VertexLocation(mappedModel), mappedModel, action, req.Stream)
- if err != nil {
- return nil, "", err
- }
- restGeminiReq := normalizeGeminiRequestForAIStudio(geminiReq)
- upstreamReq, err := http.NewRequestWithContext(ctx, http.MethodPost, fullURL, bytes.NewReader(restGeminiReq))
- if err != nil {
- return nil, "", err
- }
- upstreamReq.Header.Set("Content-Type", "application/json")
- upstreamReq.Header.Set("Authorization", "Bearer "+accessToken)
- return upstreamReq, "x-request-id", nil
- }
- requestIDHeader = "x-request-id"
- default:
- return nil, fmt.Errorf("unsupported account type: %s", account.Type)
- }
- var resp *http.Response
- signatureRetryStage := 0
- for attempt := 1; attempt <= geminiMaxRetries; attempt++ {
- upstreamReq, idHeader, err := buildReq(ctx)
- if err != nil {
- if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
- return nil, err
- }
- // Local build error: don't retry.
- if strings.Contains(err.Error(), "missing project_id") {
- return nil, s.writeClaudeError(c, http.StatusBadRequest, "invalid_request_error", err.Error())
- }
- return nil, s.writeClaudeError(c, http.StatusBadGateway, "upstream_error", err.Error())
- }
- requestIDHeader = idHeader
- resp, err = s.httpUpstream.Do(WithAccountTrafficRequest(upstreamReq, account), proxyURL, account.ID, account.Concurrency)
- if err != nil {
- if local := AccountTrafficFailover(err); local != nil {
- return nil, local
- }
- safeErr := sanitizeUpstreamErrorMessage(err.Error())
- appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
- ProxyID: opsUpstreamProxyID(account),
- ProxyName: opsUpstreamProxyName(account),
- Platform: account.Platform,
- AccountID: account.ID,
- AccountName: account.Name,
- UpstreamStatusCode: 0,
- Kind: "request_error",
- Message: safeErr,
- })
- if attempt < geminiMaxRetries {
- logger.LegacyPrintf("service.gemini_messages_compat", "Gemini account %d: upstream request failed, retry %d/%d: %v", account.ID, attempt, geminiMaxRetries, err)
- sleepGeminiBackoff(attempt)
- continue
- }
- setOpsUpstreamError(c, 0, safeErr, "")
- return nil, s.writeClaudeError(c, http.StatusBadGateway, "upstream_error", "Upstream request failed after retries: "+safeErr)
- }
- // Special-case: signature/thought_signature validation errors are not transient, but may be fixed by
- // downgrading Claude thinking/tool history to plain text (conservative two-stage retry).
- if resp.StatusCode == http.StatusBadRequest && signatureRetryStage < 2 {
- respBody := s.readUpstreamErrorBody(resp)
- _ = resp.Body.Close()
- if isGeminiSignatureRelatedError(respBody) {
- upstreamReqID := resp.Header.Get(requestIDHeader)
- if upstreamReqID == "" {
- upstreamReqID = resp.Header.Get("x-goog-request-id")
- }
- upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody))
- upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
- upstreamDetail := ""
- if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody {
- maxBytes := s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes
- if maxBytes <= 0 {
- maxBytes = 2048
- }
- upstreamDetail = truncateString(string(respBody), maxBytes)
- }
- appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
- ProxyID: opsUpstreamProxyID(account),
- ProxyName: opsUpstreamProxyName(account),
- Platform: account.Platform,
- AccountID: account.ID,
- AccountName: account.Name,
- UpstreamStatusCode: resp.StatusCode,
- UpstreamRequestID: upstreamReqID,
- Kind: "signature_error",
- Message: upstreamMsg,
- Detail: upstreamDetail,
- })
- var strippedClaudeBody []byte
- stageName := ""
- // 路径说明:本处上游是 Gemini,但被剥离的 body 是 Anthropic 格式。传 originalModel
- // (客户端原 Anthropic model)而非 mappedModel(上游 Gemini model),让剥离逻辑按
- // 客户端请求的 Anthropic 子协议族判定(详见 ResolveThinkingProtocol 文档)。
- switch signatureRetryStage {
- case 0:
- // Stage 1: disable thinking + thinking->text
- strippedClaudeBody = FilterThinkingBlocksForRetry(originalClaudeBody, originalModel)
- stageName = "thinking-only"
- signatureRetryStage = 1
- default:
- // Stage 2: additionally downgrade tool_use/tool_result blocks to text
- strippedClaudeBody = FilterSignatureSensitiveBlocksForRetry(originalClaudeBody, originalModel)
- stageName = "thinking+tools"
- signatureRetryStage = 2
- }
- retryGeminiReq, txErr := convertClaudeMessagesToGeminiGenerateContent(strippedClaudeBody)
- if txErr == nil {
- logger.LegacyPrintf("service.gemini_messages_compat", "Gemini account %d: detected signature-related 400, retrying with downgraded Claude blocks (%s)", account.ID, stageName)
- geminiReq = retryGeminiReq
- // Consume one retry budget attempt and continue with the updated request payload.
- sleepGeminiBackoff(1)
- continue
- }
- }
- // Restore body for downstream error handling.
- resp = &http.Response{
- StatusCode: http.StatusBadRequest,
- Header: resp.Header.Clone(),
- Body: io.NopCloser(bytes.NewReader(respBody)),
- }
- break
- }
- // 错误策略优先:匹配则跳过重试直接处理。
- if matched, rebuilt := s.checkErrorPolicyInLoop(ctx, account, resp, mappedModel); matched {
- resp = rebuilt
- break
- } else {
- resp = rebuilt
- }
- if resp.StatusCode >= 400 && s.shouldRetryGeminiUpstreamError(account, resp.StatusCode) {
- respBody := s.readUpstreamErrorBody(resp)
- _ = resp.Body.Close()
- // Don't treat insufficient-scope as transient.
- if resp.StatusCode == 403 && isGeminiInsufficientScope(resp.Header, respBody) {
- resp = &http.Response{
- StatusCode: resp.StatusCode,
- Header: resp.Header.Clone(),
- Body: io.NopCloser(bytes.NewReader(respBody)),
- }
- break
- }
- if resp.StatusCode == 429 {
- // Mark as rate-limited early so concurrent requests avoid this account.
- s.handleGeminiUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
- }
- if attempt < geminiMaxRetries {
- upstreamReqID := resp.Header.Get(requestIDHeader)
- if upstreamReqID == "" {
- upstreamReqID = resp.Header.Get("x-goog-request-id")
- }
- upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody))
- upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
- upstreamDetail := ""
- if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody {
- maxBytes := s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes
- if maxBytes <= 0 {
- maxBytes = 2048
- }
- upstreamDetail = truncateString(string(respBody), maxBytes)
- }
- appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
- ProxyID: opsUpstreamProxyID(account),
- ProxyName: opsUpstreamProxyName(account),
- Platform: account.Platform,
- AccountID: account.ID,
- AccountName: account.Name,
- UpstreamStatusCode: resp.StatusCode,
- UpstreamRequestID: upstreamReqID,
- Kind: "retry",
- Message: upstreamMsg,
- Detail: upstreamDetail,
- })
- logger.LegacyPrintf("service.gemini_messages_compat", "Gemini account %d: upstream status %d, retry %d/%d", account.ID, resp.StatusCode, attempt, geminiMaxRetries)
- sleepGeminiBackoff(attempt)
- continue
- }
- // Final attempt: surface the upstream error body (mapped below) instead of a generic retry error.
- resp = &http.Response{
- StatusCode: resp.StatusCode,
- Header: resp.Header.Clone(),
- Body: io.NopCloser(bytes.NewReader(respBody)),
- }
- break
- }
- break
- }
- defer func() { _ = resp.Body.Close() }()
- if resp.StatusCode >= 400 {
- respBody := s.readUpstreamErrorBody(resp)
- // 统一错误策略:自定义错误码 + 临时不可调度
- if s.rateLimitService != nil {
- policy := s.rateLimitService.CheckErrorPolicy(ctx, account, resp.StatusCode, respBody, mappedModel)
- switch policy {
- case ErrorPolicySkipped:
- upstreamReqID := resp.Header.Get(requestIDHeader)
- if upstreamReqID == "" {
- upstreamReqID = resp.Header.Get("x-goog-request-id")
- }
- if failoverErr := s.skippedErrorPolicyFailoverError(c, account, resp.StatusCode, respBody, upstreamReqID); failoverErr != nil {
- return nil, failoverErr
- }
- if account.IsCustomErrorCodesEnabled() {
- return nil, s.writeGeminiCustomCodeSkippedError(c, account, resp.StatusCode, upstreamReqID, respBody, func() {
- _ = s.writeClaudeError(c, http.StatusInternalServerError, "api_error", geminiCustomCodeSkippedClientMessage)
- })
- }
- // 池模式:客户端写出与 ErrorPolicyNone 相同(按上游真实状态码映射),仅跳过账号状态标记。
- return nil, s.writeGeminiMappedError(c, account, resp.StatusCode, upstreamReqID, respBody)
- case ErrorPolicyMatched, ErrorPolicyTempUnscheduled:
- if policy == ErrorPolicyMatched {
- s.handleGeminiUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
- }
- upstreamReqID := resp.Header.Get(requestIDHeader)
- if upstreamReqID == "" {
- upstreamReqID = resp.Header.Get("x-goog-request-id")
- }
- upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody))
- upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
- upstreamDetail := ""
- if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody {
- maxBytes := s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes
- if maxBytes <= 0 {
- maxBytes = 2048
- }
- upstreamDetail = truncateString(string(respBody), maxBytes)
- }
- appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
- ProxyID: opsUpstreamProxyID(account),
- ProxyName: opsUpstreamProxyName(account),
- Platform: account.Platform,
- AccountID: account.ID,
- AccountName: account.Name,
- UpstreamStatusCode: resp.StatusCode,
- UpstreamRequestID: upstreamReqID,
- Kind: "failover",
- Message: upstreamMsg,
- Detail: upstreamDetail,
- })
- return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: respBody}
- }
- }
- // ErrorPolicyNone → 原有逻辑
- s.handleGeminiUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
- // 精确匹配服务端配置类 400 错误,触发 failover + 临时封禁
- if resp.StatusCode == http.StatusBadRequest {
- msg400 := strings.ToLower(strings.TrimSpace(extractUpstreamErrorMessage(respBody)))
- if isGoogleProjectConfigError(msg400) {
- upstreamReqID := resp.Header.Get(requestIDHeader)
- if upstreamReqID == "" {
- upstreamReqID = resp.Header.Get("x-goog-request-id")
- }
- upstreamMsg := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractUpstreamErrorMessage(respBody)))
- upstreamDetail := ""
- if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody {
- maxBytes := s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes
- if maxBytes <= 0 {
- maxBytes = 2048
- }
- upstreamDetail = truncateString(string(respBody), maxBytes)
- }
- log.Printf("[Gemini] status=400 google_config_error failover=true upstream_message=%q account=%d", upstreamMsg, account.ID)
- appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
- ProxyID: opsUpstreamProxyID(account),
- ProxyName: opsUpstreamProxyName(account),
- Platform: account.Platform,
- AccountID: account.ID,
- AccountName: account.Name,
- UpstreamStatusCode: resp.StatusCode,
- UpstreamRequestID: upstreamReqID,
- Kind: "failover",
- Message: upstreamMsg,
- Detail: upstreamDetail,
- })
- return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: respBody, RetryableOnSameAccount: true}
- }
- }
- if s.shouldFailoverGeminiUpstreamError(resp.StatusCode) {
- upstreamReqID := resp.Header.Get(requestIDHeader)
- if upstreamReqID == "" {
- upstreamReqID = resp.Header.Get("x-goog-request-id")
- }
- upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody))
- upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
- upstreamDetail := ""
- if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody {
- maxBytes := s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes
- if maxBytes <= 0 {
- maxBytes = 2048
- }
- upstreamDetail = truncateString(string(respBody), maxBytes)
- }
- appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
- ProxyID: opsUpstreamProxyID(account),
- ProxyName: opsUpstreamProxyName(account),
- Platform: account.Platform,
- AccountID: account.ID,
- AccountName: account.Name,
- UpstreamStatusCode: resp.StatusCode,
- UpstreamRequestID: upstreamReqID,
- Kind: "failover",
- Message: upstreamMsg,
- Detail: upstreamDetail,
- })
- return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: respBody}
- }
- upstreamReqID := resp.Header.Get(requestIDHeader)
- if upstreamReqID == "" {
- upstreamReqID = resp.Header.Get("x-goog-request-id")
- }
- return nil, s.writeGeminiMappedError(c, account, resp.StatusCode, upstreamReqID, respBody)
- }
- requestID := resp.Header.Get(requestIDHeader)
- if requestID == "" {
- requestID = resp.Header.Get("x-goog-request-id")
- }
- if requestID != "" {
- c.Header("x-request-id", requestID)
- }
- var usage *ClaudeUsage
- var firstTokenMs *int
- if req.Stream {
- streamRes, err := s.handleStreamingResponse(c, resp, startTime, originalModel)
- if err != nil {
- return nil, err
- }
- usage = streamRes.usage
- firstTokenMs = streamRes.firstTokenMs
- } else {
- if useUpstreamStream {
- collected, usageObj, err := collectGeminiSSE(resp.Body, true)
- if err != nil {
- return nil, s.writeClaudeError(c, http.StatusBadGateway, "upstream_error", "Failed to read upstream stream")
- }
- collectedBytes, _ := json.Marshal(collected)
- upstreamResponseModelObserverFromContext(c).ObserveGemini(collectedBytes)
- observeGeminiImageOutputs(c, collectedBytes)
- claudeResp, usageObj2 := convertGeminiToClaudeMessage(collected, originalModel, collectedBytes, false)
- c.JSON(http.StatusOK, claudeResp)
- usage = usageObj2
- if usageObj != nil && (usageObj.InputTokens > 0 || usageObj.OutputTokens > 0) {
- usage = usageObj
- }
- } else {
- usage, err = s.handleNonStreamingResponse(c, resp, originalModel)
- if err != nil {
- return nil, err
- }
- }
- }
- // 图片生成计费
- imageInputSize := s.extractImageInputSize(body)
- imageSize := normalizeOpenAIImageSizeTier(imageInputSize)
- imageCount := resolveGeminiImageCount(c, originalModel, mappedModel)
- return &ForwardResult{
- RequestID: requestID,
- UpstreamHeaders: resp.Header,
- Usage: *usage,
- Model: originalModel,
- UpstreamModel: mappedModel,
- UpstreamResponseModel: observedUpstreamResponseModel(c),
- UpstreamResponseModelConflict: observedUpstreamResponseModelConflict(c),
- Stream: req.Stream,
- Duration: time.Since(startTime),
- FirstTokenMs: firstTokenMs,
- ImageCount: imageCount,
- ImageSize: imageSize,
- ImageInputSize: imageInputSize,
- }, nil
- }
- func isGeminiSignatureRelatedError(respBody []byte) bool {
- msg := strings.ToLower(strings.TrimSpace(extractAntigravityErrorMessage(respBody)))
- if msg == "" {
- msg = strings.ToLower(string(respBody))
- }
- return strings.Contains(msg, "thought_signature") || strings.Contains(msg, "signature")
- }
- func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.Context, account *Account, originalModel string, action string, stream bool, body []byte) (*ForwardResult, error) {
- beginUpstreamResponseModelObservation(c)
- beginGeminiImageOutputObservation(c)
- startTime := time.Now()
- if strings.TrimSpace(originalModel) == "" {
- return nil, s.writeGoogleError(c, http.StatusBadRequest, "Missing model in URL")
- }
- if strings.TrimSpace(action) == "" {
- return nil, s.writeGoogleError(c, http.StatusBadRequest, "Missing action in URL")
- }
- if len(body) == 0 {
- return nil, s.writeGoogleError(c, http.StatusBadRequest, "Request body is empty")
- }
- // 过滤掉 parts 为空的消息(Gemini API 不接受空 parts)
- if filteredBody, err := filterEmptyPartsFromGeminiRequest(body); err == nil {
- body = filteredBody
- }
- switch action {
- case "generateContent", "streamGenerateContent", "countTokens":
- // ok
- default:
- return nil, s.writeGoogleError(c, http.StatusNotFound, "Unsupported action: "+action)
- }
- // Some Gemini upstreams validate tool call parts strictly; ensure any `functionCall` part includes a
- // `thoughtSignature` to avoid frequent INVALID_ARGUMENT 400s.
- body = ensureGeminiFunctionCallThoughtSignatures(body)
- mappedModel := originalModel
- if account.Type == AccountTypeAPIKey || account.Type == AccountTypeServiceAccount {
- mappedModel = account.GetMappedModel(originalModel)
- }
- proxyURL := ""
- if account.ProxyID != nil && account.Proxy != nil {
- proxyURL = account.Proxy.URL()
- }
- useUpstreamStream := stream
- upstreamAction := action
- if account.Type == AccountTypeOAuth && !stream && action == "generateContent" && strings.TrimSpace(account.GetCredential("project_id")) != "" {
- // Code Assist's non-streaming generateContent may return no content; use streaming upstream and aggregate.
- useUpstreamStream = true
- upstreamAction = "streamGenerateContent"
- }
- forceAIStudio := action == "countTokens"
- var requestIDHeader string
- var buildReq func(ctx context.Context) (*http.Request, string, error)
- switch account.Type {
- case AccountTypeAPIKey:
- buildReq = func(ctx context.Context) (*http.Request, string, error) {
- apiKey := account.GetCredential("api_key")
- if strings.TrimSpace(apiKey) == "" {
- return nil, "", errors.New("gemini api_key not configured")
- }
- baseURL := account.GetGeminiBaseURL(geminicli.AIStudioBaseURL)
- normalizedBaseURL, err := s.validateUpstreamBaseURL(baseURL)
- if err != nil {
- return nil, "", err
- }
- fullURL, err := buildGeminiAIStudioModelActionURL(normalizedBaseURL, mappedModel, upstreamAction, useUpstreamStream)
- if err != nil {
- return nil, "", err
- }
- upstreamReq, err := http.NewRequestWithContext(ctx, http.MethodPost, fullURL, bytes.NewReader(body))
- if err != nil {
- return nil, "", err
- }
- upstreamReq.Header.Set("Content-Type", "application/json")
- upstreamReq.Header.Set("x-goog-api-key", apiKey)
- return upstreamReq, "x-request-id", nil
- }
- requestIDHeader = "x-request-id"
- case AccountTypeOAuth:
- buildReq = func(ctx context.Context) (*http.Request, string, error) {
- if s.tokenProvider == nil {
- return nil, "", errors.New("gemini token provider not configured")
- }
- accessToken, err := s.tokenProvider.GetAccessToken(ctx, account)
- if err != nil {
- return nil, "", err
- }
- projectID := strings.TrimSpace(account.GetCredential("project_id"))
- // Two modes for OAuth:
- // 1. With project_id -> Code Assist API (wrapped request)
- // 2. Without project_id -> AI Studio API (direct OAuth, like API key but with Bearer token)
- if projectID != "" && !forceAIStudio {
- // Mode 1: Code Assist API
- baseURL, err := s.validateUpstreamBaseURL(geminicli.GeminiCliBaseURL)
- if err != nil {
- return nil, "", err
- }
- fullURL := fmt.Sprintf("%s/v1internal:%s", strings.TrimRight(baseURL, "/"), upstreamAction)
- if useUpstreamStream {
- fullURL += "?alt=sse"
- }
- wrapped := map[string]any{
- "model": mappedModel,
- "project": projectID,
- }
- var inner any
- if err := json.Unmarshal(body, &inner); err != nil {
- return nil, "", fmt.Errorf("failed to parse gemini request: %w", err)
- }
- wrapped["request"] = inner
- wrappedBytes, _ := json.Marshal(wrapped)
- upstreamReq, err := http.NewRequestWithContext(ctx, http.MethodPost, fullURL, bytes.NewReader(wrappedBytes))
- if err != nil {
- return nil, "", err
- }
- upstreamReq.Header.Set("Content-Type", "application/json")
- upstreamReq.Header.Set("Authorization", "Bearer "+accessToken)
- upstreamReq.Header.Set("User-Agent", geminicli.GeminiCLIUserAgent)
- return upstreamReq, "x-request-id", nil
- } else {
- // Mode 2: AI Studio API with OAuth (like API key mode, but using Bearer token)
- baseURL := account.GetGeminiBaseURL(geminicli.AIStudioBaseURL)
- normalizedBaseURL, err := s.validateUpstreamBaseURL(baseURL)
- if err != nil {
- return nil, "", err
- }
- fullURL, err := buildGeminiAIStudioModelActionURL(normalizedBaseURL, mappedModel, upstreamAction, useUpstreamStream)
- if err != nil {
- return nil, "", err
- }
- upstreamReq, err := http.NewRequestWithContext(ctx, http.MethodPost, fullURL, bytes.NewReader(body))
- if err != nil {
- return nil, "", err
- }
- upstreamReq.Header.Set("Content-Type", "application/json")
- upstreamReq.Header.Set("Authorization", "Bearer "+accessToken)
- return upstreamReq, "x-request-id", nil
- }
- }
- requestIDHeader = "x-request-id"
- case AccountTypeServiceAccount:
- buildReq = func(ctx context.Context) (*http.Request, string, error) {
- if s.tokenProvider == nil {
- return nil, "", errors.New("gemini token provider not configured")
- }
- accessToken, err := s.tokenProvider.GetAccessToken(ctx, account)
- if err != nil {
- return nil, "", err
- }
- fullURL, err := buildVertexGeminiURL(account.VertexProjectID(), account.VertexLocation(mappedModel), mappedModel, upstreamAction, useUpstreamStream)
- if err != nil {
- return nil, "", err
- }
- upstreamReq, err := http.NewRequestWithContext(ctx, http.MethodPost, fullURL, bytes.NewReader(body))
- if err != nil {
- return nil, "", err
- }
- upstreamReq.Header.Set("Content-Type", "application/json")
- upstreamReq.Header.Set("Authorization", "Bearer "+accessToken)
- return upstreamReq, "x-request-id", nil
- }
- requestIDHeader = "x-request-id"
- default:
- return nil, s.writeGoogleError(c, http.StatusBadGateway, "Unsupported account type: "+account.Type)
- }
- var resp *http.Response
- for attempt := 1; attempt <= geminiMaxRetries; attempt++ {
- upstreamReq, idHeader, err := buildReq(ctx)
- if err != nil {
- if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
- return nil, err
- }
- // Local build error: don't retry.
- if strings.Contains(err.Error(), "missing project_id") {
- return nil, s.writeGoogleError(c, http.StatusBadRequest, err.Error())
- }
- return nil, s.writeGoogleError(c, http.StatusBadGateway, err.Error())
- }
- requestIDHeader = idHeader
- resp, err = s.httpUpstream.Do(WithAccountTrafficRequest(upstreamReq, account), proxyURL, account.ID, account.Concurrency)
- if err != nil {
- if local := AccountTrafficFailover(err); local != nil {
- return nil, local
- }
- safeErr := sanitizeUpstreamErrorMessage(err.Error())
- appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
- ProxyID: opsUpstreamProxyID(account),
- ProxyName: opsUpstreamProxyName(account),
- Platform: account.Platform,
- AccountID: account.ID,
- AccountName: account.Name,
- UpstreamStatusCode: 0,
- Kind: "request_error",
- Message: safeErr,
- })
- if attempt < geminiMaxRetries {
- logger.LegacyPrintf("service.gemini_messages_compat", "Gemini account %d: upstream request failed, retry %d/%d: %v", account.ID, attempt, geminiMaxRetries, err)
- sleepGeminiBackoff(attempt)
- continue
- }
- if action == "countTokens" {
- estimated := estimateGeminiCountTokens(body)
- c.JSON(http.StatusOK, map[string]any{"totalTokens": estimated})
- return &ForwardResult{
- RequestID: "",
- Usage: ClaudeUsage{},
- Model: originalModel,
- UpstreamModel: mappedModel,
- Stream: false,
- Duration: time.Since(startTime),
- FirstTokenMs: nil,
- }, nil
- }
- setOpsUpstreamError(c, 0, safeErr, "")
- return nil, s.writeGoogleError(c, http.StatusBadGateway, "Upstream request failed after retries: "+safeErr)
- }
- // 错误策略优先:匹配则跳过重试直接处理。
- if matched, rebuilt := s.checkErrorPolicyInLoop(ctx, account, resp, mappedModel); matched {
- resp = rebuilt
- break
- } else {
- resp = rebuilt
- }
- if resp.StatusCode >= 400 && s.shouldRetryGeminiUpstreamError(account, resp.StatusCode) {
- respBody := s.readUpstreamErrorBody(resp)
- _ = resp.Body.Close()
- // Don't treat insufficient-scope as transient.
- if resp.StatusCode == 403 && isGeminiInsufficientScope(resp.Header, respBody) {
- resp = &http.Response{
- StatusCode: resp.StatusCode,
- Header: resp.Header.Clone(),
- Body: io.NopCloser(bytes.NewReader(respBody)),
- }
- break
- }
- if resp.StatusCode == 429 {
- s.handleGeminiUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
- }
- if attempt < geminiMaxRetries {
- upstreamReqID := resp.Header.Get(requestIDHeader)
- if upstreamReqID == "" {
- upstreamReqID = resp.Header.Get("x-goog-request-id")
- }
- upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody))
- upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
- upstreamDetail := ""
- if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody {
- maxBytes := s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes
- if maxBytes <= 0 {
- maxBytes = 2048
- }
- upstreamDetail = truncateString(string(respBody), maxBytes)
- }
- appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
- ProxyID: opsUpstreamProxyID(account),
- ProxyName: opsUpstreamProxyName(account),
- Platform: account.Platform,
- AccountID: account.ID,
- AccountName: account.Name,
- UpstreamStatusCode: resp.StatusCode,
- UpstreamRequestID: upstreamReqID,
- Kind: "retry",
- Message: upstreamMsg,
- Detail: upstreamDetail,
- })
- logger.LegacyPrintf("service.gemini_messages_compat", "Gemini account %d: upstream status %d, retry %d/%d", account.ID, resp.StatusCode, attempt, geminiMaxRetries)
- sleepGeminiBackoff(attempt)
- continue
- }
- if action == "countTokens" {
- estimated := estimateGeminiCountTokens(body)
- c.JSON(http.StatusOK, map[string]any{"totalTokens": estimated})
- return &ForwardResult{
- RequestID: "",
- Usage: ClaudeUsage{},
- Model: originalModel,
- UpstreamModel: mappedModel,
- Stream: false,
- Duration: time.Since(startTime),
- FirstTokenMs: nil,
- }, nil
- }
- // Final attempt: surface the upstream error body (passed through below) instead of a generic retry error.
- resp = &http.Response{
- StatusCode: resp.StatusCode,
- Header: resp.Header.Clone(),
- Body: io.NopCloser(bytes.NewReader(respBody)),
- }
- break
- }
- break
- }
- defer func() { _ = resp.Body.Close() }()
- requestID := resp.Header.Get(requestIDHeader)
- if requestID == "" {
- requestID = resp.Header.Get("x-goog-request-id")
- }
- if requestID != "" {
- c.Header("x-request-id", requestID)
- }
- isOAuth := account.Type == AccountTypeOAuth
- if resp.StatusCode >= 400 {
- respBody := s.readUpstreamErrorBody(resp)
- // Best-effort fallback for OAuth tokens missing AI Studio scopes when calling countTokens.
- // This avoids Gemini SDKs failing hard during preflight token counting.
- // Checked before error policy so it always works regardless of custom error codes.
- if action == "countTokens" && isOAuth && isGeminiInsufficientScope(resp.Header, respBody) {
- estimated := estimateGeminiCountTokens(body)
- c.JSON(http.StatusOK, map[string]any{"totalTokens": estimated})
- return &ForwardResult{
- RequestID: requestID,
- UpstreamHeaders: resp.Header,
- Usage: ClaudeUsage{},
- Model: originalModel,
- UpstreamModel: mappedModel,
- Stream: false,
- Duration: time.Since(startTime),
- FirstTokenMs: nil,
- }, nil
- }
- // 统一错误策略:自定义错误码 + 临时不可调度
- if s.rateLimitService != nil {
- policy := s.rateLimitService.CheckErrorPolicy(ctx, account, resp.StatusCode, respBody, mappedModel)
- switch policy {
- case ErrorPolicySkipped:
- if failoverErr := s.skippedErrorPolicyFailoverError(c, account, resp.StatusCode, respBody, requestID); failoverErr != nil {
- return nil, failoverErr
- }
- if account.IsCustomErrorCodesEnabled() {
- return nil, s.writeGeminiCustomCodeSkippedError(c, account, resp.StatusCode, requestID, respBody, func() {
- _ = s.writeGoogleError(c, http.StatusInternalServerError, geminiCustomCodeSkippedClientMessage)
- })
- }
- // 池模式:客户端写出与 ErrorPolicyNone 相同(状态码/响应体保真),仅跳过账号状态标记。
- return nil, s.writeGeminiNativeUpstreamError(c, account, resp, respBody, requestID, isOAuth)
- case ErrorPolicyMatched, ErrorPolicyTempUnscheduled:
- if policy == ErrorPolicyMatched {
- s.handleGeminiUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
- }
- evBody := unwrapIfNeeded(isOAuth, respBody)
- upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(evBody))
- upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
- upstreamDetail := ""
- if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody {
- maxBytes := s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes
- if maxBytes <= 0 {
- maxBytes = 2048
- }
- upstreamDetail = truncateString(string(evBody), maxBytes)
- }
- appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
- ProxyID: opsUpstreamProxyID(account),
- ProxyName: opsUpstreamProxyName(account),
- Platform: account.Platform,
- AccountID: account.ID,
- AccountName: account.Name,
- UpstreamStatusCode: resp.StatusCode,
- UpstreamRequestID: requestID,
- Kind: "failover",
- Message: upstreamMsg,
- Detail: upstreamDetail,
- })
- return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: respBody}
- }
- }
- // ErrorPolicyNone → 原有逻辑
- s.handleGeminiUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
- // 精确匹配服务端配置类 400 错误,触发 failover + 临时封禁
- if resp.StatusCode == http.StatusBadRequest {
- msg400 := strings.ToLower(strings.TrimSpace(extractUpstreamErrorMessage(respBody)))
- if isGoogleProjectConfigError(msg400) {
- evBody := unwrapIfNeeded(isOAuth, respBody)
- upstreamMsg := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractUpstreamErrorMessage(evBody)))
- upstreamDetail := ""
- if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody {
- maxBytes := s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes
- if maxBytes <= 0 {
- maxBytes = 2048
- }
- upstreamDetail = truncateString(string(evBody), maxBytes)
- }
- log.Printf("[Gemini] status=400 google_config_error failover=true upstream_message=%q account=%d", upstreamMsg, account.ID)
- appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
- ProxyID: opsUpstreamProxyID(account),
- ProxyName: opsUpstreamProxyName(account),
- Platform: account.Platform,
- AccountID: account.ID,
- AccountName: account.Name,
- UpstreamStatusCode: resp.StatusCode,
- UpstreamRequestID: requestID,
- Kind: "failover",
- Message: upstreamMsg,
- Detail: upstreamDetail,
- })
- return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: evBody, RetryableOnSameAccount: true}
- }
- }
- if s.shouldFailoverGeminiUpstreamError(resp.StatusCode) {
- evBody := unwrapIfNeeded(isOAuth, respBody)
- upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(evBody))
- upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
- upstreamDetail := ""
- if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody {
- maxBytes := s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes
- if maxBytes <= 0 {
- maxBytes = 2048
- }
- upstreamDetail = truncateString(string(evBody), maxBytes)
- }
- appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
- ProxyID: opsUpstreamProxyID(account),
- ProxyName: opsUpstreamProxyName(account),
- Platform: account.Platform,
- AccountID: account.ID,
- AccountName: account.Name,
- UpstreamStatusCode: resp.StatusCode,
- UpstreamRequestID: requestID,
- Kind: "failover",
- Message: upstreamMsg,
- Detail: upstreamDetail,
- })
- return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: evBody}
- }
- return nil, s.writeGeminiNativeUpstreamError(c, account, resp, respBody, requestID, isOAuth)
- }
- var usage *ClaudeUsage
- var firstTokenMs *int
- if stream {
- streamRes, err := s.handleNativeStreamingResponse(c, resp, startTime, isOAuth)
- if err != nil {
- return nil, err
- }
- usage = streamRes.usage
- firstTokenMs = streamRes.firstTokenMs
- } else {
- if useUpstreamStream {
- collected, usageObj, err := collectGeminiSSE(resp.Body, isOAuth)
- if err != nil {
- return nil, s.writeGoogleError(c, http.StatusBadGateway, "Failed to read upstream stream")
- }
- b, _ := json.Marshal(collected)
- upstreamResponseModelObserverFromContext(c).ObserveGemini(b)
- observeGeminiImageOutputs(c, b)
- c.Data(http.StatusOK, "application/json", b)
- usage = usageObj
- } else {
- usageResp, err := s.handleNativeNonStreamingResponse(c, resp, isOAuth)
- if err != nil {
- return nil, err
- }
- usage = usageResp
- }
- }
- if usage == nil {
- usage = &ClaudeUsage{}
- }
- // 图片生成计费
- imageInputSize := s.extractImageInputSize(body)
- imageSize := normalizeOpenAIImageSizeTier(imageInputSize)
- imageCount := resolveGeminiImageCount(c, originalModel, mappedModel)
- return &ForwardResult{
- RequestID: requestID,
- UpstreamHeaders: resp.Header,
- Usage: *usage,
- Model: originalModel,
- UpstreamModel: mappedModel,
- UpstreamResponseModel: observedUpstreamResponseModel(c),
- UpstreamResponseModelConflict: observedUpstreamResponseModelConflict(c),
- Stream: stream,
- Duration: time.Since(startTime),
- FirstTokenMs: firstTokenMs,
- ImageCount: imageCount,
- ImageSize: imageSize,
- ImageInputSize: imageInputSize,
- }, nil
- }
- // checkErrorPolicyInLoop 在重试循环内预检查错误策略。
- // 返回 true 表示策略已匹配(调用者应 break),resp 已重建可直接使用。
- // 返回 false 表示 ErrorPolicyNone,resp 已重建,调用者继续走重试逻辑。
- func (s *GeminiMessagesCompatService) checkErrorPolicyInLoop(
- ctx context.Context, account *Account, resp *http.Response, mappedModel string,
- ) (matched bool, rebuilt *http.Response) {
- if resp.StatusCode < 400 || s.rateLimitService == nil {
- return false, resp
- }
- body := s.readUpstreamErrorBody(resp)
- _ = resp.Body.Close()
- rebuilt = &http.Response{
- StatusCode: resp.StatusCode,
- Header: resp.Header.Clone(),
- Body: io.NopCloser(bytes.NewReader(body)),
- }
- policy := s.rateLimitService.CheckErrorPolicy(ctx, account, resp.StatusCode, body, mappedModel)
- return policy != ErrorPolicyNone, rebuilt
- }
- func (s *GeminiMessagesCompatService) shouldRetryGeminiUpstreamError(account *Account, statusCode int) bool {
- switch statusCode {
- case 429, 500, 502, 503, 504, 529:
- return true
- case 403:
- // GeminiCli OAuth occasionally returns 403 transiently (activation/quota propagation); allow retry.
- if account == nil || account.Type != AccountTypeOAuth {
- return false
- }
- oauthType := strings.ToLower(strings.TrimSpace(account.GetCredential("oauth_type")))
- if oauthType == "" && strings.TrimSpace(account.GetCredential("project_id")) != "" {
- // Legacy/implicit Code Assist OAuth accounts.
- oauthType = "code_assist"
- }
- return oauthType == "code_assist"
- default:
- return false
- }
- }
- func (s *GeminiMessagesCompatService) shouldFailoverGeminiUpstreamError(statusCode int) bool {
- switch statusCode {
- case 401, 403, 429, 529:
- return true
- default:
- return statusCode >= 500
- }
- }
- // skippedErrorPolicyFailoverError 命中 ErrorPolicySkipped(池模式、或自定义错误码未命中)
- // 时构造 failover 错误:可 failover 的状态码返回 UpstreamFailoverError,交给 handler 层换号
- // (池模式账号按 pool_mode_retry_count 先同账号重试);返回 nil 表示状态码不可 failover,
- // 由调用方决定客户端写出。Skipped 只豁免账号状态标记,不豁免换号,与 OpenAI 网关路径一致。
- func (s *GeminiMessagesCompatService) skippedErrorPolicyFailoverError(c *gin.Context, account *Account, statusCode int, respBody []byte, upstreamRequestID string) *UpstreamFailoverError {
- if !s.shouldFailoverGeminiUpstreamError(statusCode) {
- return nil
- }
- upstreamMsg := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractUpstreamErrorMessage(respBody)))
- upstreamDetail := s.upstreamErrorDetail(respBody)
- appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
- ProxyID: opsUpstreamProxyID(account),
- ProxyName: opsUpstreamProxyName(account),
- Platform: account.Platform,
- AccountID: account.ID,
- AccountName: account.Name,
- UpstreamStatusCode: statusCode,
- UpstreamRequestID: upstreamRequestID,
- Kind: "failover",
- Message: upstreamMsg,
- Detail: upstreamDetail,
- })
- return &UpstreamFailoverError{
- StatusCode: statusCode,
- ResponseBody: respBody,
- RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(statusCode),
- }
- }
- // geminiCustomCodeSkippedClientMessage 自定义错误码未命中时对客户端隐藏上游细节的固定文案,
- // 与 OpenAI 网关路径同场景的文案一致。
- const geminiCustomCodeSkippedClientMessage = "Upstream gateway error"
- // upstreamErrorDetail 按配置截断上游错误响应体,用于 ops 错误日志的 Detail 字段;
- // 未开启 LogUpstreamErrorBody 时返回空。
- func (s *GeminiMessagesCompatService) upstreamErrorDetail(body []byte) string {
- if s.cfg == nil || !s.cfg.Gateway.LogUpstreamErrorBody {
- return ""
- }
- maxBytes := s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes
- if maxBytes <= 0 {
- maxBytes = 2048
- }
- return truncateString(string(body), maxBytes)
- }
- // writeGeminiCustomCodeSkippedError 处理自定义错误码未命中且不可 failover 的上游错误:
- // 客户端统一收到 500 + 固定文案(由 write 按端点格式写出),不透传上游细节;
- // 上游真实状态码与错误信息仅记录到 ops 错误日志。
- func (s *GeminiMessagesCompatService) writeGeminiCustomCodeSkippedError(c *gin.Context, account *Account, upstreamStatus int, upstreamRequestID string, body []byte, write func()) error {
- upstreamMsg := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractUpstreamErrorMessage(body)))
- upstreamDetail := s.upstreamErrorDetail(body)
- setOpsUpstreamError(c, upstreamStatus, upstreamMsg, upstreamDetail)
- appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
- ProxyID: opsUpstreamProxyID(account),
- ProxyName: opsUpstreamProxyName(account),
- Platform: account.Platform,
- AccountID: account.ID,
- AccountName: account.Name,
- UpstreamStatusCode: upstreamStatus,
- UpstreamRequestID: upstreamRequestID,
- Kind: "http_error",
- Message: upstreamMsg,
- Detail: upstreamDetail,
- })
- write()
- if upstreamMsg == "" {
- return fmt.Errorf("gemini upstream error: %d (not in custom error codes)", upstreamStatus)
- }
- return fmt.Errorf("gemini upstream error: %d (not in custom error codes) message=%s", upstreamStatus, upstreamMsg)
- }
- // writeGeminiNativeUpstreamError 将不可 failover 的上游错误按原始状态码与响应体透传给客户端,
- // 并记录 ops 错误事件。状态码保真:下游据此区分请求级错误与可重试的链路故障。
- func (s *GeminiMessagesCompatService) writeGeminiNativeUpstreamError(c *gin.Context, account *Account, resp *http.Response, respBody []byte, requestID string, isOAuth bool) error {
- respBody = unwrapIfNeeded(isOAuth, respBody)
- upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody))
- upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
- upstreamDetail := s.upstreamErrorDetail(respBody)
- if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody {
- logger.LegacyPrintf("service.gemini_messages_compat", "[Gemini] native upstream error %d: %s", resp.StatusCode, truncateForLog(respBody, s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes))
- }
- setOpsUpstreamError(c, resp.StatusCode, upstreamMsg, upstreamDetail)
- appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
- ProxyID: opsUpstreamProxyID(account),
- ProxyName: opsUpstreamProxyName(account),
- Platform: account.Platform,
- AccountID: account.ID,
- AccountName: account.Name,
- UpstreamStatusCode: resp.StatusCode,
- UpstreamRequestID: requestID,
- Kind: "http_error",
- Message: upstreamMsg,
- Detail: upstreamDetail,
- })
- contentType := resp.Header.Get("Content-Type")
- if contentType == "" {
- contentType = "application/json"
- }
- MarkResponseCommitted(c)
- c.Data(resp.StatusCode, contentType, respBody)
- if upstreamMsg == "" {
- return fmt.Errorf("gemini upstream error: %d", resp.StatusCode)
- }
- return fmt.Errorf("gemini upstream error: %d message=%s", resp.StatusCode, upstreamMsg)
- }
- func sleepGeminiBackoff(attempt int) {
- delay := geminiRetryBaseDelay * time.Duration(1<<uint(attempt-1))
- if delay > geminiRetryMaxDelay {
- delay = geminiRetryMaxDelay
- }
- // +/- 20% jitter
- r := mathrand.New(mathrand.NewSource(time.Now().UnixNano()))
- jitter := time.Duration(float64(delay) * 0.2 * (r.Float64()*2 - 1))
- sleepFor := delay + jitter
- if sleepFor < 0 {
- sleepFor = 0
- }
- time.Sleep(sleepFor)
- }
- var (
- sensitiveQueryParamRegex = regexp.MustCompile(`(?i)([?&](?:key|client_secret|access_token|refresh_token)=)[^&"\s]+`)
- retryInRegex = regexp.MustCompile(`Please retry in ([0-9.]+)s`)
- )
- func sanitizeUpstreamErrorMessage(msg string) string {
- if msg == "" {
- return msg
- }
- return sensitiveQueryParamRegex.ReplaceAllString(msg, `$1***`)
- }
- func (s *GeminiMessagesCompatService) writeGeminiMappedError(c *gin.Context, account *Account, upstreamStatus int, upstreamRequestID string, body []byte) error {
- MarkResponseCommitted(c)
- upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(body))
- upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
- upstreamDetail := ""
- if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody {
- maxBytes := s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes
- if maxBytes <= 0 {
- maxBytes = 2048
- }
- upstreamDetail = truncateString(string(body), maxBytes)
- }
- setOpsUpstreamError(c, upstreamStatus, upstreamMsg, upstreamDetail)
- appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
- ProxyID: opsUpstreamProxyID(account),
- ProxyName: opsUpstreamProxyName(account),
- Platform: account.Platform,
- AccountID: account.ID,
- AccountName: account.Name,
- UpstreamStatusCode: upstreamStatus,
- UpstreamRequestID: upstreamRequestID,
- Kind: "http_error",
- Message: upstreamMsg,
- Detail: upstreamDetail,
- })
- if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody {
- logger.LegacyPrintf("service.gemini_messages_compat", "[Gemini] upstream error %d: %s", upstreamStatus, truncateForLog(body, s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes))
- }
- if status, errType, errMsg, matched := applyErrorPassthroughRule(
- c,
- PlatformGemini,
- upstreamStatus,
- body,
- http.StatusBadGateway,
- "upstream_error",
- "Upstream request failed",
- ); matched {
- c.JSON(status, gin.H{
- "type": "error",
- "error": gin.H{"type": errType, "message": errMsg},
- })
- if upstreamMsg == "" {
- upstreamMsg = errMsg
- }
- if upstreamMsg == "" {
- return fmt.Errorf("upstream error: %d (passthrough rule matched)", upstreamStatus)
- }
- return fmt.Errorf("upstream error: %d (passthrough rule matched) message=%s", upstreamStatus, upstreamMsg)
- }
- var statusCode int
- var errType, errMsg string
- if mapped := mapGeminiErrorBodyToClaudeError(body); mapped != nil {
- errType = mapped.Type
- if mapped.Message != "" {
- errMsg = mapped.Message
- }
- if mapped.StatusCode > 0 {
- statusCode = mapped.StatusCode
- }
- }
- switch upstreamStatus {
- case 400:
- if statusCode == 0 {
- statusCode = http.StatusBadRequest
- }
- if errType == "" {
- errType = "invalid_request_error"
- }
- // 400 是确定性的请求错误:回传上游 message(已脱敏),客户端据此定位非法字段。
- if errMsg == "" {
- errMsg = upstreamMsg
- }
- if errMsg == "" {
- errMsg = "Invalid request"
- }
- case 401:
- if statusCode == 0 {
- statusCode = http.StatusBadGateway
- }
- if errType == "" {
- errType = "authentication_error"
- }
- if errMsg == "" {
- errMsg = "Upstream authentication failed, please contact administrator"
- }
- case 403:
- if statusCode == 0 {
- statusCode = http.StatusBadGateway
- }
- if errType == "" {
- errType = "permission_error"
- }
- if errMsg == "" {
- errMsg = "Upstream access forbidden, please contact administrator"
- }
- case 404:
- if statusCode == 0 {
- statusCode = http.StatusNotFound
- }
- if errType == "" {
- errType = "not_found_error"
- }
- if errMsg == "" {
- errMsg = "Resource not found"
- }
- case 429:
- if statusCode == 0 {
- statusCode = http.StatusTooManyRequests
- }
- if errType == "" {
- errType = "rate_limit_error"
- }
- if errMsg == "" {
- errMsg = "Upstream rate limit exceeded, please retry later"
- }
- case 529:
- if statusCode == 0 {
- statusCode = http.StatusServiceUnavailable
- }
- if errType == "" {
- errType = "overloaded_error"
- }
- if errMsg == "" {
- errMsg = "Upstream service overloaded, please retry later"
- }
- case 500, 502, 503, 504:
- if statusCode == 0 {
- statusCode = http.StatusBadGateway
- }
- if errType == "" {
- switch upstreamStatus {
- case 504:
- errType = "timeout_error"
- case 503:
- errType = "overloaded_error"
- default:
- errType = "api_error"
- }
- }
- if errMsg == "" {
- errMsg = "Upstream service temporarily unavailable"
- }
- default:
- if statusCode == 0 {
- statusCode = http.StatusBadGateway
- }
- if errType == "" {
- errType = "upstream_error"
- }
- if errMsg == "" {
- errMsg = "Upstream request failed"
- }
- }
- c.JSON(statusCode, gin.H{
- "type": "error",
- "error": gin.H{"type": errType, "message": errMsg},
- })
- if upstreamMsg == "" {
- return fmt.Errorf("upstream error: %d", upstreamStatus)
- }
- return fmt.Errorf("upstream error: %d message=%s", upstreamStatus, upstreamMsg)
- }
- type claudeErrorMapping struct {
- Type string
- Message string
- StatusCode int
- }
- func mapGeminiErrorBodyToClaudeError(body []byte) *claudeErrorMapping {
- if len(body) == 0 {
- return nil
- }
- var parsed struct {
- Error struct {
- Code int `json:"code"`
- Message string `json:"message"`
- Status string `json:"status"`
- } `json:"error"`
- }
- if err := json.Unmarshal(body, &parsed); err != nil {
- return nil
- }
- if strings.TrimSpace(parsed.Error.Status) == "" && parsed.Error.Code == 0 && strings.TrimSpace(parsed.Error.Message) == "" {
- return nil
- }
- mapped := &claudeErrorMapping{
- Type: mapGeminiStatusToClaudeErrorType(parsed.Error.Status),
- Message: "",
- }
- if mapped.Type == "" {
- mapped.Type = "upstream_error"
- }
- switch strings.ToUpper(strings.TrimSpace(parsed.Error.Status)) {
- case "INVALID_ARGUMENT":
- mapped.StatusCode = http.StatusBadRequest
- case "NOT_FOUND":
- mapped.StatusCode = http.StatusNotFound
- case "RESOURCE_EXHAUSTED":
- mapped.StatusCode = http.StatusTooManyRequests
- default:
- // Keep StatusCode unset and let HTTP status mapping decide.
- }
- // Keep messages generic by default; upstream error message can be long or include sensitive fragments.
- return mapped
- }
- func mapGeminiStatusToClaudeErrorType(status string) string {
- switch strings.ToUpper(strings.TrimSpace(status)) {
- case "INVALID_ARGUMENT":
- return "invalid_request_error"
- case "PERMISSION_DENIED":
- return "permission_error"
- case "NOT_FOUND":
- return "not_found_error"
- case "RESOURCE_EXHAUSTED":
- return "rate_limit_error"
- case "UNAUTHENTICATED":
- return "authentication_error"
- case "UNAVAILABLE":
- return "overloaded_error"
- case "INTERNAL":
- return "api_error"
- case "DEADLINE_EXCEEDED":
- return "timeout_error"
- default:
- return ""
- }
- }
- type geminiStreamResult struct {
- usage *ClaudeUsage
- firstTokenMs *int
- }
- func (s *GeminiMessagesCompatService) handleNonStreamingResponse(c *gin.Context, resp *http.Response, originalModel string) (*ClaudeUsage, error) {
- body, err := io.ReadAll(io.LimitReader(resp.Body, 8<<20))
- if err != nil {
- return nil, s.writeClaudeError(c, http.StatusBadGateway, "upstream_error", "Failed to read upstream response")
- }
- unwrappedBody, err := unwrapGeminiResponse(body)
- if err != nil {
- return nil, s.writeClaudeError(c, http.StatusBadGateway, "upstream_error", "Failed to parse upstream response")
- }
- observer := upstreamResponseModelObserverFromContext(c)
- if observer == nil {
- observer = beginUpstreamResponseModelObservation(c)
- }
- observer.ObserveGemini(unwrappedBody)
- observeGeminiImageOutputs(c, unwrappedBody)
- var geminiResp map[string]any
- if err := json.Unmarshal(unwrappedBody, &geminiResp); err != nil {
- return nil, s.writeClaudeError(c, http.StatusBadGateway, "upstream_error", "Failed to parse upstream response")
- }
- claudeResp, usage := convertGeminiToClaudeMessage(geminiResp, originalModel, unwrappedBody, false)
- c.JSON(http.StatusOK, claudeResp)
- return usage, nil
- }
- func (s *GeminiMessagesCompatService) handleStreamingResponse(c *gin.Context, resp *http.Response, startTime time.Time, originalModel string) (*geminiStreamResult, error) {
- c.Header("Content-Type", "text/event-stream")
- c.Header("Cache-Control", "no-cache")
- c.Header("Connection", "keep-alive")
- c.Header("X-Accel-Buffering", "no")
- c.Status(http.StatusOK)
- flusher, ok := c.Writer.(http.Flusher)
- if !ok {
- return nil, errors.New("streaming not supported")
- }
- messageID := generateAnthropicMsgID()
- messageStart := map[string]any{
- "type": "message_start",
- "message": map[string]any{
- "id": messageID,
- "type": "message",
- "role": "assistant",
- "model": originalModel,
- "content": []any{},
- "stop_reason": nil,
- "stop_sequence": nil,
- "usage": map[string]any{
- "input_tokens": 0,
- "output_tokens": 0,
- },
- },
- }
- writeSSE(c.Writer, "message_start", messageStart)
- flusher.Flush()
- var firstTokenMs *int
- var usage ClaudeUsage
- finishReason := ""
- sawToolUse := false
- nextBlockIndex := 0
- openBlockIndex := -1
- openBlockType := ""
- seenText := ""
- openToolIndex := -1
- openToolID := ""
- openToolName := ""
- seenToolJSON := ""
- reader := bufio.NewReader(resp.Body)
- for {
- line, err := reader.ReadString('\n')
- if err != nil && !errors.Is(err, io.EOF) {
- return nil, fmt.Errorf("stream read error: %w", err)
- }
- if !strings.HasPrefix(line, "data:") {
- if errors.Is(err, io.EOF) {
- break
- }
- continue
- }
- payload := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
- if payload == "" || payload == "[DONE]" {
- if errors.Is(err, io.EOF) {
- break
- }
- continue
- }
- unwrappedBytes, err := unwrapGeminiResponse([]byte(payload))
- if err != nil {
- continue
- }
- observer := upstreamResponseModelObserverFromContext(c)
- if observer == nil {
- observer = beginUpstreamResponseModelObservation(c)
- }
- observer.ObserveGemini(unwrappedBytes)
- observeGeminiImageOutputs(c, unwrappedBytes)
- var geminiResp map[string]any
- if err := json.Unmarshal(unwrappedBytes, &geminiResp); err != nil {
- continue
- }
- if fr := extractGeminiFinishReason(geminiResp); fr != "" {
- finishReason = fr
- }
- parts := extractGeminiParts(geminiResp)
- for _, part := range parts {
- if text, ok := part["text"].(string); ok && text != "" {
- // Close an open tool_use block before starting text, mirroring
- // the functionCall branch (which closes open text blocks) and
- // the chat-completions sibling's closeOpenTool(). Otherwise a
- // tool→text sequence keeps the tool_use block open while the
- // text block starts, emitting overlapping Anthropic content
- // blocks that violate the SSE contract.
- if openToolIndex >= 0 {
- writeSSE(c.Writer, "content_block_stop", map[string]any{
- "type": "content_block_stop",
- "index": openToolIndex,
- })
- openToolIndex = -1
- openToolName = ""
- seenToolJSON = ""
- }
- delta, newSeen := computeGeminiTextDelta(seenText, text)
- seenText = newSeen
- if delta == "" {
- continue
- }
- if openBlockType != "text" {
- if openBlockIndex >= 0 {
- writeSSE(c.Writer, "content_block_stop", map[string]any{
- "type": "content_block_stop",
- "index": openBlockIndex,
- })
- }
- openBlockType = "text"
- openBlockIndex = nextBlockIndex
- nextBlockIndex++
- writeSSE(c.Writer, "content_block_start", map[string]any{
- "type": "content_block_start",
- "index": openBlockIndex,
- "content_block": map[string]any{
- "type": "text",
- "text": "",
- },
- })
- }
- if firstTokenMs == nil {
- ms := int(time.Since(startTime).Milliseconds())
- firstTokenMs = &ms
- }
- writeSSE(c.Writer, "content_block_delta", map[string]any{
- "type": "content_block_delta",
- "index": openBlockIndex,
- "delta": map[string]any{
- "type": "text_delta",
- "text": delta,
- },
- })
- flusher.Flush()
- continue
- }
- if fc, ok := part["functionCall"].(map[string]any); ok && fc != nil {
- name, _ := fc["name"].(string)
- args := fc["args"]
- if strings.TrimSpace(name) == "" {
- name = "tool"
- }
- // Close any open text block before tool_use.
- if openBlockIndex >= 0 {
- writeSSE(c.Writer, "content_block_stop", map[string]any{
- "type": "content_block_stop",
- "index": openBlockIndex,
- })
- openBlockIndex = -1
- openBlockType = ""
- }
- // If we receive streamed tool args in pieces, keep a single tool block open and emit deltas.
- if openToolIndex >= 0 && openToolName != name {
- writeSSE(c.Writer, "content_block_stop", map[string]any{
- "type": "content_block_stop",
- "index": openToolIndex,
- })
- openToolIndex = -1
- openToolName = ""
- seenToolJSON = ""
- }
- if openToolIndex < 0 {
- openToolID = "toolu_" + randomHex(8)
- openToolIndex = nextBlockIndex
- openToolName = name
- nextBlockIndex++
- sawToolUse = true
- writeSSE(c.Writer, "content_block_start", map[string]any{
- "type": "content_block_start",
- "index": openToolIndex,
- "content_block": map[string]any{
- "type": "tool_use",
- "id": openToolID,
- "name": name,
- "input": map[string]any{},
- },
- })
- }
- argsJSONText := "{}"
- switch v := args.(type) {
- case nil:
- // keep default "{}"
- case string:
- if strings.TrimSpace(v) != "" {
- argsJSONText = v
- }
- default:
- if b, err := json.Marshal(args); err == nil && len(b) > 0 {
- argsJSONText = string(b)
- }
- }
- delta, newSeen := computeGeminiTextDelta(seenToolJSON, argsJSONText)
- seenToolJSON = newSeen
- if delta != "" {
- writeSSE(c.Writer, "content_block_delta", map[string]any{
- "type": "content_block_delta",
- "index": openToolIndex,
- "delta": map[string]any{
- "type": "input_json_delta",
- "partial_json": delta,
- },
- })
- }
- flusher.Flush()
- }
- }
- if u := extractGeminiUsage(unwrappedBytes); u != nil {
- usage = *u
- }
- // Process the final unterminated line at EOF as well.
- if errors.Is(err, io.EOF) {
- break
- }
- }
- if openBlockIndex >= 0 {
- writeSSE(c.Writer, "content_block_stop", map[string]any{
- "type": "content_block_stop",
- "index": openBlockIndex,
- })
- }
- if openToolIndex >= 0 {
- writeSSE(c.Writer, "content_block_stop", map[string]any{
- "type": "content_block_stop",
- "index": openToolIndex,
- })
- }
- stopReason := mapGeminiFinishReasonToClaudeStopReason(finishReason)
- if sawToolUse {
- stopReason = "tool_use"
- }
- usageObj := map[string]any{
- "output_tokens": usage.OutputTokens,
- }
- if usage.InputTokens > 0 {
- usageObj["input_tokens"] = usage.InputTokens
- }
- writeSSE(c.Writer, "message_delta", map[string]any{
- "type": "message_delta",
- "delta": map[string]any{
- "stop_reason": stopReason,
- "stop_sequence": nil,
- },
- "usage": usageObj,
- })
- writeSSE(c.Writer, "message_stop", map[string]any{
- "type": "message_stop",
- })
- flusher.Flush()
- return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, nil
- }
- func writeSSE(w io.Writer, event string, data any) {
- if event != "" {
- _, _ = fmt.Fprintf(w, "event: %s\n", event)
- }
- b, _ := json.Marshal(data)
- _, _ = fmt.Fprintf(w, "data: %s\n\n", string(b))
- }
- func randomHex(nBytes int) string {
- b := make([]byte, nBytes)
- _, _ = rand.Read(b)
- return hex.EncodeToString(b)
- }
- // generateAnthropicMsgID 生成 Anthropic 官方格式的 message ID:msg_01 + 22 位 Base62
- func generateAnthropicMsgID() string {
- const charset = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
- const idLen = 22
- randomBytes := make([]byte, idLen)
- if _, err := rand.Read(randomBytes); err != nil {
- return "msg_01" + randomHex(11)
- }
- b := make([]byte, idLen)
- for i := range b {
- b[i] = charset[int(randomBytes[i])%len(charset)]
- }
- return "msg_01" + string(b)
- }
- func (s *GeminiMessagesCompatService) writeClaudeError(c *gin.Context, status int, errType, message string) error {
- MarkResponseCommitted(c)
- c.JSON(status, gin.H{
- "type": "error",
- "error": gin.H{"type": errType, "message": message},
- })
- return fmt.Errorf("%s", message)
- }
- func (s *GeminiMessagesCompatService) writeGoogleError(c *gin.Context, status int, message string) error {
- MarkResponseCommitted(c)
- c.JSON(status, gin.H{
- "error": gin.H{
- "code": status,
- "message": message,
- "status": googleapi.HTTPStatusToGoogleStatus(status),
- },
- })
- return fmt.Errorf("%s", message)
- }
- func unwrapIfNeeded(isOAuth bool, raw []byte) []byte {
- if !isOAuth {
- return raw
- }
- inner, err := unwrapGeminiResponse(raw)
- if err != nil {
- return raw
- }
- return inner
- }
- func collectGeminiSSE(body io.Reader, isOAuth bool) (map[string]any, *ClaudeUsage, error) {
- reader := bufio.NewReader(body)
- var last map[string]any
- var lastWithParts map[string]any
- var collectedTextParts []string // Collect all text parts for aggregation
- usage := &ClaudeUsage{}
- for {
- line, err := reader.ReadString('\n')
- if len(line) > 0 {
- trimmed := strings.TrimRight(line, "\r\n")
- if strings.HasPrefix(trimmed, "data:") {
- payload := strings.TrimSpace(strings.TrimPrefix(trimmed, "data:"))
- switch payload {
- case "", "[DONE]":
- if payload == "[DONE]" {
- return mergeCollectedTextParts(pickGeminiCollectResult(last, lastWithParts), collectedTextParts), usage, nil
- }
- default:
- var parsed map[string]any
- var rawBytes []byte
- if isOAuth {
- innerBytes, err := unwrapGeminiResponse([]byte(payload))
- if err == nil {
- rawBytes = innerBytes
- _ = json.Unmarshal(innerBytes, &parsed)
- }
- } else {
- rawBytes = []byte(payload)
- _ = json.Unmarshal(rawBytes, &parsed)
- }
- if parsed != nil {
- last = parsed
- if u := extractGeminiUsage(rawBytes); u != nil {
- usage = u
- }
- if parts := extractGeminiParts(parsed); len(parts) > 0 {
- lastWithParts = parsed
- // Collect text from each part for aggregation
- for _, part := range parts {
- if text, ok := part["text"].(string); ok && text != "" {
- collectedTextParts = append(collectedTextParts, text)
- }
- }
- }
- }
- }
- }
- }
- if errors.Is(err, io.EOF) {
- break
- }
- if err != nil {
- return nil, nil, err
- }
- }
- return mergeCollectedTextParts(pickGeminiCollectResult(last, lastWithParts), collectedTextParts), usage, nil
- }
- func pickGeminiCollectResult(last map[string]any, lastWithParts map[string]any) map[string]any {
- if lastWithParts != nil {
- return lastWithParts
- }
- if last != nil {
- return last
- }
- return map[string]any{}
- }
- // mergeCollectedTextParts merges all collected text chunks into the final response.
- // This fixes the issue where non-streaming responses only returned the last chunk
- // instead of the complete aggregated text.
- func mergeCollectedTextParts(response map[string]any, textParts []string) map[string]any {
- if len(textParts) == 0 {
- return response
- }
- // Join all text parts
- mergedText := strings.Join(textParts, "")
- // Deep copy response
- result := make(map[string]any)
- for k, v := range response {
- result[k] = v
- }
- // Get or create candidates
- candidates, ok := result["candidates"].([]any)
- if !ok || len(candidates) == 0 {
- candidates = []any{map[string]any{}}
- }
- // Get first candidate
- candidate, ok := candidates[0].(map[string]any)
- if !ok {
- candidate = make(map[string]any)
- candidates[0] = candidate
- }
- // Get or create content
- content, ok := candidate["content"].(map[string]any)
- if !ok {
- content = map[string]any{"role": "model"}
- candidate["content"] = content
- }
- // Get existing parts
- existingParts, ok := content["parts"].([]any)
- if !ok {
- existingParts = []any{}
- }
- // Find and update first text part, or create new one
- 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 {
- // Replace with merged text
- 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...)
- }
- content["parts"] = newParts
- result["candidates"] = candidates
- return result
- }
- type geminiNativeStreamResult struct {
- usage *ClaudeUsage
- firstTokenMs *int
- }
- func isGeminiInsufficientScope(headers http.Header, body []byte) bool {
- if strings.Contains(strings.ToLower(headers.Get("Www-Authenticate")), "insufficient_scope") {
- return true
- }
- lower := strings.ToLower(string(body))
- return strings.Contains(lower, "insufficient authentication scopes") || strings.Contains(lower, "access_token_scope_insufficient")
- }
- func estimateGeminiCountTokens(reqBody []byte) int {
- total := 0
- // systemInstruction.parts[].text
- gjson.GetBytes(reqBody, "systemInstruction.parts").ForEach(func(_, part gjson.Result) bool {
- if t := strings.TrimSpace(part.Get("text").String()); t != "" {
- total += estimateTokensForText(t)
- }
- return true
- })
- // contents[].parts[].text
- gjson.GetBytes(reqBody, "contents").ForEach(func(_, content gjson.Result) bool {
- content.Get("parts").ForEach(func(_, part gjson.Result) bool {
- if t := strings.TrimSpace(part.Get("text").String()); t != "" {
- total += estimateTokensForText(t)
- }
- return true
- })
- return true
- })
- if total < 0 {
- return 0
- }
- return total
- }
- func estimateTokensForText(s string) int {
- s = strings.TrimSpace(s)
- if s == "" {
- return 0
- }
- runes := []rune(s)
- if len(runes) == 0 {
- return 0
- }
- ascii := 0
- for _, r := range runes {
- if r <= 0x7f {
- ascii++
- }
- }
- asciiRatio := float64(ascii) / float64(len(runes))
- if asciiRatio >= 0.8 {
- // Roughly 4 chars per token for English-like text.
- return (len(runes) + 3) / 4
- }
- // For CJK-heavy text, approximate 1 rune per token.
- return len(runes)
- }
- type UpstreamHTTPResult struct {
- StatusCode int
- Headers http.Header
- Body []byte
- }
- func (s *GeminiMessagesCompatService) handleNativeNonStreamingResponse(c *gin.Context, resp *http.Response, isOAuth bool) (*ClaudeUsage, error) {
- if s.cfg != nil && s.cfg.Gateway.GeminiDebugResponseHeaders {
- logger.LegacyPrintf("service.gemini_messages_compat", "[GeminiAPI] ========== Response Headers ==========")
- for key, values := range resp.Header {
- if strings.HasPrefix(strings.ToLower(key), "x-ratelimit") {
- logger.LegacyPrintf("service.gemini_messages_compat", "[GeminiAPI] %s: %v", key, values)
- }
- }
- logger.LegacyPrintf("service.gemini_messages_compat", "[GeminiAPI] ========================================")
- }
- respBody, err := ReadUpstreamResponseBody(resp.Body, s.cfg, c, openAITooLargeError)
- if err != nil {
- return nil, err
- }
- if isOAuth {
- unwrappedBody, uwErr := unwrapGeminiResponse(respBody)
- if uwErr == nil {
- respBody = unwrappedBody
- }
- }
- observer := upstreamResponseModelObserverFromContext(c)
- if observer == nil {
- observer = beginUpstreamResponseModelObservation(c)
- }
- observer.ObserveGemini(respBody)
- observeGeminiImageOutputs(c, respBody)
- responseheaders.WriteFilteredHeaders(c.Writer.Header(), resp.Header, s.responseHeaderFilter)
- contentType := resp.Header.Get("Content-Type")
- if contentType == "" {
- contentType = "application/json"
- }
- c.Data(resp.StatusCode, contentType, respBody)
- if u := extractGeminiUsage(respBody); u != nil {
- return u, nil
- }
- return &ClaudeUsage{}, nil
- }
- func (s *GeminiMessagesCompatService) handleNativeStreamingResponse(c *gin.Context, resp *http.Response, startTime time.Time, isOAuth bool) (*geminiNativeStreamResult, error) {
- if s.cfg != nil && s.cfg.Gateway.GeminiDebugResponseHeaders {
- logger.LegacyPrintf("service.gemini_messages_compat", "[GeminiAPI] ========== Streaming Response Headers ==========")
- for key, values := range resp.Header {
- if strings.HasPrefix(strings.ToLower(key), "x-ratelimit") {
- logger.LegacyPrintf("service.gemini_messages_compat", "[GeminiAPI] %s: %v", key, values)
- }
- }
- logger.LegacyPrintf("service.gemini_messages_compat", "[GeminiAPI] ====================================================")
- }
- if s.responseHeaderFilter != nil {
- responseheaders.WriteFilteredHeaders(c.Writer.Header(), resp.Header, s.responseHeaderFilter)
- }
- c.Status(resp.StatusCode)
- c.Header("Cache-Control", "no-cache")
- c.Header("Connection", "keep-alive")
- c.Header("X-Accel-Buffering", "no")
- contentType := resp.Header.Get("Content-Type")
- if contentType == "" {
- contentType = "text/event-stream; charset=utf-8"
- }
- c.Header("Content-Type", contentType)
- flusher, ok := c.Writer.(http.Flusher)
- if !ok {
- return nil, errors.New("streaming not supported")
- }
- reader := bufio.NewReader(resp.Body)
- usage := &ClaudeUsage{}
- observer := upstreamResponseModelObserverFromContext(c)
- if observer == nil {
- observer = beginUpstreamResponseModelObservation(c)
- }
- var firstTokenMs *int
- for {
- line, err := reader.ReadString('\n')
- if len(line) > 0 {
- trimmed := strings.TrimRight(line, "\r\n")
- if strings.HasPrefix(trimmed, "data:") {
- payload := strings.TrimSpace(strings.TrimPrefix(trimmed, "data:"))
- // Keepalive / done markers
- if payload == "" || payload == "[DONE]" {
- _, _ = io.WriteString(c.Writer, line)
- flusher.Flush()
- } else {
- var rawToWrite string
- rawToWrite = payload
- var rawBytes []byte
- if isOAuth {
- innerBytes, err := unwrapGeminiResponse([]byte(payload))
- if err == nil {
- rawToWrite = string(innerBytes)
- rawBytes = innerBytes
- }
- } else {
- rawBytes = []byte(payload)
- }
- if u := extractGeminiUsage(rawBytes); u != nil {
- usage = u
- }
- observer.ObserveGemini(rawBytes)
- observeGeminiImageOutputs(c, rawBytes)
- if firstTokenMs == nil {
- ms := int(time.Since(startTime).Milliseconds())
- firstTokenMs = &ms
- }
- if isOAuth {
- // SSE format requires double newline (\n\n) to separate events
- _, _ = fmt.Fprintf(c.Writer, "data: %s\n\n", rawToWrite)
- } else {
- // Pass-through for AI Studio responses.
- _, _ = io.WriteString(c.Writer, line)
- }
- flusher.Flush()
- }
- } else {
- _, _ = io.WriteString(c.Writer, line)
- flusher.Flush()
- }
- }
- if errors.Is(err, io.EOF) {
- break
- }
- if err != nil {
- return nil, err
- }
- }
- return &geminiNativeStreamResult{usage: usage, firstTokenMs: firstTokenMs}, nil
- }
- // ForwardAIStudioGET forwards a GET request to AI Studio (generativelanguage.googleapis.com) for
- // endpoints like /v1beta/models and /v1beta/models/{model}.
- //
- // This is used to support Gemini SDKs that call models listing endpoints before generation.
- func (s *GeminiMessagesCompatService) ForwardAIStudioGET(ctx context.Context, account *Account, path string) (*UpstreamHTTPResult, error) {
- if account == nil {
- return nil, errors.New("account is nil")
- }
- // path 会被直接拼到上游 base URL 后面,因此按路径护栏逐片段校验,
- // 见 upstream_path_guard.go。
- sanitizedPath, ok := sanitizedUpstreamPathSuffix(path)
- if !ok || sanitizedPath == "" {
- return nil, errors.New("invalid path")
- }
- path = sanitizedPath
- baseURL := account.GetGeminiBaseURL(geminicli.AIStudioBaseURL)
- normalizedBaseURL, err := s.validateUpstreamBaseURL(baseURL)
- if err != nil {
- return nil, err
- }
- fullURL := strings.TrimRight(normalizedBaseURL, "/") + path
- var proxyURL string
- if account.ProxyID != nil && account.Proxy != nil {
- proxyURL = account.Proxy.URL()
- }
- req, err := http.NewRequestWithContext(ctx, http.MethodGet, fullURL, nil)
- if err != nil {
- return nil, err
- }
- switch account.Type {
- case AccountTypeAPIKey:
- apiKey := strings.TrimSpace(account.GetCredential("api_key"))
- if apiKey == "" {
- return nil, errors.New("gemini api_key not configured")
- }
- req.Header.Set("x-goog-api-key", apiKey)
- case AccountTypeOAuth:
- if s.tokenProvider == nil {
- return nil, errors.New("gemini token provider not configured")
- }
- accessToken, err := s.tokenProvider.GetAccessToken(ctx, account)
- if err != nil {
- return nil, err
- }
- req.Header.Set("Authorization", "Bearer "+accessToken)
- default:
- return nil, fmt.Errorf("unsupported account type: %s", account.Type)
- }
- resp, err := s.httpUpstream.Do(WithAccountTrafficRequest(req, account), proxyURL, account.ID, account.Concurrency)
- if err != nil {
- return nil, err
- }
- defer func() { _ = resp.Body.Close() }()
- body, _ := io.ReadAll(io.LimitReader(resp.Body, 8<<20))
- wwwAuthenticate := resp.Header.Get("Www-Authenticate")
- filteredHeaders := responseheaders.FilterHeaders(resp.Header, s.responseHeaderFilter)
- if wwwAuthenticate != "" {
- filteredHeaders.Set("Www-Authenticate", wwwAuthenticate)
- }
- return &UpstreamHTTPResult{
- StatusCode: resp.StatusCode,
- Headers: filteredHeaders,
- Body: body,
- }, nil
- }
- // unwrapGeminiResponse 解包 Gemini OAuth 响应中的 response 字段
- // 使用 gjson 零拷贝提取,避免完整 Unmarshal+Marshal
- func unwrapGeminiResponse(raw []byte) ([]byte, error) {
- result := gjson.GetBytes(raw, "response")
- if result.Exists() && result.Type == gjson.JSON {
- return []byte(result.Raw), nil
- }
- return raw, nil
- }
- func convertGeminiToClaudeMessage(geminiResp map[string]any, originalModel string, rawData []byte, includeInlineData bool) (map[string]any, *ClaudeUsage) {
- usage := extractGeminiUsage(rawData)
- if usage == nil {
- usage = &ClaudeUsage{}
- }
- contentBlocks := make([]any, 0)
- sawToolUse := false
- if candidates, ok := geminiResp["candidates"].([]any); ok && len(candidates) > 0 {
- if cand, ok := candidates[0].(map[string]any); ok {
- if content, ok := cand["content"].(map[string]any); ok {
- if parts, ok := content["parts"].([]any); ok {
- for _, part := range parts {
- pm, ok := part.(map[string]any)
- if !ok {
- continue
- }
- if text, ok := pm["text"].(string); ok && text != "" {
- contentBlocks = append(contentBlocks, map[string]any{
- "type": "text",
- "text": text,
- })
- }
- if inlineData, ok := pm["inlineData"].(map[string]any); includeInlineData && ok {
- mimeType, _ := inlineData["mimeType"].(string)
- data, _ := inlineData["data"].(string)
- if isGeminiInlineImageMIMEType(mimeType) && isValidBase64(data) {
- contentBlocks = append(contentBlocks, map[string]any{
- "type": "text",
- "text": fmt.Sprintf("", mimeType, data),
- })
- }
- }
- if fc, ok := pm["functionCall"].(map[string]any); ok {
- name, _ := fc["name"].(string)
- if strings.TrimSpace(name) == "" {
- name = "tool"
- }
- args := fc["args"]
- sawToolUse = true
- contentBlocks = append(contentBlocks, map[string]any{
- "type": "tool_use",
- "id": "toolu_" + randomHex(8),
- "name": name,
- "input": args,
- })
- }
- }
- }
- }
- }
- }
- stopReason := mapGeminiFinishReasonToClaudeStopReason(extractGeminiFinishReason(geminiResp))
- if sawToolUse {
- stopReason = "tool_use"
- }
- resp := map[string]any{
- "id": generateAnthropicMsgID(),
- "type": "message",
- "role": "assistant",
- "model": originalModel,
- "content": contentBlocks,
- "stop_reason": stopReason,
- "stop_sequence": nil,
- "usage": map[string]any{
- "input_tokens": usage.InputTokens,
- "output_tokens": usage.OutputTokens,
- },
- }
- return resp, usage
- }
- func isGeminiInlineImageMIMEType(mimeType string) bool {
- switch mimeType {
- case "image/gif", "image/jpeg", "image/png", "image/webp":
- return true
- default:
- return false
- }
- }
- func isValidBase64(data string) bool {
- if data == "" {
- return false
- }
- _, err := base64.StdEncoding.DecodeString(data)
- return err == nil
- }
- func extractGeminiUsage(data []byte) *ClaudeUsage {
- usage := gjson.GetBytes(data, "usageMetadata")
- if !usage.Exists() {
- return nil
- }
- prompt := int(usage.Get("promptTokenCount").Int())
- cand := int(usage.Get("candidatesTokenCount").Int())
- cached := int(usage.Get("cachedContentTokenCount").Int())
- thoughts := int(usage.Get("thoughtsTokenCount").Int())
- // 从 candidatesTokensDetails 提取 IMAGE 模态 token 数
- imageTokens := 0
- candidateDetails := usage.Get("candidatesTokensDetails")
- if candidateDetails.Exists() {
- candidateDetails.ForEach(func(_, detail gjson.Result) bool {
- if detail.Get("modality").String() == "IMAGE" {
- imageTokens = int(detail.Get("tokenCount").Int())
- return false
- }
- return true
- })
- }
- // 注意:Gemini 的 promptTokenCount 包含 cachedContentTokenCount,
- // 但 Claude 的 input_tokens 不包含 cache_read_input_tokens,需要减去
- return &ClaudeUsage{
- InputTokens: prompt - cached,
- OutputTokens: cand + thoughts,
- CacheReadInputTokens: cached,
- ImageOutputTokens: imageTokens,
- }
- }
- func asInt(v any) (int, bool) {
- switch t := v.(type) {
- case float64:
- return int(t), true
- case int:
- return t, true
- case int64:
- return int(t), true
- case json.Number:
- i, err := t.Int64()
- if err != nil {
- return 0, false
- }
- return int(i), true
- default:
- return 0, false
- }
- }
- func (s *GeminiMessagesCompatService) handleGeminiUpstreamError(ctx context.Context, account *Account, statusCode int, headers http.Header, body []byte) {
- // 遵守自定义错误码策略:未命中则跳过所有限流处理
- if !account.ShouldHandleErrorCode(statusCode) {
- return
- }
- if s.rateLimitService != nil && (statusCode == 401 || statusCode == 403 || statusCode == 529) {
- s.rateLimitService.HandleUpstreamError(ctx, account, statusCode, headers, body)
- return
- }
- if statusCode != 429 {
- return
- }
- // 池模式账号不写账号级限流:账号留在池内,由 failover / 同号重试消化 429。
- // 自定义错误码优先级高于池模式,开启后仍按其命中结果标记。
- if account.IsPoolMode() && !account.IsCustomErrorCodesEnabled() {
- return
- }
- oauthType := account.GeminiOAuthType()
- tierID := account.GeminiTierID()
- projectID := strings.TrimSpace(account.GetCredential("project_id"))
- isCodeAssist := account.IsGeminiCodeAssist()
- resetAt := ParseGeminiRateLimitResetTime(body)
- if resetAt == nil {
- // 根据账号类型使用不同的默认重置时间
- var ra time.Time
- if isCodeAssist || oauthType == "google_one" {
- // Gemini CLI / Google One: fallback cooldown by tier
- cooldown := geminiCooldownForTier(tierID)
- if s.rateLimitService != nil {
- cooldown = s.rateLimitService.GeminiCooldown(ctx, account)
- }
- ra = time.Now().Add(cooldown)
- if isCodeAssist {
- logger.LegacyPrintf("service.gemini_messages_compat", "[Gemini 429] Account %d (Code Assist, tier=%s, project=%s) rate limited, cooldown=%v", account.ID, tierID, projectID, time.Until(ra).Truncate(time.Second))
- } else {
- logger.LegacyPrintf("service.gemini_messages_compat", "[Gemini 429] Account %d (Google One OAuth, tier=%s, project=%s) rate limited, cooldown=%v", account.ID, tierID, projectID, time.Until(ra).Truncate(time.Second))
- }
- } else {
- // API Key / AI Studio OAuth: PST 午夜
- if ts := nextGeminiDailyResetUnix(); ts != nil {
- ra = time.Unix(*ts, 0)
- logger.LegacyPrintf("service.gemini_messages_compat", "[Gemini 429] Account %d (API Key/AI Studio, type=%s) rate limited, reset at PST midnight (%v)", account.ID, account.Type, ra)
- } else {
- // 兜底:5 分钟
- ra = time.Now().Add(5 * time.Minute)
- logger.LegacyPrintf("service.gemini_messages_compat", "[Gemini 429] Account %d rate limited, fallback to 5min", account.ID)
- }
- }
- _ = s.accountRepo.SetRateLimited(ctx, account.ID, ra)
- return
- }
- // 使用解析到的重置时间
- resetTime := time.Unix(*resetAt, 0)
- _ = s.accountRepo.SetRateLimited(ctx, account.ID, resetTime)
- logger.LegacyPrintf("service.gemini_messages_compat", "[Gemini 429] Account %d rate limited until %v (oauth_type=%s, tier=%s)",
- account.ID, resetTime, oauthType, tierID)
- }
- // ParseGeminiRateLimitResetTime 解析 Gemini 格式的 429 响应,返回重置时间的 Unix 时间戳
- func ParseGeminiRateLimitResetTime(body []byte) *int64 {
- // 第一阶段:gjson 结构化提取
- errMsg := gjson.GetBytes(body, "error.message").String()
- if looksLikeGeminiDailyQuota(errMsg) {
- if ts := nextGeminiDailyResetUnix(); ts != nil {
- return ts
- }
- }
- // 遍历 error.details 查找 quotaResetDelay
- var found *int64
- gjson.GetBytes(body, "error.details").ForEach(func(_, detail gjson.Result) bool {
- v := detail.Get("metadata.quotaResetDelay").String()
- if v == "" {
- return true
- }
- if dur, err := time.ParseDuration(v); err == nil {
- // Use ceil to avoid undercounting fractional seconds (e.g. 10.1s should not become 10s),
- // which can affect scheduling decisions around thresholds (like 10s).
- ts := time.Now().Unix() + int64(math.Ceil(dur.Seconds()))
- found = &ts
- return false
- }
- return true
- })
- if found != nil {
- return found
- }
- // 第二阶段:regex 回退匹配 "Please retry in Xs"
- matches := retryInRegex.FindStringSubmatch(string(body))
- if len(matches) == 2 {
- if dur, err := time.ParseDuration(matches[1] + "s"); err == nil {
- ts := time.Now().Unix() + int64(math.Ceil(dur.Seconds()))
- return &ts
- }
- }
- return nil
- }
- func looksLikeGeminiDailyQuota(message string) bool {
- m := strings.ToLower(message)
- if strings.Contains(m, "per day") || strings.Contains(m, "requests per day") || strings.Contains(m, "quota") && strings.Contains(m, "per day") {
- return true
- }
- return false
- }
- func nextGeminiDailyResetUnix() *int64 {
- reset := geminiDailyResetTime(time.Now())
- ts := reset.Unix()
- return &ts
- }
- func ensureGeminiFunctionCallThoughtSignatures(body []byte) []byte {
- // Fast path: only run when functionCall is present.
- if !bytes.Contains(body, []byte(`"functionCall"`)) {
- return body
- }
- var payload map[string]any
- if err := json.Unmarshal(body, &payload); err != nil {
- return body
- }
- contentsAny, ok := payload["contents"].([]any)
- if !ok || len(contentsAny) == 0 {
- return body
- }
- modified := false
- for _, c := range contentsAny {
- cm, ok := c.(map[string]any)
- if !ok {
- continue
- }
- partsAny, ok := cm["parts"].([]any)
- if !ok || len(partsAny) == 0 {
- continue
- }
- for _, p := range partsAny {
- pm, ok := p.(map[string]any)
- if !ok || pm == nil {
- continue
- }
- if fc, ok := pm["functionCall"].(map[string]any); !ok || fc == nil {
- continue
- }
- ts, _ := pm["thoughtSignature"].(string)
- if strings.TrimSpace(ts) == "" {
- pm["thoughtSignature"] = geminiDummyThoughtSignature
- modified = true
- }
- }
- }
- if !modified {
- return body
- }
- b, err := json.Marshal(payload)
- if err != nil {
- return body
- }
- return b
- }
- func extractGeminiFinishReason(geminiResp map[string]any) string {
- if candidates, ok := geminiResp["candidates"].([]any); ok && len(candidates) > 0 {
- if cand, ok := candidates[0].(map[string]any); ok {
- if fr, ok := cand["finishReason"].(string); ok {
- return fr
- }
- }
- }
- return ""
- }
- func extractGeminiParts(geminiResp map[string]any) []map[string]any {
- if candidates, ok := geminiResp["candidates"].([]any); ok && len(candidates) > 0 {
- if cand, ok := candidates[0].(map[string]any); ok {
- if content, ok := cand["content"].(map[string]any); ok {
- if partsAny, ok := content["parts"].([]any); ok && len(partsAny) > 0 {
- out := make([]map[string]any, 0, len(partsAny))
- for _, p := range partsAny {
- pm, ok := p.(map[string]any)
- if !ok {
- continue
- }
- out = append(out, pm)
- }
- return out
- }
- }
- }
- }
- return nil
- }
- func computeGeminiTextDelta(seen, incoming string) (delta, newSeen string) {
- incoming = strings.TrimSuffix(incoming, "\u0000")
- if incoming == "" {
- return "", seen
- }
- // Cumulative mode: incoming contains full text so far.
- if strings.HasPrefix(incoming, seen) {
- return strings.TrimPrefix(incoming, seen), incoming
- }
- // Duplicate/rewind: ignore.
- if strings.HasPrefix(seen, incoming) {
- return "", seen
- }
- // Delta mode: treat incoming as incremental chunk.
- return incoming, seen + incoming
- }
- func mapGeminiFinishReasonToClaudeStopReason(finishReason string) string {
- switch strings.ToUpper(strings.TrimSpace(finishReason)) {
- case "MAX_TOKENS":
- return "max_tokens"
- case "STOP":
- return "end_turn"
- default:
- return "end_turn"
- }
- }
- func convertClaudeMessagesToGeminiGenerateContent(body []byte) ([]byte, error) {
- var req map[string]any
- if err := json.Unmarshal(body, &req); err != nil {
- return nil, err
- }
- toolUseIDToName := make(map[string]string)
- systemText := extractClaudeSystemText(req["system"])
- contents, err := convertClaudeMessagesToGeminiContents(req["messages"], toolUseIDToName)
- if err != nil {
- return nil, err
- }
- out := make(map[string]any)
- if systemText != "" {
- out["systemInstruction"] = map[string]any{
- "parts": []any{map[string]any{"text": systemText}},
- }
- }
- out["contents"] = contents
- if tools := convertClaudeToolsToGeminiTools(req["tools"]); tools != nil {
- out["tools"] = tools
- }
- generationConfig := convertClaudeGenerationConfig(req)
- if generationConfig != nil {
- out["generationConfig"] = generationConfig
- }
- stripGeminiFunctionIDs(out)
- return json.Marshal(out)
- }
- func stripGeminiFunctionIDs(req map[string]any) {
- // Defensive cleanup: some upstreams reject unexpected `id` fields in functionCall/functionResponse.
- contents, ok := req["contents"].([]any)
- if !ok {
- return
- }
- for _, c := range contents {
- cm, ok := c.(map[string]any)
- if !ok {
- continue
- }
- contentParts, ok := cm["parts"].([]any)
- if !ok {
- continue
- }
- for _, p := range contentParts {
- pm, ok := p.(map[string]any)
- if !ok {
- continue
- }
- if fc, ok := pm["functionCall"].(map[string]any); ok && fc != nil {
- delete(fc, "id")
- }
- if fr, ok := pm["functionResponse"].(map[string]any); ok && fr != nil {
- delete(fr, "id")
- }
- }
- }
- }
- func extractClaudeSystemText(system any) string {
- switch v := system.(type) {
- case string:
- return strings.TrimSpace(v)
- case []any:
- var parts []string
- for _, p := range v {
- pm, ok := p.(map[string]any)
- if !ok {
- continue
- }
- if t, _ := pm["type"].(string); t != "text" {
- continue
- }
- if text, ok := pm["text"].(string); ok && strings.TrimSpace(text) != "" {
- parts = append(parts, text)
- }
- }
- return strings.TrimSpace(strings.Join(parts, "\n"))
- default:
- return ""
- }
- }
- func convertClaudeMessagesToGeminiContents(messages any, toolUseIDToName map[string]string) ([]any, error) {
- arr, ok := messages.([]any)
- if !ok {
- return nil, errors.New("messages must be an array")
- }
- out := make([]any, 0, len(arr))
- for _, m := range arr {
- mm, ok := m.(map[string]any)
- if !ok {
- continue
- }
- role, _ := mm["role"].(string)
- role = strings.ToLower(strings.TrimSpace(role))
- gRole := "user"
- if role == "assistant" {
- gRole = "model"
- }
- parts := make([]any, 0)
- switch content := mm["content"].(type) {
- case string:
- // 字符串形式的 content,保留所有内容(包括空白)
- parts = append(parts, map[string]any{"text": content})
- case []any:
- // 如果只有一个 block,不过滤空白(让上游 API 报错)
- singleBlock := len(content) == 1
- for _, block := range content {
- bm, ok := block.(map[string]any)
- if !ok {
- continue
- }
- bt, _ := bm["type"].(string)
- switch bt {
- case "text":
- if text, ok := bm["text"].(string); ok {
- // 单个 block 时保留所有内容(包括空白)
- // 多个 blocks 时过滤掉空白
- if singleBlock || strings.TrimSpace(text) != "" {
- parts = append(parts, map[string]any{"text": text})
- }
- }
- case "tool_use":
- id, _ := bm["id"].(string)
- name, _ := bm["name"].(string)
- if strings.TrimSpace(id) != "" && strings.TrimSpace(name) != "" {
- toolUseIDToName[id] = name
- }
- signature, _ := bm["signature"].(string)
- signature = strings.TrimSpace(signature)
- if signature == "" {
- signature = geminiDummyThoughtSignature
- }
- parts = append(parts, map[string]any{
- "thoughtSignature": signature,
- "functionCall": map[string]any{
- "name": name,
- "args": bm["input"],
- },
- })
- case "tool_result":
- toolUseID, _ := bm["tool_use_id"].(string)
- name := toolUseIDToName[toolUseID]
- if name == "" {
- name = "tool"
- }
- parts = append(parts, map[string]any{
- "functionResponse": map[string]any{
- "name": name,
- "response": map[string]any{
- "content": extractClaudeContentText(bm["content"]),
- },
- },
- })
- case "image":
- if src, ok := bm["source"].(map[string]any); ok {
- if srcType, _ := src["type"].(string); srcType == "base64" {
- mediaType, _ := src["media_type"].(string)
- data, _ := src["data"].(string)
- if mediaType != "" && data != "" {
- parts = append(parts, map[string]any{
- "inlineData": map[string]any{
- "mimeType": mediaType,
- "data": data,
- },
- })
- }
- }
- }
- default:
- // best-effort: preserve unknown blocks as text
- if b, err := json.Marshal(bm); err == nil {
- parts = append(parts, map[string]any{"text": string(b)})
- }
- }
- }
- default:
- // ignore
- }
- out = append(out, map[string]any{
- "role": gRole,
- "parts": parts,
- })
- }
- return out, nil
- }
- func extractClaudeContentText(v any) string {
- switch t := v.(type) {
- case string:
- return t
- case []any:
- var sb strings.Builder
- for _, part := range t {
- pm, ok := part.(map[string]any)
- if !ok {
- continue
- }
- if pm["type"] == "text" {
- if text, ok := pm["text"].(string); ok {
- _, _ = sb.WriteString(text)
- }
- }
- }
- return sb.String()
- default:
- b, _ := json.Marshal(t)
- return string(b)
- }
- }
- func convertClaudeToolsToGeminiTools(tools any) []any {
- arr, ok := tools.([]any)
- if !ok || len(arr) == 0 {
- return nil
- }
- hasWebSearch := false
- funcDecls := make([]any, 0, len(arr))
- for _, t := range arr {
- tm, ok := t.(map[string]any)
- if !ok {
- continue
- }
- if isClaudeWebSearchToolMap(tm) {
- hasWebSearch = true
- continue
- }
- var name, desc string
- var params any
- // 检查是否为 custom 类型工具 (MCP)
- toolType, _ := tm["type"].(string)
- if toolType == "custom" {
- // Custom 格式: 从 custom 字段获取 description 和 input_schema
- custom, ok := tm["custom"].(map[string]any)
- if !ok {
- continue
- }
- name, _ = tm["name"].(string)
- desc, _ = custom["description"].(string)
- params = custom["input_schema"]
- } else {
- // 标准格式: 从顶层字段获取
- name, _ = tm["name"].(string)
- desc, _ = tm["description"].(string)
- params = tm["input_schema"]
- }
- if name == "" {
- continue
- }
- // 为 nil params 提供默认值
- if params == nil {
- params = map[string]any{
- "type": "object",
- "properties": map[string]any{},
- }
- }
- // 清理 JSON Schema
- cleanedParams := cleanToolSchema(params)
- funcDecls = append(funcDecls, map[string]any{
- "name": name,
- "description": desc,
- "parameters": cleanedParams,
- })
- }
- out := make([]any, 0, 2)
- if len(funcDecls) > 0 {
- out = append(out, map[string]any{
- "functionDeclarations": funcDecls,
- })
- }
- if hasWebSearch {
- out = append(out, map[string]any{
- "googleSearch": map[string]any{},
- })
- }
- if len(out) == 0 {
- return nil
- }
- return out
- }
- func normalizeGeminiRequestForAIStudio(body []byte) []byte {
- var payload map[string]any
- if err := json.Unmarshal(body, &payload); err != nil {
- return body
- }
- tools, ok := payload["tools"].([]any)
- if !ok || len(tools) == 0 {
- return body
- }
- modified := false
- for _, rawTool := range tools {
- tool, ok := rawTool.(map[string]any)
- if !ok {
- continue
- }
- googleSearch, ok := tool["googleSearch"]
- if !ok {
- continue
- }
- if _, exists := tool["google_search"]; exists {
- continue
- }
- tool["google_search"] = googleSearch
- delete(tool, "googleSearch")
- modified = true
- }
- if !modified {
- return body
- }
- normalized, err := json.Marshal(payload)
- if err != nil {
- return body
- }
- return normalized
- }
- func isClaudeWebSearchToolMap(tool map[string]any) bool {
- toolType, _ := tool["type"].(string)
- // A function named web_search is still a client-side function. This is
- // especially important for Chat Completions clients such as Hermes, whose
- // built-in runtime tools are represented as ordinary function tools.
- // Promote only explicitly typed server-side search tools to Gemini's
- // built-in googleSearch tool.
- return strings.HasPrefix(toolType, "web_search") || toolType == "google_search"
- }
- // cleanToolSchema 清理工具的 JSON Schema,移除 Gemini 不支持的字段
- func cleanToolSchema(schema any) any {
- if schema == nil {
- return nil
- }
- switch v := schema.(type) {
- case map[string]any:
- cleaned := make(map[string]any)
- for key, value := range v {
- // 跳过不支持的字段
- if key == "$schema" || key == "$id" || key == "$ref" ||
- key == "$defs" || key == "definitions" ||
- key == "additionalProperties" || key == "patternProperties" || key == "minLength" ||
- key == "maxLength" || key == "minItems" || key == "maxItems" || key == "exclusiveMinimum" ||
- key == "deprecated" {
- continue
- }
- // 递归清理嵌套对象
- cleaned[key] = cleanToolSchema(value)
- }
- if enum, ok := cleaned["enum"].([]any); ok {
- if normalized, ok := normalizeGeminiEnum(enum); ok {
- cleaned["enum"] = normalized
- } else {
- delete(cleaned, "enum")
- }
- }
- // 规范化 type 字段为大写
- if typeVal, ok := cleaned["type"].(string); ok {
- cleaned["type"] = strings.ToUpper(typeVal)
- } else if typeValues, ok := cleaned["type"].([]any); ok {
- for _, typeValue := range typeValues {
- typeName, ok := typeValue.(string)
- if ok && !strings.EqualFold(typeName, "null") {
- cleaned["type"] = strings.ToUpper(typeName)
- break
- }
- }
- if _, ok := cleaned["type"].([]any); ok {
- delete(cleaned, "type")
- }
- }
- if cleaned["type"] == "INTEGER" {
- if minimum, ok := incrementIntegralSchemaBound(v["exclusiveMinimum"]); ok {
- if existing, exists := cleaned["minimum"]; !exists || schemaNumberLess(existing, minimum) {
- cleaned["minimum"] = minimum
- }
- }
- }
- return cleaned
- case []any:
- cleaned := make([]any, len(v))
- for i, item := range v {
- cleaned[i] = cleanToolSchema(item)
- }
- return cleaned
- default:
- return v
- }
- }
- func normalizeGeminiEnum(values []any) ([]any, bool) {
- normalized := make([]any, len(values))
- for i, value := range values {
- if stringValue, ok := value.(string); ok {
- normalized[i] = stringValue
- continue
- }
- switch value.(type) {
- case nil, bool, float32, float64, int, int8, int16, int32, int64, uint, uint8, uint16, uint32, uint64, json.Number:
- encoded, err := json.Marshal(value)
- if err != nil {
- return nil, false
- }
- normalized[i] = string(encoded)
- default:
- return nil, false
- }
- }
- return normalized, true
- }
- func incrementIntegralSchemaBound(value any) (any, bool) {
- switch v := value.(type) {
- case float64:
- if math.IsNaN(v) || math.IsInf(v, 0) || v != math.Trunc(v) || v+1 <= v {
- return nil, false
- }
- return v + 1, true
- case int:
- if v == math.MaxInt {
- return nil, false
- }
- return v + 1, true
- case int64:
- if v == math.MaxInt64 {
- return nil, false
- }
- return v + 1, true
- case json.Number:
- i, err := v.Int64()
- if err != nil || i == math.MaxInt64 {
- return nil, false
- }
- return json.Number(fmt.Sprintf("%d", i+1)), true
- default:
- return nil, false
- }
- }
- func schemaNumberLess(left, right any) bool {
- leftNumber, leftOK := schemaNumberFloat64(left)
- rightNumber, rightOK := schemaNumberFloat64(right)
- return leftOK && rightOK && leftNumber < rightNumber
- }
- func schemaNumberFloat64(value any) (float64, bool) {
- switch v := value.(type) {
- case float64:
- return v, !math.IsNaN(v) && !math.IsInf(v, 0)
- case int:
- return float64(v), true
- case int64:
- return float64(v), true
- case json.Number:
- n, err := v.Float64()
- return n, err == nil && !math.IsInf(n, 0)
- default:
- return 0, false
- }
- }
- func convertClaudeGenerationConfig(req map[string]any) map[string]any {
- out := make(map[string]any)
- if mt, ok := asInt(req["max_tokens"]); ok && mt > 0 {
- out["maxOutputTokens"] = mt
- }
- if temp, ok := req["temperature"].(float64); ok {
- out["temperature"] = temp
- }
- if topP, ok := req["top_p"].(float64); ok {
- out["topP"] = topP
- }
- if stopSeq, ok := req["stop_sequences"].([]any); ok && len(stopSeq) > 0 {
- out["stopSequences"] = stopSeq
- }
- if len(out) == 0 {
- return nil
- }
- return out
- }
- func (s *GeminiMessagesCompatService) extractImageInputSize(body []byte) string {
- var req struct {
- GenerationConfig *struct {
- ImageConfig *struct {
- ImageSize string `json:"imageSize"`
- } `json:"imageConfig"`
- } `json:"generationConfig"`
- }
- if err := json.Unmarshal(body, &req); err != nil {
- return ""
- }
- if req.GenerationConfig != nil && req.GenerationConfig.ImageConfig != nil {
- return strings.TrimSpace(req.GenerationConfig.ImageConfig.ImageSize)
- }
- return ""
- }
|