package service import ( "bufio" "bytes" "crypto/rand" "encoding/hex" "encoding/json" "fmt" "io" "log" "net/http" "strconv" "strings" "time" "github.com/Wei-Shaw/sub2api/internal/model" "github.com/Wei-Shaw/sub2api/internal/pkg/claude" "github.com/Wei-Shaw/sub2api/internal/pkg/openai" "github.com/Wei-Shaw/sub2api/internal/service/ports" "github.com/gin-gonic/gin" "github.com/google/uuid" ) const ( testClaudeAPIURL = "https://api.anthropic.com/v1/messages" testOpenAIAPIURL = "https://api.openai.com/v1/responses" 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"` Success bool `json:"success,omitempty"` Error string `json:"error,omitempty"` } // AccountTestService handles account testing operations type AccountTestService struct { accountRepo ports.AccountRepository oauthService *OAuthService openaiOAuthService *OpenAIOAuthService httpUpstream ports.HTTPUpstream } // NewAccountTestService creates a new AccountTestService func NewAccountTestService(accountRepo ports.AccountRepository, oauthService *OAuthService, openaiOAuthService *OpenAIOAuthService, httpUpstream ports.HTTPUpstream) *AccountTestService { return &AccountTestService{ accountRepo: accountRepo, oauthService: oauthService, openaiOAuthService: openaiOAuthService, httpUpstream: httpUpstream, } } // generateSessionString generates a Claude Code style session string func generateSessionString() (string, error) { bytes := make([]byte, 32) if _, err := rand.Read(bytes); err != nil { return "", err } hex64 := hex.EncodeToString(bytes) sessionUUID := uuid.New().String() return fmt.Sprintf("user_%s_account__session_%s", hex64, sessionUUID), 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": "You are Claude Code, Anthropic's official CLI for Claude.", "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 func (s *AccountTestService) TestAccountConnection(c *gin.Context, accountID int64, modelID string) error { ctx := c.Request.Context() // Get account account, err := s.accountRepo.GetByID(ctx, accountID) if err != nil { return s.sendErrorAndEnd(c, "Account not found") } // Route to platform-specific test method if account.IsOpenAI() { return s.testOpenAIAccountConnection(c, account, modelID) } return s.testClaudeAccountConnection(c, account, modelID) } // testClaudeAccountConnection tests an Anthropic Claude account's connection func (s *AccountTestService) testClaudeAccountConnection(c *gin.Context, account *model.Account, modelID string) error { 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 var authToken string var useBearer bool var apiURL string if account.IsOAuth() { // OAuth or Setup Token - use Bearer token useBearer = true 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 expiresAtStr := account.GetCredential("expires_at"); expiresAtStr != "" { expiresAt, err := strconv.ParseInt(expiresAtStr, 10, 64) if err == nil && time.Now().Unix()+300 > expiresAt { 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 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") } 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") req.Header.Set("anthropic-beta", claude.DefaultBetaHeader) // Apply Claude Code client headers for key, value := range claude.DefaultHeaders { req.Header.Set(key, value) } // Set authentication header if useBearer { req.Header.Set("Authorization", "Bearer "+authToken) } else { req.Header.Set("x-api-key", authToken) } // Get proxy URL proxyURL := "" if account.ProxyID != nil && account.Proxy != nil { proxyURL = account.Proxy.URL() } resp, err := s.httpUpstream.Do(req, proxyURL) 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 *model.Account, modelID string) error { 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 } } } // Determine authentication method and API URL var authToken string var apiURL string var isOAuth bool var chatgptAccountID string if account.IsOAuth() { isOAuth = true // OAuth - use Bearer token with ChatGPT internal API authToken = account.GetOpenAIAccessToken() if authToken == "" { return s.sendErrorAndEnd(c, "No access token available") } // 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())) } authToken = tokenInfo.AccessToken } // OAuth uses ChatGPT internal API apiURL = chatgptCodexAPIURL chatgptAccountID = account.GetChatGPTAccountID() } else if account.Type == "apikey" { // API Key - use Platform API authToken = account.GetOpenAIApiKey() if authToken == "" { return s.sendErrorAndEnd(c, "No API key available") } baseURL := account.GetOpenAIBaseURL() if baseURL == "" { baseURL = "https://api.openai.com" } apiURL = strings.TrimSuffix(baseURL, "/") + "/v1/responses" } 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 payload := createOpenAITestPayload(testModelID, isOAuth) 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("Authorization", "Bearer "+authToken) // Set OAuth-specific headers for ChatGPT internal API if isOAuth { req.Host = "chatgpt.com" req.Header.Set("accept", "text/event-stream") if chatgptAccountID != "" { req.Header.Set("chatgpt-account-id", chatgptAccountID) } } // Get proxy URL proxyURL := "" if account.ProxyID != nil && account.Proxy != nil { proxyURL = account.Proxy.URL() } resp, err := s.httpUpstream.Do(req, proxyURL) 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.processOpenAIStream(c, resp.Body) } // 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 and instructions if isOAuth { payload["store"] = false payload["instructions"] = openai.DefaultInstructions } return payload } // 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 == "" || !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 } 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) } } } // 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 == "" || !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 } 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) } } } // 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 } 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) }