opencode-cursor-agent/internal/providers/geminiweb/pool.go

package geminiweb

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"sync"
	"time"
)

type GeminiSession struct {
	Name         string `json:"name"`
	CookieFile   string `json:"cookie_file"`
	LastUsed     int64  `json:"last_used"`
	ActiveCount  int    `json:"active_count"`
	RateLimitEnd int64  `json:"rate_limit_end"`
}

type SessionPool struct {
	mu       sync.Mutex
	sessions []*GeminiSession
	dir      string
	maxCount int
}

func NewSessionPool(dir string, maxSessions int) (*SessionPool, error) {
	if err := os.MkdirAll(dir, 0755); err != nil {
		return nil, fmt.Errorf("failed to create session dir: %w", err)
	}

	sessions, err := loadSessions(dir)
	if err != nil {
		return nil, fmt.Errorf("failed to load sessions: %w", err)
	}

	return &SessionPool{
		sessions: sessions,
		dir:      dir,
		maxCount: maxSessions,
	}, nil
}

func loadSessions(dir string) ([]*GeminiSession, error) {
	entries, err := os.ReadDir(dir)
	if err != nil {
		return nil, err
	}

	var sessions []*GeminiSession
	for _, entry := range entries {
		if !entry.IsDir() {
			continue
		}
		name := entry.Name()
		metaPath := filepath.Join(dir, name, "session.json")
		data, err := os.ReadFile(metaPath)
		if err != nil {
			continue
		}

		var s GeminiSession
		if err := json.Unmarshal(data, &s); err != nil {
			continue
		}
		sessions = append(sessions, &s)
	}

	return sessions, nil
}

func (p *SessionPool) Count() int {
	p.mu.Lock()
	defer p.mu.Unlock()
	return len(p.sessions)
}

func (p *SessionPool) GetAvailable() *GeminiSession {
	p.mu.Lock()
	defer p.mu.Unlock()

	now := time.Now().UnixMilli()

	var available []*GeminiSession
	for _, s := range p.sessions {
		if s.RateLimitEnd < now {
			available = append(available, s)
		}
	}

	if len(available) == 0 {
		return nil
	}

	var best *GeminiSession
	for _, s := range available {
		if best == nil || s.ActiveCount < best.ActiveCount {
			best = s
		} else if s.ActiveCount == best.ActiveCount && s.LastUsed < best.LastUsed {
			best = s
		}
	}

	return best
}

func (p *SessionPool) StartSession(s *GeminiSession) {
	p.mu.Lock()
	defer p.mu.Unlock()
	s.ActiveCount++
	s.LastUsed = time.Now().UnixMilli()
	p.saveSession(s)
}

func (p *SessionPool) EndSession(s *GeminiSession) {
	p.mu.Lock()
	defer p.mu.Unlock()
	if s.ActiveCount > 0 {
		s.ActiveCount--
	}
	p.saveSession(s)
}

func (p *SessionPool) RateLimitSession(s *GeminiSession, durationMs int64) {
	p.mu.Lock()
	defer p.mu.Unlock()
	s.RateLimitEnd = time.Now().UnixMilli() + durationMs
	p.saveSession(s)
}

func (p *SessionPool) saveSession(s *GeminiSession) {
	metaPath := filepath.Join(p.dir, s.Name, "session.json")
	data, err := json.MarshalIndent(s, "", "  ")
	if err != nil {
		return
	}
	_ = os.WriteFile(metaPath, data, 0644)
}

func (p *SessionPool) CreateSession(name string) (*GeminiSession, error) {
	p.mu.Lock()
	defer p.mu.Unlock()

	sessionDir := filepath.Join(p.dir, name)
	if err := os.MkdirAll(sessionDir, 0755); err != nil {
		return nil, fmt.Errorf("failed to create session dir: %w", err)
	}

	s := &GeminiSession{
		Name:       name,
		CookieFile: filepath.Join(sessionDir, "cookies.json"),
		LastUsed:   time.Now().UnixMilli(),
	}

	p.sessions = append(p.sessions, s)
	p.saveSession(s)

	return s, nil
}

func (p *SessionPool) GetSessionNames() []string {
	p.mu.Lock()
	defer p.mu.Unlock()
	names := make([]string, len(p.sessions))
	for i, s := range p.sessions {
		names[i] = s.Name
	}
	return names
}