feat(repository): 实现 Gemini OAuth 和 Token 缓存客户端
- 添加 Gemini OAuth 客户端实现 - 实现 Redis 基础的 Token 缓存 - 添加 gemini-cli Code Assist 客户端封装
This commit is contained in:
84
backend/internal/repository/gemini_oauth_client.go
Normal file
84
backend/internal/repository/gemini_oauth_client.go
Normal file
@@ -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
|
||||||
|
}
|
||||||
44
backend/internal/repository/gemini_token_cache.go
Normal file
44
backend/internal/repository/gemini_token_cache.go
Normal file
@@ -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()
|
||||||
|
}
|
||||||
95
backend/internal/repository/geminicli_codeassist_client.go
Normal file
95
backend/internal/repository/geminicli_codeassist_client.go
Normal file
@@ -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",
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user