李建琦
6d655c9903
Release / update-version (push) Has been cancelled
Release / build-frontend (push) Has been cancelled
Release / release (push) Has been cancelled
Release / sync-version-file (push) Has been cancelled
CI / shell (push) Canceled after 0s
CI / test (push) Canceled after 0s
CI / frontend (push) Canceled after 0s
CI / golangci-lint (push) Canceled after 0s
Security Scan / backend-security (push) Canceled after 0s
Security Scan / frontend-security (push) Canceled after 0s
3151 lines
111 KiB
Go
3151 lines
111 KiB
Go
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"
|
||
)
|
||
|
||
// 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 = 12 * time.Second
|
||
)
|
||
|
||
// 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 {
|
||
accountRepo AccountRepository
|
||
geminiTokenProvider *GeminiTokenProvider
|
||
claudeTokenProvider *ClaudeTokenProvider
|
||
grokTokenProvider *GrokTokenProvider
|
||
antigravityGatewayService *AntigravityGatewayService
|
||
httpUpstream HTTPUpstream
|
||
cfg *config.Config
|
||
settingService *SettingService
|
||
tlsFPProfileService *TLSFingerprintProfileService
|
||
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
|
||
}
|
||
}
|
||
|
||
// 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")
|
||
}
|
||
|
||
// 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
|
||
}
|
||
|
||
// Route to platform-specific test method
|
||
if account.IsCNProvider() {
|
||
switch account.GetAPIProtocol() {
|
||
case APIProtocolAdaptive:
|
||
return s.testCNProviderAdaptiveConnection(c, account, modelID, prompt)
|
||
case APIProtocolChatCompletions:
|
||
return s.testCNProviderChatCompletionsConnection(c, account, modelID, prompt)
|
||
}
|
||
}
|
||
|
||
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")
|
||
}
|
||
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)
|
||
setAnthropicAPIKeyAuthHeader(req.Header, account, authToken)
|
||
}
|
||
|
||
// 账号级请求头覆写：测试请求与真实转发保持一致的最终头
|
||
account.ApplyHeaderOverrides(req.Header)
|
||
|
||
// Get proxy URL
|
||
proxyURL := ""
|
||
if account.ProxyID != nil && account.Proxy != nil {
|
||
proxyURL = account.Proxy.URL()
|
||
}
|
||
|
||
resp, err := s.httpUpstream.DoWithTLS(req, 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")
|
||
}
|
||
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(req, 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,
|
||
}
|
||
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(req, 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.GetOpenAIApiKey()
|
||
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 = buildOpenAIResponsesURL(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, testModelID)
|
||
}
|
||
payload := createOpenAITestPayload(upstreamTestModelID, isOAuth)
|
||
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) {
|
||
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))
|
||
|
||
// 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)
|
||
|
||
// Get proxy URL
|
||
proxyURL := ""
|
||
if account.ProxyID != nil && account.Proxy != nil {
|
||
proxyURL = account.Proxy.URL()
|
||
}
|
||
|
||
resp, err := s.httpUpstream.DoWithTLS(req, 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 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 !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(req, 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(req, 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"})
|
||
}
|
||
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(req, 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(statusReq, 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(req, 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,
|
||
}
|
||
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(req, 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 {
|
||
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(req, 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(req, 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) error {
|
||
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)
|
||
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(req, 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.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 = buildOpenAIResponsesURL(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.httpUpstream.DoWithTLS(req, proxyURL, account.ID, account.Concurrency, s.tlsFPProfileService.ResolveTLSProfile(account))
|
||
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(req, 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()
|
||
|
||
// 默认模型：Claude 使用 claude-sonnet-4-5，Gemini 使用 gemini-3-pro-preview
|
||
testModelID := modelID
|
||
if testModelID == "" {
|
||
testModelID = "claude-sonnet-4-5"
|
||
}
|
||
|
||
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
|
||
}
|
||
|
||
// 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
|
||
|
||
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 != "" {
|
||
s.sendEvent(c, TestEvent{Type: "content", Text: delta})
|
||
}
|
||
case "response.completed", "response.done":
|
||
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",
|
||
}
|
||
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(req, 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)
|
||
|
||
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.httpUpstream.Do(req, proxyURL, account.ID, account.Concurrency)
|
||
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 {
|
||
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 {
|
||
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
|
||
}
|