From 96cebd2a35daac81902944c2e53c5dd83a5acf68 Mon Sep 17 00:00:00 2001
From: hkfires <10558748+hkfires@users.noreply.github.com>
Date: Thu, 25 Sep 2025 20:39:15 +0800
Subject: [PATCH] feat(auth): add interactive prompts to Gemini web auth flow

---
 internal/cmd/gemini-web_auth.go        | 71 ++++++++++++++++----------
 internal/provider/gemini-web/client.go | 10 ++--
 internal/provider/gemini-web/media.go  |  4 +-
 3 files changed, 51 insertions(+), 34 deletions(-)

diff --git a/internal/cmd/gemini-web_auth.go b/internal/cmd/gemini-web_auth.go
index 80f227f7..40dda308 100644
--- a/internal/cmd/gemini-web_auth.go
+++ b/internal/cmd/gemini-web_auth.go
@@ -56,8 +56,19 @@ func DoGeminiWebAuth(cfg *config.Config) {
 
 	secure1psid := strings.TrimSpace(cookieMap["__Secure-1PSID"])
 	secure1psidts := strings.TrimSpace(cookieMap["__Secure-1PSIDTS"])
+	// Fallback: prompt user to input missing values
+	if secure1psid == "" {
+		fmt.Print("Cookie missing __Secure-1PSID. Enter __Secure-1PSID: ")
+		v, _ := reader.ReadString('\n')
+		secure1psid = strings.TrimSpace(v)
+	}
+	if secure1psidts == "" {
+		fmt.Print("Cookie missing __Secure-1PSIDTS. Enter __Secure-1PSIDTS: ")
+		v, _ := reader.ReadString('\n')
+		secure1psidts = strings.TrimSpace(v)
+	}
 	if secure1psid == "" || secure1psidts == "" {
-		fmt.Println("Cookie does not contain __Secure-1PSID or __Secure-1PSIDTS")
+		log.Fatal("__Secure-1PSID and __Secure-1PSIDTS cannot be empty")
 		return
 	}
 
@@ -78,36 +89,34 @@ func DoGeminiWebAuth(cfg *config.Config) {
 	req.Header.Set("Content-Type", "application/x-www-form-urlencoded;charset=UTF-8")
 
 	resp, err := httpClient.Do(req)
+	email := ""
 	if err != nil {
 		fmt.Printf("Request to ListAccounts failed: %v\n", err)
-		return
-	}
-	defer resp.Body.Close()
-	if resp.StatusCode != http.StatusOK {
-		fmt.Printf("ListAccounts returned status code: %d\n", resp.StatusCode)
-		return
-	}
-
-	var payload []any
-	if err := json.NewDecoder(resp.Body).Decode(&payload); err != nil {
-		fmt.Printf("Failed to parse ListAccounts response: %v\n", err)
-		return
-	}
-
-	// Expected structure like: ["gaia.l.a.r", [["gaia.l.a",1,"Name","email@example.com", ... ]]]
-	email := ""
-	if len(payload) >= 2 {
-		if accounts, ok := payload[1].([]any); ok && len(accounts) >= 1 {
-			if first, ok := accounts[0].([]any); ok && len(first) >= 4 {
-				if em, ok := first[3].(string); ok {
-					email = strings.TrimSpace(em)
+	} else {
+		defer resp.Body.Close()
+		if resp.StatusCode != http.StatusOK {
+			fmt.Printf("ListAccounts returned status code: %d\n", resp.StatusCode)
+		} else {
+			var payload []any
+			if err := json.NewDecoder(resp.Body).Decode(&payload); err != nil {
+				fmt.Printf("Failed to parse ListAccounts response: %v\n", err)
+			} else {
+				// Expected structure like: ["gaia.l.a.r", [["gaia.l.a",1,"Name","email@example.com", ... ]]]
+				if len(payload) >= 2 {
+					if accounts, ok := payload[1].([]any); ok && len(accounts) >= 1 {
+						if first, ok := accounts[0].([]any); ok && len(first) >= 4 {
+							if em, ok := first[3].(string); ok {
+								email = strings.TrimSpace(em)
+							}
+						}
+					}
+				}
+				if email == "" {
+					fmt.Println("Failed to parse email from ListAccounts response")
 				}
 			}
 		}
 	}
-	if email == "" {
-		fmt.Println("Failed to parse email from ListAccounts response; fallback to filename-based label")
-	}
 
 	// Generate a filename based on the SHA256 hash of the PSID
 	hasher := sha256.New()
@@ -115,10 +124,18 @@ func DoGeminiWebAuth(cfg *config.Config) {
 	hash := hex.EncodeToString(hasher.Sum(nil))
 	fileName := fmt.Sprintf("gemini-web-%s.json", hash[:16])
 
-	// Decide label: prefer email; fallback to file name without .json
+	// Decide label: prefer email; fallback prompt then file name without .json
+	defaultLabel := strings.TrimSuffix(fileName, ".json")
 	label := email
 	if label == "" {
-		label = strings.TrimSuffix(fileName, ".json")
+		fmt.Printf("Enter label for this auth (default: %s): ", defaultLabel)
+		v, _ := reader.ReadString('\n')
+		v = strings.TrimSpace(v)
+		if v != "" {
+			label = v
+		} else {
+			label = defaultLabel
+		}
 	}
 
 	tokenStorage := &gemini.GeminiWebTokenStorage{
diff --git a/internal/provider/gemini-web/client.go b/internal/provider/gemini-web/client.go
index 8dd28e3e..4ce9a263 100644
--- a/internal/provider/gemini-web/client.go
+++ b/internal/provider/gemini-web/client.go
@@ -292,7 +292,7 @@ func (c *GeminiClient) Close(delaySec float64) {
 	c.Running = false
 }
 
-// ensureRunning mirrors the Python decorator behavior and retries on APIError.
+// ensureRunning mirrors the decorator behavior and retries on APIError.
 func (c *GeminiClient) ensureRunning() error {
 	if c.Running {
 		return nil
@@ -419,7 +419,7 @@ func (c *GeminiClient) generateOnce(prompt string, files []string, model Model,
 	}()
 
 	if resp.StatusCode == 429 {
-		// Surface 429 as TemporarilyBlocked to match Python behavior
+		// Surface 429 as TemporarilyBlocked to match reference behavior
 		c.Close(0)
 		return empty, &TemporarilyBlocked{GeminiError{Msg: "Too many requests. IP temporarily blocked."}}
 	}
@@ -499,7 +499,7 @@ func (c *GeminiClient) generateOnce(prompt string, files []string, model Model,
 				}
 			}
 		}
-		// Parse nested error code to align with Python mapping
+		// Parse nested error code to align with error mapping
 		var top []any
 		// Prefer lastTop from fallback scan; otherwise try parts[2]
 		if len(lastTop) > 0 {
@@ -522,7 +522,7 @@ func (c *GeminiClient) generateOnce(prompt string, files []string, model Model,
 			}
 		}
 		// Debug("Invalid response: control frames only; no body found")
-		// Close the client to force re-initialization on next request (parity with Python client behavior)
+		// Close the client to force re-initialization on next request (parity with reference client behavior)
 		c.Close(0)
 		return empty, &APIError{Msg: "Failed to generate contents. Invalid response data received."}
 	}
@@ -745,7 +745,7 @@ func (c *GeminiClient) generateOnce(prompt string, files []string, model Model,
 }
 
 // extractErrorCode attempts to navigate the known nested error structure and fetch the integer code.
-// Mirrors Python path: response_json[0][5][2][0][1][0]
+// Mirrors reference path: response_json[0][5][2][0][1][0]
 func extractErrorCode(top []any) (int, bool) {
 	if len(top) == 0 {
 		return 0, false
diff --git a/internal/provider/gemini-web/media.go b/internal/provider/gemini-web/media.go
index e9dcecde..585eff90 100644
--- a/internal/provider/gemini-web/media.go
+++ b/internal/provider/gemini-web/media.go
@@ -52,7 +52,7 @@ func (i Image) Save(path string, filename string, cookies map[string]string, ver
 			filename = q[0]
 		}
 	}
-	// Regex validation (align with Python: ^(.*\.\w+)) to extract name with extension.
+	// Regex validation (pattern: ^(.*\.\w+)) to extract name with extension.
 	if filename != "" {
 		re := regexp.MustCompile(`^(.*\.\w+)`)
 		if m := re.FindStringSubmatch(filename); len(m) >= 2 {
@@ -70,7 +70,7 @@ func (i Image) Save(path string, filename string, cookies map[string]string, ver
 	client := newHTTPClient(httpOptions{ProxyURL: i.Proxy, Insecure: insecure, FollowRedirects: true})
 	client.Timeout = 120 * time.Second
 
-	// Helper to set raw Cookie header using provided cookies (to mirror Python client behavior).
+	// Helper to set raw Cookie header using provided cookies (parity with the reference client behavior).
 	buildCookieHeader := func(m map[string]string) string {
 		if len(m) == 0 {
 			return ""