package executor import ( "context" "encoding/json" "fmt" "io" "net/http" "net/url" "strconv" "strings" "time" "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" sdkAuth "github.com/router-for-me/CLIProxyAPI/v7/sdk/auth" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" log "github.com/sirupsen/logrus" ) // Refresh refreshes the authentication credentials using the refresh token. func (e *AntigravityExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) { if refreshed, handled, err := helps.RefreshAuthViaHome(ctx, e.cfg, auth); handled { return refreshed, err } if auth == nil { return auth, nil } updated, errRefresh := e.refreshToken(ctx, auth.Clone()) if errRefresh != nil { return nil, errRefresh } return updated, nil } func (e *AntigravityExecutor) ShouldPrepareRequestAuth(auth *cliproxyauth.Auth) bool { return antigravityProjectIDFromAuth(auth) == "" } func (e *AntigravityExecutor) PrepareRequestAuth(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) { if auth == nil || !e.ShouldPrepareRequestAuth(auth) { return nil, nil } updated := auth.Clone() token, refreshedAuth, errToken := e.ensureAccessToken(ctx, updated) if errToken != nil { return nil, errToken } if refreshedAuth != nil { updated = refreshedAuth } if antigravityProjectIDFromAuth(updated) != "" { return updated, nil } projectID, errProject := e.fetchAntigravityProjectID(ctx, updated, token) if errProject != nil { return nil, missingAntigravityProjectIDError(errProject) } if projectID == "" { return nil, missingAntigravityProjectIDError(nil) } if updated.Metadata == nil { updated.Metadata = make(map[string]any) } updated.Metadata["project_id"] = projectID return updated, nil } func (e *AntigravityExecutor) ensureAccessToken(ctx context.Context, auth *cliproxyauth.Auth) (string, *cliproxyauth.Auth, error) { if auth == nil { return "", nil, statusErr{code: http.StatusUnauthorized, msg: "missing auth"} } accessToken := metaStringValue(auth.Metadata, "access_token") expiry := tokenExpiry(auth.Metadata) if accessToken != "" && expiry.After(time.Now().Add(refreshSkew)) { e.maybeRefreshAntigravityCreditsHint(ctx, auth, accessToken) return accessToken, nil, nil } refreshCtx := context.Background() if ctx != nil { if rt, ok := ctx.Value("cliproxy.roundtripper").(http.RoundTripper); ok && rt != nil { refreshCtx = context.WithValue(refreshCtx, "cliproxy.roundtripper", rt) } } if refreshed, handled, err := helps.RefreshAuthViaHome(refreshCtx, e.cfg, auth); handled { if err != nil { return "", nil, err } token := metaStringValue(refreshed.Metadata, "access_token") if strings.TrimSpace(token) == "" { return "", nil, statusErr{code: http.StatusUnauthorized, msg: "missing access token"} } e.maybeRefreshAntigravityCreditsHint(ctx, refreshed, token) return token, refreshed, nil } updated, errRefresh := e.refreshToken(refreshCtx, auth.Clone()) if errRefresh != nil { return "", nil, errRefresh } return metaStringValue(updated.Metadata, "access_token"), updated, nil } func (e *AntigravityExecutor) refreshToken(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) { if auth == nil { return nil, statusErr{code: http.StatusUnauthorized, msg: "missing auth"} } refreshToken := metaStringValue(auth.Metadata, "refresh_token") if refreshToken == "" { return auth, statusErr{code: http.StatusUnauthorized, msg: "missing refresh token"} } if ctx == nil { ctx = context.Background() } refreshToken = strings.TrimSpace(refreshToken) result, errRefresh, _ := antigravityRefreshGroup.Do(refreshToken, func() (interface{}, error) { return e.refreshTokenSingleFlight(context.WithoutCancel(ctx), auth, refreshToken) }) if errRefresh != nil { return auth, errRefresh } tokenResp, ok := result.(*antigravityTokenRefreshData) if !ok || tokenResp == nil { return auth, fmt.Errorf("antigravity token refresh failed: invalid single-flight result") } if auth.Metadata == nil { auth.Metadata = make(map[string]any) } auth.Metadata["access_token"] = tokenResp.AccessToken if tokenResp.RefreshToken != "" { auth.Metadata["refresh_token"] = tokenResp.RefreshToken } auth.Metadata["expires_in"] = tokenResp.ExpiresIn now := time.Now() auth.Metadata["timestamp"] = now.UnixMilli() auth.Metadata["expired"] = now.Add(time.Duration(tokenResp.ExpiresIn) * time.Second).Format(time.RFC3339) auth.Metadata["type"] = antigravityAuthType if errProject := e.ensureAntigravityProjectID(ctx, auth, tokenResp.AccessToken); errProject != nil { log.Warnf("antigravity executor: ensure project id failed: %v", errProject) } e.updateAntigravityCreditsBalance(ctx, auth, tokenResp.AccessToken) return auth, nil } func (e *AntigravityExecutor) refreshTokenSingleFlight(ctx context.Context, auth *cliproxyauth.Auth, refreshToken string) (*antigravityTokenRefreshData, error) { form := url.Values{} form.Set("client_id", antigravityClientID) form.Set("client_secret", antigravityClientSecret) form.Set("grant_type", "refresh_token") form.Set("refresh_token", refreshToken) httpReq, errReq := http.NewRequestWithContext(ctx, http.MethodPost, "https://oauth2.googleapis.com/token", strings.NewReader(form.Encode())) if errReq != nil { return nil, errReq } httpReq.Header.Set("Host", "oauth2.googleapis.com") httpReq.Header.Set("Content-Type", "application/x-www-form-urlencoded") // Real Antigravity uses Go's default User-Agent for OAuth token refresh httpReq.Header.Set("User-Agent", "Go-http-client/2.0") httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0) httpResp, errDo := httpClient.Do(httpReq) if errDo != nil { return nil, errDo } defer func() { if errClose := httpResp.Body.Close(); errClose != nil { log.Errorf("antigravity executor: close response body error: %v", errClose) } }() bodyBytes, errRead := io.ReadAll(httpResp.Body) if errRead != nil { return nil, errRead } if httpResp.StatusCode < http.StatusOK || httpResp.StatusCode >= http.StatusMultipleChoices { sErr := statusErr{code: httpResp.StatusCode, msg: string(bodyBytes)} if httpResp.StatusCode == http.StatusTooManyRequests { if retryAfter, parseErr := helps.ParseRetryDelay(bodyBytes); parseErr == nil && retryAfter != nil { sErr.retryAfter = retryAfter } } return nil, sErr } var tokenResp antigravityTokenRefreshData if errUnmarshal := json.Unmarshal(bodyBytes, &tokenResp); errUnmarshal != nil { return nil, errUnmarshal } return &tokenResp, nil } func (e *AntigravityExecutor) ensureAntigravityProjectID(ctx context.Context, auth *cliproxyauth.Auth, accessToken string) error { if auth == nil { return nil } if antigravityProjectIDFromAuth(auth) != "" { return nil } projectID, errFetch := e.fetchAntigravityProjectID(ctx, auth, accessToken) if errFetch != nil { return errFetch } if projectID == "" { return nil } if auth.Metadata == nil { auth.Metadata = make(map[string]any) } auth.Metadata["project_id"] = projectID return nil } func (e *AntigravityExecutor) fetchAntigravityProjectID(ctx context.Context, auth *cliproxyauth.Auth, accessToken string) (string, error) { token := strings.TrimSpace(accessToken) if token == "" { token = metaStringValue(auth.Metadata, "access_token") } if token == "" { return "", nil } httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0) projectID, errFetch := sdkAuth.FetchAntigravityProjectID(ctx, token, httpClient) if errFetch != nil { return "", errFetch } return strings.TrimSpace(projectID), nil } func (e *AntigravityExecutor) projectIDForRequest(_ context.Context, auth *cliproxyauth.Auth, _ string) (string, error) { if projectID := antigravityProjectIDFromAuth(auth); projectID != "" { return projectID, nil } return "", missingAntigravityProjectIDError(nil) } func antigravityProjectIDFromAuth(auth *cliproxyauth.Auth) string { if auth == nil || auth.Metadata == nil { return "" } if pid, ok := auth.Metadata["project_id"].(string); ok { return strings.TrimSpace(pid) } return "" } func missingAntigravityProjectIDError(cause error) statusErr { msg := "antigravity auth missing project_id" if cause != nil { msg = fmt.Sprintf("%s: %v", msg, cause) } return statusErr{code: http.StatusBadRequest, msg: msg} } func tokenExpiry(metadata map[string]any) time.Time { if metadata == nil { return time.Time{} } if expStr, ok := metadata["expired"].(string); ok { expStr = strings.TrimSpace(expStr) if expStr != "" { if parsed, errParse := time.Parse(time.RFC3339, expStr); errParse == nil { return parsed } } } expiresIn, hasExpires := int64Value(metadata["expires_in"]) tsMs, hasTimestamp := int64Value(metadata["timestamp"]) if hasExpires && hasTimestamp { return time.Unix(0, tsMs*int64(time.Millisecond)).Add(time.Duration(expiresIn) * time.Second) } return time.Time{} } func metaStringValue(metadata map[string]any, key string) string { if metadata == nil { return "" } if v, ok := metadata[key]; ok { switch typed := v.(type) { case string: return strings.TrimSpace(typed) case []byte: return strings.TrimSpace(string(typed)) } } return "" } func int64Value(value any) (int64, bool) { switch typed := value.(type) { case int: return int64(typed), true case int64: return typed, true case float64: return int64(typed), true case json.Number: if i, errParse := typed.Int64(); errParse == nil { return i, true } case string: if strings.TrimSpace(typed) == "" { return 0, false } if i, errParse := strconv.ParseInt(strings.TrimSpace(typed), 10, 64); errParse == nil { return i, true } } return 0, false }