feat(gemini-web): Namespace conversation index by account label

This commit is contained in:
hkfires
2025-09-29 21:05:20 +08:00
parent 82187bffba
commit 6080527e9e
3 changed files with 182 additions and 91 deletions

View File

@@ -1,6 +1,7 @@
package conversation package conversation
import ( import (
"bytes"
"encoding/json" "encoding/json"
"errors" "errors"
"os" "os"
@@ -82,38 +83,79 @@ func StoreMatch(hash string, record MatchRecord) error {
if err != nil { if err != nil {
return err return err
} }
return bucket.Put([]byte(hash), payload) // Namespace by account label to avoid cross-account collisions.
label := strings.ToLower(strings.TrimSpace(record.AccountLabel))
if label == "" {
return errors.New("gemini-web conversation: empty account label")
}
key := []byte(hash + ":" + label)
if err := bucket.Put(key, payload); err != nil {
return err
}
// Best-effort cleanup of legacy single-key format (hash -> MatchRecord).
// We do not know its label; leave it for lookup fallback/cleanup elsewhere.
return nil
}) })
} }
// LookupMatch retrieves a stored mapping. // LookupMatch retrieves a stored mapping.
// It prefers namespaced entries (hash:label). If multiple labels exist for the same
// hash, it returns not found to avoid redirecting to the wrong credential.
// Falls back to legacy single-key entries if present.
func LookupMatch(hash string) (MatchRecord, bool, error) { func LookupMatch(hash string) (MatchRecord, bool, error) {
db, err := openIndex() db, err := openIndex()
if err != nil { if err != nil {
return MatchRecord{}, false, err return MatchRecord{}, false, err
} }
var record MatchRecord var foundOne bool
var single MatchRecord
err = db.View(func(tx *bolt.Tx) error { err = db.View(func(tx *bolt.Tx) error {
bucket := tx.Bucket([]byte(bucketMatches)) bucket := tx.Bucket([]byte(bucketMatches))
if bucket == nil { if bucket == nil {
return nil return nil
} }
// Scan namespaced keys with prefix "hash:"
prefix := []byte(hash + ":")
c := bucket.Cursor()
for k, v := c.Seek(prefix); k != nil && bytes.HasPrefix(k, prefix); k, v = c.Next() {
if len(v) == 0 {
continue
}
var rec MatchRecord
if err := json.Unmarshal(v, &rec); err != nil {
// Ignore malformed; removal is handled elsewhere.
continue
}
if strings.TrimSpace(rec.AccountLabel) == "" || rec.PrefixLen <= 0 {
continue
}
if foundOne {
// More than one distinct label exists for this hash; ambiguous.
return nil
}
single = rec
foundOne = true
}
if foundOne {
return nil
}
// Fallback to legacy single-key format
raw := bucket.Get([]byte(hash)) raw := bucket.Get([]byte(hash))
if len(raw) == 0 { if len(raw) == 0 {
return nil return nil
} }
return json.Unmarshal(raw, &record) return json.Unmarshal(raw, &single)
}) })
if err != nil { if err != nil {
return MatchRecord{}, false, err return MatchRecord{}, false, err
} }
if record.AccountLabel == "" || record.PrefixLen <= 0 { if strings.TrimSpace(single.AccountLabel) == "" || single.PrefixLen <= 0 {
return MatchRecord{}, false, nil return MatchRecord{}, false, nil
} }
return record, true, nil return single, true, nil
} }
// RemoveMatch deletes a mapping for the given hash. // RemoveMatch deletes all mappings for the given hash (all labels and legacy key).
func RemoveMatch(hash string) error { func RemoveMatch(hash string) error {
db, err := openIndex() db, err := openIndex()
if err != nil { if err != nil {
@@ -124,7 +166,47 @@ func RemoveMatch(hash string) error {
if bucket == nil { if bucket == nil {
return nil return nil
} }
return bucket.Delete([]byte(hash)) // Delete namespaced entries
prefix := []byte(hash + ":")
c := bucket.Cursor()
for k, _ := c.Seek(prefix); k != nil && bytes.HasPrefix(k, prefix); k, _ = c.Next() {
if err := bucket.Delete(k); err != nil {
return err
}
}
// Delete legacy entry
_ = bucket.Delete([]byte(hash))
return nil
})
}
// RemoveMatchForLabel deletes the mapping for the given hash and label only.
func RemoveMatchForLabel(hash, label string) error {
label = strings.ToLower(strings.TrimSpace(label))
if strings.TrimSpace(hash) == "" || label == "" {
return nil
}
db, err := openIndex()
if err != nil {
return err
}
return db.Update(func(tx *bolt.Tx) error {
bucket := tx.Bucket([]byte(bucketMatches))
if bucket == nil {
return nil
}
// Remove namespaced key
_ = bucket.Delete([]byte(hash + ":" + label))
// If legacy single-key exists and matches label, remove it as well.
if raw := bucket.Get([]byte(hash)); len(raw) > 0 {
var rec MatchRecord
if err := json.Unmarshal(raw, &rec); err == nil {
if strings.EqualFold(strings.TrimSpace(rec.AccountLabel), label) {
_ = bucket.Delete([]byte(hash))
}
}
}
return nil
}) })
} }

View File

@@ -75,7 +75,7 @@ func (s *geminiWebStickySelector) Pick(ctx context.Context, provider, model stri
} }
return auth, nil return auth, nil
} }
_ = conversation.RemoveMatch(candidate.Hash) _ = conversation.RemoveMatchForLabel(candidate.Hash, label)
} }
} }

View File

@@ -208,7 +208,16 @@ func (s *Service) applyCoreAuthRemoval(ctx context.Context, id string) {
GlobalModelRegistry().UnregisterClient(id) GlobalModelRegistry().UnregisterClient(id)
if existing, ok := s.coreManager.GetByID(id); ok && existing != nil { if existing, ok := s.coreManager.GetByID(id); ok && existing != nil {
if strings.EqualFold(existing.Provider, "gemini-web") { if strings.EqualFold(existing.Provider, "gemini-web") {
label := strings.TrimSpace(existing.Label) // Prefer the stable cookie label stored in metadata when available.
var label string
if existing.Metadata != nil {
if v, ok := existing.Metadata["label"].(string); ok {
label = strings.TrimSpace(v)
}
}
if label == "" {
label = strings.TrimSpace(existing.Label)
}
if label != "" { if label != "" {
if err := conversation.RemoveMatchesByLabel(label); err != nil { if err := conversation.RemoveMatchesByLabel(label); err != nil {
log.Debugf("failed to remove gemini web sticky entries for %s: %v", label, err) log.Debugf("failed to remove gemini web sticky entries for %s: %v", label, err)