From f3d58fa0ce63ead7e30ca8d0358ffe7d9cabf26b Mon Sep 17 00:00:00 2001
From: hkfires <10558748+hkfires@users.noreply.github.com>
Date: Sat, 24 Jan 2026 07:36:52 +0800
Subject: [PATCH] fix(auth): correct antigravity oauth redirect and expiry

---
 .../api/handlers/management/auth_files.go     |  6 ++----
 internal/auth/antigravity/auth.go             | 21 +++++++++++++------
 sdk/auth/antigravity.go                       |  9 +++-----
 3 files changed, 20 insertions(+), 16 deletions(-)

diff --git a/internal/api/handlers/management/auth_files.go b/internal/api/handlers/management/auth_files.go
index 791c4094..a36dbe20 100644
--- a/internal/api/handlers/management/auth_files.go
+++ b/internal/api/handlers/management/auth_files.go
@@ -1398,7 +1398,7 @@ func (h *Handler) RequestAntigravityToken(c *gin.Context) {
 
 	fmt.Println("Initializing Antigravity authentication...")
 
-	authSvc := antigravity.NewAntigravityAuth(h.cfg)
+	authSvc := antigravity.NewAntigravityAuth(h.cfg, nil)
 
 	state, errState := misc.GenerateRandomState()
 	if errState != nil {
@@ -1408,9 +1408,7 @@ func (h *Handler) RequestAntigravityToken(c *gin.Context) {
 	}
 
 	redirectURI := fmt.Sprintf("http://localhost:%d/oauth-callback", antigravity.CallbackPort)
-	authURL := authSvc.BuildAuthURL(state)
-	// Override redirect URI if needed (BuildAuthURL hardcodes it)
-	authURL = strings.ReplaceAll(authURL, fmt.Sprintf("http://localhost:%d/oauth-callback", antigravity.CallbackPort), redirectURI)
+	authURL := authSvc.BuildAuthURL(state, redirectURI)
 
 	RegisterOAuthSession(state, "antigravity")
 
diff --git a/internal/auth/antigravity/auth.go b/internal/auth/antigravity/auth.go
index 80dab17d..a85aa5e9 100644
--- a/internal/auth/antigravity/auth.go
+++ b/internal/auth/antigravity/auth.go
@@ -20,7 +20,7 @@ import (
 type TokenResponse struct {
 	AccessToken  string `json:"access_token"`
 	RefreshToken string `json:"refresh_token"`
-	ExpiresIn    int64  `json:"expires_token"`
+	ExpiresIn    int64  `json:"expires_in"`
 	TokenType    string `json:"token_type"`
 }
 
@@ -34,20 +34,29 @@ type AntigravityAuth struct {
 	httpClient *http.Client
 }
 
-// NewAntigravityAuth creates a new Antigravity auth service
-func NewAntigravityAuth(cfg *config.Config) *AntigravityAuth {
+// NewAntigravityAuth creates a new Antigravity auth service.
+func NewAntigravityAuth(cfg *config.Config, httpClient *http.Client) *AntigravityAuth {
+	if httpClient != nil {
+		return &AntigravityAuth{httpClient: httpClient}
+	}
+	if cfg == nil {
+		cfg = &config.Config{}
+	}
 	return &AntigravityAuth{
 		httpClient: util.SetProxy(&cfg.SDKConfig, &http.Client{}),
 	}
 }
 
-// BuildAuthURL generates the OAuth authorization URL
-func (o *AntigravityAuth) BuildAuthURL(state string) string {
+// BuildAuthURL generates the OAuth authorization URL.
+func (o *AntigravityAuth) BuildAuthURL(state, redirectURI string) string {
+	if strings.TrimSpace(redirectURI) == "" {
+		redirectURI = fmt.Sprintf("http://localhost:%d/oauth-callback", CallbackPort)
+	}
 	params := url.Values{}
 	params.Set("access_type", "offline")
 	params.Set("client_id", ClientID)
 	params.Set("prompt", "consent")
-	params.Set("redirect_uri", fmt.Sprintf("http://localhost:%d/oauth-callback", CallbackPort))
+	params.Set("redirect_uri", redirectURI)
 	params.Set("response_type", "code")
 	params.Set("scope", strings.Join(Scopes, " "))
 	params.Set("state", state)
diff --git a/sdk/auth/antigravity.go b/sdk/auth/antigravity.go
index de182eb3..7281fdb8 100644
--- a/sdk/auth/antigravity.go
+++ b/sdk/auth/antigravity.go
@@ -49,7 +49,7 @@ func (AntigravityAuthenticator) Login(ctx context.Context, cfg *config.Config, o
 		callbackPort = opts.CallbackPort
 	}
 
-	authSvc := antigravity.NewAntigravityAuth(cfg)
+	authSvc := antigravity.NewAntigravityAuth(cfg, nil)
 
 	state, err := misc.GenerateRandomState()
 	if err != nil {
@@ -67,9 +67,7 @@ func (AntigravityAuthenticator) Login(ctx context.Context, cfg *config.Config, o
 	}()
 
 	redirectURI := fmt.Sprintf("http://localhost:%d/oauth-callback", port)
-	authURL := authSvc.BuildAuthURL(state)
-	// Override redirect URI in authURL
-	authURL = strings.ReplaceAll(authURL, fmt.Sprintf("http://localhost:%d/oauth-callback", antigravity.CallbackPort), redirectURI)
+	authURL := authSvc.BuildAuthURL(state, redirectURI)
 
 	if !opts.NoBrowser {
 		fmt.Println("Opening browser for antigravity authentication")
@@ -256,7 +254,6 @@ func startAntigravityCallbackServer(port int) (*http.Server, int, <-chan callbac
 // FetchAntigravityProjectID exposes project discovery for external callers.
 func FetchAntigravityProjectID(ctx context.Context, accessToken string, httpClient *http.Client) (string, error) {
 	cfg := &config.Config{}
-	// Set the httpClient if provided (for proxy support)
-	authSvc := antigravity.NewAntigravityAuth(cfg)
+	authSvc := antigravity.NewAntigravityAuth(cfg, httpClient)
 	return authSvc.FetchProjectID(ctx, accessToken)
 }