diff --git a/backend/internal/repository/gemini_oauth_client.go b/backend/internal/repository/gemini_oauth_client.go new file mode 100644 index 00000000..07364368 --- /dev/null +++ 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 +} diff --git a/backend/internal/repository/gemini_token_cache.go b/backend/internal/repository/gemini_token_cache.go new file mode 100644 index 00000000..9d294605 --- /dev/null +++ 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() +} diff --git a/backend/internal/repository/geminicli_codeassist_client.go b/backend/internal/repository/geminicli_codeassist_client.go new file mode 100644 index 00000000..63f1719c --- /dev/null +++ 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", + }, + } +}