sub2api/backend/internal/service/antigravity_heartbeat.go

package service

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"math/rand"
	"net/http"
	"sync"
	"time"

	"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
)

// AntigravityHeartbeat 模拟真实 Antigravity IDE 的心跳行为
// 真实 IDE 每 5 分钟发送 loadCodeAssist + fetchAvailableModels
type AntigravityHeartbeat struct {
	mu       sync.Mutex
	sessions map[int64]*heartbeatSession // accountID -> session
	stopCh   chan struct{}
}

type heartbeatSession struct {
	accountID   int64
	accessToken string
	projectID   string
	proxyURL    string
	lastBeat    time.Time
}

// NewAntigravityHeartbeat 创建心跳管理器
func NewAntigravityHeartbeat() *AntigravityHeartbeat {
	hb := &AntigravityHeartbeat{
		sessions: make(map[int64]*heartbeatSession),
		stopCh:   make(chan struct{}),
	}
	go hb.loop()
	return hb
}

// Register 注册账号心跳（首次 API 调用时调用）
func (h *AntigravityHeartbeat) Register(accountID int64, accessToken, projectID, proxyURL string) {
	h.mu.Lock()
	defer h.mu.Unlock()

	if _, exists := h.sessions[accountID]; exists {
		// 更新 token（可能已刷新）
		h.sessions[accountID].accessToken = accessToken
		return
	}

	h.sessions[accountID] = &heartbeatSession{
		accountID:   accountID,
		accessToken: accessToken,
		projectID:   projectID,
		proxyURL:    proxyURL,
		lastBeat:    time.Now(),
	}
	log.Printf("[antigravity-heartbeat] registered account %d (project: %s)", accountID, projectID)
}

// UpdateToken 更新账号的 access token（token 刷新后调用）
func (h *AntigravityHeartbeat) UpdateToken(accountID int64, accessToken string) {
	h.mu.Lock()
	defer h.mu.Unlock()
	if s, ok := h.sessions[accountID]; ok {
		s.accessToken = accessToken
	}
}

// Unregister 移除账号心跳
func (h *AntigravityHeartbeat) Unregister(accountID int64) {
	h.mu.Lock()
	defer h.mu.Unlock()
	delete(h.sessions, accountID)
}

// Stop 停止心跳
func (h *AntigravityHeartbeat) Stop() {
	select {
	case <-h.stopCh:
	default:
		close(h.stopCh)
	}
}

func (h *AntigravityHeartbeat) loop() {
	ticker := time.NewTicker(1 * time.Minute)
	defer ticker.Stop()

	for {
		select {
		case <-h.stopCh:
			return
		case <-ticker.C:
			h.tick()
		}
	}
}

func (h *AntigravityHeartbeat) tick() {
	h.mu.Lock()
	// 收集需要心跳的 session
	var toSend []*heartbeatSession
	now := time.Now()
	for _, s := range h.sessions {
		if now.Sub(s.lastBeat) >= 5*time.Minute {
			s.lastBeat = now
			// 复制一份避免持锁时发请求
			cp := *s
			toSend = append(toSend, &cp)
		}
	}
	h.mu.Unlock()

	for _, s := range toSend {
		go h.sendHeartbeat(s)
	}
}

func (h *AntigravityHeartbeat) sendHeartbeat(s *heartbeatSession) {
	client, err := antigravity.NewClient(s.proxyURL)
	if err != nil {
		log.Printf("[antigravity-heartbeat] account %d: client error: %v", s.accountID, err)
		return
	}

	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
	defer cancel()

	// 1. loadCodeAssist
	h.doLoadCodeAssist(ctx, client, s)

	// 模拟真实 IDE 的延迟（~500ms）
	time.Sleep(time.Duration(400+rand.Intn(200)) * time.Millisecond)

	// 2. fetchAvailableModels
	h.doFetchAvailableModels(ctx, client, s)
}

func (h *AntigravityHeartbeat) doLoadCodeAssist(ctx context.Context, client *antigravity.Client, s *heartbeatSession) {
	reqBody := map[string]any{
		"metadata": map[string]string{
			"ideType": "ANTIGRAVITY",
		},
	}
	body, _ := json.Marshal(reqBody)

	for _, baseURL := range antigravity.BaseURLs {
		apiURL := fmt.Sprintf("%s/v1internal:loadCodeAssist", baseURL)
		req, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, bytes.NewReader(body))
		if err != nil {
			continue
		}
		req.Header.Set("Content-Type", "application/json")
		req.Header.Set("Authorization", "Bearer "+s.accessToken)
		req.Header.Set("User-Agent", antigravity.GetUserAgent())
		req.Header.Set("X-Goog-Api-Client", antigravity.GetGoogAPIClient())

		resp, err := client.DoRaw(req)
		if err != nil {
			continue
		}
		io.Copy(io.Discard, resp.Body)
		resp.Body.Close()

		if resp.StatusCode == http.StatusOK {
			return
		}
	}
}

func (h *AntigravityHeartbeat) doFetchAvailableModels(ctx context.Context, client *antigravity.Client, s *heartbeatSession) {
	reqBody := map[string]string{
		"project": s.projectID,
	}
	body, _ := json.Marshal(reqBody)

	for _, baseURL := range antigravity.BaseURLs {
		apiURL := fmt.Sprintf("%s/v1internal:fetchAvailableModels", baseURL)
		req, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, bytes.NewReader(body))
		if err != nil {
			continue
		}
		req.Header.Set("Content-Type", "application/json")
		req.Header.Set("Authorization", "Bearer "+s.accessToken)
		req.Header.Set("User-Agent", antigravity.GetUserAgent())
		req.Header.Set("X-Goog-Api-Client", antigravity.GetGoogAPIClient())

		resp, err := client.DoRaw(req)
		if err != nil {
			continue
		}
		io.Copy(io.Discard, resp.Body)
		resp.Body.Close()

		if resp.StatusCode == http.StatusOK {
			return
		}
	}
}