* fix(ops): 修复运维监控系统的关键安全和稳定性问题
## 修复内容
### P0 严重问题
1. **DNS Rebinding防护** (ops_alert_service.go)
- 实现IP钉住机制防止验证后的DNS rebinding攻击
- 自定义Transport.DialContext强制只允许拨号到验证过的公网IP
- 扩展IP黑名单,包括云metadata地址(169.254.169.254)
- 添加完整的单元测试覆盖
2. **OpsAlertService生命周期管理** (wire.go)
- 在ProvideOpsMetricsCollector中添加opsAlertService.Start()调用
- 确保stopCtx正确初始化,避免nil指针问题
- 实现防御式启动,保证服务启动顺序
3. **数据库查询排序** (ops_repo.go)
- 在ListRecentSystemMetrics中添加显式ORDER BY updated_at DESC, id DESC
- 在GetLatestSystemMetric中添加排序保证
- 避免数据库返回顺序不确定导致告警误判
### P1 重要问题
4. **并发安全** (ops_metrics_collector.go)
- 为lastGCPauseTotal字段添加sync.Mutex保护
- 防止数据竞争
5. **Goroutine泄漏** (ops_error_logger.go)
- 实现worker pool模式限制并发goroutine数量
- 使用256容量缓冲队列和10个固定worker
- 非阻塞投递,队列满时丢弃任务
6. **生命周期控制** (ops_alert_service.go)
- 添加Start/Stop方法实现优雅关闭
- 使用context控制goroutine生命周期
- 实现WaitGroup等待后台任务完成
7. **Webhook URL验证** (ops_alert_service.go)
- 防止SSRF攻击:验证scheme、禁止内网IP
- DNS解析验证,拒绝解析到私有IP的域名
- 添加8个单元测试覆盖各种攻击场景
8. **资源泄漏** (ops_repo.go)
- 修复多处defer rows.Close()问题
- 简化冗余的defer func()包装
9. **HTTP超时控制** (ops_alert_service.go)
- 创建带10秒超时的http.Client
- 添加buildWebhookHTTPClient辅助函数
- 防止HTTP请求无限期挂起
10. **数据库查询优化** (ops_repo.go)
- 将GetWindowStats的4次独立查询合并为1次CTE查询
- 减少网络往返和表扫描次数
- 显著提升性能
11. **重试机制** (ops_alert_service.go)
- 实现邮件发送重试:最多3次,指数退避(1s/2s/4s)
- 添加webhook备用通道
- 实现完整的错误处理和日志记录
12. **魔法数字** (ops_repo.go, ops_metrics_collector.go)
- 提取硬编码数字为有意义的常量
- 提高代码可读性和可维护性
## 测试验证
- ✅ go test ./internal/service -tags opsalert_unit 通过
- ✅ 所有webhook验证测试通过
- ✅ 重试机制测试通过
## 影响范围
- 运维监控系统安全性显著提升
- 系统稳定性和性能优化
- 无破坏性变更,向后兼容
* feat(ops): 运维监控系统V2 - 完整实现
## 核心功能
- 运维监控仪表盘V2(实时监控、历史趋势、告警管理)
- WebSocket实时QPS/TPS监控(30s心跳,自动重连)
- 系统指标采集(CPU、内存、延迟、错误率等)
- 多维度统计分析(按provider、model、user等维度)
- 告警规则管理(阈值配置、通知渠道)
- 错误日志追踪(详细错误信息、堆栈跟踪)
## 数据库Schema (Migration 025)
### 扩展现有表
- ops_system_metrics: 新增RED指标、错误分类、延迟指标、资源指标、业务指标
- ops_alert_rules: 新增JSONB字段(dimension_filters, notify_channels, notify_config)
### 新增表
- ops_dimension_stats: 多维度统计数据
- ops_data_retention_config: 数据保留策略配置
### 新增视图和函数
- ops_latest_metrics: 最新1分钟窗口指标(已修复字段名和window过滤)
- ops_active_alerts: 当前活跃告警(已修复字段名和状态值)
- calculate_health_score: 健康分数计算函数
## 一致性修复(98/100分)
### P0级别(阻塞Migration)
- ✅ 修复ops_latest_metrics视图字段名(latency_p99→p99_latency_ms, cpu_usage→cpu_usage_percent)
- ✅ 修复ops_active_alerts视图字段名(metric→metric_type, triggered_at→fired_at, trigger_value→metric_value, threshold→threshold_value)
- ✅ 统一告警历史表名(删除ops_alert_history,使用ops_alert_events)
- ✅ 统一API参数限制(ListMetricsHistory和ListErrorLogs的limit改为5000)
### P1级别(功能完整性)
- ✅ 修复ops_latest_metrics视图未过滤window_minutes(添加WHERE m.window_minutes = 1)
- ✅ 修复数据回填UPDATE逻辑(QPS计算改为request_count/(window_minutes*60.0))
- ✅ 添加ops_alert_rules JSONB字段后端支持(Go结构体+序列化)
### P2级别(优化)
- ✅ 前端WebSocket自动重连(指数退避1s→2s→4s→8s→16s,最大5次)
- ✅ 后端WebSocket心跳检测(30s ping,60s pong超时)
## 技术实现
### 后端 (Go)
- Handler层: ops_handler.go(REST API), ops_ws_handler.go(WebSocket)
- Service层: ops_service.go(核心逻辑), ops_cache.go(缓存), ops_alerts.go(告警)
- Repository层: ops_repo.go(数据访问), ops.go(模型定义)
- 路由: admin.go(新增ops相关路由)
- 依赖注入: wire_gen.go(自动生成)
### 前端 (Vue3 + TypeScript)
- 组件: OpsDashboardV2.vue(仪表盘主组件)
- API: ops.ts(REST API + WebSocket封装)
- 路由: index.ts(新增/admin/ops路由)
- 国际化: en.ts, zh.ts(中英文支持)
## 测试验证
- ✅ 所有Go测试通过
- ✅ Migration可正常执行
- ✅ WebSocket连接稳定
- ✅ 前后端数据结构对齐
* refactor: 代码清理和测试优化
## 测试文件优化
- 简化integration test fixtures和断言
- 优化test helper函数
- 统一测试数据格式
## 代码清理
- 移除未使用的代码和注释
- 简化concurrency_cache实现
- 优化middleware错误处理
## 小修复
- 修复gateway_handler和openai_gateway_handler的小问题
- 统一代码风格和格式
变更统计: 27个文件,292行新增,322行删除(净减少30行)
* fix(ops): 运维监控系统安全加固和功能优化
## 安全增强
- feat(security): WebSocket日志脱敏机制,防止token/api_key泄露
- feat(security): X-Forwarded-Host白名单验证,防止CSRF绕过
- feat(security): Origin策略配置化,支持strict/permissive模式
- feat(auth): WebSocket认证支持query参数传递token
## 配置优化
- feat(config): 支持环境变量配置代理信任和Origin策略
- OPS_WS_TRUST_PROXY
- OPS_WS_TRUSTED_PROXIES
- OPS_WS_ORIGIN_POLICY
- fix(ops): 错误日志查询限流从5000降至500,优化内存使用
## 架构改进
- refactor(ops): 告警服务解耦,独立运行评估定时器
- refactor(ops): OpsDashboard统一版本,移除V2分离
## 测试和文档
- test(ops): 添加WebSocket安全验证单元测试(8个测试用例)
- test(ops): 添加告警服务集成测试
- docs(api): 更新API文档,标注限流变更
- docs: 添加CHANGELOG记录breaking changes
## 修复文件
Backend:
- backend/internal/server/middleware/logger.go
- backend/internal/handler/admin/ops_handler.go
- backend/internal/handler/admin/ops_ws_handler.go
- backend/internal/server/middleware/admin_auth.go
- backend/internal/service/ops_alert_service.go
- backend/internal/service/ops_metrics_collector.go
- backend/internal/service/wire.go
Frontend:
- frontend/src/views/admin/ops/OpsDashboard.vue
- frontend/src/router/index.ts
- frontend/src/api/admin/ops.ts
Tests:
- backend/internal/handler/admin/ops_ws_handler_test.go (新增)
- backend/internal/service/ops_alert_service_integration_test.go (新增)
Docs:
- CHANGELOG.md (新增)
- docs/API-运维监控中心2.0.md (更新)
* fix(migrations): 修复calculate_health_score函数类型匹配问题
在ops_latest_metrics视图中添加显式类型转换,确保参数类型与函数签名匹配
* fix(lint): 修复golangci-lint检查发现的所有问题
- 将Redis依赖从service层移到repository层
- 添加错误检查(WebSocket连接和读取超时)
- 运行gofmt格式化代码
- 添加nil指针检查
- 删除未使用的alertService字段
修复问题:
- depguard: 3个(service层不应直接import redis)
- errcheck: 3个(未检查错误返回值)
- gofmt: 2个(代码格式问题)
- staticcheck: 4个(nil指针解引用)
- unused: 1个(未使用字段)
代码统计:
- 修改文件:11个
- 删除代码:490行
- 新增代码:105行
- 净减少:385行
837 lines
24 KiB
Go
837 lines
24 KiB
Go
package service
|
||
|
||
import (
|
||
"bufio"
|
||
"bytes"
|
||
"context"
|
||
"crypto/rand"
|
||
"encoding/hex"
|
||
"encoding/json"
|
||
"fmt"
|
||
"io"
|
||
"log"
|
||
"net/http"
|
||
"regexp"
|
||
"strings"
|
||
"time"
|
||
|
||
"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/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*`)
|
||
|
||
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 AccountRepository
|
||
oauthService *OAuthService
|
||
openaiOAuthService *OpenAIOAuthService
|
||
geminiTokenProvider *GeminiTokenProvider
|
||
antigravityGatewayService *AntigravityGatewayService
|
||
httpUpstream HTTPUpstream
|
||
}
|
||
|
||
// NewAccountTestService creates a new AccountTestService
|
||
func NewAccountTestService(
|
||
accountRepo AccountRepository,
|
||
oauthService *OAuthService,
|
||
openaiOAuthService *OpenAIOAuthService,
|
||
geminiTokenProvider *GeminiTokenProvider,
|
||
antigravityGatewayService *AntigravityGatewayService,
|
||
httpUpstream HTTPUpstream,
|
||
) *AccountTestService {
|
||
return &AccountTestService{
|
||
accountRepo: accountRepo,
|
||
oauthService: oauthService,
|
||
openaiOAuthService: openaiOAuthService,
|
||
geminiTokenProvider: geminiTokenProvider,
|
||
antigravityGatewayService: antigravityGatewayService,
|
||
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)
|
||
}
|
||
|
||
if account.IsGemini() {
|
||
return s.testGeminiAccountConnection(c, account, modelID)
|
||
}
|
||
|
||
if account.Platform == PlatformAntigravity {
|
||
return s.testAntigravityAccountConnection(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 *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 expiresAt := account.GetCredentialAsTime("expires_at"); expiresAt != nil {
|
||
if time.Now().Add(5 * time.Minute).After(*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, 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.processClaudeStream(c, resp.Body)
|
||
}
|
||
|
||
// testOpenAIAccountConnection tests an OpenAI account's connection
|
||
func (s *AccountTestService) testOpenAIAccountConnection(c *gin.Context, account *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, "/") + "/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, 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.processOpenAIStream(c, resp.Body)
|
||
}
|
||
|
||
// 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-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
|
||
}
|
||
|
||
// 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 == "" || !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 {
|
||
// Check for completion
|
||
if finishReason, ok := candidate["finishReason"].(string); ok && finishReason != "" {
|
||
s.sendEvent(c, TestEvent{Type: "test_complete", Success: true})
|
||
return nil
|
||
}
|
||
|
||
// Extract content
|
||
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})
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// 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
|
||
}
|
||
|
||
// 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)
|
||
}
|
||
}
|
||
}
|
||
|
||
// 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) {
|
||
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 "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)
|
||
}
|