| 12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640164116421643164416451646164716481649165016511652165316541655165616571658165916601661166216631664166516661667166816691670167116721673167416751676167716781679168016811682168316841685168616871688168916901691169216931694169516961697169816991700170117021703170417051706170717081709171017111712171317141715171617171718171917201721172217231724172517261727172817291730173117321733173417351736173717381739174017411742174317441745174617471748174917501751175217531754175517561757175817591760176117621763176417651766176717681769177017711772177317741775177617771778177917801781178217831784178517861787178817891790179117921793179417951796179717981799180018011802180318041805180618071808180918101811181218131814181518161817181818191820182118221823182418251826182718281829183018311832183318341835183618371838183918401841184218431844184518461847184818491850185118521853185418551856185718581859186018611862186318641865186618671868186918701871187218731874187518761877187818791880188118821883188418851886188718881889189018911892189318941895189618971898189919001901190219031904190519061907190819091910191119121913191419151916191719181919192019211922192319241925192619271928192919301931193219331934193519361937193819391940194119421943194419451946194719481949195019511952195319541955195619571958195919601961196219631964196519661967196819691970197119721973197419751976197719781979198019811982198319841985198619871988198919901991199219931994199519961997199819992000200120022003200420052006200720082009201020112012201320142015201620172018201920202021202220232024202520262027202820292030203120322033203420352036203720382039204020412042204320442045204620472048204920502051205220532054205520562057205820592060206120622063206420652066206720682069207020712072207320742075207620772078207920802081208220832084208520862087208820892090209120922093209420952096209720982099210021012102210321042105210621072108210921102111211221132114211521162117211821192120212121222123212421252126212721282129213021312132213321342135213621372138213921402141214221432144214521462147214821492150215121522153215421552156215721582159216021612162216321642165216621672168216921702171217221732174217521762177217821792180218121822183218421852186218721882189219021912192219321942195219621972198219922002201220222032204220522062207220822092210221122122213221422152216221722182219222022212222222322242225222622272228222922302231223222332234223522362237223822392240224122422243224422452246224722482249225022512252225322542255225622572258225922602261226222632264226522662267226822692270227122722273227422752276227722782279228022812282228322842285228622872288228922902291229222932294229522962297229822992300230123022303230423052306230723082309231023112312231323142315231623172318231923202321232223232324232523262327232823292330233123322333233423352336233723382339234023412342234323442345234623472348234923502351235223532354235523562357235823592360236123622363236423652366236723682369237023712372237323742375237623772378237923802381238223832384238523862387238823892390239123922393239423952396239723982399240024012402240324042405240624072408240924102411241224132414241524162417241824192420242124222423242424252426242724282429243024312432243324342435243624372438243924402441244224432444244524462447244824492450245124522453245424552456245724582459246024612462246324642465246624672468246924702471247224732474247524762477247824792480248124822483248424852486248724882489249024912492249324942495249624972498249925002501250225032504250525062507250825092510251125122513251425152516251725182519252025212522252325242525252625272528252925302531253225332534253525362537253825392540254125422543254425452546254725482549255025512552255325542555255625572558255925602561256225632564256525662567256825692570257125722573257425752576257725782579258025812582258325842585258625872588258925902591259225932594259525962597259825992600260126022603260426052606260726082609261026112612261326142615261626172618261926202621262226232624262526262627262826292630263126322633263426352636263726382639264026412642264326442645264626472648264926502651265226532654265526562657265826592660266126622663266426652666266726682669267026712672267326742675267626772678267926802681268226832684268526862687268826892690269126922693269426952696269726982699270027012702270327042705270627072708270927102711271227132714271527162717271827192720272127222723272427252726272727282729273027312732273327342735273627372738273927402741274227432744274527462747274827492750275127522753275427552756275727582759276027612762276327642765276627672768276927702771277227732774277527762777277827792780278127822783278427852786278727882789279027912792279327942795279627972798279928002801280228032804280528062807280828092810281128122813281428152816281728182819282028212822282328242825282628272828282928302831283228332834283528362837283828392840284128422843284428452846284728482849285028512852285328542855285628572858285928602861286228632864286528662867286828692870287128722873287428752876287728782879288028812882288328842885288628872888288928902891289228932894289528962897289828992900290129022903290429052906290729082909291029112912291329142915291629172918291929202921292229232924292529262927292829292930293129322933293429352936293729382939294029412942294329442945294629472948294929502951295229532954295529562957295829592960296129622963296429652966296729682969297029712972297329742975297629772978297929802981298229832984298529862987298829892990299129922993299429952996299729982999300030013002300330043005300630073008300930103011301230133014301530163017301830193020302130223023302430253026302730283029303030313032303330343035303630373038303930403041304230433044304530463047304830493050305130523053305430553056305730583059306030613062306330643065306630673068306930703071307230733074307530763077307830793080308130823083308430853086308730883089309030913092309330943095309630973098309931003101310231033104310531063107310831093110311131123113311431153116311731183119312031213122312331243125312631273128312931303131313231333134313531363137313831393140314131423143314431453146314731483149315031513152315331543155315631573158315931603161316231633164316531663167316831693170317131723173317431753176317731783179318031813182318331843185318631873188318931903191319231933194319531963197319831993200320132023203320432053206320732083209321032113212321332143215321632173218321932203221322232233224322532263227322832293230323132323233323432353236323732383239324032413242324332443245324632473248324932503251325232533254325532563257325832593260326132623263326432653266326732683269327032713272327332743275327632773278327932803281328232833284328532863287328832893290329132923293329432953296329732983299330033013302330333043305 |
- package service
- import (
- "bufio"
- "bytes"
- "context"
- "crypto/rand"
- "encoding/base64"
- "encoding/binary"
- "encoding/hex"
- "encoding/json"
- "errors"
- "fmt"
- "image"
- _ "image/gif"
- _ "image/jpeg"
- _ "image/png"
- "io"
- "log"
- "mime/multipart"
- "net/http"
- "net/http/httptest"
- "net/url"
- "regexp"
- "strings"
- "sync"
- "time"
- "github.com/Wei-Shaw/sub2api/internal/config"
- "github.com/Wei-Shaw/sub2api/internal/pkg/claude"
- "github.com/Wei-Shaw/sub2api/internal/pkg/geminicli"
- "github.com/Wei-Shaw/sub2api/internal/pkg/openai"
- "github.com/Wei-Shaw/sub2api/internal/pkg/openai_compat"
- "github.com/Wei-Shaw/sub2api/internal/pkg/xai"
- "github.com/Wei-Shaw/sub2api/internal/util/urlvalidator"
- "github.com/gin-gonic/gin"
- "github.com/google/uuid"
- "github.com/tidwall/gjson"
- )
- // sseDataPrefix matches SSE data lines with optional whitespace after colon.
- // Some upstream APIs return non-standard "data:" without space (should be "data: ").
- var sseDataPrefix = regexp.MustCompile(`^data:\s*`)
- const (
- testClaudeAPIURL = "https://api.anthropic.com/v1/messages?beta=true"
- chatgptCodexAPIURL = "https://chatgpt.com/backend-api/codex/responses"
- defaultAntigravityTestModel = "claude-sonnet-4-6"
- )
- // TestEvent represents a SSE event for account testing
- type TestEvent struct {
- Type string `json:"type"`
- Text string `json:"text,omitempty"`
- Model string `json:"model,omitempty"`
- Status string `json:"status,omitempty"`
- Code string `json:"code,omitempty"`
- ImageURL string `json:"image_url,omitempty"`
- // AudioURL / VideoURL are data: or https URLs for in-browser media players.
- AudioURL string `json:"audio_url,omitempty"`
- VideoURL string `json:"video_url,omitempty"`
- MimeType string `json:"mime_type,omitempty"`
- Data any `json:"data,omitempty"`
- Success bool `json:"success,omitempty"`
- Error string `json:"error,omitempty"`
- }
- // AccountTestOptions carries optional media for admin connectivity tests.
- // ImageDataURL / AudioDataURL are full data URLs (data:<mime>;base64,...).
- type AccountTestOptions struct {
- ImageDataURL string
- AudioDataURL string
- }
- func firstAccountTestOptions(opts []AccountTestOptions) AccountTestOptions {
- if len(opts) == 0 {
- return AccountTestOptions{}
- }
- return opts[0]
- }
- // maxAccountTestMediaBytes caps inbound data-URL payloads for admin tests (~8 MiB).
- const maxAccountTestMediaBytes = 8 << 20
- const (
- defaultGeminiTextTestPrompt = "hi"
- defaultGeminiImageTestPrompt = "Generate a cute orange cat astronaut sticker on a clean pastel background."
- defaultOpenAIImageTestPrompt = "Generate a cute orange cat astronaut sticker on a clean pastel background."
- defaultGrokImageTestPrompt = "Generate a cute orange cat astronaut sticker on a clean pastel background."
- defaultGrokVideoTestPrompt = "A red ball bouncing once on a white floor, short simple motion."
- defaultGrokSearchTestQuery = "xAI Grok"
- defaultGrokTTSTestText = "Hello from Sub2API account connectivity test."
- // Grok account-test modes (admin UI). Empty / default / text = Responses probe.
- // image/video may also be inferred from model_id when mode is default.
- AccountTestModeGrokText = "text"
- AccountTestModeGrokImage = "image"
- AccountTestModeGrokVideo = "video"
- AccountTestModeGrokSearch = "search"
- AccountTestModeGrokTTS = "tts"
- AccountTestModeGrokSTT = "stt"
- AccountTestModeGrokRealtime = "realtime"
- defaultGrokRealtimeTestModel = "grok-voice-latest"
- grokRealtimeProbeTimeout = DefaultGrokRealtimeDialTimeout
- )
- // isOpenAIImageModel checks if the model is an OpenAI image generation model (e.g. gpt-image-2).
- func isOpenAIImageModel(model string) bool {
- return strings.HasPrefix(strings.ToLower(model), "gpt-image-")
- }
- func isGrokVideoGenerationModel(model string) bool {
- return isGrokVideoBillingModel(model) ||
- strings.HasPrefix(strings.ToLower(strings.TrimSpace(model)), "grok-video")
- }
- func normalizeGrokAccountTestMode(mode string) string {
- switch strings.ToLower(strings.TrimSpace(mode)) {
- case AccountTestModeGrokText:
- return AccountTestModeGrokText
- case AccountTestModeGrokImage:
- return AccountTestModeGrokImage
- case AccountTestModeGrokVideo:
- return AccountTestModeGrokVideo
- case AccountTestModeGrokSearch:
- return AccountTestModeGrokSearch
- case AccountTestModeGrokTTS:
- return AccountTestModeGrokTTS
- case AccountTestModeGrokSTT:
- return AccountTestModeGrokSTT
- case AccountTestModeGrokRealtime:
- return AccountTestModeGrokRealtime
- default:
- return AccountTestModeDefault
- }
- }
- // AccountTestService handles account testing operations
- type AccountTestService struct {
- concurrencyService *ConcurrencyService
- accountRepo AccountRepository
- geminiTokenProvider *GeminiTokenProvider
- claudeTokenProvider *ClaudeTokenProvider
- grokTokenProvider *GrokTokenProvider
- antigravityGatewayService *AntigravityGatewayService
- httpUpstream HTTPUpstream
- cfg *config.Config
- settingService *SettingService
- tlsFPProfileService *TLSFingerprintProfileService
- modelMetadataRegistryMu sync.Mutex
- modelMetadataRegistry map[string]modelsDevProvider
- modelMetadataRegistryAt time.Time
- pluginManager *PluginManager
- openaiGatewayService *OpenAIGatewayService
- agentIdentityTaskMu sync.Mutex
- agentIdentityWS agentIdentityWSConnectionInvalidator
- // grokWSDialer is optional; realtime account tests use the default OpenAI-style
- // WS dialer when nil (supports proxy + coder/websocket handshake).
- grokWSDialer openAIWSClientDialer
- }
- func (s *AccountTestService) SetSettingService(settingService *SettingService) {
- if s != nil {
- s.settingService = settingService
- }
- }
- func (s *AccountTestService) SetPluginManager(pluginManager *PluginManager) {
- if s != nil {
- s.pluginManager = pluginManager
- }
- }
- func (s *AccountTestService) SetOpenAIGatewayService(gateway *OpenAIGatewayService) {
- if s != nil {
- s.openaiGatewayService = gateway
- }
- }
- // FetchOpenAIAccountModels uses the shared cached discovery path for the test picker.
- func (s *AccountTestService) FetchOpenAIAccountModels(ctx context.Context, account *Account) ([]openai.Model, error) {
- if s == nil || s.openaiGatewayService == nil {
- return nil, errors.New("OpenAI model discovery service is unavailable")
- }
- response, err := s.openaiGatewayService.FetchOpenAIModelsList(ctx, account)
- if err != nil {
- return nil, err
- }
- var payload struct {
- Data []openai.Model `json:"data"`
- }
- if err := json.Unmarshal(response.Body, &payload); err != nil {
- return nil, fmt.Errorf("decode OpenAI account models: %w", err)
- }
- // Standard model catalogs do not require the fields used by the admin picker.
- // Populate them here without changing the shared discovery response or cache.
- for i := range payload.Data {
- model := &payload.Data[i]
- if strings.TrimSpace(model.DisplayName) == "" {
- model.DisplayName = model.ID
- }
- if strings.TrimSpace(model.Type) == "" {
- model.Type = "model"
- }
- }
- // Codex discovery lists Responses drivers, not image_generation tool models.
- // Add locally supported image choices only to the OAuth test picker; keep the
- // shared upstream catalog and API-key discovery authoritative.
- if account != nil && account.IsOpenAIOAuthLike() {
- seen := make(map[string]bool, len(payload.Data))
- for _, model := range payload.Data {
- seen[model.ID] = true
- }
- for _, model := range openai.DefaultModels {
- if IsGPTImageGenerationModel(model.ID) && account.IsModelSupported(model.ID) && !seen[model.ID] {
- payload.Data = append(payload.Data, model)
- seen[model.ID] = true
- }
- }
- for model := range account.GetModelMapping() {
- if IsGPTImageGenerationModel(model) && !strings.Contains(model, "*") && !seen[model] {
- payload.Data = append(payload.Data, openai.Model{ID: model, Object: "model", Type: "model", OwnedBy: "openai", DisplayName: model})
- }
- }
- }
- return payload.Data, nil
- }
- // NewAccountTestService creates a new AccountTestService
- func NewAccountTestService(
- accountRepo AccountRepository,
- geminiTokenProvider *GeminiTokenProvider,
- claudeTokenProvider *ClaudeTokenProvider,
- grokTokenProvider *GrokTokenProvider,
- antigravityGatewayService *AntigravityGatewayService,
- httpUpstream HTTPUpstream,
- cfg *config.Config,
- tlsFPProfileService *TLSFingerprintProfileService,
- ) *AccountTestService {
- return &AccountTestService{
- accountRepo: accountRepo,
- geminiTokenProvider: geminiTokenProvider,
- claudeTokenProvider: claudeTokenProvider,
- grokTokenProvider: grokTokenProvider,
- antigravityGatewayService: antigravityGatewayService,
- httpUpstream: httpUpstream,
- cfg: cfg,
- tlsFPProfileService: tlsFPProfileService,
- }
- }
- func (s *AccountTestService) validateUpstreamBaseURL(raw string) (string, error) {
- if s.cfg == nil {
- return "", errors.New("config is not available")
- }
- if !s.cfg.Security.URLAllowlist.Enabled {
- return urlvalidator.ValidateURLFormat(raw, s.cfg.Security.URLAllowlist.AllowInsecureHTTP)
- }
- 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 "", err
- }
- return normalized, nil
- }
- // generateSessionString generates a Claude Code style session string.
- // The output format is determined by the UA version in claude.DefaultHeaders,
- // ensuring consistency between the user_id format and the UA sent to upstream.
- func generateSessionString() (string, error) {
- b := make([]byte, 32)
- if _, err := rand.Read(b); err != nil {
- return "", err
- }
- hex64 := hex.EncodeToString(b)
- sessionUUID := uuid.New().String()
- uaVersion := ExtractCLIVersion(claude.DefaultHeaders["User-Agent"])
- return FormatMetadataUserID(hex64, "", sessionUUID, uaVersion), nil
- }
- // createTestPayload creates a Claude Code style test request payload
- func createTestPayload(modelID string) (map[string]any, error) {
- sessionID, err := generateSessionString()
- if err != nil {
- return nil, err
- }
- return map[string]any{
- "model": modelID,
- "messages": []map[string]any{
- {
- "role": "user",
- "content": []map[string]any{
- {
- "type": "text",
- "text": "hi",
- "cache_control": map[string]string{
- "type": "ephemeral",
- },
- },
- },
- },
- },
- "system": []map[string]any{
- {
- "type": "text",
- "text": claudeCodeSystemPrompt,
- "cache_control": map[string]string{
- "type": "ephemeral",
- },
- },
- },
- "metadata": map[string]string{
- "user_id": sessionID,
- },
- "max_tokens": 1024,
- "temperature": 1,
- "stream": true,
- }, nil
- }
- // TestAccountConnection tests an account's connection by sending a test request
- // All account types use full Claude Code client characteristics, only auth header differs
- // modelID is optional - if empty, defaults to claude.DefaultTestModel
- // mode is optional - "compact" routes OpenAI accounts to the /responses/compact probe path
- // opts is optional media (image/audio data URLs for real generation / STT).
- func (s *AccountTestService) TestAccountConnection(c *gin.Context, accountID int64, modelID string, prompt string, mode string, opts ...AccountTestOptions) error {
- ctx := c.Request.Context()
- testOpts := firstAccountTestOptions(opts)
- // Get account
- account, err := s.accountRepo.GetByID(ctx, accountID)
- if err != nil {
- return s.sendErrorAndEnd(c, "Account not found")
- }
- if err := accountTestCooldown(ctx, account, resolveIntelligentTestModel(account, modelID), time.Now()); err != nil {
- var wait *TestAdmissionWaitError
- if errors.As(err, &wait) {
- return s.sendErrorAndEnd(c, fmt.Sprintf("%s;请在 %s 后重新测试", wait.Reason, wait.Until.Local().Format("2006-01-02 15:04:05")))
- }
- return s.sendErrorAndEnd(c, err.Error())
- }
- release, err := s.acquireTestAccountSlot(ctx, account)
- if err != nil {
- return s.sendErrorAndEnd(c, err.Error())
- }
- defer release()
- // Synthetic UI load-test accounts exercise the real SSE parsing and modal
- // interactions, but intentionally do not send their placeholder credentials
- // to an upstream provider.
- if account.IsSyntheticUITest() {
- testModelID := modelID
- if testModelID == "" {
- testModelID = claude.DefaultTestModel
- }
- s.sendEvent(c, TestEvent{Type: "test_start", Model: testModelID})
- s.sendEvent(c, TestEvent{Type: "content", Text: "Synthetic Anthropic OAuth account is healthy and interactive."})
- s.sendEvent(c, TestEvent{Type: "test_complete", Success: true})
- return nil
- }
- return s.testAccountConnectionForAccount(c, account, modelID, prompt, mode, testOpts)
- }
- // testAccountConnectionForAccount always invokes the real provider adapter.
- // Capability jobs use this path after authorization and persistence.
- func (s *AccountTestService) testAccountConnectionForAccount(c *gin.Context, account *Account, modelID, prompt, mode string, testOpts AccountTestOptions) error {
- // Route to platform-specific test method
- if account.IsCNProvider() {
- switch account.GetAPIProtocol() {
- case APIProtocolAdaptive:
- return s.testCNProviderAdaptiveConnection(c, account, modelID, prompt)
- case APIProtocolResponses:
- return s.testOpenAIAccountConnection(c, account, modelID, prompt, normalizeAccountTestMode(mode))
- case APIProtocolChatCompletions:
- return s.testCNProviderChatCompletionsConnection(c, account, modelID, prompt)
- case APIProtocolAnthropic:
- return s.testCNProviderAnthropicConnection(c, account, modelID)
- }
- }
- if account.IsOpenAI() {
- return s.testOpenAIAccountConnection(c, account, modelID, prompt, normalizeAccountTestMode(mode))
- }
- if account.IsGemini() {
- return s.testGeminiAccountConnection(c, account, modelID, prompt)
- }
- if account.Platform == PlatformGrok {
- return s.testGrokAccountConnection(c, account, modelID, prompt, mode, testOpts)
- }
- if account.Platform == PlatformAntigravity {
- return s.routeAntigravityTest(c, account, modelID, prompt)
- }
- return s.testClaudeAccountConnection(c, account, modelID)
- }
- func (s *AccountTestService) testCNProviderChatCompletionsConnection(c *gin.Context, account *Account, modelID string, prompt string) error {
- testModelID := strings.TrimSpace(modelID)
- if testModelID == "" {
- testModelID = openai.DefaultTestModel
- }
- testModelID = account.GetMappedModel(testModelID)
- authToken := strings.TrimSpace(account.GetOpenAIProtocolAPIKey())
- if authToken == "" {
- return s.sendErrorAndEnd(c, "No API key available")
- }
- baseURL := account.GetOpenAIBaseURL()
- normalizedBaseURL, err := s.validateUpstreamBaseURL(baseURL)
- if err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Invalid base URL: %s", err.Error()))
- }
- return s.testOpenAIChatCompletionsConnection(c, account, testModelID, prompt, normalizedBaseURL, authToken)
- }
- // testClaudeAccountConnection tests an Anthropic Claude account's connection
- func (s *AccountTestService) testClaudeAccountConnection(c *gin.Context, account *Account, modelID string) error {
- ctx := c.Request.Context()
- // Determine the model to use
- testModelID := modelID
- if testModelID == "" {
- testModelID = claude.DefaultTestModel
- }
- // API Key 账号测试连接时也需要应用通配符模型映射。
- if account.Type == "apikey" {
- testModelID = account.GetMappedModel(testModelID)
- }
- // Bedrock accounts use a separate test path
- if account.IsBedrock() {
- return s.testBedrockAccountConnection(c, ctx, account, testModelID)
- }
- if account.Type == AccountTypeServiceAccount {
- return s.testClaudeVertexServiceAccountConnection(c, ctx, account, testModelID)
- }
- // Determine authentication method and API URL
- var authToken string
- var apiURL string
- if account.IsOAuth() {
- apiURL = testClaudeAPIURL
- authToken = account.GetCredential("access_token")
- if authToken == "" {
- return s.sendErrorAndEnd(c, "No access token available")
- }
- } else if account.Type == "apikey" {
- authToken = account.GetCredential("api_key")
- if authToken == "" {
- return s.sendErrorAndEnd(c, "No API key available")
- }
- baseURL := account.GetBaseURL()
- if baseURL == "" {
- baseURL = "https://api.anthropic.com"
- }
- normalizedBaseURL, err := s.validateUpstreamBaseURL(baseURL)
- if err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Invalid base URL: %s", err.Error()))
- }
- apiURL = strings.TrimSuffix(normalizedBaseURL, "/") + "/v1/messages?beta=true"
- } else {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Unsupported account type: %s", account.Type))
- }
- // Set SSE headers
- c.Writer.Header().Set("Content-Type", "text/event-stream")
- c.Writer.Header().Set("Cache-Control", "no-cache")
- c.Writer.Header().Set("Connection", "keep-alive")
- c.Writer.Header().Set("X-Accel-Buffering", "no")
- c.Writer.Flush()
- // Create Claude Code style payload (same for all account types)
- payload, err := createTestPayload(testModelID)
- if err != nil {
- return s.sendErrorAndEnd(c, "Failed to create test payload")
- }
- applyIntelligentPayloadPrompt(ctx, payload)
- payloadBytes, _ := json.Marshal(payload)
- // Send test_start event
- s.sendEvent(c, TestEvent{Type: "test_start", Model: testModelID})
- req, err := http.NewRequestWithContext(ctx, "POST", apiURL, bytes.NewReader(payloadBytes))
- if err != nil {
- return s.sendErrorAndEnd(c, "Failed to create request")
- }
- // Set common headers
- req.Header.Set("Content-Type", "application/json")
- req.Header.Set("anthropic-version", "2023-06-01")
- // Apply Claude Code client headers
- for key, value := range claude.DefaultHeaders {
- req.Header.Set(key, value)
- }
- // Set authentication header
- if account.IsOAuth() {
- req.Header.Set("anthropic-beta", claude.DefaultBetaHeader)
- req.Header.Set("Authorization", "Bearer "+authToken)
- } else {
- req.Header.Set("anthropic-beta", claude.APIKeyBetaHeader)
- // Ollama Cloud Anthropic 兼容端点按实际 base_url 强制 Bearer,
- // 其余保持 extra/default 行为。
- setAnthropicAPIKeyAuthHeader(req.Header, account, authToken, account.GetBaseURL())
- }
- // 账号级请求头覆写:测试请求与真实转发保持一致的最终头
- account.ApplyHeaderOverrides(req.Header)
- // Get proxy URL
- proxyURL := ""
- if account.ProxyID != nil && account.Proxy != nil {
- proxyURL = account.Proxy.URL()
- }
- resp, err := s.httpUpstream.DoWithTLS(WithAccountTrafficRequest(req, account), proxyURL, account.ID, account.Concurrency, s.tlsFPProfileService.ResolveTLSProfile(account))
- if err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Request failed: %s", err.Error()))
- }
- defer func() { _ = resp.Body.Close() }()
- if resp.StatusCode != http.StatusOK {
- body, _ := io.ReadAll(resp.Body)
- errMsg := fmt.Sprintf("API returned %d: %s", resp.StatusCode, string(body))
- // 403 表示账号被上游封禁,标记为 error 状态
- if resp.StatusCode == http.StatusForbidden {
- _ = s.accountRepo.SetError(ctx, account.ID, errMsg)
- }
- return s.sendErrorAndEnd(c, errMsg)
- }
- // Process SSE stream
- return s.processClaudeStream(c, resp.Body)
- }
- func (s *AccountTestService) testClaudeVertexServiceAccountConnection(c *gin.Context, ctx context.Context, account *Account, testModelID string) error {
- if mappedModel, matched := account.ResolveMappedModel(testModelID); matched {
- testModelID = mappedModel
- } else {
- testModelID = normalizeVertexAnthropicModelID(claude.NormalizeModelID(testModelID))
- }
- c.Writer.Header().Set("Content-Type", "text/event-stream")
- c.Writer.Header().Set("Cache-Control", "no-cache")
- c.Writer.Header().Set("Connection", "keep-alive")
- c.Writer.Header().Set("X-Accel-Buffering", "no")
- c.Writer.Flush()
- payload, err := createTestPayload(testModelID)
- if err != nil {
- return s.sendErrorAndEnd(c, "Failed to create test payload")
- }
- applyIntelligentPayloadPrompt(ctx, payload)
- payloadBytes, _ := json.Marshal(payload)
- vertexBody, err := buildVertexAnthropicRequestBody(payloadBytes)
- if err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Failed to create Vertex request body: %s", err.Error()))
- }
- if s.claudeTokenProvider == nil {
- return s.sendErrorAndEnd(c, "Claude token provider not configured")
- }
- accessToken, err := s.claudeTokenProvider.GetAccessToken(ctx, account)
- if err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Failed to get service account access token: %s", err.Error()))
- }
- fullURL, err := buildVertexAnthropicURL(account.VertexProjectID(), account.VertexLocation(testModelID), testModelID, true)
- if err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Failed to build Vertex URL: %s", err.Error()))
- }
- s.sendEvent(c, TestEvent{Type: "test_start", Model: testModelID})
- req, err := http.NewRequestWithContext(ctx, http.MethodPost, fullURL, bytes.NewReader(vertexBody))
- if err != nil {
- return s.sendErrorAndEnd(c, "Failed to create request")
- }
- req.Header.Set("Content-Type", "application/json")
- req.Header.Set("Authorization", "Bearer "+accessToken)
- proxyURL := ""
- if account.ProxyID != nil && account.Proxy != nil {
- proxyURL = account.Proxy.URL()
- }
- resp, err := s.httpUpstream.DoWithTLS(WithAccountTrafficRequest(req, account), proxyURL, account.ID, account.Concurrency, s.tlsFPProfileService.ResolveTLSProfile(account))
- if err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Request failed: %s", err.Error()))
- }
- defer func() { _ = resp.Body.Close() }()
- if resp.StatusCode != http.StatusOK {
- body, _ := io.ReadAll(resp.Body)
- errMsg := fmt.Sprintf("API returned %d: %s", resp.StatusCode, string(body))
- if resp.StatusCode == http.StatusForbidden {
- _ = s.accountRepo.SetError(ctx, account.ID, errMsg)
- }
- return s.sendErrorAndEnd(c, errMsg)
- }
- return s.processClaudeStream(c, resp.Body)
- }
- // testBedrockAccountConnection tests a Bedrock (SigV4 or API Key) account using non-streaming invoke
- func (s *AccountTestService) testBedrockAccountConnection(c *gin.Context, ctx context.Context, account *Account, testModelID string) error {
- region := bedrockRuntimeRegion(account)
- resolvedModelID, ok := ResolveBedrockModelID(account, testModelID)
- if !ok {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Unsupported Bedrock model: %s", testModelID))
- }
- testModelID = resolvedModelID
- // Set SSE headers (test UI expects SSE)
- c.Writer.Header().Set("Content-Type", "text/event-stream")
- c.Writer.Header().Set("Cache-Control", "no-cache")
- c.Writer.Header().Set("Connection", "keep-alive")
- c.Writer.Header().Set("X-Accel-Buffering", "no")
- c.Writer.Flush()
- // Create a minimal Bedrock-compatible payload (no stream, no cache_control)
- bedrockPayload := map[string]any{
- "anthropic_version": "bedrock-2023-05-31",
- "messages": []map[string]any{
- {
- "role": "user",
- "content": []map[string]any{
- {
- "type": "text",
- "text": "hi",
- },
- },
- },
- },
- "max_tokens": 256,
- "temperature": 1,
- }
- applyIntelligentPayloadPrompt(ctx, bedrockPayload)
- bedrockBody, _ := json.Marshal(bedrockPayload)
- // Use non-streaming endpoint (response is standard Claude JSON)
- apiURL := BuildBedrockURL(region, testModelID, false)
- s.sendEvent(c, TestEvent{Type: "test_start", Model: testModelID})
- req, err := http.NewRequestWithContext(ctx, "POST", apiURL, bytes.NewReader(bedrockBody))
- if err != nil {
- return s.sendErrorAndEnd(c, "Failed to create request")
- }
- req.Header.Set("Content-Type", "application/json")
- // Sign or set auth based on account type
- if account.IsBedrockAPIKey() {
- apiKey := account.GetCredential("api_key")
- if apiKey == "" {
- return s.sendErrorAndEnd(c, "No API key available")
- }
- req.Header.Set("Authorization", "Bearer "+apiKey)
- } else {
- signer, err := NewBedrockSignerFromAccount(account)
- if err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Failed to create Bedrock signer: %s", err.Error()))
- }
- if err := signer.SignRequest(ctx, req, bedrockBody); err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Failed to sign request: %s", err.Error()))
- }
- }
- proxyURL := ""
- if account.ProxyID != nil && account.Proxy != nil {
- proxyURL = account.Proxy.URL()
- }
- resp, err := s.httpUpstream.DoWithTLS(WithAccountTrafficRequest(req, account), proxyURL, account.ID, account.Concurrency, nil)
- if err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Request failed: %s", err.Error()))
- }
- defer func() { _ = resp.Body.Close() }()
- body, _ := io.ReadAll(resp.Body)
- if resp.StatusCode != http.StatusOK {
- return s.sendErrorAndEnd(c, fmt.Sprintf("API returned %d: %s", resp.StatusCode, string(body)))
- }
- // Bedrock non-streaming response is standard Claude JSON, extract the text
- var result struct {
- Content []struct {
- Text string `json:"text"`
- } `json:"content"`
- }
- if err := json.Unmarshal(body, &result); err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Failed to parse response: %s", err.Error()))
- }
- text := ""
- if len(result.Content) > 0 {
- text = result.Content[0].Text
- }
- if text == "" {
- text = "(empty response)"
- }
- s.sendEvent(c, TestEvent{Type: "content", Text: text})
- s.sendEvent(c, TestEvent{Type: "test_complete", Success: true})
- return nil
- }
- // testOpenAIAccountConnection tests an OpenAI account's connection
- func (s *AccountTestService) testOpenAIAccountConnection(c *gin.Context, account *Account, modelID string, prompt string, mode string) error {
- ctx := c.Request.Context()
- mode = normalizeAccountTestMode(mode)
- // Default to openai.DefaultTestModel for OpenAI testing
- testModelID := modelID
- if testModelID == "" {
- testModelID = openai.DefaultTestModel
- }
- // Align test routing with gateway behavior: OpenAI accounts apply normal
- // account model mapping. Native remote compaction v2 rides the ordinary
- // /responses wire and does NOT apply the legacy compact-only mapping
- // (post-#5641 semantics: compact_model_mapping is /responses/compact-only).
- testModelID = account.GetMappedModel(testModelID)
- if mode == AccountTestModeCompact {
- return s.testOpenAICompactConnection(c, account, testModelID)
- }
- // Route to image generation test if an image model is selected
- if isOpenAIImageModel(testModelID) {
- imagePrompt := strings.TrimSpace(prompt)
- if imagePrompt == "" {
- imagePrompt = defaultOpenAIImageTestPrompt
- }
- if account.Type == "apikey" {
- return s.testOpenAIImageAPIKey(c, ctx, account, testModelID, imagePrompt)
- }
- return s.testOpenAIImageOAuth(c, ctx, account, testModelID, imagePrompt)
- }
- credentialAccount := account
- if account.IsCredentialShadow() {
- resolved, err := resolveCredentialAccount(ctx, s.accountRepo, account)
- if err != nil {
- return s.sendErrorAndEnd(c, err.Error())
- }
- credentialAccount = resolved
- }
- // Determine authentication method and API URL
- var authToken string
- var apiURL string
- var isOAuth bool
- if credentialAccount.IsOAuth() {
- isOAuth = true
- // Agent Identity signs each request and does not retain the OAuth token.
- if !credentialAccount.IsOpenAIAgentIdentity() {
- authToken = credentialAccount.GetOpenAIAccessToken()
- }
- if authToken == "" && !credentialAccount.IsOpenAIAgentIdentity() {
- return s.sendErrorAndEnd(c, "No access token available")
- }
- // OAuth uses ChatGPT internal API
- apiURL = chatgptCodexAPIURL
- } else if credentialAccount.Type == "apikey" {
- // API Key - use Platform API
- authToken = credentialAccount.GetOpenAIProtocolAPIKey()
- if authToken == "" {
- return s.sendErrorAndEnd(c, "No API key available")
- }
- baseURL := credentialAccount.GetOpenAIBaseURL()
- if baseURL == "" {
- baseURL = "https://api.openai.com"
- }
- normalizedBaseURL, err := s.validateUpstreamBaseURL(baseURL)
- if err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Invalid base URL: %s", err.Error()))
- }
- if !openai_compat.ShouldUseResponsesAPI(account.Extra) {
- return s.testOpenAIChatCompletionsConnection(c, account, testModelID, prompt, normalizedBaseURL, authToken)
- }
- apiURL = buildOpenAIResponsesURLForPlatform(credentialAccount.Platform, normalizedBaseURL)
- } else {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Unsupported account type: %s", account.Type))
- }
- // Set SSE headers
- c.Writer.Header().Set("Content-Type", "text/event-stream")
- c.Writer.Header().Set("Cache-Control", "no-cache")
- c.Writer.Header().Set("Connection", "keep-alive")
- c.Writer.Header().Set("X-Accel-Buffering", "no")
- c.Writer.Flush()
- // Create OpenAI Responses API payload. OAuth accounts use ChatGPT Codex
- // upstream and must apply the same model normalization as real forwarding.
- upstreamTestModelID := testModelID
- if isOAuth {
- upstreamTestModelID = normalizeOpenAIModelForUpstream(credentialAccount, resolveIntelligentTestModel(credentialAccount, testModelID))
- }
- payload := createOpenAITestPayload(upstreamTestModelID, isOAuth)
- applyIntelligentPayloadPrompt(ctx, payload)
- if err := prepareIntelligentTestProtection(c, account, payload); err != nil {
- return s.sendErrorAndEnd(c, err.Error())
- }
- payloadBytes, _ := json.Marshal(payload)
- // Send test_start event once. A task-invalid Agent Identity response may
- // restart this probe after registering a replacement task.
- if !agentIdentityTaskRecoveryWasTried(ctx) {
- // Report the model sent to the upstream, not a display alias that the
- // ChatGPT/Codex endpoint cannot accept.
- s.sendEvent(c, TestEvent{Type: "test_start", Model: upstreamTestModelID})
- }
- req, err := http.NewRequestWithContext(ctx, "POST", apiURL, bytes.NewReader(payloadBytes))
- if err != nil {
- return s.sendErrorAndEnd(c, "Failed to create request")
- }
- req = req.WithContext(WithHTTPUpstreamProfile(req.Context(), HTTPUpstreamProfileOpenAI))
- // Set common headers
- req.Header.Set("Content-Type", "application/json")
- if !isOAuth {
- applyOpenAICodexProbeHeaders(req.Header)
- }
- if credentialAccount.IsOpenAIAgentIdentity() {
- authHeaders, authErr := buildAgentIdentityAuthenticationHeaders(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, credentialAccount)
- if authErr != nil {
- return s.sendErrorAndEnd(c, "Failed to build Agent Identity authentication")
- }
- for key, values := range authHeaders {
- for _, value := range values {
- req.Header.Add(key, value)
- }
- }
- } else {
- req.Header.Set("Authorization", "Bearer "+authToken)
- }
- // Set OAuth-specific headers for ChatGPT internal API
- if isOAuth {
- req.Host = "chatgpt.com"
- req.Header.Set("accept", "text/event-stream")
- req.Header.Set("OpenAI-Beta", "responses=experimental")
- canonical := resolveCodexOutboundIdentity("")
- req.Header.Set("Originator", canonical.originator)
- if customUA := strings.TrimSpace(credentialAccount.GetOpenAIUserAgent()); customUA != "" {
- req.Header.Set("User-Agent", customUA)
- } else {
- req.Header.Set("User-Agent", canonical.userAgent)
- }
- setOpenAIChatGPTAccountHeaders(req.Header, credentialAccount)
- // 与真实转发一致:账号级自定义 UA 同样作为管理员显式配置传入,否则测试用的身份
- // 与该账号真实出站的身份不是同一个(issue #3901 的配对不变式由收口保证)。
- enforceCodexIdentityHeadersWithUA(req.Header, credentialAccount.GetOpenAIUserAgent())
- }
- // 账号级请求头覆写:测试请求与真实转发保持一致的最终头
- credentialAccount.ApplyHeaderOverrides(req.Header)
- if err := applyIntelligentTestProtection(c, account, req.Header, payloadBytes); err != nil {
- return s.sendErrorAndEnd(c, err.Error())
- }
- // Get proxy URL
- proxyURL := ""
- if account.ProxyID != nil && account.Proxy != nil {
- proxyURL = account.Proxy.URL()
- }
- resp, err := s.doOpenAIAccountTestUpstream(req, proxyURL, account, true)
- if err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Request failed: %s", err.Error()))
- }
- defer func() { _ = resp.Body.Close() }()
- if isOAuth && s.accountRepo != nil {
- if updates, err := extractOpenAICodexProbeUpdates(resp); err == nil && len(updates) > 0 {
- _ = s.accountRepo.UpdateExtra(ctx, account.ID, updates)
- mergeAccountExtra(account, updates)
- }
- }
- if resp.StatusCode != http.StatusOK {
- body, _ := io.ReadAll(resp.Body)
- body = redactAgentIdentitySensitiveBodyForAccount(ctx, s.accountRepo, credentialAccount, body)
- if !agentIdentityTaskRecoveryWasTried(ctx) && credentialAccount.IsOpenAIAgentIdentity() && isAgentIdentityTaskInvalidHTTPResponse(resp.StatusCode, body) {
- expectedTaskID := credentialAccount.GetCredential("task_id")
- if err := ensureAgentIdentityTaskForAccount(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, credentialAccount, expectedTaskID); err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Agent Identity task recovery failed: %s", err.Error()))
- }
- c.Request = c.Request.WithContext(markAgentIdentityTaskRecoveryTried(ctx))
- return s.testOpenAIAccountConnection(c, account, modelID, prompt, mode)
- }
- if resp.StatusCode == http.StatusTooManyRequests {
- s.reconcileOpenAI429State(ctx, account, resp.Header, body)
- }
- // 401 Unauthorized: 标记账号为永久错误
- if resp.StatusCode == http.StatusUnauthorized && s.accountRepo != nil {
- errMsg := fmt.Sprintf("Authentication failed (401): %s", string(body))
- _ = s.accountRepo.SetError(ctx, account.ID, errMsg)
- }
- return s.sendErrorAndEnd(c, fmt.Sprintf("API returned %d: %s", resp.StatusCode, string(body)))
- }
- // Process SSE stream
- return s.processOpenAIStream(c, resp.Body)
- }
- // testGrokAccountConnection routes Grok admin connectivity tests by explicit mode first,
- // then by selected model family for media. Standalone modes (search/tts/stt) never share
- // the text Responses path; image/video never hit Responses either.
- //
- // Modes:
- // - default/text → Responses (optional model)
- // - image → /v1/images/generations (model optional; defaults to grok-imagine-image)
- // - video → /v1/videos/generations (model optional; defaults to grok-imagine-video)
- // - search → standalone web-search probe (gateway /v1/web_search semantics)
- // - tts → HTTP /v1/tts
- // - stt → HTTP /v1/stt (synthetic tiny wav probe)
- // - realtime → WS /v1/realtime dial + optional first server event
- //
- // When mode is default, image/video can still be inferred from model_id for backward compat.
- func (s *AccountTestService) testGrokAccountConnection(c *gin.Context, account *Account, modelID, prompt, mode string, opts AccountTestOptions) error {
- ctx := c.Request.Context()
- // Realtime is WebSocket-only and does not need HTTP upstream.
- mode = normalizeGrokAccountTestMode(mode)
- if mode != AccountTestModeGrokRealtime && s.httpUpstream == nil {
- return s.sendErrorAndEnd(c, "HTTP upstream not configured")
- }
- authToken, err := s.grokTestAccessToken(ctx, account)
- if err != nil {
- return s.sendErrorAndEnd(c, err.Error())
- }
- // Explicit standalone / media modes always win over model id.
- switch mode {
- case AccountTestModeGrokSearch:
- return s.testGrokWebSearch(c, ctx, account, authToken, prompt)
- case AccountTestModeGrokTTS:
- return s.testGrokTTS(c, ctx, account, authToken, prompt)
- case AccountTestModeGrokSTT:
- return s.testGrokSTT(c, ctx, account, authToken, opts.AudioDataURL)
- case AccountTestModeGrokRealtime:
- return s.testGrokRealtime(c, ctx, account, authToken, modelID)
- case AccountTestModeGrokImage:
- return s.testGrokImageGeneration(c, ctx, account, authToken, resolveGrokImageTestModel(account, modelID), resolveGrokImagePrompt(prompt), opts.ImageDataURL)
- case AccountTestModeGrokVideo:
- return s.testGrokVideoGeneration(c, ctx, account, authToken, resolveGrokVideoTestModel(account, modelID), resolveGrokVideoPrompt(prompt), opts)
- case AccountTestModeGrokText:
- // Force text Responses even if model_id looks like media.
- testModelID := strings.TrimSpace(modelID)
- if testModelID == "" {
- testModelID = grokDefaultResponsesModel
- }
- if mapped := strings.TrimSpace(account.GetMappedModel(testModelID)); mapped != "" {
- testModelID = mapped
- }
- return s.testGrokResponsesConnection(c, ctx, account, authToken, testModelID)
- }
- // mode == default: infer from model family (legacy UI / API clients).
- testModelID := strings.TrimSpace(modelID)
- if testModelID == "" {
- testModelID = grokDefaultResponsesModel
- }
- if mapped := strings.TrimSpace(account.GetMappedModel(testModelID)); mapped != "" {
- testModelID = mapped
- }
- switch {
- case isGrokImageGenerationModel(testModelID):
- return s.testGrokImageGeneration(c, ctx, account, authToken, testModelID, resolveGrokImagePrompt(prompt), opts.ImageDataURL)
- case isGrokVideoGenerationModel(testModelID):
- return s.testGrokVideoGeneration(c, ctx, account, authToken, testModelID, resolveGrokVideoPrompt(prompt), opts)
- default:
- return s.testGrokResponsesConnection(c, ctx, account, authToken, testModelID)
- }
- }
- func resolveGrokImagePrompt(prompt string) string {
- if strings.TrimSpace(prompt) == "" {
- return defaultGrokImageTestPrompt
- }
- return strings.TrimSpace(prompt)
- }
- func resolveGrokVideoPrompt(prompt string) string {
- if strings.TrimSpace(prompt) == "" {
- return defaultGrokVideoTestPrompt
- }
- return strings.TrimSpace(prompt)
- }
- func resolveGrokImageTestModel(account *Account, modelID string) string {
- testModelID := strings.TrimSpace(modelID)
- if testModelID == "" {
- testModelID = "grok-imagine-image"
- }
- if mapped := strings.TrimSpace(account.GetMappedModel(testModelID)); mapped != "" {
- return mapped
- }
- return testModelID
- }
- func resolveGrokVideoTestModel(account *Account, modelID string) string {
- testModelID := strings.TrimSpace(modelID)
- if testModelID == "" {
- testModelID = "grok-imagine-video"
- }
- if mapped := strings.TrimSpace(account.GetMappedModel(testModelID)); mapped != "" {
- return mapped
- }
- return testModelID
- }
- func (s *AccountTestService) grokTestAccessToken(ctx context.Context, account *Account) (string, error) {
- switch account.Type {
- case AccountTypeOAuth:
- if s.grokTokenProvider == nil {
- return "", fmt.Errorf("grok token provider not configured")
- }
- // Manual tests skip production scheduling eligibility so paused/rate-limited
- // accounts can still be probed by admins (same as Codex/OpenAI tests).
- token, err := s.grokTokenProvider.GetAccessTokenForManualTest(ctx, account)
- if err != nil {
- return "", fmt.Errorf("failed to get grok access token: %s", err.Error())
- }
- return token, nil
- case AccountTypeAPIKey:
- authToken := strings.TrimSpace(account.GetCredential("api_key"))
- if authToken == "" {
- return "", fmt.Errorf("grok api key is missing")
- }
- return authToken, nil
- default:
- return "", fmt.Errorf("unsupported grok account type: %s", account.Type)
- }
- }
- func (s *AccountTestService) grokTestProxyURL(account *Account) string {
- if account.ProxyID != nil && account.Proxy != nil {
- return account.Proxy.URL()
- }
- return ""
- }
- func (s *AccountTestService) prepareGrokTestSSE(c *gin.Context) {
- c.Writer.Header().Set("Content-Type", "text/event-stream")
- c.Writer.Header().Set("Cache-Control", "no-cache")
- c.Writer.Header().Set("Connection", "keep-alive")
- c.Writer.Header().Set("X-Accel-Buffering", "no")
- c.Writer.Flush()
- }
- func (s *AccountTestService) applyGrokTestRequestHeaders(req *http.Request, account *Account, authToken string, accept string) {
- req.Header.Set("Content-Type", "application/json")
- if accept != "" {
- req.Header.Set("Accept", accept)
- }
- req.Header.Set("Authorization", "Bearer "+authToken)
- // Match gateway media/voice: CLI identity headers only on the CLI chat proxy.
- // api.x.ai media (images/videos) rejects or mistreats OAuth when CLI headers
- // are stamped on the official API host (e.g. ZDR upload_url false positives).
- if account.IsGrokOAuth() && req.URL != nil && isGrokCLIProxyTarget(req.URL.String()) {
- applyGrokCLIHeaders(req.Header)
- }
- account.ApplyHeaderOverrides(req.Header)
- }
- func (s *AccountTestService) observeGrokTestResponse(ctx context.Context, account *Account, resp *http.Response) {
- if resp == nil {
- return
- }
- now := time.Now()
- // Error bodies carry Grok's free-usage, billing, and content-policy
- // classifications when quota headers are absent. Read only non-success
- // responses here, then restore the body because the caller still needs it
- // for the user-facing test result.
- var responseBody []byte
- if resp.StatusCode >= http.StatusBadRequest && resp.Body != nil {
- responseBody, _ = io.ReadAll(resp.Body)
- _ = resp.Body.Close()
- resp.Body = io.NopCloser(bytes.NewReader(responseBody))
- }
- snapshot := parseGrokQuotaSnapshot(resp.Header, resp.StatusCode, now)
- stampGrokQuotaSnapshotForPlan(account, snapshot, grokRequestedModelFromCtx(ctx))
- if snapshot != nil && s.accountRepo != nil {
- resetAt, limited := grokRateLimitResetAtForAccount(account, snapshot, now)
- if limited {
- normalizeGrokExhaustedWindowResets(snapshot, resetAt, now)
- }
- _ = s.accountRepo.UpdateExtra(ctx, account.ID, map[string]any{
- grokQuotaSnapshotExtraKey: snapshot,
- })
- if limited {
- persistGrokRateLimit(ctx, s.accountRepo, account, resetAt)
- } else if isSuccessfulGrokRateLimitRecovery(account, snapshot) {
- clearGrokRateLimitAfterRecovery(ctx, s.accountRepo, account)
- }
- } else if s.accountRepo != nil && isSuccessfulGrokRateLimitRecovery(account, &xai.QuotaSnapshot{StatusCode: resp.StatusCode}) {
- clearGrokRateLimitAfterRecovery(ctx, s.accountRepo, account)
- }
- if s.accountRepo == nil || len(responseBody) == 0 {
- if resp.StatusCode == http.StatusPaymentRequired && s.accountRepo != nil {
- stateCtx, cancel := openAIAccountStateContext(ctx)
- defer cancel()
- _ = s.accountRepo.SetTempUnschedulable(stateCtx, account.ID, now.Add(30*time.Minute), "grok payment required")
- }
- return
- }
- if isGrokContentPolicyRejection(resp.StatusCode, responseBody) {
- return
- }
- decision := classifyGrokUpstreamFailure(resp.StatusCode, responseBody, "")
- if decision.Class == GrokFailureFreeUsage {
- if resetAt, limited := grokRateLimitResetAtForAccount(account, snapshot, now); limited && resetAt.After(now) {
- persistGrokRateLimit(ctx, s.accountRepo, account, resetAt)
- } else {
- stateCtx, cancel := openAIAccountStateContext(ctx)
- _ = s.accountRepo.SetTempUnschedulable(stateCtx, account.ID, now.Add(grokFreeUsageProbeCooldown), "grok free usage exhausted")
- cancel()
- }
- return
- }
- if decision.Class == GrokFailureBilling && (isGrokSpendingLimitError(responseBody) || strings.Contains(strings.ToLower(decision.Reason), "credit")) {
- persistGrokRateLimit(ctx, s.accountRepo, account, grokSpendingLimitResetAt(account, now))
- return
- }
- cooldown := time.Duration(0)
- reason := ""
- switch resp.StatusCode {
- case http.StatusUnauthorized:
- cooldown, reason = 10*time.Minute, "grok oauth token unauthorized"
- case http.StatusPaymentRequired:
- cooldown, reason = 30*time.Minute, "grok payment required"
- case http.StatusForbidden:
- cooldown, reason = 30*time.Minute, "grok entitlement or subscription tier denied"
- default:
- if resp.StatusCode >= 500 {
- cooldown, reason = 2*time.Minute, "grok upstream temporary error"
- }
- }
- if decision.Class == GrokFailureBilling && cooldown == 0 {
- cooldown, reason = 30*time.Minute, "grok payment required"
- }
- if cooldown > 0 {
- stateCtx, cancel := openAIAccountStateContext(ctx)
- defer cancel()
- until := now.Add(cooldown)
- if account.TempUnschedulableUntil != nil && account.TempUnschedulableUntil.After(until) {
- until = *account.TempUnschedulableUntil
- }
- _ = s.accountRepo.SetTempUnschedulable(
- stateCtx,
- account.ID,
- until,
- reason,
- )
- }
- }
- func (s *AccountTestService) testGrokResponsesConnection(c *gin.Context, ctx context.Context, account *Account, authToken, testModelID string) error {
- apiURL, err := buildGrokResponsesURL(account, s.cfg, s.settingService)
- if err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Invalid Grok base URL: %s", err.Error()))
- }
- s.prepareGrokTestSSE(c)
- payloadBytes, err := buildGrokQuotaProbeBody(testModelID)
- if err != nil {
- return s.sendErrorAndEnd(c, "Failed to create Grok test payload")
- }
- if intelligentPrompt(ctx) != "" {
- payload := createOpenAITestPayload(testModelID, false)
- applyIntelligentPayloadPrompt(ctx, payload)
- payloadBytes, _ = json.Marshal(payload)
- }
- if !agentIdentityTaskRecoveryWasTried(ctx) {
- s.sendEvent(c, TestEvent{Type: "test_start", Model: testModelID})
- }
- req, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, bytes.NewReader(payloadBytes))
- if err != nil {
- return s.sendErrorAndEnd(c, "Failed to create Grok request")
- }
- s.applyGrokTestRequestHeaders(req, account, authToken, "application/json, text/event-stream")
- resp, err := s.httpUpstream.Do(WithAccountTrafficRequest(req, account), s.grokTestProxyURL(account), account.ID, account.Concurrency)
- if err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Grok Responses API request failed: %s", err.Error()))
- }
- defer func() { _ = resp.Body.Close() }()
- s.observeGrokTestResponse(withGrokTeamRateLimitModel(ctx, testModelID), account, resp)
- if resp.StatusCode != http.StatusOK {
- body, _ := io.ReadAll(resp.Body)
- return s.sendErrorAndEnd(c, fmt.Sprintf("Grok Responses API returned %d: %s", resp.StatusCode, string(body)))
- }
- return s.processOpenAIStream(c, resp.Body)
- }
- func (s *AccountTestService) testGrokImageGeneration(c *gin.Context, ctx context.Context, account *Account, authToken, modelID, prompt, imageDataURL string) error {
- // With a source image, prefer /images/edits; otherwise /images/generations.
- endpoint := GrokMediaEndpointImagesGenerations
- imageDataURL = strings.TrimSpace(imageDataURL)
- hasSourceImage := imageDataURL != ""
- if hasSourceImage {
- endpoint = GrokMediaEndpointImagesEdits
- }
- // Align model aliases with gateway (e.g. grok-imagine → grok-imagine-image-quality).
- modelID = NormalizeGrokMediaModelForEndpoint(endpoint, modelID, hasSourceImage)
- if modelID == "" {
- modelID = "grok-imagine-image-quality"
- }
- apiURL, err := buildGrokMediaURL(account, s.cfg, endpoint, "")
- if err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Invalid Grok media base URL: %s", err.Error()))
- }
- s.prepareGrokTestSSE(c)
- s.sendEvent(c, TestEvent{Type: "test_start", Model: modelID})
- if endpoint == GrokMediaEndpointImagesEdits {
- s.sendEvent(c, TestEvent{Type: "status", Text: "Calling Grok /v1/images/edits with uploaded source image..."})
- } else {
- s.sendEvent(c, TestEvent{Type: "status", Text: "Calling Grok /v1/images/generations..."})
- }
- // Zero-data-retention teams reject URL format; always request base64 for admin tests.
- payload := map[string]any{
- "model": modelID,
- "prompt": prompt,
- "n": 1,
- "response_format": "b64_json",
- }
- if hasSourceImage {
- normalized, err := normalizeAccountTestImageDataURL(imageDataURL)
- if err != nil {
- return s.sendErrorAndEnd(c, err.Error())
- }
- // Match gateway prepareGrokMediaForwardBody shape: {url, type:image_url}.
- payload["image"] = grokMediaImageObject(normalized)
- s.sendEvent(c, TestEvent{Type: "content", Text: fmt.Sprintf("source image ready (%d chars data URL)\n", len(normalized))})
- }
- payloadBytes, err := json.Marshal(payload)
- if err != nil {
- return s.sendErrorAndEnd(c, "Failed to marshal Grok image request")
- }
- req, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, bytes.NewReader(payloadBytes))
- if err != nil {
- return s.sendErrorAndEnd(c, "Failed to create Grok image request")
- }
- s.applyGrokTestRequestHeaders(req, account, authToken, "application/json")
- req.ContentLength = int64(len(payloadBytes))
- req.GetBody = func() (io.ReadCloser, error) {
- return io.NopCloser(bytes.NewReader(payloadBytes)), nil
- }
- // One retry on transport EOF (proxies occasionally drop large edit payloads).
- var resp *http.Response
- var doErr error
- for attempt := 0; attempt < 2; attempt++ {
- if attempt > 0 {
- s.sendEvent(c, TestEvent{Type: "status", Text: "Retrying Grok image request after transport error..."})
- req, err = http.NewRequestWithContext(ctx, http.MethodPost, apiURL, bytes.NewReader(payloadBytes))
- if err != nil {
- return s.sendErrorAndEnd(c, "Failed to create Grok image retry request")
- }
- s.applyGrokTestRequestHeaders(req, account, authToken, "application/json")
- req.ContentLength = int64(len(payloadBytes))
- }
- resp, doErr = s.httpUpstream.Do(WithAccountTrafficRequest(req, account), s.grokTestProxyURL(account), account.ID, account.Concurrency)
- if doErr == nil {
- break
- }
- if !isTransientGrokTransportError(doErr) || attempt == 1 {
- return s.sendErrorAndEnd(c, formatGrokImageTransportError(doErr, hasSourceImage, len(payloadBytes)))
- }
- }
- defer func() { _ = resp.Body.Close() }()
- s.observeGrokTestResponse(ctx, account, resp)
- body, err := io.ReadAll(resp.Body)
- if err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Failed to read Grok image response: %s", err.Error()))
- }
- if resp.StatusCode != http.StatusOK {
- return s.sendErrorAndEnd(c, formatGrokImagesAPIError(resp.StatusCode, body, hasSourceImage))
- }
- var result struct {
- Data []struct {
- URL string `json:"url"`
- B64JSON string `json:"b64_json"`
- RevisedPrompt string `json:"revised_prompt"`
- MimeType string `json:"mime_type"`
- } `json:"data"`
- }
- if err := json.Unmarshal(body, &result); err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Failed to parse Grok image response: %s", err.Error()))
- }
- if len(result.Data) == 0 {
- return s.sendErrorAndEnd(c, "No images returned from Grok API")
- }
- for _, item := range result.Data {
- if item.RevisedPrompt != "" {
- s.sendEvent(c, TestEvent{Type: "content", Text: item.RevisedPrompt})
- }
- mimeType := strings.TrimSpace(item.MimeType)
- if mimeType == "" {
- mimeType = "image/jpeg"
- }
- switch {
- case strings.TrimSpace(item.B64JSON) != "":
- s.sendEvent(c, TestEvent{
- Type: "image",
- ImageURL: "data:" + mimeType + ";base64," + item.B64JSON,
- MimeType: mimeType,
- })
- case strings.TrimSpace(item.URL) != "":
- s.sendEvent(c, TestEvent{Type: "image", ImageURL: item.URL, MimeType: mimeType})
- }
- }
- s.sendEvent(c, TestEvent{Type: "test_complete", Success: true})
- return nil
- }
- func (s *AccountTestService) testGrokVideoGeneration(c *gin.Context, ctx context.Context, account *Account, authToken, modelID, prompt string, opts AccountTestOptions) error {
- apiURL, err := buildGrokMediaURL(account, s.cfg, GrokMediaEndpointVideosGenerations, "")
- if err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Invalid Grok media base URL: %s", err.Error()))
- }
- s.prepareGrokTestSSE(c)
- s.sendEvent(c, TestEvent{Type: "test_start", Model: modelID})
- s.sendEvent(c, TestEvent{Type: "status", Text: "Calling Grok /v1/videos/generations..."})
- payload := map[string]any{
- "model": modelID,
- "prompt": prompt,
- "duration": 6,
- "aspect_ratio": "16:9",
- "resolution": "480p",
- }
- if img := strings.TrimSpace(opts.ImageDataURL); img != "" {
- normalized, err := normalizeAccountTestImageDataURL(img)
- if err != nil {
- return s.sendErrorAndEnd(c, err.Error())
- }
- // First-frame / image-to-video input (xAI image field).
- payload["image"] = grokMediaImageObject(normalized)
- s.sendEvent(c, TestEvent{Type: "content", Text: "using uploaded first-frame / reference image\n"})
- }
- applyIntelligentPayloadPrompt(ctx, payload)
- payloadBytes, _ := json.Marshal(payload)
- req, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, bytes.NewReader(payloadBytes))
- if err != nil {
- return s.sendErrorAndEnd(c, "Failed to create Grok video request")
- }
- s.applyGrokTestRequestHeaders(req, account, authToken, "application/json")
- resp, err := s.httpUpstream.Do(WithAccountTrafficRequest(req, account), s.grokTestProxyURL(account), account.ID, account.Concurrency)
- if err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Grok video request failed: %s", err.Error()))
- }
- defer func() { _ = resp.Body.Close() }()
- s.observeGrokTestResponse(ctx, account, resp)
- body, err := io.ReadAll(resp.Body)
- if err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Failed to read Grok video response: %s", err.Error()))
- }
- if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusAccepted && resp.StatusCode != http.StatusCreated {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Grok videos API returned %d: %s", resp.StatusCode, string(body)))
- }
- requestID := strings.TrimSpace(gjson.GetBytes(body, "request_id").String())
- if requestID == "" {
- requestID = strings.TrimSpace(gjson.GetBytes(body, "id").String())
- }
- if requestID == "" {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Grok video create response missing request_id: %s", string(body)))
- }
- s.sendEvent(c, TestEvent{Type: "content", Text: fmt.Sprintf("video request accepted: %s\n", requestID)})
- s.sendEvent(c, TestEvent{Type: "status", Text: "Polling video status until done (max ~60s)..."})
- statusURL, err := buildGrokMediaURL(account, s.cfg, GrokMediaEndpointVideoStatus, requestID)
- if err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Invalid Grok video status URL: %s", err.Error()))
- }
- deadline := time.Now().Add(60 * time.Second)
- for time.Now().Before(deadline) {
- if ctx.Err() != nil {
- return s.sendErrorAndEnd(c, "Grok video poll canceled")
- }
- statusReq, err := http.NewRequestWithContext(ctx, http.MethodGet, statusURL, nil)
- if err != nil {
- return s.sendErrorAndEnd(c, "Failed to create Grok video status request")
- }
- s.applyGrokTestRequestHeaders(statusReq, account, authToken, "application/json")
- statusResp, err := s.httpUpstream.Do(WithAccountTrafficRequest(statusReq, account), s.grokTestProxyURL(account), account.ID, account.Concurrency)
- if err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Grok video status failed: %s", err.Error()))
- }
- statusBody, _ := io.ReadAll(statusResp.Body)
- _ = statusResp.Body.Close()
- if statusResp.StatusCode != http.StatusOK && statusResp.StatusCode != http.StatusAccepted {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Grok video status returned %d: %s", statusResp.StatusCode, string(statusBody)))
- }
- st := strings.ToLower(strings.TrimSpace(gjson.GetBytes(statusBody, "status").String()))
- progress := gjson.GetBytes(statusBody, "progress")
- if progress.Exists() {
- s.sendEvent(c, TestEvent{Type: "status", Text: fmt.Sprintf("status=%s progress=%v", st, progress.Value())})
- } else {
- s.sendEvent(c, TestEvent{Type: "status", Text: "status=" + st})
- }
- switch st {
- case "done", "completed", "succeeded", "success":
- return s.emitGrokVideoResult(c, ctx, account, authToken, requestID, statusBody)
- case "failed", "error", "canceled", "cancelled":
- return s.sendErrorAndEnd(c, fmt.Sprintf("Grok video failed: %s", string(statusBody)))
- }
- select {
- case <-ctx.Done():
- return s.sendErrorAndEnd(c, "Grok video poll canceled")
- case <-time.After(3 * time.Second):
- }
- }
- return s.sendErrorAndEnd(c, "Grok video still processing after 60s (request_id="+requestID+")")
- }
- // emitGrokVideoResult surfaces a playable video URL or downloads /content as data URL.
- func (s *AccountTestService) emitGrokVideoResult(c *gin.Context, ctx context.Context, account *Account, authToken, requestID string, statusBody []byte) error {
- videoURL := firstNonEmpty(
- strings.TrimSpace(gjson.GetBytes(statusBody, "video.url").String()),
- strings.TrimSpace(gjson.GetBytes(statusBody, "url").String()),
- strings.TrimSpace(gjson.GetBytes(statusBody, "video_url").String()),
- strings.TrimSpace(gjson.GetBytes(statusBody, "download_url").String()),
- )
- if videoURL != "" && (strings.HasPrefix(videoURL, "http://") || strings.HasPrefix(videoURL, "https://") || strings.HasPrefix(videoURL, "data:")) {
- s.sendEvent(c, TestEvent{Type: "content", Text: "video ready: " + videoURL + "\n"})
- s.sendEvent(c, TestEvent{Type: "video", VideoURL: videoURL, MimeType: "video/mp4"})
- s.sendEvent(c, TestEvent{Type: "test_complete", Success: true})
- return nil
- }
- // Fetch binary content via official /videos/{id}/content (Bearer-authenticated).
- contentURL, err := buildGrokMediaURL(account, s.cfg, GrokMediaEndpointVideoContent, requestID)
- if err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Invalid Grok video content URL: %s", err.Error()))
- }
- s.sendEvent(c, TestEvent{Type: "status", Text: "Downloading video content for preview..."})
- req, err := http.NewRequestWithContext(ctx, http.MethodGet, contentURL, nil)
- if err != nil {
- return s.sendErrorAndEnd(c, "Failed to create Grok video content request")
- }
- s.applyGrokTestRequestHeaders(req, account, authToken, "video/*, application/octet-stream, */*")
- resp, err := s.httpUpstream.Do(WithAccountTrafficRequest(req, account), s.grokTestProxyURL(account), account.ID, account.Concurrency)
- if err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Grok video content download failed: %s", err.Error()))
- }
- defer func() { _ = resp.Body.Close() }()
- body, _ := io.ReadAll(io.LimitReader(resp.Body, 64<<20)) // 64 MiB cap for admin preview
- if resp.StatusCode != http.StatusOK {
- // Fall back to status URL when binary content is unavailable.
- if videoURL != "" {
- s.sendEvent(c, TestEvent{Type: "content", Text: "video completed; content download unavailable, reported url=" + videoURL + "\n"})
- s.sendEvent(c, TestEvent{Type: "video", VideoURL: videoURL, MimeType: "video/mp4"})
- s.sendEvent(c, TestEvent{Type: "test_complete", Success: true})
- return nil
- }
- return s.sendErrorAndEnd(c, fmt.Sprintf("Grok video content returned %d: %s", resp.StatusCode, truncateString(string(body), 300)))
- }
- ct := resp.Header.Get("Content-Type")
- if ct == "" || strings.HasPrefix(ct, "application/octet-stream") {
- ct = "video/mp4"
- }
- // Keep only type/subtype for data URL.
- if i := strings.Index(ct, ";"); i >= 0 {
- ct = strings.TrimSpace(ct[:i])
- }
- dataURL := "data:" + ct + ";base64," + base64.StdEncoding.EncodeToString(body)
- s.sendEvent(c, TestEvent{Type: "content", Text: fmt.Sprintf("video content downloaded: content-type=%s bytes=%d\n", ct, len(body))})
- s.sendEvent(c, TestEvent{Type: "video", VideoURL: dataURL, MimeType: ct})
- s.sendEvent(c, TestEvent{Type: "test_complete", Success: true})
- return nil
- }
- func (s *AccountTestService) testGrokWebSearch(c *gin.Context, ctx context.Context, account *Account, authToken, query string) error {
- query = strings.TrimSpace(query)
- if query == "" {
- query = defaultGrokSearchTestQuery
- }
- // Account-test "web_search" mode mirrors the standalone gateway endpoint
- // POST /v1/web_search (not a free-form chat with tools). Implementation still
- // uses the same DoGrokNativeResponsesJSON helper as the gateway handler so
- // results match production search.
- s.prepareGrokTestSSE(c)
- s.sendEvent(c, TestEvent{Type: "test_start", Model: "grok-web-search"})
- s.sendEvent(c, TestEvent{Type: "status", Text: "Calling standalone web_search probe (same as gateway /v1/web_search)..."})
- // Keep parity with handler.buildGrokWebSearchPrompt / include sources.
- const maxResults = 5
- prompt := fmt.Sprintf(
- `Search the web for the user query below. Return ONLY valid JSON with this exact shape: {"results":[{"url":"https://...","title":"page title","snippet":"concise factual summary"}]}. Return at most %d unique results. Every URL must be an actual web_search source. Populate a non-empty title and snippet for every result. Do not wrap the JSON in markdown.
- User query:
- %s`, maxResults, query)
- payload := map[string]any{
- "model": grokDefaultResponsesModel,
- "input": prompt,
- "tools": []map[string]any{{"type": "web_search"}},
- "include": []string{"web_search_call.action.sources"},
- "store": false,
- "stream": false,
- }
- applyIntelligentPayloadPrompt(ctx, payload)
- payloadBytes, _ := json.Marshal(payload)
- apiURL, err := buildGrokResponsesURL(account, s.cfg, s.settingService)
- if err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Invalid Grok base URL: %s", err.Error()))
- }
- req, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, bytes.NewReader(payloadBytes))
- if err != nil {
- return s.sendErrorAndEnd(c, "Failed to create standalone web_search probe request")
- }
- s.applyGrokTestRequestHeaders(req, account, authToken, "application/json")
- resp, err := s.httpUpstream.Do(WithAccountTrafficRequest(req, account), s.grokTestProxyURL(account), account.ID, account.Concurrency)
- if err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("standalone web_search probe failed: %s", err.Error()))
- }
- defer func() { _ = resp.Body.Close() }()
- s.observeGrokTestResponse(withGrokTeamRateLimitModel(ctx, grokDefaultResponsesModel), account, resp)
- body, _ := io.ReadAll(resp.Body)
- if resp.StatusCode != http.StatusOK {
- return s.sendErrorAndEnd(c, fmt.Sprintf("standalone web_search probe returned %d: %s", resp.StatusCode, string(body)))
- }
- // Normalize like gateway extractGrokWebSearchSources (URL-only sources are enough for connectivity).
- sourceCount := 0
- gjson.GetBytes(body, "output").ForEach(func(_, item gjson.Result) bool {
- if item.Get("type").String() != "web_search_call" {
- return true
- }
- sources := item.Get("action.sources")
- if sources.IsArray() {
- sourceCount += len(sources.Array())
- }
- return true
- })
- searchCount := countGrokNativeSearchCallsFromJSONBytes(body)
- s.sendEvent(c, TestEvent{Type: "content", Text: fmt.Sprintf("web_search ok: query=%q tool_calls=%d sources=%d\n", query, searchCount, sourceCount)})
- // Optional: first structured result title if model returned JSON text.
- gjson.GetBytes(body, "output").ForEach(func(_, item gjson.Result) bool {
- if item.Get("type").String() != "message" {
- return true
- }
- for _, part := range item.Get("content").Array() {
- text := strings.TrimSpace(part.Get("text").String())
- if text == "" {
- continue
- }
- if len(text) > 300 {
- text = text[:300] + "..."
- }
- s.sendEvent(c, TestEvent{Type: "content", Text: text + "\n"})
- return false
- }
- return true
- })
- if searchCount == 0 && sourceCount == 0 {
- return s.sendErrorAndEnd(c, "standalone web_search probe completed but no search sources/tool calls were observed")
- }
- s.sendEvent(c, TestEvent{Type: "test_complete", Success: true})
- return nil
- }
- func (s *AccountTestService) testGrokTTS(c *gin.Context, ctx context.Context, account *Account, authToken, text string) error {
- text = strings.TrimSpace(text)
- if text == "" {
- text = defaultGrokTTSTestText
- }
- apiURL, err := buildGrokVoiceURL(account, s.cfg, "tts")
- if err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Invalid Grok TTS URL: %s", err.Error()))
- }
- s.prepareGrokTestSSE(c)
- s.sendEvent(c, TestEvent{Type: "test_start", Model: "grok-voice-tts"})
- s.sendEvent(c, TestEvent{Type: "status", Text: "Calling standalone /v1/tts..."})
- // xAI requires `language`; optional voice_id. Prefer the shape that matches
- // live gateway probes (text + language [+ voice_id]).
- payloads := []map[string]any{
- {"text": text, "language": "en", "voice_id": "Ara"},
- {"text": text, "language": "en"},
- {"text": text, "language": "English", "voice_id": "Ara"},
- }
- var lastBody string
- var lastCode int
- for _, payload := range payloads {
- applyIntelligentPayloadPrompt(ctx, payload)
- payloadBytes, _ := json.Marshal(payload)
- req, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, bytes.NewReader(payloadBytes))
- if err != nil {
- return s.sendErrorAndEnd(c, "Failed to create Grok TTS request")
- }
- s.applyGrokTestRequestHeaders(req, account, authToken, "audio/*, application/json, */*")
- resp, err := s.httpUpstream.Do(WithAccountTrafficRequest(req, account), s.grokTestProxyURL(account), account.ID, account.Concurrency)
- if err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Grok TTS failed: %s", err.Error()))
- }
- body, _ := io.ReadAll(resp.Body)
- _ = resp.Body.Close()
- s.observeGrokTestResponse(ctx, account, resp)
- lastCode = resp.StatusCode
- lastBody = string(body)
- if resp.StatusCode == http.StatusOK {
- ct := resp.Header.Get("Content-Type")
- if ct == "" {
- ct = "audio/mpeg"
- }
- if i := strings.Index(ct, ";"); i >= 0 {
- ct = strings.TrimSpace(ct[:i])
- }
- // Cap preview size so SSE stays manageable (~4 MiB audio).
- if len(body) > 4<<20 {
- body = body[:4<<20]
- }
- audioURL := "data:" + ct + ";base64," + base64.StdEncoding.EncodeToString(body)
- s.sendEvent(c, TestEvent{Type: "content", Text: fmt.Sprintf("tts ok: content-type=%s bytes=%d\n", ct, len(body))})
- s.sendEvent(c, TestEvent{Type: "audio", AudioURL: audioURL, MimeType: ct})
- s.sendEvent(c, TestEvent{Type: "test_complete", Success: true})
- return nil
- }
- if resp.StatusCode < 400 || resp.StatusCode >= 500 {
- break
- }
- }
- return s.sendErrorAndEnd(c, fmt.Sprintf("Grok TTS returned %d: %s", lastCode, lastBody))
- }
- // testGrokSTT posts audio to /v1/stt. When audioDataURL is set, uses the
- // uploaded file; otherwise a tiny synthetic silent WAV for connectivity only.
- func (s *AccountTestService) testGrokSTT(c *gin.Context, ctx context.Context, account *Account, authToken, audioDataURL string) error {
- apiURL, err := buildGrokVoiceURL(account, s.cfg, "stt")
- if err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Invalid Grok STT URL: %s", err.Error()))
- }
- s.prepareGrokTestSSE(c)
- s.sendEvent(c, TestEvent{Type: "test_start", Model: "grok-voice-stt"})
- var audioBytes []byte
- filename := "probe.wav"
- if audioDataURL = strings.TrimSpace(audioDataURL); audioDataURL != "" {
- if err := validateAccountTestDataURL(audioDataURL, "audio/"); err != nil {
- return s.sendErrorAndEnd(c, err.Error())
- }
- raw, mime, err := decodeAccountTestDataURL(audioDataURL)
- if err != nil {
- return s.sendErrorAndEnd(c, "Invalid audio data URL: "+err.Error())
- }
- audioBytes = raw
- filename = sttFilenameForMIME(mime)
- s.sendEvent(c, TestEvent{Type: "status", Text: "Calling standalone /v1/stt with uploaded audio..."})
- s.sendEvent(c, TestEvent{Type: "content", Text: fmt.Sprintf("uploaded audio: mime=%s bytes=%d\n", mime, len(audioBytes))})
- } else {
- audioBytes = minimalSilentWAV()
- s.sendEvent(c, TestEvent{Type: "status", Text: "Calling standalone /v1/stt with a synthetic silent WAV..."})
- }
- var bodyBuf bytes.Buffer
- w := multipart.NewWriter(&bodyBuf)
- part, err := w.CreateFormFile("file", filename)
- if err != nil {
- return s.sendErrorAndEnd(c, "Failed to build STT multipart body")
- }
- if _, err := part.Write(audioBytes); err != nil {
- return s.sendErrorAndEnd(c, "Failed to write STT audio part")
- }
- _ = w.WriteField("model", "grok-stt")
- _ = w.WriteField("language", "en")
- if err := w.Close(); err != nil {
- return s.sendErrorAndEnd(c, "Failed to finalize STT multipart body")
- }
- req, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, &bodyBuf)
- if err != nil {
- return s.sendErrorAndEnd(c, "Failed to create Grok STT request")
- }
- req.Header.Set("Content-Type", w.FormDataContentType())
- req.Header.Set("Accept", "application/json")
- req.Header.Set("Authorization", "Bearer "+authToken)
- if account.IsGrokOAuth() {
- applyGrokCLIHeaders(req.Header)
- }
- account.ApplyHeaderOverrides(req.Header)
- resp, err := s.httpUpstream.Do(WithAccountTrafficRequest(req, account), s.grokTestProxyURL(account), account.ID, account.Concurrency)
- if err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Grok STT failed: %s", err.Error()))
- }
- defer func() { _ = resp.Body.Close() }()
- s.observeGrokTestResponse(ctx, account, resp)
- respBody, _ := io.ReadAll(resp.Body)
- if resp.StatusCode != http.StatusOK {
- // 4xx on synthetic audio still proves the STT endpoint is wired; report clearly.
- return s.sendErrorAndEnd(c, fmt.Sprintf("Grok STT returned %d: %s", resp.StatusCode, string(respBody)))
- }
- text := strings.TrimSpace(gjson.GetBytes(respBody, "text").String())
- if text == "" {
- text = strings.TrimSpace(string(respBody))
- if len(text) > 200 {
- text = text[:200] + "..."
- }
- }
- s.sendEvent(c, TestEvent{Type: "content", Text: fmt.Sprintf("stt ok: %s\n", text)})
- s.sendEvent(c, TestEvent{Type: "test_complete", Success: true})
- return nil
- }
- // testGrokRealtime dials the standalone xAI Voice Realtime WebSocket
- // (wss://api.x.ai/v1/realtime?model=...) to verify auth + endpoint reachability.
- // It does not run a full audio session — success is WS handshake, optionally
- // enriched with the first server event type when one arrives quickly.
- func (s *AccountTestService) testGrokRealtime(c *gin.Context, ctx context.Context, account *Account, authToken, modelID string) (trafficErr error) {
- ctx, permit, admissionErr := beginAccountTrafficTurn(ctx, s.httpUpstream, account)
- if admissionErr != nil {
- return s.sendErrorAndEnd(c, admissionErr.Error())
- }
- defer func() { finishAccountTrafficTurn(permit, trafficErr) }()
- model := strings.TrimSpace(modelID)
- if model == "" {
- model = defaultGrokRealtimeTestModel
- }
- if mapped := strings.TrimSpace(account.GetMappedModel(model)); mapped != "" {
- model = mapped
- }
- base, err := buildGrokVoiceURL(account, s.cfg, "realtime")
- if err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Invalid Grok Realtime URL: %s", err.Error()))
- }
- u, err := url.Parse(base)
- if err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Invalid Grok Realtime URL: %s", err.Error()))
- }
- switch strings.ToLower(u.Scheme) {
- case "https":
- u.Scheme = "wss"
- case "http":
- u.Scheme = "ws"
- case "wss", "ws":
- // already websocket
- default:
- return s.sendErrorAndEnd(c, "Invalid Grok Realtime URL scheme")
- }
- q := u.Query()
- if q.Get("model") == "" {
- q.Set("model", model)
- }
- u.RawQuery = q.Encode()
- wsURL := u.String()
- s.prepareGrokTestSSE(c)
- s.sendEvent(c, TestEvent{Type: "test_start", Model: model})
- s.sendEvent(c, TestEvent{Type: "status", Text: "Dialing standalone wss /v1/realtime (connectivity probe)..."})
- s.sendEvent(c, TestEvent{Type: "content", Text: fmt.Sprintf("realtime target: %s\n", redactGrokRealtimeURLForLog(wsURL))})
- headers := http.Header{}
- headers.Set("Authorization", "Bearer "+authToken)
- if account.IsGrokOAuth() {
- applyGrokCLIHeaders(headers)
- }
- account.ApplyHeaderOverrides(headers)
- dialer := s.grokWSDialer
- if dialer == nil {
- dialer = newDefaultOpenAIWSClientDialer()
- }
- dialCtx, cancel := context.WithTimeout(ctx, grokRealtimeProbeTimeout)
- defer cancel()
- conn, status, _, dialErr := dialer.Dial(dialCtx, wsURL, headers, s.grokTestProxyURL(account))
- if dialErr != nil {
- detail := dialErr.Error()
- var hs *openAIWSHandshakeError
- if errors.As(dialErr, &hs) && len(hs.Body) > 0 {
- body := strings.TrimSpace(string(hs.Body))
- if len(body) > 300 {
- body = body[:300] + "..."
- }
- detail = fmt.Sprintf("%s body=%s", detail, body)
- }
- if status > 0 {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Grok Realtime WS handshake failed (HTTP %d): %s", status, detail))
- }
- return s.sendErrorAndEnd(c, fmt.Sprintf("Grok Realtime WS dial failed: %s", detail))
- }
- defer func() { _ = conn.Close() }()
- s.sendEvent(c, TestEvent{Type: "content", Text: "realtime ws handshake ok\n"})
- // Best-effort: read one server event if it arrives quickly (session.created etc.).
- // Handshake alone is enough for connectivity; missing first event is not a failure.
- readCtx, readCancel := context.WithTimeout(ctx, 3*time.Second)
- defer readCancel()
- if msg, readErr := conn.ReadMessage(readCtx); readErr == nil && len(msg) > 0 {
- eventType := strings.TrimSpace(gjson.GetBytes(msg, "type").String())
- if eventType == "" {
- eventType = "unknown"
- }
- preview := strings.TrimSpace(string(msg))
- if len(preview) > 240 {
- preview = preview[:240] + "..."
- }
- s.sendEvent(c, TestEvent{
- Type: "content",
- Text: fmt.Sprintf("realtime first event: type=%s payload=%s\n", eventType, preview),
- })
- } else {
- s.sendEvent(c, TestEvent{
- Type: "content",
- Text: "realtime handshake succeeded (no server event within 3s; still connectivity OK)\n",
- })
- }
- s.sendEvent(c, TestEvent{Type: "test_complete", Success: true})
- return nil
- }
- // validateAccountTestDataURL ensures data URLs are well-formed and size-bounded.
- func validateAccountTestDataURL(raw, requiredPrefix string) error {
- raw = strings.TrimSpace(raw)
- if raw == "" {
- return fmt.Errorf("media data URL is empty")
- }
- if !strings.HasPrefix(raw, "data:") {
- return fmt.Errorf("media must be a data: URL (data:<mime>;base64,...)")
- }
- // Rough size check before decode (base64 expands ~4/3).
- if len(raw) > maxAccountTestMediaBytes*2 {
- return fmt.Errorf("media data URL exceeds size limit")
- }
- _, mime, err := decodeAccountTestDataURL(raw)
- if err != nil {
- return err
- }
- if requiredPrefix != "" && !strings.HasPrefix(strings.ToLower(mime), strings.ToLower(requiredPrefix)) {
- return fmt.Errorf("expected media type prefix %q, got %q", requiredPrefix, mime)
- }
- return nil
- }
- // normalizeAccountTestImageDataURL validates an image data URL, enforces xAI
- // minimum dimensions (8x8), and rewrites to a clean data:image/<type>;base64,... form.
- func normalizeAccountTestImageDataURL(raw string) (string, error) {
- if err := validateAccountTestDataURL(raw, "image/"); err != nil {
- return "", err
- }
- data, mime, err := decodeAccountTestDataURL(raw)
- if err != nil {
- return "", err
- }
- // Soft cap decoded bytes (~4 MiB) for edit payloads to avoid upstream/proxy EOF.
- const maxDecodedImage = 4 << 20
- if len(data) > maxDecodedImage {
- return "", fmt.Errorf(
- "source image is too large (%d bytes decoded). Please use a smaller image (under ~4 MB) for admin edit tests",
- len(data),
- )
- }
- cfg, _, err := image.DecodeConfig(bytes.NewReader(data))
- if err != nil {
- // Keep raw data URL if decoder does not understand the codec (e.g. webp
- // without golang.org/x/image/webp); still send upstream and let xAI validate.
- return "data:" + mime + ";base64," + base64.StdEncoding.EncodeToString(data), nil
- }
- if cfg.Width < 8 || cfg.Height < 8 {
- return "", fmt.Errorf(
- "source image is too small (%dx%d). xAI requires both width and height to be at least 8 pixels",
- cfg.Width, cfg.Height,
- )
- }
- // Prefer a stable mime from config when known.
- if mime == "" || mime == "application/octet-stream" {
- mime = "image/png"
- }
- return "data:" + mime + ";base64," + base64.StdEncoding.EncodeToString(data), nil
- }
- func isTransientGrokTransportError(err error) bool {
- if err == nil {
- return false
- }
- msg := strings.ToLower(err.Error())
- return strings.Contains(msg, "unexpected eof") ||
- strings.Contains(msg, "connection reset") ||
- strings.Contains(msg, "broken pipe") ||
- strings.Contains(msg, "i/o timeout") ||
- strings.Contains(msg, "timeout awaiting response")
- }
- func formatGrokImageTransportError(err error, hasSourceImage bool, payloadBytes int) string {
- base := fmt.Sprintf("Grok image request failed: %s", err.Error())
- if !hasSourceImage {
- return base
- }
- return base + fmt.Sprintf(
- " (edit payload ~%d bytes). Tips: use a smaller source image (<4 MB / lower resolution), ensure the account proxy is stable, and retry. xAI /images/edits expects image as {\"url\":\"data:image/...;base64,...\",\"type\":\"image_url\"}.",
- payloadBytes,
- )
- }
- func formatGrokImagesAPIError(status int, body []byte, hasSourceImage bool) string {
- msg := strings.TrimSpace(string(body))
- if len(msg) > 800 {
- msg = msg[:800] + "..."
- }
- prefix := fmt.Sprintf("Grok images API returned %d: %s", status, msg)
- lower := strings.ToLower(msg)
- if hasSourceImage && (strings.Contains(lower, "too small") || strings.Contains(lower, "at least 8")) {
- return prefix + " — upload a source image with both width and height ≥ 8 px."
- }
- return prefix
- }
- func decodeAccountTestDataURL(raw string) (data []byte, mime string, err error) {
- raw = strings.TrimSpace(raw)
- if !strings.HasPrefix(raw, "data:") {
- return nil, "", fmt.Errorf("not a data URL")
- }
- rest := strings.TrimPrefix(raw, "data:")
- comma := strings.Index(rest, ",")
- if comma < 0 {
- return nil, "", fmt.Errorf("invalid data URL (missing comma)")
- }
- meta := rest[:comma]
- payload := rest[comma+1:]
- mime = "application/octet-stream"
- if semi := strings.Index(meta, ";"); semi >= 0 {
- if t := strings.TrimSpace(meta[:semi]); t != "" {
- mime = t
- }
- } else if t := strings.TrimSpace(meta); t != "" {
- mime = t
- }
- if !strings.Contains(strings.ToLower(meta), ";base64") {
- return nil, "", fmt.Errorf("only base64 data URLs are supported")
- }
- decoded, err := base64.StdEncoding.DecodeString(payload)
- if err != nil {
- // Some browsers emit URL-safe base64 without padding.
- decoded, err = base64.RawStdEncoding.DecodeString(strings.TrimRight(payload, "="))
- if err != nil {
- return nil, "", fmt.Errorf("base64 decode failed: %w", err)
- }
- }
- if len(decoded) == 0 {
- return nil, "", fmt.Errorf("decoded media is empty")
- }
- if len(decoded) > maxAccountTestMediaBytes {
- return nil, "", fmt.Errorf("media exceeds %d byte limit", maxAccountTestMediaBytes)
- }
- return decoded, mime, nil
- }
- func sttFilenameForMIME(mime string) string {
- switch strings.ToLower(strings.TrimSpace(mime)) {
- case "audio/mpeg", "audio/mp3":
- return "upload.mp3"
- case "audio/wav", "audio/x-wav", "audio/wave":
- return "upload.wav"
- case "audio/webm":
- return "upload.webm"
- case "audio/ogg", "audio/opus":
- return "upload.ogg"
- case "audio/mp4", "audio/m4a", "audio/x-m4a":
- return "upload.m4a"
- default:
- return "upload.bin"
- }
- }
- // redactGrokRealtimeURLForLog strips query secrets while keeping model for diagnostics.
- func redactGrokRealtimeURLForLog(raw string) string {
- u, err := url.Parse(strings.TrimSpace(raw))
- if err != nil || u == nil {
- return raw
- }
- // Keep model query only.
- model := u.Query().Get("model")
- u.RawQuery = ""
- if model != "" {
- u.RawQuery = "model=" + url.QueryEscape(model)
- }
- // Never log bearer in fragment/userinfo.
- u.User = nil
- u.Fragment = ""
- return u.String()
- }
- // minimalSilentWAV returns a valid tiny mono 8kHz 16-bit PCM WAV (~0.05s silence).
- func minimalSilentWAV() []byte {
- // 400 samples * 2 bytes = 800 data bytes
- const sampleRate = 8000
- const numSamples = 400
- dataSize := numSamples * 2
- buf := make([]byte, 44+dataSize)
- copy(buf[0:], []byte("RIFF"))
- binary.LittleEndian.PutUint32(buf[4:], uint32(36+dataSize))
- copy(buf[8:], []byte("WAVE"))
- copy(buf[12:], []byte("fmt "))
- binary.LittleEndian.PutUint32(buf[16:], 16) // PCM chunk size
- binary.LittleEndian.PutUint16(buf[20:], 1) // PCM
- binary.LittleEndian.PutUint16(buf[22:], 1) // mono
- binary.LittleEndian.PutUint32(buf[24:], sampleRate)
- binary.LittleEndian.PutUint32(buf[28:], sampleRate*2) // byte rate
- binary.LittleEndian.PutUint16(buf[32:], 2) // block align
- binary.LittleEndian.PutUint16(buf[34:], 16) // bits
- copy(buf[36:], []byte("data"))
- binary.LittleEndian.PutUint32(buf[40:], uint32(dataSize))
- // samples already zero (silence)
- return buf
- }
- // testOpenAIChatCompletionsConnection tests an OpenAI-compatible APIKey account
- // through the raw /v1/chat/completions endpoint.
- func (s *AccountTestService) testOpenAIChatCompletionsConnection(
- c *gin.Context,
- account *Account,
- testModelID string,
- prompt string,
- normalizedBaseURL string,
- authToken string,
- ) error {
- ctx := c.Request.Context()
- apiURL := buildOpenAIChatCompletionsURL(normalizedBaseURL)
- c.Writer.Header().Set("Content-Type", "text/event-stream")
- c.Writer.Header().Set("Cache-Control", "no-cache")
- c.Writer.Header().Set("Connection", "keep-alive")
- c.Writer.Header().Set("X-Accel-Buffering", "no")
- c.Writer.Flush()
- payload := createOpenAIChatCompletionsTestPayload(testModelID, prompt)
- applyIntelligentPayloadPrompt(ctx, payload)
- payloadBytes, _ := json.Marshal(payload)
- s.sendEvent(c, TestEvent{Type: "test_start", Model: testModelID})
- s.sendEvent(c, TestEvent{Type: "status", Text: "正在通过 /v1/chat/completions 测试连接"})
- req, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, bytes.NewReader(payloadBytes))
- if err != nil {
- return s.sendErrorAndEnd(c, "Failed to create Chat Completions request")
- }
- req = req.WithContext(WithHTTPUpstreamProfile(req.Context(), HTTPUpstreamProfileOpenAI))
- req.Header.Set("Content-Type", "application/json")
- req.Header.Set("Accept", "text/event-stream")
- req.Header.Set("Authorization", "Bearer "+authToken)
- // 账号级请求头覆写:测试请求与真实转发保持一致的最终头
- account.ApplyHeaderOverrides(req.Header)
- proxyURL := ""
- if account.ProxyID != nil && account.Proxy != nil {
- proxyURL = account.Proxy.URL()
- }
- resp, err := s.httpUpstream.DoWithTLS(WithAccountTrafficRequest(req, account), proxyURL, account.ID, account.Concurrency, s.tlsFPProfileService.ResolveTLSProfile(account))
- if err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Chat Completions API (/v1/chat/completions) request failed: %s", err.Error()))
- }
- defer func() { _ = resp.Body.Close() }()
- if resp.StatusCode != http.StatusOK {
- body, _ := io.ReadAll(resp.Body)
- if resp.StatusCode == http.StatusTooManyRequests {
- s.reconcileOpenAI429State(ctx, account, resp.Header, body)
- }
- if resp.StatusCode == http.StatusUnauthorized && s.accountRepo != nil {
- errMsg := fmt.Sprintf("Chat Completions authentication failed (401): %s", string(body))
- _ = s.accountRepo.SetError(ctx, account.ID, errMsg)
- }
- return s.sendErrorAndEnd(c, fmt.Sprintf("Chat Completions API (/v1/chat/completions) returned %d: %s", resp.StatusCode, string(body)))
- }
- return s.processOpenAIChatCompletionsStream(c, resp.Body)
- }
- // testOpenAICompactConnection probes native remote compaction v2 (streaming
- // /responses with a compaction_trigger input item) and persists the resulting
- // capability state on the account. The legacy unary /responses/compact
- // endpoint has been sunset upstream (404, #5598/#5624) and is no longer probed.
- func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account *Account, testModelID string) error {
- ctx := c.Request.Context()
- credentialAccount := account
- if account.IsShadow() {
- resolved, err := resolveCredentialAccount(ctx, s.accountRepo, account)
- if err != nil {
- return s.sendErrorAndEnd(c, "Failed to resolve account credentials")
- }
- credentialAccount = resolved
- }
- authToken := ""
- apiURL := ""
- isOAuth := false
- switch {
- case credentialAccount.IsOAuth():
- isOAuth = true
- if !credentialAccount.IsOpenAIAgentIdentity() {
- authToken = credentialAccount.GetOpenAIAccessToken()
- }
- if authToken == "" && !credentialAccount.IsOpenAIAgentIdentity() {
- return s.sendErrorAndEnd(c, "No access token available")
- }
- apiURL = chatgptCodexAPIURL
- case account.Type == AccountTypeAPIKey:
- authToken = account.GetOpenAIProtocolAPIKey()
- if authToken == "" {
- return s.sendErrorAndEnd(c, "No API key available")
- }
- baseURL := account.GetOpenAIBaseURL()
- if baseURL == "" {
- baseURL = "https://api.openai.com"
- }
- normalizedBaseURL, err := s.validateUpstreamBaseURL(baseURL)
- if err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Invalid base URL: %s", err.Error()))
- }
- apiURL = buildOpenAIResponsesURLForPlatform(account.Platform, normalizedBaseURL)
- default:
- return s.sendErrorAndEnd(c, fmt.Sprintf("Unsupported account type: %s", account.Type))
- }
- c.Writer.Header().Set("Content-Type", "text/event-stream")
- c.Writer.Header().Set("Cache-Control", "no-cache")
- c.Writer.Header().Set("Connection", "keep-alive")
- c.Writer.Header().Set("X-Accel-Buffering", "no")
- c.Writer.Flush()
- // 原生 v2 走普通 /responses 线:OAuth 与真实转发一致做上游模型归一化。
- if isOAuth {
- testModelID = normalizeOpenAIModelForUpstream(credentialAccount, testModelID)
- }
- payloadBytes, _ := json.Marshal(createOpenAICompactProbePayload(testModelID, isOAuth))
- if !agentIdentityTaskRecoveryWasTried(ctx) {
- s.sendEvent(c, TestEvent{Type: "test_start", Model: testModelID})
- }
- req, err := http.NewRequestWithContext(ctx, "POST", apiURL, bytes.NewReader(payloadBytes))
- if err != nil {
- return s.sendErrorAndEnd(c, "Failed to create request")
- }
- req = req.WithContext(WithHTTPUpstreamProfile(req.Context(), HTTPUpstreamProfileOpenAI))
- req.Header.Set("Content-Type", "application/json")
- // v2 探测是流式请求;同时补注协商头,与真实 codex 出站线型一致。
- req.Header.Set("Accept", "text/event-stream")
- ensureOpenAIRemoteCompactionV2BetaFeature(req.Header)
- if credentialAccount.IsOpenAIAgentIdentity() {
- authHeaders, authErr := buildAgentIdentityAuthenticationHeaders(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, credentialAccount)
- if authErr != nil {
- return s.sendErrorAndEnd(c, "Failed to build Agent Identity authentication")
- }
- for key, values := range authHeaders {
- for _, value := range values {
- req.Header.Add(key, value)
- }
- }
- } else {
- req.Header.Set("Authorization", "Bearer "+authToken)
- }
- applyOpenAICodexProbeHeaders(req.Header)
- if isOAuth {
- enforceCodexIdentityHeadersWithUA(req.Header, credentialAccount.GetOpenAIUserAgent())
- }
- probeSessionID := compactProbeSessionID(account.ID)
- req.Header.Set("Session_ID", probeSessionID)
- req.Header.Set("Conversation_ID", probeSessionID)
- if isOAuth {
- req.Host = "chatgpt.com"
- setOpenAIChatGPTAccountHeaders(req.Header, credentialAccount)
- // 指纹收敛:探测与真实转发走同一个 /responses 端点,身份也必须同构,
- // 否则探测流量会以「缺 x-codex-installation-id + 非收敛 session」的
- // 形态暴露在上游眼里。账号关闭收敛(off)时返回 nil,探测保持原样。
- if fpIDs := resolveCodexFingerprintIDsFromRequest(account, req.Header); fpIDs != nil {
- applyCodexFingerprintHeaders(req.Header, fpIDs)
- }
- }
- // 账号级请求头覆写:测试请求与真实转发保持一致的最终头
- account.ApplyHeaderOverrides(req.Header)
- proxyURL := ""
- if account.ProxyID != nil && account.Proxy != nil {
- proxyURL = account.Proxy.URL()
- }
- resp, err := s.doOpenAIAccountTestUpstream(req, proxyURL, account, true)
- if err != nil {
- if s.accountRepo != nil {
- updates := buildOpenAICompactProbeExtraUpdates(nil, nil, err, false, time.Now())
- _ = s.accountRepo.UpdateExtra(ctx, account.ID, updates)
- mergeAccountExtra(account, updates)
- }
- return s.sendErrorAndEnd(c, fmt.Sprintf("Request failed: %s", err.Error()))
- }
- defer func() { _ = resp.Body.Close() }()
- body, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
- body = redactAgentIdentitySensitiveBodyForAccount(ctx, s.accountRepo, credentialAccount, body)
- if !agentIdentityTaskRecoveryWasTried(ctx) && credentialAccount.IsOpenAIAgentIdentity() && isAgentIdentityTaskInvalidHTTPResponse(resp.StatusCode, body) {
- expectedTaskID := credentialAccount.GetCredential("task_id")
- if err := ensureAgentIdentityTaskForAccount(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, credentialAccount, expectedTaskID); err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Agent Identity task recovery failed: %s", err.Error()))
- }
- c.Request = c.Request.WithContext(markAgentIdentityTaskRecoveryTried(ctx))
- return s.testOpenAICompactConnection(c, account, testModelID)
- }
- compactionFound := openAICompactProbeFoundCompactionItem(body)
- if s.accountRepo != nil {
- updates := buildOpenAICompactProbeExtraUpdates(resp, body, nil, compactionFound, time.Now())
- if codexUpdates, err := extractOpenAICodexProbeUpdates(resp); err == nil && len(codexUpdates) > 0 {
- updates = mergeExtraUpdates(updates, codexUpdates)
- }
- if len(updates) > 0 {
- _ = s.accountRepo.UpdateExtra(ctx, account.ID, updates)
- mergeAccountExtra(account, updates)
- }
- // 探测如返回 429,主动同步限流状态,避免后续短时间内继续选中。
- if resp.StatusCode == http.StatusTooManyRequests {
- s.reconcileOpenAI429State(ctx, account, resp.Header, body)
- }
- }
- if resp.StatusCode != http.StatusOK {
- if resp.StatusCode == http.StatusUnauthorized && s.accountRepo != nil {
- errMsg := fmt.Sprintf("Authentication failed (401): %s", string(body))
- _ = s.accountRepo.SetError(ctx, account.ID, errMsg)
- }
- return s.sendErrorAndEnd(c, fmt.Sprintf("API returned %d: %s", resp.StatusCode, string(body)))
- }
- if !compactionFound {
- return s.sendErrorAndEnd(c, "Upstream returned 2xx without a compaction output item (native remote compaction v2 unsupported on this chain)")
- }
- s.sendEvent(c, TestEvent{Type: "content", Text: "Compact probe succeeded (native remote compaction v2)"})
- s.sendEvent(c, TestEvent{Type: "test_complete", Success: true})
- return nil
- }
- func (s *AccountTestService) reconcileOpenAI429State(ctx context.Context, account *Account, headers http.Header, body []byte) {
- if s == nil || s.accountRepo == nil || account == nil {
- return
- }
- persistOpenAI429PlanType(ctx, s.accountRepo, account, body)
- var resetAt *time.Time
- if calculated := calculateOpenAI429ResetTime(headers); calculated != nil {
- resetAt = calculated
- } else if unixTs := parseOpenAIRateLimitResetTime(body); unixTs != nil {
- t := time.Unix(*unixTs, 0)
- resetAt = &t
- }
- if resetAt == nil {
- return
- }
- if err := s.accountRepo.SetRateLimited(ctx, account.ID, *resetAt); err != nil {
- return
- }
- now := time.Now()
- account.RateLimitedAt = &now
- account.RateLimitResetAt = resetAt
- if account.Status == StatusError {
- if err := s.accountRepo.ClearError(ctx, account.ID); err != nil {
- return
- }
- account.Status = StatusActive
- account.ErrorMessage = ""
- }
- }
- // testGeminiAccountConnection tests a Gemini account's connection
- func (s *AccountTestService) testGeminiAccountConnection(c *gin.Context, account *Account, modelID string, prompt string) error {
- ctx := c.Request.Context()
- // Determine the model to use
- testModelID := modelID
- if testModelID == "" {
- testModelID = geminicli.DefaultTestModel
- }
- // For static upstream credentials with model mapping, map the model
- if account.Type == AccountTypeAPIKey || account.Type == AccountTypeServiceAccount {
- mapping := account.GetModelMapping()
- if len(mapping) > 0 {
- if mappedModel, exists := mapping[testModelID]; exists {
- testModelID = mappedModel
- }
- }
- }
- // Set SSE headers
- c.Writer.Header().Set("Content-Type", "text/event-stream")
- c.Writer.Header().Set("Cache-Control", "no-cache")
- c.Writer.Header().Set("Connection", "keep-alive")
- c.Writer.Header().Set("X-Accel-Buffering", "no")
- c.Writer.Flush()
- // Create test payload (Gemini format)
- payload := createGeminiTestPayload(testModelID, prompt)
- // Build request based on account type
- var req *http.Request
- var err error
- switch account.Type {
- case AccountTypeAPIKey:
- req, err = s.buildGeminiAPIKeyRequest(ctx, account, testModelID, payload)
- case AccountTypeOAuth:
- req, err = s.buildGeminiOAuthRequest(ctx, account, testModelID, payload)
- case AccountTypeServiceAccount:
- req, err = s.buildGeminiServiceAccountRequest(ctx, account, testModelID, payload)
- default:
- return s.sendErrorAndEnd(c, fmt.Sprintf("Unsupported account type: %s", account.Type))
- }
- if err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Failed to build request: %s", err.Error()))
- }
- // Send test_start event
- s.sendEvent(c, TestEvent{Type: "test_start", Model: testModelID})
- // Get proxy and execute request
- proxyURL := ""
- if account.ProxyID != nil && account.Proxy != nil {
- proxyURL = account.Proxy.URL()
- }
- resp, err := s.httpUpstream.DoWithTLS(WithAccountTrafficRequest(req, account), proxyURL, account.ID, account.Concurrency, s.tlsFPProfileService.ResolveTLSProfile(account))
- if err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Request failed: %s", err.Error()))
- }
- defer func() { _ = resp.Body.Close() }()
- if resp.StatusCode != http.StatusOK {
- body, _ := io.ReadAll(resp.Body)
- return s.sendErrorAndEnd(c, fmt.Sprintf("API returned %d: %s", resp.StatusCode, string(body)))
- }
- // Process SSE stream
- return s.processGeminiStream(c, resp.Body)
- }
- // routeAntigravityTest 路由 Antigravity 账号的测试请求。
- // APIKey 类型走原生协议(与 gateway_handler 路由一致),OAuth/Upstream 走 CRS 中转。
- func (s *AccountTestService) routeAntigravityTest(c *gin.Context, account *Account, modelID string, prompt string) error {
- if account.Type == AccountTypeAPIKey {
- if strings.HasPrefix(modelID, "gemini-") {
- return s.testGeminiAccountConnection(c, account, modelID, prompt)
- }
- return s.testClaudeAccountConnection(c, account, modelID)
- }
- return s.testAntigravityAccountConnection(c, account, modelID)
- }
- // testAntigravityAccountConnection tests an Antigravity account's connection
- // 支持 Claude 和 Gemini 两种协议,使用非流式请求
- func (s *AccountTestService) testAntigravityAccountConnection(c *gin.Context, account *Account, modelID string) error {
- ctx := c.Request.Context()
- testModelID := antigravityConnectionTestModel(modelID)
- if s.antigravityGatewayService == nil {
- return s.sendErrorAndEnd(c, "Antigravity gateway service not configured")
- }
- // Set SSE headers
- c.Writer.Header().Set("Content-Type", "text/event-stream")
- c.Writer.Header().Set("Cache-Control", "no-cache")
- c.Writer.Header().Set("Connection", "keep-alive")
- c.Writer.Header().Set("X-Accel-Buffering", "no")
- c.Writer.Flush()
- // Send test_start event
- s.sendEvent(c, TestEvent{Type: "test_start", Model: testModelID})
- // 调用 AntigravityGatewayService.TestConnection(复用协议转换逻辑)
- result, err := s.antigravityGatewayService.TestConnection(ctx, account, testModelID)
- if err != nil {
- return s.sendErrorAndEnd(c, err.Error())
- }
- // 发送响应内容
- if result.Text != "" {
- s.sendEvent(c, TestEvent{Type: "content", Text: result.Text})
- }
- s.sendEvent(c, TestEvent{Type: "test_complete", Success: true})
- return nil
- }
- func antigravityConnectionTestModel(modelID string) string {
- if modelID == "" {
- return defaultAntigravityTestModel
- }
- return modelID
- }
- // buildGeminiAPIKeyRequest builds request for Gemini API Key accounts
- func (s *AccountTestService) buildGeminiAPIKeyRequest(ctx context.Context, account *Account, modelID string, payload []byte) (*http.Request, error) {
- apiKey := account.GetCredential("api_key")
- if strings.TrimSpace(apiKey) == "" {
- return nil, fmt.Errorf("no API key available")
- }
- baseURL := account.GetCredential("base_url")
- if baseURL == "" {
- baseURL = geminicli.AIStudioBaseURL
- }
- normalizedBaseURL, err := s.validateUpstreamBaseURL(baseURL)
- if err != nil {
- return nil, err
- }
- // Use streamGenerateContent for real-time feedback
- fullURL, err := buildGeminiAIStudioModelActionURL(normalizedBaseURL, modelID, "streamGenerateContent", true)
- if err != nil {
- return nil, err
- }
- req, err := http.NewRequestWithContext(ctx, "POST", fullURL, bytes.NewReader(payload))
- if err != nil {
- return nil, err
- }
- req.Header.Set("Content-Type", "application/json")
- req.Header.Set("x-goog-api-key", apiKey)
- return req, nil
- }
- // buildGeminiOAuthRequest builds request for Gemini OAuth accounts
- func (s *AccountTestService) buildGeminiOAuthRequest(ctx context.Context, account *Account, modelID string, payload []byte) (*http.Request, error) {
- if s.geminiTokenProvider == nil {
- return nil, fmt.Errorf("gemini token provider not configured")
- }
- // Get access token (auto-refreshes if needed)
- accessToken, err := s.geminiTokenProvider.GetAccessToken(ctx, account)
- if err != nil {
- return nil, fmt.Errorf("failed to get access token: %w", err)
- }
- projectID := strings.TrimSpace(account.GetCredential("project_id"))
- if projectID == "" {
- // AI Studio OAuth mode (no project_id): call generativelanguage API directly with Bearer token.
- baseURL := account.GetCredential("base_url")
- if strings.TrimSpace(baseURL) == "" {
- baseURL = geminicli.AIStudioBaseURL
- }
- normalizedBaseURL, err := s.validateUpstreamBaseURL(baseURL)
- if err != nil {
- return nil, err
- }
- fullURL, err := buildGeminiAIStudioModelActionURL(normalizedBaseURL, modelID, "streamGenerateContent", true)
- if err != nil {
- return nil, err
- }
- req, err := http.NewRequestWithContext(ctx, http.MethodPost, fullURL, bytes.NewReader(payload))
- if err != nil {
- return nil, err
- }
- req.Header.Set("Content-Type", "application/json")
- req.Header.Set("Authorization", "Bearer "+accessToken)
- return req, nil
- }
- // Code Assist mode (with project_id)
- return s.buildCodeAssistRequest(ctx, accessToken, projectID, modelID, payload)
- }
- func (s *AccountTestService) buildGeminiServiceAccountRequest(ctx context.Context, account *Account, modelID string, payload []byte) (*http.Request, error) {
- if s.geminiTokenProvider == nil {
- return nil, fmt.Errorf("gemini token provider not configured")
- }
- accessToken, err := s.geminiTokenProvider.GetAccessToken(ctx, account)
- if err != nil {
- return nil, fmt.Errorf("failed to get service account access token: %w", err)
- }
- fullURL, err := buildVertexGeminiURL(account.VertexProjectID(), account.VertexLocation(modelID), modelID, "streamGenerateContent", true)
- if err != nil {
- return nil, err
- }
- req, err := http.NewRequestWithContext(ctx, http.MethodPost, fullURL, bytes.NewReader(payload))
- if err != nil {
- return nil, err
- }
- req.Header.Set("Content-Type", "application/json")
- req.Header.Set("Authorization", "Bearer "+accessToken)
- return req, nil
- }
- // buildCodeAssistRequest builds request for Google Code Assist API (used by Gemini CLI and Antigravity)
- func (s *AccountTestService) buildCodeAssistRequest(ctx context.Context, accessToken, projectID, modelID string, payload []byte) (*http.Request, error) {
- var inner map[string]any
- if err := json.Unmarshal(payload, &inner); err != nil {
- return nil, err
- }
- wrapped := map[string]any{
- "model": modelID,
- "project": projectID,
- "request": inner,
- }
- wrappedBytes, _ := json.Marshal(wrapped)
- normalizedBaseURL, err := s.validateUpstreamBaseURL(geminicli.GeminiCliBaseURL)
- if err != nil {
- return nil, err
- }
- fullURL := fmt.Sprintf("%s/v1internal:streamGenerateContent?alt=sse", normalizedBaseURL)
- req, err := http.NewRequestWithContext(ctx, "POST", fullURL, bytes.NewReader(wrappedBytes))
- if err != nil {
- return nil, err
- }
- req.Header.Set("Content-Type", "application/json")
- req.Header.Set("Authorization", "Bearer "+accessToken)
- req.Header.Set("User-Agent", geminicli.GeminiCLIUserAgent)
- return req, nil
- }
- // createGeminiTestPayload creates a minimal test payload for Gemini API.
- // Image models use the image-generation path so the frontend can preview the returned image.
- func createGeminiTestPayload(modelID string, prompt string) []byte {
- if isImageGenerationModel(modelID) {
- imagePrompt := strings.TrimSpace(prompt)
- if imagePrompt == "" {
- imagePrompt = defaultGeminiImageTestPrompt
- }
- payload := map[string]any{
- "contents": []map[string]any{
- {
- "role": "user",
- "parts": []map[string]any{
- {"text": imagePrompt},
- },
- },
- },
- "generationConfig": map[string]any{
- "responseModalities": []string{"TEXT", "IMAGE"},
- "imageConfig": map[string]any{
- "aspectRatio": "1:1",
- },
- },
- }
- bytes, _ := json.Marshal(payload)
- return bytes
- }
- textPrompt := strings.TrimSpace(prompt)
- if textPrompt == "" {
- textPrompt = defaultGeminiTextTestPrompt
- }
- payload := map[string]any{
- "contents": []map[string]any{
- {
- "role": "user",
- "parts": []map[string]any{
- {"text": textPrompt},
- },
- },
- },
- "systemInstruction": map[string]any{
- "parts": []map[string]any{
- {"text": "You are a helpful AI assistant."},
- },
- },
- }
- bytes, _ := json.Marshal(payload)
- return bytes
- }
- // processGeminiStream processes SSE stream from Gemini API
- func (s *AccountTestService) processGeminiStream(c *gin.Context, body io.Reader) error {
- reader := bufio.NewReader(body)
- for {
- line, err := reader.ReadString('\n')
- if err != nil {
- if err == io.EOF {
- s.sendEvent(c, TestEvent{Type: "test_complete", Success: true})
- return nil
- }
- return s.sendErrorAndEnd(c, fmt.Sprintf("Stream read error: %s", err.Error()))
- }
- line = strings.TrimSpace(line)
- if line == "" || !strings.HasPrefix(line, "data: ") {
- continue
- }
- jsonStr := strings.TrimPrefix(line, "data: ")
- if jsonStr == "[DONE]" {
- s.sendEvent(c, TestEvent{Type: "test_complete", Success: true})
- return nil
- }
- var data map[string]any
- if err := json.Unmarshal([]byte(jsonStr), &data); err != nil {
- continue
- }
- // Support two Gemini response formats:
- // - AI Studio: {"candidates": [...]}
- // - Gemini CLI: {"response": {"candidates": [...]}}
- if resp, ok := data["response"].(map[string]any); ok && resp != nil {
- data = resp
- }
- if candidates, ok := data["candidates"].([]any); ok && len(candidates) > 0 {
- if candidate, ok := candidates[0].(map[string]any); ok {
- // Extract content first (before checking completion)
- if content, ok := candidate["content"].(map[string]any); ok {
- if parts, ok := content["parts"].([]any); ok {
- for _, part := range parts {
- if partMap, ok := part.(map[string]any); ok {
- if text, ok := partMap["text"].(string); ok && text != "" {
- s.sendEvent(c, TestEvent{Type: "content", Text: text})
- }
- if inlineData, ok := partMap["inlineData"].(map[string]any); ok {
- mimeType, _ := inlineData["mimeType"].(string)
- data, _ := inlineData["data"].(string)
- if strings.HasPrefix(strings.ToLower(mimeType), "image/") && data != "" {
- s.sendEvent(c, TestEvent{
- Type: "image",
- ImageURL: fmt.Sprintf("data:%s;base64,%s", mimeType, data),
- MimeType: mimeType,
- })
- }
- }
- }
- }
- }
- }
- // Check for completion after extracting content
- if finishReason, ok := candidate["finishReason"].(string); ok && finishReason != "" {
- s.sendEvent(c, TestEvent{Type: "test_complete", Success: true})
- return nil
- }
- }
- }
- // Handle errors
- if errData, ok := data["error"].(map[string]any); ok {
- errorMsg := "Unknown error"
- if msg, ok := errData["message"].(string); ok {
- errorMsg = msg
- }
- return s.sendErrorAndEnd(c, errorMsg)
- }
- }
- }
- // createOpenAITestPayload creates a test payload for OpenAI Responses API
- func createOpenAITestPayload(modelID string, isOAuth bool) map[string]any {
- payload := map[string]any{
- "model": modelID,
- "input": []map[string]any{
- {
- "role": "user",
- "content": []map[string]any{
- {
- "type": "input_text",
- "text": "hi",
- },
- },
- },
- },
- "stream": true,
- }
- // OAuth accounts using ChatGPT internal API require store: false
- if isOAuth {
- payload["store"] = false
- }
- // All accounts require instructions for Responses API
- payload["instructions"] = openai.DefaultInstructions
- return payload
- }
- func createOpenAIChatCompletionsTestPayload(modelID string, prompt string) map[string]any {
- testPrompt := strings.TrimSpace(prompt)
- if testPrompt == "" {
- testPrompt = "hi"
- }
- return map[string]any{
- "model": modelID,
- "messages": []map[string]any{
- {
- "role": "user",
- "content": testPrompt,
- },
- },
- "stream": true,
- }
- }
- // processClaudeStream processes the SSE stream from Claude API
- func (s *AccountTestService) processClaudeStream(c *gin.Context, body io.Reader) error {
- reader := bufio.NewReader(body)
- for {
- line, err := reader.ReadString('\n')
- if err != nil {
- if err == io.EOF {
- s.sendEvent(c, TestEvent{Type: "test_complete", Success: true})
- return nil
- }
- return s.sendErrorAndEnd(c, fmt.Sprintf("Stream read error: %s", err.Error()))
- }
- line = strings.TrimSpace(line)
- if line == "" || !sseDataPrefix.MatchString(line) {
- continue
- }
- jsonStr := sseDataPrefix.ReplaceAllString(line, "")
- if jsonStr == "[DONE]" {
- s.sendEvent(c, TestEvent{Type: "test_complete", Success: true})
- return nil
- }
- var data map[string]any
- if err := json.Unmarshal([]byte(jsonStr), &data); err != nil {
- continue
- }
- eventType, _ := data["type"].(string)
- switch eventType {
- case "content_block_delta":
- if delta, ok := data["delta"].(map[string]any); ok {
- if text, ok := delta["text"].(string); ok {
- s.sendEvent(c, TestEvent{Type: "content", Text: text})
- }
- }
- case "message_stop":
- s.sendEvent(c, TestEvent{Type: "test_complete", Success: true})
- return nil
- case "error":
- errorMsg := "Unknown error"
- if errData, ok := data["error"].(map[string]any); ok {
- if msg, ok := errData["message"].(string); ok {
- errorMsg = msg
- }
- }
- return s.sendErrorAndEnd(c, errorMsg)
- }
- }
- }
- // processOpenAIChatCompletionsStream processes SSE chunks from the
- // OpenAI-compatible Chat Completions API.
- func (s *AccountTestService) processOpenAIChatCompletionsStream(c *gin.Context, body io.Reader) error {
- reader := bufio.NewReader(body)
- seenJSON := false
- seenFinish := false
- for {
- line, err := reader.ReadString('\n')
- if err != nil {
- if err == io.EOF {
- if seenFinish {
- s.sendEvent(c, TestEvent{Type: "status", Text: "已通过 /v1/chat/completions 验证"})
- s.sendEvent(c, TestEvent{Type: "test_complete", Success: true})
- return nil
- }
- if seenJSON {
- return s.sendErrorAndEnd(c, "Chat Completions stream from /v1/chat/completions ended before [DONE]")
- }
- return s.sendErrorAndEnd(c, "Invalid Chat Completions response from /v1/chat/completions: expected SSE JSON data")
- }
- return s.sendErrorAndEnd(c, fmt.Sprintf("Chat Completions stream read error from /v1/chat/completions: %s", err.Error()))
- }
- line = strings.TrimSpace(line)
- if line == "" || !sseDataPrefix.MatchString(line) {
- continue
- }
- jsonStr := sseDataPrefix.ReplaceAllString(line, "")
- if jsonStr == "[DONE]" {
- s.sendEvent(c, TestEvent{Type: "status", Text: "已通过 /v1/chat/completions 验证"})
- s.sendEvent(c, TestEvent{Type: "test_complete", Success: true})
- return nil
- }
- var data map[string]any
- if err := json.Unmarshal([]byte(jsonStr), &data); err != nil {
- return s.sendErrorAndEnd(c, "Invalid Chat Completions response from /v1/chat/completions: expected JSON data")
- }
- seenJSON = true
- if errData, ok := data["error"].(map[string]any); ok {
- errorMsg := "Chat Completions API (/v1/chat/completions) returned an error"
- if msg, ok := errData["message"].(string); ok && msg != "" {
- errorMsg = msg
- }
- return s.sendErrorAndEnd(c, fmt.Sprintf("Chat Completions API (/v1/chat/completions) error: %s", errorMsg))
- }
- choices, ok := data["choices"].([]any)
- if !ok {
- continue
- }
- for _, choiceValue := range choices {
- choice, ok := choiceValue.(map[string]any)
- if !ok {
- continue
- }
- if delta, ok := choice["delta"].(map[string]any); ok {
- if text, ok := delta["content"].(string); ok && text != "" {
- s.sendEvent(c, TestEvent{Type: "content", Text: text})
- }
- }
- if message, ok := choice["message"].(map[string]any); ok {
- if text, ok := message["content"].(string); ok && text != "" {
- s.sendEvent(c, TestEvent{Type: "content", Text: text})
- }
- }
- if finishReason, ok := choice["finish_reason"].(string); ok && finishReason != "" {
- seenFinish = true
- }
- }
- }
- }
- // processOpenAIStream processes the SSE stream from OpenAI Responses API
- func (s *AccountTestService) processOpenAIStream(c *gin.Context, body io.Reader) error {
- reader := bufio.NewReader(body)
- seenCompleted := false
- seenText := false
- for {
- line, err := reader.ReadString('\n')
- if err != nil {
- if err == io.EOF {
- if seenCompleted {
- s.sendEvent(c, TestEvent{Type: "test_complete", Success: true})
- return nil
- }
- return s.sendErrorAndEnd(c, "Stream ended before response.completed")
- }
- return s.sendErrorAndEnd(c, fmt.Sprintf("Stream read error: %s", err.Error()))
- }
- line = strings.TrimSpace(line)
- if line == "" || !sseDataPrefix.MatchString(line) {
- continue
- }
- jsonStr := sseDataPrefix.ReplaceAllString(line, "")
- if jsonStr == "[DONE]" {
- if seenCompleted {
- s.sendEvent(c, TestEvent{Type: "test_complete", Success: true})
- return nil
- }
- return s.sendErrorAndEnd(c, "Stream ended before response.completed")
- }
- var data map[string]any
- if err := json.Unmarshal([]byte(jsonStr), &data); err != nil {
- continue
- }
- eventType, _ := data["type"].(string)
- switch eventType {
- case "response.output_text.delta":
- // OpenAI Responses API uses "delta" field for text content
- if delta, ok := data["delta"].(string); ok && delta != "" {
- seenText = true
- s.sendEvent(c, TestEvent{Type: "content", Text: delta})
- }
- case "response.completed", "response.done":
- if !seenText {
- response, _ := data["response"].(map[string]any)
- if output, ok := response["output"].([]any); ok {
- for _, raw := range output {
- if item, ok := raw.(map[string]any); ok {
- if content, ok := item["content"].([]any); ok {
- for _, part := range content {
- if text, ok := part.(map[string]any); ok && text["type"] == "output_text" {
- if value, ok := text["text"].(string); ok && value != "" {
- s.sendEvent(c, TestEvent{Type: "content", Text: value})
- }
- }
- }
- }
- }
- }
- }
- }
- s.sendEvent(c, TestEvent{Type: "test_complete", Success: true})
- return nil
- case "response.failed":
- errorMsg := "OpenAI response failed"
- if responseData, ok := data["response"].(map[string]any); ok {
- if errData, ok := responseData["error"].(map[string]any); ok {
- if msg, ok := errData["message"].(string); ok && msg != "" {
- errorMsg = msg
- }
- }
- }
- return s.sendErrorAndEnd(c, errorMsg)
- case "error":
- errorMsg := "Unknown error"
- if errData, ok := data["error"].(map[string]any); ok {
- if msg, ok := errData["message"].(string); ok {
- errorMsg = msg
- }
- }
- return s.sendErrorAndEnd(c, errorMsg)
- }
- }
- }
- // testOpenAIImageAPIKey tests OpenAI image generation using an API Key account.
- func (s *AccountTestService) testOpenAIImageAPIKey(c *gin.Context, ctx context.Context, account *Account, modelID, prompt string) error {
- authToken := account.GetOpenAIApiKey()
- if authToken == "" {
- return s.sendErrorAndEnd(c, "No API key available")
- }
- baseURL := account.GetOpenAIBaseURL()
- if baseURL == "" {
- baseURL = "https://api.openai.com"
- }
- normalizedBaseURL, err := s.validateUpstreamBaseURL(baseURL)
- if err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Invalid base URL: %s", err.Error()))
- }
- apiURL := buildOpenAIImagesURL(normalizedBaseURL, openAIImagesGenerationsEndpoint)
- // Set SSE headers
- c.Writer.Header().Set("Content-Type", "text/event-stream")
- c.Writer.Header().Set("Cache-Control", "no-cache")
- c.Writer.Header().Set("Connection", "keep-alive")
- c.Writer.Header().Set("X-Accel-Buffering", "no")
- c.Writer.Flush()
- s.sendEvent(c, TestEvent{Type: "test_start", Model: modelID})
- payload := map[string]any{
- "model": modelID,
- "prompt": prompt,
- "n": 1,
- "response_format": "b64_json",
- }
- applyIntelligentPayloadPrompt(ctx, payload)
- payloadBytes, _ := json.Marshal(payload)
- req, err := http.NewRequestWithContext(ctx, "POST", apiURL, bytes.NewReader(payloadBytes))
- if err != nil {
- return s.sendErrorAndEnd(c, "Failed to create request")
- }
- req = req.WithContext(WithHTTPUpstreamProfile(req.Context(), HTTPUpstreamProfileOpenAI))
- req.Header.Set("Content-Type", "application/json")
- req.Header.Set("Authorization", "Bearer "+authToken)
- // 账号级请求头覆写:测试请求与真实转发保持一致的最终头
- account.ApplyHeaderOverrides(req.Header)
- proxyURL := ""
- if account.ProxyID != nil && account.Proxy != nil {
- proxyURL = account.Proxy.URL()
- }
- resp, err := s.httpUpstream.DoWithTLS(WithAccountTrafficRequest(req, account), proxyURL, account.ID, account.Concurrency, s.tlsFPProfileService.ResolveTLSProfile(account))
- if err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Request failed: %s", err.Error()))
- }
- defer func() { _ = resp.Body.Close() }()
- body, err := io.ReadAll(resp.Body)
- if err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Failed to read response: %s", err.Error()))
- }
- if resp.StatusCode != http.StatusOK {
- return s.sendErrorAndEnd(c, fmt.Sprintf("API returned %d: %s", resp.StatusCode, string(body)))
- }
- // Parse {"data": [{"b64_json": "...", "revised_prompt": "..."}]}
- var result struct {
- Data []struct {
- B64JSON string `json:"b64_json"`
- RevisedPrompt string `json:"revised_prompt"`
- } `json:"data"`
- }
- if err := json.Unmarshal(body, &result); err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Failed to parse response: %s", err.Error()))
- }
- if len(result.Data) == 0 {
- return s.sendErrorAndEnd(c, "No images returned from API")
- }
- for _, item := range result.Data {
- if item.RevisedPrompt != "" {
- s.sendEvent(c, TestEvent{Type: "content", Text: item.RevisedPrompt})
- }
- if item.B64JSON != "" {
- s.sendEvent(c, TestEvent{
- Type: "image",
- ImageURL: "data:image/png;base64," + item.B64JSON,
- MimeType: "image/png",
- })
- }
- }
- s.sendEvent(c, TestEvent{Type: "test_complete", Success: true})
- return nil
- }
- // testOpenAIImageOAuth tests OpenAI image generation using an OAuth account via Codex /responses API.
- func (s *AccountTestService) testOpenAIImageOAuth(c *gin.Context, ctx context.Context, account *Account, modelID, prompt string) error {
- credentialAccount := account
- if account.IsShadow() {
- resolved, err := resolveCredentialAccount(ctx, s.accountRepo, account)
- if err != nil {
- return s.sendErrorAndEnd(c, "Failed to resolve account credentials")
- }
- credentialAccount = resolved
- }
- authToken := ""
- if !credentialAccount.IsOpenAIAgentIdentity() {
- authToken = credentialAccount.GetOpenAIAccessToken()
- }
- if authToken == "" && !credentialAccount.IsOpenAIAgentIdentity() {
- return s.sendErrorAndEnd(c, "No access token available")
- }
- // Set SSE headers
- c.Writer.Header().Set("Content-Type", "text/event-stream")
- c.Writer.Header().Set("Cache-Control", "no-cache")
- c.Writer.Header().Set("Connection", "keep-alive")
- c.Writer.Header().Set("X-Accel-Buffering", "no")
- c.Writer.Flush()
- s.sendEvent(c, TestEvent{Type: "test_start", Model: modelID})
- s.sendEvent(c, TestEvent{Type: "content", Text: "Calling Codex /responses image tool...\n"})
- parsed := &OpenAIImagesRequest{
- Endpoint: openAIImagesGenerationsEndpoint,
- Model: strings.TrimSpace(modelID),
- Prompt: prompt,
- }
- applyOpenAIImagesDefaults(parsed)
- s.sendEvent(c, TestEvent{Type: "content", Text: fmt.Sprintf("Responses driver: %s; image model: %s\n", openAIImagesResponsesMainModelValue(), parsed.Model)})
- responsesBody, err := buildOpenAIImagesResponsesRequest(parsed, parsed.Model)
- if err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Failed to build image request: %s", err.Error()))
- }
- req, err := http.NewRequestWithContext(ctx, http.MethodPost, chatgptCodexAPIURL, bytes.NewReader(responsesBody))
- if err != nil {
- return s.sendErrorAndEnd(c, "Failed to create request")
- }
- req = req.WithContext(WithHTTPUpstreamProfile(req.Context(), HTTPUpstreamProfileOpenAI))
- req.Host = "chatgpt.com"
- if credentialAccount.IsOpenAIAgentIdentity() {
- authHeaders, authErr := buildAgentIdentityAuthenticationHeaders(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, credentialAccount)
- if authErr != nil {
- return s.sendErrorAndEnd(c, "Failed to build Agent Identity authentication")
- }
- for key, values := range authHeaders {
- for _, value := range values {
- req.Header.Add(key, value)
- }
- }
- } else {
- req.Header.Set("Authorization", "Bearer "+authToken)
- }
- req.Header.Set("Content-Type", "application/json")
- req.Header.Set("Accept", "text/event-stream")
- req.Header.Set("OpenAI-Beta", "responses=experimental")
- canonical := resolveCodexOutboundIdentity("")
- req.Header.Set("originator", canonical.originator)
- if customUA := strings.TrimSpace(credentialAccount.GetOpenAIUserAgent()); customUA != "" {
- req.Header.Set("User-Agent", customUA)
- } else {
- req.Header.Set("User-Agent", canonical.userAgent)
- }
- setOpenAIChatGPTAccountHeaders(req.Header, credentialAccount)
- // 与真实转发一致:账号级自定义 UA 同样作为管理员显式配置传入,否则测试用的身份
- // 与该账号真实出站的身份不是同一个(issue #3901 的配对不变式由收口保证)。
- enforceCodexIdentityHeadersWithUA(req.Header, credentialAccount.GetOpenAIUserAgent())
- proxyURL := ""
- if account.ProxyID != nil && account.Proxy != nil {
- proxyURL = account.Proxy.URL()
- }
- resp, err := s.doOpenAIAccountTestUpstream(req, proxyURL, account, false)
- if err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Responses API request failed: %s", err.Error()))
- }
- defer func() {
- if resp != nil && resp.Body != nil {
- _ = resp.Body.Close()
- }
- }()
- if resp.StatusCode >= 400 {
- body, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
- body = redactAgentIdentitySensitiveBodyForAccount(ctx, s.accountRepo, credentialAccount, body)
- message := strings.TrimSpace(extractUpstreamErrorMessage(body))
- if message == "" {
- message = fmt.Sprintf("Responses API returned %d", resp.StatusCode)
- }
- return s.sendErrorAndEnd(c, message)
- }
- body, err := io.ReadAll(resp.Body)
- if err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Failed to read image response: %s", err.Error()))
- }
- body = redactAgentIdentitySensitiveBodyForAccount(ctx, s.accountRepo, credentialAccount, body)
- results, _, _, _, _, err := collectOpenAIImagesFromResponsesBody(body)
- if err != nil {
- return s.sendErrorAndEnd(c, fmt.Sprintf("Failed to parse image response: %s", err.Error()))
- }
- if len(results) == 0 {
- if upstreamErr := extractOpenAIImagesUpstreamError(body); upstreamErr != nil {
- return s.sendErrorAndEnd(c, upstreamErr.clientMessage())
- }
- if textErr := openAIImagesTextFallbackError(body); textErr != nil {
- return s.sendErrorAndEnd(c, textErr.clientMessage())
- }
- return s.sendErrorAndEnd(c, "No images returned from responses API")
- }
- for _, item := range results {
- if item.RevisedPrompt != "" {
- s.sendEvent(c, TestEvent{Type: "content", Text: item.RevisedPrompt})
- }
- mimeType := openAIImageOutputMIMEType(item.OutputFormat)
- s.sendEvent(c, TestEvent{
- Type: "image",
- ImageURL: "data:" + mimeType + ";base64," + item.Result,
- MimeType: mimeType,
- })
- }
- s.sendEvent(c, TestEvent{Type: "test_complete", Success: true})
- return nil
- }
- func (s *AccountTestService) sendEvent(c *gin.Context, event TestEvent) {
- if event.Type == "test_complete" {
- if suppress, ok := c.Get(accountTestSuppressCompletionContextKey); ok {
- if suppressCompletion, _ := suppress.(bool); suppressCompletion {
- return
- }
- }
- }
- eventJSON, _ := json.Marshal(event)
- if _, err := fmt.Fprintf(c.Writer, "data: %s\n\n", eventJSON); err != nil {
- log.Printf("failed to write SSE event: %v", err)
- return
- }
- c.Writer.Flush()
- }
- // sendErrorAndEnd sends an error event and ends the stream
- func (s *AccountTestService) sendErrorAndEnd(c *gin.Context, errorMsg string) error {
- if c.Request != nil {
- if run := intelligentContext(c.Request.Context()); run != nil {
- errorMsg = run.capture.redact(errorMsg)
- }
- }
- log.Printf("Account test error: %s", errorMsg)
- s.sendEvent(c, TestEvent{Type: "error", Error: errorMsg})
- return fmt.Errorf("%s", errorMsg)
- }
- // RunTestBackground executes an account test in-memory (no real HTTP client),
- // capturing SSE output via httptest.NewRecorder, then parses the result.
- func (s *AccountTestService) RunTestBackground(ctx context.Context, accountID int64, modelID string) (*ScheduledTestResult, error) {
- startedAt := time.Now()
- w := httptest.NewRecorder()
- ginCtx, _ := gin.CreateTestContext(w)
- ginCtx.Request = (&http.Request{}).WithContext(ctx)
- testErr := s.TestAccountConnection(ginCtx, accountID, modelID, "", AccountTestModeDefault)
- finishedAt := time.Now()
- body := w.Body.String()
- responseText, errMsg := parseTestSSEOutput(body)
- status := "success"
- if testErr != nil || errMsg != "" {
- status = "failed"
- if errMsg == "" && testErr != nil {
- errMsg = testErr.Error()
- }
- }
- return &ScheduledTestResult{
- Status: status,
- ResponseText: responseText,
- ErrorMessage: errMsg,
- LatencyMs: finishedAt.Sub(startedAt).Milliseconds(),
- StartedAt: startedAt,
- FinishedAt: finishedAt,
- }, nil
- }
- // parseTestSSEOutput extracts response text and error message from captured SSE output.
- func parseTestSSEOutput(body string) (responseText, errMsg string) {
- var texts []string
- for _, line := range strings.Split(body, "\n") {
- line = strings.TrimSpace(line)
- if !strings.HasPrefix(line, "data: ") {
- continue
- }
- jsonStr := strings.TrimPrefix(line, "data: ")
- var event TestEvent
- if err := json.Unmarshal([]byte(jsonStr), &event); err != nil {
- continue
- }
- switch event.Type {
- case "content":
- if event.Text != "" {
- texts = append(texts, event.Text)
- }
- case "error":
- errMsg = event.Error
- }
- }
- responseText = strings.Join(texts, "")
- return
- }
|