Files
sub2api/backend/internal/service/account_test_service.go

958 lines
29 KiB
Go
Raw Normal View History

2025-12-18 13:50:39 +08:00
package service
import (
"bufio"
"bytes"
"context"
2025-12-18 13:50:39 +08:00
"crypto/rand"
"encoding/hex"
"encoding/json"
"fmt"
"io"
"log"
"net/http"
"regexp"
2025-12-18 13:50:39 +08:00
"strings"
"time"
2025-12-24 21:07:21 +08:00
"github.com/Wei-Shaw/sub2api/internal/pkg/claude"
"github.com/Wei-Shaw/sub2api/internal/pkg/geminicli"
2025-12-24 21:07:21 +08:00
"github.com/Wei-Shaw/sub2api/internal/pkg/openai"
2025-12-18 13:50:39 +08:00
"github.com/gin-gonic/gin"
"github.com/google/uuid"
)
// 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*`)
2025-12-18 13:50:39 +08:00
const (
2025-12-22 22:58:31 +08:00
testClaudeAPIURL = "https://api.anthropic.com/v1/messages"
testOpenAIAPIURL = "https://api.openai.com/v1/responses"
chatgptCodexAPIURL = "https://chatgpt.com/backend-api/codex/responses"
2025-12-18 13:50:39 +08:00
)
// 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"`
Success bool `json:"success,omitempty"`
Error string `json:"error,omitempty"`
}
// AccountTestService handles account testing operations
type AccountTestService struct {
accountRepo AccountRepository
oauthService *OAuthService
openaiOAuthService *OpenAIOAuthService
geminiTokenProvider *GeminiTokenProvider
antigravityGatewayService *AntigravityGatewayService
httpUpstream HTTPUpstream
2025-12-18 13:50:39 +08:00
}
// NewAccountTestService creates a new AccountTestService
func NewAccountTestService(
accountRepo AccountRepository,
oauthService *OAuthService,
openaiOAuthService *OpenAIOAuthService,
geminiTokenProvider *GeminiTokenProvider,
antigravityGatewayService *AntigravityGatewayService,
httpUpstream HTTPUpstream,
) *AccountTestService {
2025-12-18 13:50:39 +08:00
return &AccountTestService{
accountRepo: accountRepo,
oauthService: oauthService,
openaiOAuthService: openaiOAuthService,
geminiTokenProvider: geminiTokenProvider,
antigravityGatewayService: antigravityGatewayService,
httpUpstream: httpUpstream,
2025-12-18 13:50:39 +08:00
}
}
// generateSessionString generates a Claude Code style session string
func generateSessionString() (string, error) {
2025-12-18 13:50:39 +08:00
bytes := make([]byte, 32)
if _, err := rand.Read(bytes); err != nil {
return "", err
}
2025-12-18 13:50:39 +08:00
hex64 := hex.EncodeToString(bytes)
sessionUUID := uuid.New().String()
return fmt.Sprintf("user_%s_account__session_%s", hex64, sessionUUID), nil
2025-12-18 13:50:39 +08:00
}
// createTestPayload creates a Claude Code style test request payload
2025-12-20 16:19:40 +08:00
func createTestPayload(modelID string) (map[string]any, error) {
sessionID, err := generateSessionString()
if err != nil {
return nil, err
}
2025-12-20 16:19:40 +08:00
return map[string]any{
"model": modelID,
2025-12-20 16:19:40 +08:00
"messages": []map[string]any{
2025-12-18 13:50:39 +08:00
{
"role": "user",
2025-12-20 16:19:40 +08:00
"content": []map[string]any{
2025-12-18 13:50:39 +08:00
{
"type": "text",
"text": "hi",
"cache_control": map[string]string{
"type": "ephemeral",
},
},
},
},
},
2025-12-20 16:19:40 +08:00
"system": []map[string]any{
2025-12-18 13:50:39 +08:00
{
"type": "text",
"text": "You are Claude Code, Anthropic's official CLI for Claude.",
"cache_control": map[string]string{
"type": "ephemeral",
},
},
},
"metadata": map[string]string{
"user_id": sessionID,
2025-12-18 13:50:39 +08:00
},
"max_tokens": 1024,
2025-12-18 13:50:39 +08:00
"temperature": 1,
"stream": true,
}, nil
2025-12-18 13:50:39 +08:00
}
// 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
func (s *AccountTestService) TestAccountConnection(c *gin.Context, accountID int64, modelID string) error {
2025-12-18 13:50:39 +08:00
ctx := c.Request.Context()
// Get account
account, err := s.accountRepo.GetByID(ctx, accountID)
2025-12-18 13:50:39 +08:00
if err != nil {
return s.sendErrorAndEnd(c, "Account not found")
}
2025-12-22 22:58:31 +08:00
// Route to platform-specific test method
if account.IsOpenAI() {
return s.testOpenAIAccountConnection(c, account, modelID)
}
if account.IsGemini() {
return s.testGeminiAccountConnection(c, account, modelID)
}
if account.Platform == PlatformAntigravity {
return s.testAntigravityAccountConnection(c, account, modelID)
}
2025-12-22 22:58:31 +08:00
return s.testClaudeAccountConnection(c, account, modelID)
}
// testClaudeAccountConnection tests an Anthropic Claude account's connection
func (s *AccountTestService) testClaudeAccountConnection(c *gin.Context, account *Account, modelID string) error {
2025-12-22 22:58:31 +08:00
ctx := c.Request.Context()
// Determine the model to use
testModelID := modelID
if testModelID == "" {
testModelID = claude.DefaultTestModel
}
// For API Key accounts with model mapping, map the model
if account.Type == "apikey" {
mapping := account.GetModelMapping()
if len(mapping) > 0 {
if mappedModel, exists := mapping[testModelID]; exists {
testModelID = mappedModel
}
}
}
// Determine authentication method and API URL
2025-12-18 13:50:39 +08:00
var authToken string
var useBearer bool
2025-12-18 13:50:39 +08:00
var apiURL string
if account.IsOAuth() {
// OAuth or Setup Token - use Bearer token
useBearer = true
2025-12-18 13:50:39 +08:00
apiURL = testClaudeAPIURL
authToken = account.GetCredential("access_token")
if authToken == "" {
return s.sendErrorAndEnd(c, "No access token available")
}
// Check if token needs refresh
needRefresh := false
if expiresAt := account.GetCredentialAsTime("expires_at"); expiresAt != nil {
if time.Now().Add(5 * time.Minute).After(*expiresAt) {
2025-12-18 13:50:39 +08:00
needRefresh = true
}
}
if needRefresh && s.oauthService != nil {
tokenInfo, err := s.oauthService.RefreshAccountToken(ctx, account)
if err != nil {
return s.sendErrorAndEnd(c, fmt.Sprintf("Failed to refresh token: %s", err.Error()))
}
authToken = tokenInfo.AccessToken
}
} else if account.Type == "apikey" {
// API Key - use x-api-key header
useBearer = false
2025-12-18 13:50:39 +08:00
authToken = account.GetCredential("api_key")
if authToken == "" {
return s.sendErrorAndEnd(c, "No API key available")
}
apiURL = account.GetBaseURL()
if apiURL == "" {
apiURL = "https://api.anthropic.com"
}
apiURL = strings.TrimSuffix(apiURL, "/") + "/v1/messages"
} 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")
}
2025-12-18 13:50:39 +08:00
payloadBytes, _ := json.Marshal(payload)
// Send test_start event
s.sendEvent(c, TestEvent{Type: "test_start", Model: testModelID})
2025-12-18 13:50:39 +08:00
req, err := http.NewRequestWithContext(ctx, "POST", apiURL, bytes.NewReader(payloadBytes))
if err != nil {
return s.sendErrorAndEnd(c, "Failed to create request")
}
// Set common headers
2025-12-18 13:50:39 +08:00
req.Header.Set("Content-Type", "application/json")
req.Header.Set("anthropic-version", "2023-06-01")
req.Header.Set("anthropic-beta", claude.DefaultBetaHeader)
// Apply Claude Code client headers
for key, value := range claude.DefaultHeaders {
req.Header.Set(key, value)
}
2025-12-18 13:50:39 +08:00
// Set authentication header
if useBearer {
2025-12-18 13:50:39 +08:00
req.Header.Set("Authorization", "Bearer "+authToken)
} else {
req.Header.Set("x-api-key", authToken)
}
2025-12-20 11:56:11 +08:00
// Get proxy URL
proxyURL := ""
2025-12-18 13:50:39 +08:00
if account.ProxyID != nil && account.Proxy != nil {
2025-12-20 11:56:11 +08:00
proxyURL = account.Proxy.URL()
2025-12-18 13:50:39 +08:00
}
resp, err := s.httpUpstream.Do(req, proxyURL, account.ID, account.Concurrency)
2025-12-22 22:58:31 +08:00
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.processClaudeStream(c, resp.Body)
}
// testOpenAIAccountConnection tests an OpenAI account's connection
func (s *AccountTestService) testOpenAIAccountConnection(c *gin.Context, account *Account, modelID string) error {
2025-12-22 22:58:31 +08:00
ctx := c.Request.Context()
// Default to openai.DefaultTestModel for OpenAI testing
testModelID := modelID
if testModelID == "" {
testModelID = openai.DefaultTestModel
}
// For API Key accounts with model mapping, map the model
if account.Type == "apikey" {
mapping := account.GetModelMapping()
if len(mapping) > 0 {
if mappedModel, exists := mapping[testModelID]; exists {
testModelID = mappedModel
}
}
}
// Set SSE headers early
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})
// Get proxy URL
proxyURL := ""
if account.ProxyID != nil && account.Proxy != nil {
proxyURL = account.Proxy.URL()
}
2025-12-22 22:58:31 +08:00
if account.IsOAuth() {
// OAuth - use ChatGPT internal API (Responses API)
return s.testOpenAIOAuthAccount(c, ctx, account, testModelID, proxyURL)
}
2025-12-22 22:58:31 +08:00
// API Key - try Chat Completions API first, fallback to Responses API
return s.testOpenAIApiKeyAccount(c, ctx, account, testModelID, proxyURL)
}
2025-12-22 22:58:31 +08:00
// testOpenAIOAuthAccount tests OAuth account using ChatGPT internal API
func (s *AccountTestService) testOpenAIOAuthAccount(c *gin.Context, ctx context.Context, account *Account, testModelID, proxyURL string) error {
authToken := account.GetOpenAIAccessToken()
if authToken == "" {
return s.sendErrorAndEnd(c, "No access token available")
}
2025-12-22 22:58:31 +08:00
// Check if token is expired and refresh if needed
if account.IsOpenAITokenExpired() && s.openaiOAuthService != nil {
tokenInfo, err := s.openaiOAuthService.RefreshAccountToken(ctx, account)
if err != nil {
return s.sendErrorAndEnd(c, fmt.Sprintf("Failed to refresh token: %s", err.Error()))
2025-12-22 22:58:31 +08:00
}
authToken = tokenInfo.AccessToken
2025-12-22 22:58:31 +08:00
}
// Create Responses API payload
payload := createOpenAITestPayload(testModelID, true)
2025-12-22 22:58:31 +08:00
payloadBytes, _ := json.Marshal(payload)
req, err := http.NewRequestWithContext(ctx, "POST", chatgptCodexAPIURL, bytes.NewReader(payloadBytes))
2025-12-22 22:58:31 +08:00
if err != nil {
return s.sendErrorAndEnd(c, "Failed to create request")
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", "Bearer "+authToken)
req.Host = "chatgpt.com"
req.Header.Set("accept", "text/event-stream")
if chatgptAccountID := account.GetChatGPTAccountID(); chatgptAccountID != "" {
req.Header.Set("chatgpt-account-id", chatgptAccountID)
}
2025-12-22 22:58:31 +08:00
resp, err := s.httpUpstream.Do(req, proxyURL, account.ID, account.Concurrency)
if err != nil {
return s.sendErrorAndEnd(c, fmt.Sprintf("Request failed: %s", err.Error()))
2025-12-22 22:58:31 +08:00
}
defer func() { _ = resp.Body.Close() }()
2025-12-22 22:58:31 +08:00
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(resp.Body)
return s.sendErrorAndEnd(c, fmt.Sprintf("API returned %d: %s", resp.StatusCode, string(body)))
}
return s.processOpenAIStream(c, resp.Body)
}
// testOpenAIApiKeyAccount tests API Key account, trying Chat Completions first, then Responses API
func (s *AccountTestService) testOpenAIApiKeyAccount(c *gin.Context, ctx context.Context, account *Account, testModelID, proxyURL 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"
}
baseURL = strings.TrimSuffix(baseURL, "/")
// Try Chat Completions API first (more compatible with third-party proxies)
chatCompletionsURL := baseURL + "/v1/chat/completions"
chatPayload := createOpenAIChatCompletionsPayload(testModelID)
chatPayloadBytes, _ := json.Marshal(chatPayload)
req, err := http.NewRequestWithContext(ctx, "POST", chatCompletionsURL, bytes.NewReader(chatPayloadBytes))
if err != nil {
return s.sendErrorAndEnd(c, "Failed to create request")
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", "Bearer "+authToken)
resp, err := s.httpUpstream.Do(req, proxyURL, account.ID, account.Concurrency)
if err != nil {
// Network error, try Responses API
s.sendEvent(c, TestEvent{Type: "info", Text: "Chat Completions API failed, trying Responses API..."})
return s.tryOpenAIResponsesAPI(c, ctx, account, testModelID, baseURL, authToken, proxyURL)
}
defer func() { _ = resp.Body.Close() }()
if resp.StatusCode == http.StatusOK {
// Chat Completions API succeeded
return s.processOpenAIChatCompletionsStream(c, resp.Body)
}
// Chat Completions API failed, try Responses API
_ = resp.Body.Close()
s.sendEvent(c, TestEvent{Type: "info", Text: "Chat Completions API failed, trying Responses API..."})
return s.tryOpenAIResponsesAPI(c, ctx, account, testModelID, baseURL, authToken, proxyURL)
}
// tryOpenAIResponsesAPI tries the OpenAI Responses API as fallback
func (s *AccountTestService) tryOpenAIResponsesAPI(c *gin.Context, ctx context.Context, account *Account, testModelID, baseURL, authToken, proxyURL string) error {
responsesURL := baseURL + "/v1/responses"
payload := createOpenAITestPayload(testModelID, false)
payloadBytes, _ := json.Marshal(payload)
req, err := http.NewRequestWithContext(ctx, "POST", responsesURL, bytes.NewReader(payloadBytes))
if err != nil {
return s.sendErrorAndEnd(c, "Failed to create request")
2025-12-22 22:58:31 +08:00
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", "Bearer "+authToken)
2025-12-22 22:58:31 +08:00
resp, err := s.httpUpstream.Do(req, proxyURL, account.ID, account.Concurrency)
2025-12-18 13:50:39 +08:00
if err != nil {
return s.sendErrorAndEnd(c, fmt.Sprintf("Request failed: %s", err.Error()))
}
defer func() { _ = resp.Body.Close() }()
2025-12-18 13:50:39 +08:00
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(resp.Body)
return s.sendErrorAndEnd(c, fmt.Sprintf("API returned %d: %s", resp.StatusCode, string(body)))
}
2025-12-22 22:58:31 +08:00
return s.processOpenAIStream(c, resp.Body)
}
// createOpenAIChatCompletionsPayload creates a test payload for OpenAI Chat Completions API
func createOpenAIChatCompletionsPayload(modelID string) map[string]any {
return map[string]any{
"model": modelID,
"messages": []map[string]any{
{
"role": "user",
"content": "hi",
},
},
"stream": true,
"max_tokens": 100,
}
}
// processOpenAIChatCompletionsStream processes the SSE stream from OpenAI Chat Completions API
func (s *AccountTestService) processOpenAIChatCompletionsStream(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
}
// Handle Chat Completions format: choices[0].delta.content
if choices, ok := data["choices"].([]any); ok && len(choices) > 0 {
if choice, ok := choices[0].(map[string]any); ok {
// Check finish_reason
if finishReason, ok := choice["finish_reason"].(string); ok && finishReason != "" {
s.sendEvent(c, TestEvent{Type: "test_complete", Success: true})
return nil
}
// Extract content from delta
if delta, ok := choice["delta"].(map[string]any); ok {
if content, ok := delta["content"].(string); ok && content != "" {
s.sendEvent(c, TestEvent{Type: "content", Text: content})
}
}
}
}
// Handle error
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)
}
}
}
// testGeminiAccountConnection tests a Gemini account's connection
func (s *AccountTestService) testGeminiAccountConnection(c *gin.Context, account *Account, modelID string) error {
ctx := c.Request.Context()
// Determine the model to use
testModelID := modelID
if testModelID == "" {
testModelID = geminicli.DefaultTestModel
}
// For API Key accounts with model mapping, map the model
if account.Type == AccountTypeApiKey {
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()
// 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)
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.Do(req, proxyURL, account.ID, account.Concurrency)
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)
}
// 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-5Gemini 使用 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
}
// Use streamGenerateContent for real-time feedback
fullURL := fmt.Sprintf("%s/v1beta/models/%s:streamGenerateContent?alt=sse",
strings.TrimRight(baseURL, "/"), modelID)
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
}
fullURL := fmt.Sprintf("%s/v1beta/models/%s:streamGenerateContent?alt=sse", strings.TrimRight(baseURL, "/"), modelID)
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)
}
// 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)
fullURL := fmt.Sprintf("%s/v1internal:streamGenerateContent?alt=sse", geminicli.GeminiCliBaseURL)
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
func createGeminiTestPayload() []byte {
payload := map[string]any{
"contents": []map[string]any{
{
"role": "user",
"parts": []map[string]any{
{"text": "hi"},
},
},
},
"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 == "" || !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
}
// 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 finishReason)
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})
}
}
}
}
}
// 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)
}
}
}
2025-12-22 22:58:31 +08:00
// 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
2025-12-22 22:58:31 +08:00
if isOAuth {
payload["store"] = false
}
// All accounts require instructions for Responses API
payload["instructions"] = openai.DefaultInstructions
2025-12-22 22:58:31 +08:00
return payload
2025-12-18 13:50:39 +08:00
}
2025-12-22 22:58:31 +08:00
// processClaudeStream processes the SSE stream from Claude API
func (s *AccountTestService) processClaudeStream(c *gin.Context, body io.Reader) error {
2025-12-18 13:50:39 +08:00
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) {
2025-12-18 13:50:39 +08:00
continue
}
jsonStr := sseDataPrefix.ReplaceAllString(line, "")
2025-12-18 13:50:39 +08:00
if jsonStr == "[DONE]" {
s.sendEvent(c, TestEvent{Type: "test_complete", Success: true})
return nil
}
2025-12-20 16:19:40 +08:00
var data map[string]any
2025-12-18 13:50:39 +08:00
if err := json.Unmarshal([]byte(jsonStr), &data); err != nil {
continue
}
eventType, _ := data["type"].(string)
switch eventType {
case "content_block_delta":
2025-12-20 16:19:40 +08:00
if delta, ok := data["delta"].(map[string]any); ok {
2025-12-18 13:50:39 +08:00
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"
2025-12-20 16:19:40 +08:00
if errData, ok := data["error"].(map[string]any); ok {
2025-12-18 13:50:39 +08:00
if msg, ok := errData["message"].(string); ok {
errorMsg = msg
}
}
return s.sendErrorAndEnd(c, errorMsg)
}
}
}
2025-12-22 22:58:31 +08:00
// processOpenAIStream processes the SSE stream from OpenAI Responses API
func (s *AccountTestService) processOpenAIStream(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) {
2025-12-22 22:58:31 +08:00
continue
}
jsonStr := sseDataPrefix.ReplaceAllString(line, "")
2025-12-22 22:58:31 +08:00
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 "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":
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)
}
}
}
2025-12-18 13:50:39 +08:00
// sendEvent sends a SSE event to the client
func (s *AccountTestService) sendEvent(c *gin.Context, event TestEvent) {
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
}
2025-12-18 13:50:39 +08:00
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)
2025-12-18 13:50:39 +08:00
}