feat(repository): 实现 Gemini OAuth 和 Token 缓存客户端

- 添加 Gemini OAuth 客户端实现 - 实现 Redis 基础的 Token 缓存 - 添加 gemini-cli Code Assist 客户端封装
2025-12-25 06:43:35 -08:00
parent aea48ae1ab
commit 2bafc28a9b
3 changed files with 223 additions and 0 deletions
--- a/backend/internal/repository/gemini_oauth_client.go
+++ b/backend/internal/repository/gemini_oauth_client.go
@@ -0,0 +1,84 @@
+package repository
+
+import (
+	"context"
+	"fmt"
+	"net/url"
+	"time"
+
+	"github.com/Wei-Shaw/sub2api/internal/config"
+	"github.com/Wei-Shaw/sub2api/internal/pkg/geminicli"
+	"github.com/Wei-Shaw/sub2api/internal/service/ports"
+
+	"github.com/imroc/req/v3"
+)
+
+type geminiOAuthClient struct {
+	tokenURL string
+	cfg      *config.Config
+}
+
+func NewGeminiOAuthClient(cfg *config.Config) ports.GeminiOAuthClient {
+	return &geminiOAuthClient{
+		tokenURL: geminicli.TokenURL,
+		cfg:      cfg,
+	}
+}
+
+func (c *geminiOAuthClient) ExchangeCode(ctx context.Context, code, codeVerifier, redirectURI, proxyURL string) (*geminicli.TokenResponse, error) {
+	client := createGeminiReqClient(proxyURL)
+
+	formData := url.Values{}
+	formData.Set("grant_type", "authorization_code")
+	formData.Set("client_id", c.cfg.Gemini.OAuth.ClientID)
+	formData.Set("client_secret", c.cfg.Gemini.OAuth.ClientSecret)
+	formData.Set("code", code)
+	formData.Set("code_verifier", codeVerifier)
+	formData.Set("redirect_uri", redirectURI)
+
+	var tokenResp geminicli.TokenResponse
+	resp, err := client.R().
+		SetContext(ctx).
+		SetFormDataFromValues(formData).
+		SetSuccessResult(&tokenResp).
+		Post(c.tokenURL)
+	if err != nil {
+		return nil, fmt.Errorf("request failed: %w", err)
+	}
+	if !resp.IsSuccessState() {
+		return nil, fmt.Errorf("token exchange failed: status %d, body: %s", resp.StatusCode, geminicli.SanitizeBodyForLogs(resp.String()))
+	}
+	return &tokenResp, nil
+}
+
+func (c *geminiOAuthClient) RefreshToken(ctx context.Context, refreshToken, proxyURL string) (*geminicli.TokenResponse, error) {
+	client := createGeminiReqClient(proxyURL)
+
+	formData := url.Values{}
+	formData.Set("grant_type", "refresh_token")
+	formData.Set("refresh_token", refreshToken)
+	formData.Set("client_id", c.cfg.Gemini.OAuth.ClientID)
+	formData.Set("client_secret", c.cfg.Gemini.OAuth.ClientSecret)
+
+	var tokenResp geminicli.TokenResponse
+	resp, err := client.R().
+		SetContext(ctx).
+		SetFormDataFromValues(formData).
+		SetSuccessResult(&tokenResp).
+		Post(c.tokenURL)
+	if err != nil {
+		return nil, fmt.Errorf("request failed: %w", err)
+	}
+	if !resp.IsSuccessState() {
+		return nil, fmt.Errorf("token refresh failed: status %d, body: %s", resp.StatusCode, geminicli.SanitizeBodyForLogs(resp.String()))
+	}
+	return &tokenResp, nil
+}
+
+func createGeminiReqClient(proxyURL string) *req.Client {
+	client := req.C().SetTimeout(60 * time.Second)
+	if proxyURL != "" {
+		client.SetProxyURL(proxyURL)
+	}
+	return client
+}
--- a/backend/internal/repository/gemini_token_cache.go
+++ b/backend/internal/repository/gemini_token_cache.go
@@ -0,0 +1,44 @@
+package repository
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/Wei-Shaw/sub2api/internal/service/ports"
+
+	"github.com/redis/go-redis/v9"
+)
+
+const (
+	geminiTokenKeyPrefix       = "gemini:token:"
+	geminiRefreshLockKeyPrefix = "gemini:refresh_lock:"
+)
+
+type geminiTokenCache struct {
+	rdb *redis.Client
+}
+
+func NewGeminiTokenCache(rdb *redis.Client) ports.GeminiTokenCache {
+	return &geminiTokenCache{rdb: rdb}
+}
+
+func (c *geminiTokenCache) GetAccessToken(ctx context.Context, cacheKey string) (string, error) {
+	key := fmt.Sprintf("%s%s", geminiTokenKeyPrefix, cacheKey)
+	return c.rdb.Get(ctx, key).Result()
+}
+
+func (c *geminiTokenCache) SetAccessToken(ctx context.Context, cacheKey string, token string, ttl time.Duration) error {
+	key := fmt.Sprintf("%s%s", geminiTokenKeyPrefix, cacheKey)
+	return c.rdb.Set(ctx, key, token, ttl).Err()
+}
+
+func (c *geminiTokenCache) AcquireRefreshLock(ctx context.Context, cacheKey string, ttl time.Duration) (bool, error) {
+	key := fmt.Sprintf("%s%s", geminiRefreshLockKeyPrefix, cacheKey)
+	return c.rdb.SetNX(ctx, key, 1, ttl).Result()
+}
+
+func (c *geminiTokenCache) ReleaseRefreshLock(ctx context.Context, cacheKey string) error {
+	key := fmt.Sprintf("%s%s", geminiRefreshLockKeyPrefix, cacheKey)
+	return c.rdb.Del(ctx, key).Err()
+}
--- a/backend/internal/repository/geminicli_codeassist_client.go
+++ b/backend/internal/repository/geminicli_codeassist_client.go
@@ -0,0 +1,95 @@
+package repository
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/Wei-Shaw/sub2api/internal/pkg/geminicli"
+	"github.com/Wei-Shaw/sub2api/internal/service/ports"
+
+	"github.com/imroc/req/v3"
+)
+
+type geminiCliCodeAssistClient struct {
+	baseURL string
+}
+
+func NewGeminiCliCodeAssistClient() ports.GeminiCliCodeAssistClient {
+	return &geminiCliCodeAssistClient{baseURL: geminicli.GeminiCliBaseURL}
+}
+
+func (c *geminiCliCodeAssistClient) LoadCodeAssist(ctx context.Context, accessToken, proxyURL string, reqBody *geminicli.LoadCodeAssistRequest) (*geminicli.LoadCodeAssistResponse, error) {
+	if reqBody == nil {
+		reqBody = defaultLoadCodeAssistRequest()
+	}
+
+	var out geminicli.LoadCodeAssistResponse
+	resp, err := createGeminiCliReqClient(proxyURL).R().
+		SetContext(ctx).
+		SetHeader("Authorization", "Bearer "+accessToken).
+		SetHeader("Content-Type", "application/json").
+		SetHeader("User-Agent", geminicli.GeminiCLIUserAgent).
+		SetBody(reqBody).
+		SetSuccessResult(&out).
+		Post(c.baseURL + "/v1internal:loadCodeAssist")
+	if err != nil {
+		return nil, fmt.Errorf("request failed: %w", err)
+	}
+	if !resp.IsSuccessState() {
+		return nil, fmt.Errorf("loadCodeAssist failed: status %d, body: %s", resp.StatusCode, geminicli.SanitizeBodyForLogs(resp.String()))
+	}
+	return &out, nil
+}
+
+func (c *geminiCliCodeAssistClient) OnboardUser(ctx context.Context, accessToken, proxyURL string, reqBody *geminicli.OnboardUserRequest) (*geminicli.OnboardUserResponse, error) {
+	if reqBody == nil {
+		reqBody = defaultOnboardUserRequest()
+	}
+
+	var out geminicli.OnboardUserResponse
+	resp, err := createGeminiCliReqClient(proxyURL).R().
+		SetContext(ctx).
+		SetHeader("Authorization", "Bearer "+accessToken).
+		SetHeader("Content-Type", "application/json").
+		SetHeader("User-Agent", geminicli.GeminiCLIUserAgent).
+		SetBody(reqBody).
+		SetSuccessResult(&out).
+		Post(c.baseURL + "/v1internal:onboardUser")
+	if err != nil {
+		return nil, fmt.Errorf("request failed: %w", err)
+	}
+	if !resp.IsSuccessState() {
+		return nil, fmt.Errorf("onboardUser failed: status %d, body: %s", resp.StatusCode, geminicli.SanitizeBodyForLogs(resp.String()))
+	}
+	return &out, nil
+}
+
+func createGeminiCliReqClient(proxyURL string) *req.Client {
+	client := req.C().SetTimeout(30 * time.Second)
+	if proxyURL != "" {
+		client.SetProxyURL(proxyURL)
+	}
+	return client
+}
+
+func defaultLoadCodeAssistRequest() *geminicli.LoadCodeAssistRequest {
+	return &geminicli.LoadCodeAssistRequest{
+		Metadata: geminicli.LoadCodeAssistMetadata{
+			IDEType:    "ANTIGRAVITY",
+			Platform:   "PLATFORM_UNSPECIFIED",
+			PluginType: "GEMINI",
+		},
+	}
+}
+
+func defaultOnboardUserRequest() *geminicli.OnboardUserRequest {
+	return &geminicli.OnboardUserRequest{
+		TierID: "LEGACY",
+		Metadata: geminicli.LoadCodeAssistMetadata{
+			IDEType:    "ANTIGRAVITY",
+			Platform:   "PLATFORM_UNSPECIFIED",
+			PluginType: "GEMINI",
+		},
+	}
+}