package geminiwebapi import ( "crypto/sha256" "encoding/hex" "encoding/json" "fmt" "os" "path/filepath" "strings" "time" "github.com/luispater/CLIProxyAPI/v5/internal/auth/gemini" ) // StoredMessage represents a single message in a conversation record. type StoredMessage struct { Role string `json:"role"` Content string `json:"content"` Name string `json:"name,omitempty"` } // ConversationRecord stores a full conversation with its metadata for persistence. type ConversationRecord struct { Model string `json:"model"` ClientID string `json:"client_id"` Metadata []string `json:"metadata,omitempty"` Messages []StoredMessage `json:"messages"` CreatedAt time.Time `json:"created_at"` UpdatedAt time.Time `json:"updated_at"` } // Sha256Hex computes the SHA256 hash of a string and returns its hex representation. func Sha256Hex(s string) string { sum := sha256.Sum256([]byte(s)) return hex.EncodeToString(sum[:]) } // RoleText represents a turn in a conversation with a role and text content. type RoleText struct { Role string Text string } func ToStoredMessages(msgs []RoleText) []StoredMessage { out := make([]StoredMessage, 0, len(msgs)) for _, m := range msgs { out = append(out, StoredMessage{ Role: m.Role, Content: m.Text, }) } return out } func HashMessage(m StoredMessage) string { s := fmt.Sprintf(`{"content":%q,"role":%q}`, m.Content, strings.ToLower(m.Role)) return Sha256Hex(s) } func HashConversation(clientID, model string, msgs []StoredMessage) string { var b strings.Builder b.WriteString(clientID) b.WriteString("|") b.WriteString(model) for _, m := range msgs { b.WriteString("|") b.WriteString(HashMessage(m)) } return Sha256Hex(b.String()) } // ConvStorePath returns the path for account-level metadata persistence based on token file path. func ConvStorePath(tokenFilePath string) string { wd, err := os.Getwd() if err != nil || wd == "" { wd = "." } convDir := filepath.Join(wd, "conv") base := strings.TrimSuffix(filepath.Base(tokenFilePath), filepath.Ext(tokenFilePath)) return filepath.Join(convDir, base+".conv.json") } // ConvDataPath returns the path for full conversation persistence based on token file path. func ConvDataPath(tokenFilePath string) string { wd, err := os.Getwd() if err != nil || wd == "" { wd = "." } convDir := filepath.Join(wd, "conv") base := strings.TrimSuffix(filepath.Base(tokenFilePath), filepath.Ext(tokenFilePath)) return filepath.Join(convDir, base+".data.json") } // LoadConvStore reads the account-level metadata store from disk. func LoadConvStore(path string) (map[string][]string, error) { b, err := os.ReadFile(path) if err != nil { // Missing file is not an error; return empty map return map[string][]string{}, nil } var tmp map[string][]string if err := json.Unmarshal(b, &tmp); err != nil { return nil, err } if tmp == nil { tmp = map[string][]string{} } return tmp, nil } // SaveConvStore writes the account-level metadata store to disk atomically. func SaveConvStore(path string, data map[string][]string) error { if data == nil { data = map[string][]string{} } payload, err := json.MarshalIndent(data, "", " ") if err != nil { return err } // Ensure directory exists if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { return err } tmp := path + ".tmp" if err := os.WriteFile(tmp, payload, 0o644); err != nil { return err } return os.Rename(tmp, path) } // AccountMetaKey builds the key for account-level metadata map. func AccountMetaKey(email, modelName string) string { return fmt.Sprintf("account-meta|%s|%s", email, modelName) } // LoadConvData reads the full conversation data and index from disk. func LoadConvData(path string) (map[string]ConversationRecord, map[string]string, error) { b, err := os.ReadFile(path) if err != nil { // Missing file is not an error; return empty sets return map[string]ConversationRecord{}, map[string]string{}, nil } var wrapper struct { Items map[string]ConversationRecord `json:"items"` Index map[string]string `json:"index"` } if err := json.Unmarshal(b, &wrapper); err != nil { return nil, nil, err } if wrapper.Items == nil { wrapper.Items = map[string]ConversationRecord{} } if wrapper.Index == nil { wrapper.Index = map[string]string{} } return wrapper.Items, wrapper.Index, nil } // SaveConvData writes the full conversation data and index to disk atomically. func SaveConvData(path string, items map[string]ConversationRecord, index map[string]string) error { if items == nil { items = map[string]ConversationRecord{} } if index == nil { index = map[string]string{} } wrapper := struct { Items map[string]ConversationRecord `json:"items"` Index map[string]string `json:"index"` }{Items: items, Index: index} payload, err := json.MarshalIndent(wrapper, "", " ") if err != nil { return err } if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { return err } tmp := path + ".tmp" if err := os.WriteFile(tmp, payload, 0o644); err != nil { return err } return os.Rename(tmp, path) } // BuildConversationRecord constructs a ConversationRecord from history and the latest output. // Returns false when output is empty or has no candidates. func BuildConversationRecord(model, clientID string, history []RoleText, output *ModelOutput, metadata []string) (ConversationRecord, bool) { if output == nil || len(output.Candidates) == 0 { return ConversationRecord{}, false } text := "" if t := output.Candidates[0].Text; t != "" { text = RemoveThinkTags(t) } final := append([]RoleText{}, history...) final = append(final, RoleText{Role: "assistant", Text: text}) rec := ConversationRecord{ Model: model, ClientID: clientID, Metadata: metadata, Messages: ToStoredMessages(final), CreatedAt: time.Now(), UpdatedAt: time.Now(), } return rec, true } // FindByMessageListIn looks up a conversation record by hashed message list. // It attempts both the stable client ID and a legacy email-based ID. func FindByMessageListIn(items map[string]ConversationRecord, index map[string]string, stableClientID, email, model string, msgs []RoleText) (ConversationRecord, bool) { stored := ToStoredMessages(msgs) stableHash := HashConversation(stableClientID, model, stored) fallbackHash := HashConversation(email, model, stored) // Try stable hash via index indirection first if key, ok := index["hash:"+stableHash]; ok { if rec, ok2 := items[key]; ok2 { return rec, true } } if rec, ok := items[stableHash]; ok { return rec, true } // Fallback to legacy hash (email-based) if key, ok := index["hash:"+fallbackHash]; ok { if rec, ok2 := items[key]; ok2 { return rec, true } } if rec, ok := items[fallbackHash]; ok { return rec, true } return ConversationRecord{}, false } // FindConversationIn tries exact then sanitized assistant messages. func FindConversationIn(items map[string]ConversationRecord, index map[string]string, stableClientID, email, model string, msgs []RoleText) (ConversationRecord, bool) { if len(msgs) == 0 { return ConversationRecord{}, false } if rec, ok := FindByMessageListIn(items, index, stableClientID, email, model, msgs); ok { return rec, true } if rec, ok := FindByMessageListIn(items, index, stableClientID, email, model, SanitizeAssistantMessages(msgs)); ok { return rec, true } return ConversationRecord{}, false } // FindReusableSessionIn returns reusable metadata and the remaining message suffix. func FindReusableSessionIn(items map[string]ConversationRecord, index map[string]string, stableClientID, email, model string, msgs []RoleText) ([]string, []RoleText) { if len(msgs) < 2 { return nil, nil } searchEnd := len(msgs) for searchEnd >= 2 { sub := msgs[:searchEnd] tail := sub[len(sub)-1] if strings.EqualFold(tail.Role, "assistant") || strings.EqualFold(tail.Role, "system") { if rec, ok := FindConversationIn(items, index, stableClientID, email, model, sub); ok { remain := msgs[searchEnd:] return rec.Metadata, remain } } searchEnd-- } return nil, nil } // CookiesSidecarPath derives the sidecar cookie file path from the main token JSON path. func CookiesSidecarPath(mainPath string) string { if strings.HasSuffix(mainPath, ".json") { return strings.TrimSuffix(mainPath, ".json") + ".cookies" } return mainPath + ".cookies" } // FileExists reports whether the given path exists and is a regular file. func FileExists(path string) bool { if path == "" { return false } if st, err := os.Stat(path); err == nil && !st.IsDir() { return true } return false } // ApplyCookiesSidecarToTokenStorage loads cookies from sidecar into the provided token storage. // Returns true when a sidecar was found and applied. func ApplyCookiesSidecarToTokenStorage(tokenFilePath string, ts *gemini.GeminiWebTokenStorage) (bool, error) { if ts == nil { return false, nil } side := CookiesSidecarPath(tokenFilePath) if !FileExists(side) { return false, nil } data, err := os.ReadFile(side) if err != nil || len(data) == 0 { return false, err } var latest gemini.GeminiWebTokenStorage if err := json.Unmarshal(data, &latest); err != nil { return false, err } if latest.Secure1PSID != "" { ts.Secure1PSID = latest.Secure1PSID } if latest.Secure1PSIDTS != "" { ts.Secure1PSIDTS = latest.Secure1PSIDTS } return true, nil } // SaveCookiesSidecar writes the current cookies into a sidecar file next to the token file. // This keeps the main token JSON stable until an orderly flush. func SaveCookiesSidecar(tokenFilePath string, cookies map[string]string) error { side := CookiesSidecarPath(tokenFilePath) ts := &gemini.GeminiWebTokenStorage{Type: "gemini-web"} if v := cookies["__Secure-1PSID"]; v != "" { ts.Secure1PSID = v } if v := cookies["__Secure-1PSIDTS"]; v != "" { ts.Secure1PSIDTS = v } if err := os.MkdirAll(filepath.Dir(side), 0o700); err != nil { return err } return ts.SaveTokenToFile(side) } // FlushCookiesSidecarToMain merges the sidecar cookies into the main token file and removes the sidecar. // If sidecar is missing, it will combine the provided base token storage with the latest cookies. func FlushCookiesSidecarToMain(tokenFilePath string, cookies map[string]string, base *gemini.GeminiWebTokenStorage) error { if tokenFilePath == "" { return nil } side := CookiesSidecarPath(tokenFilePath) var merged gemini.GeminiWebTokenStorage var fromSidecar bool if FileExists(side) { if data, err := os.ReadFile(side); err == nil && len(data) > 0 { if err2 := json.Unmarshal(data, &merged); err2 == nil { fromSidecar = true } } } if !fromSidecar { if base != nil { merged = *base } if v := cookies["__Secure-1PSID"]; v != "" { merged.Secure1PSID = v } if v := cookies["__Secure-1PSIDTS"]; v != "" { merged.Secure1PSIDTS = v } } merged.Type = "gemini-web" if err := os.MkdirAll(filepath.Dir(tokenFilePath), 0o700); err != nil { return err } if err := merged.SaveTokenToFile(tokenFilePath); err != nil { return err } if FileExists(side) { _ = os.Remove(side) } return nil } // IsSelfPersistedToken compares provided token storage with current cookies. // Removed: IsSelfPersistedToken (client-side no longer needs self-originated write detection)