diff --git a/internal/provider/gemini-web/conversation/index.go b/internal/provider/gemini-web/conversation/index.go index 5d46717e..ab06bbf5 100644 --- a/internal/provider/gemini-web/conversation/index.go +++ b/internal/provider/gemini-web/conversation/index.go @@ -1,16 +1,16 @@ package conversation import ( - "bytes" - "encoding/json" - "errors" - "os" - "path/filepath" - "strings" - "sync" - "time" + "bytes" + "encoding/json" + "errors" + "os" + "path/filepath" + "strings" + "sync" + "time" - bolt "go.etcd.io/bbolt" + bolt "go.etcd.io/bbolt" ) const ( @@ -66,36 +66,36 @@ func indexPath() string { // StoreMatch persists or updates a conversation hash mapping. func StoreMatch(hash string, record MatchRecord) error { - if strings.TrimSpace(hash) == "" { - return errors.New("gemini-web conversation: empty hash") - } - db, err := openIndex() - if err != nil { - return err - } - record.UpdatedAt = time.Now().UTC().Unix() - payload, err := json.Marshal(record) - if err != nil { - return err - } - return db.Update(func(tx *bolt.Tx) error { - bucket, err := tx.CreateBucketIfNotExists([]byte(bucketMatches)) - if err != nil { - return err - } - // 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 - }) + if strings.TrimSpace(hash) == "" { + return errors.New("gemini-web conversation: empty hash") + } + db, err := openIndex() + if err != nil { + return err + } + record.UpdatedAt = time.Now().UTC().Unix() + payload, err := json.Marshal(record) + if err != nil { + return err + } + return db.Update(func(tx *bolt.Tx) error { + bucket, err := tx.CreateBucketIfNotExists([]byte(bucketMatches)) + if err != nil { + return err + } + // 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. @@ -103,111 +103,122 @@ func StoreMatch(hash string, record MatchRecord) error { // 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) { - db, err := openIndex() - if err != nil { - return MatchRecord{}, false, err - } - var foundOne bool - var single MatchRecord - err = db.View(func(tx *bolt.Tx) error { - bucket := tx.Bucket([]byte(bucketMatches)) - if bucket == 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)) - if len(raw) == 0 { - return nil - } - return json.Unmarshal(raw, &single) - }) - if err != nil { - return MatchRecord{}, false, err - } - if strings.TrimSpace(single.AccountLabel) == "" || single.PrefixLen <= 0 { - return MatchRecord{}, false, nil - } - return single, true, nil + db, err := openIndex() + if err != nil { + return MatchRecord{}, false, err + } + var foundOne bool + var ambiguous bool + var firstLabel string + var single MatchRecord + err = db.View(func(tx *bolt.Tx) error { + bucket := tx.Bucket([]byte(bucketMatches)) + if bucket == 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 + } + label := strings.ToLower(strings.TrimSpace(rec.AccountLabel)) + if !foundOne { + firstLabel = label + single = rec + foundOne = true + continue + } + if label != firstLabel { + ambiguous = true + // Early exit scan; ambiguity detected. + return nil + } + } + if foundOne { + return nil + } + // Fallback to legacy single-key format + raw := bucket.Get([]byte(hash)) + if len(raw) == 0 { + return nil + } + return json.Unmarshal(raw, &single) + }) + if err != nil { + return MatchRecord{}, false, err + } + if ambiguous { + return MatchRecord{}, false, nil + } + if strings.TrimSpace(single.AccountLabel) == "" || single.PrefixLen <= 0 { + return MatchRecord{}, false, nil + } + return single, true, nil } // RemoveMatch deletes all mappings for the given hash (all labels and legacy key). func RemoveMatch(hash string) error { - 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 - } - // 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 - }) + 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 + } + // 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 - }) + 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 + }) } // RemoveMatchesByLabel removes all entries associated with the specified label. diff --git a/sdk/cliproxy/auth/selector_rr.go b/sdk/cliproxy/auth/selector_rr.go index a36c9413..1dbf357d 100644 --- a/sdk/cliproxy/auth/selector_rr.go +++ b/sdk/cliproxy/auth/selector_rr.go @@ -64,20 +64,20 @@ func (s *geminiWebStickySelector) Pick(ctx context.Context, provider, model stri if label == "" { continue } - auth := findAuthByLabel(auths, label) - if auth != nil { - if opts.Metadata != nil { - opts.Metadata[conversation.MetadataMatchKey] = &conversation.MatchResult{ - Hash: candidate.Hash, - Record: record, - Model: normalizedModel, - } - } - return auth, nil - } - _ = conversation.RemoveMatchForLabel(candidate.Hash, label) - } - } + auth := findAuthByLabel(auths, label) + if auth != nil { + if opts.Metadata != nil { + opts.Metadata[conversation.MetadataMatchKey] = &conversation.MatchResult{ + Hash: candidate.Hash, + Record: record, + Model: normalizedModel, + } + } + return auth, nil + } + _ = conversation.RemoveMatchForLabel(candidate.Hash, label) + } + } return s.base.Pick(ctx, provider, model, opts, auths) } diff --git a/sdk/cliproxy/service.go b/sdk/cliproxy/service.go index b54e8bc9..32738b7d 100644 --- a/sdk/cliproxy/service.go +++ b/sdk/cliproxy/service.go @@ -206,30 +206,30 @@ func (s *Service) applyCoreAuthRemoval(ctx context.Context, id string) { return } GlobalModelRegistry().UnregisterClient(id) - if existing, ok := s.coreManager.GetByID(id); ok && existing != nil { - if strings.EqualFold(existing.Provider, "gemini-web") { - // 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 err := conversation.RemoveMatchesByLabel(label); err != nil { - log.Debugf("failed to remove gemini web sticky entries for %s: %v", label, err) - } - } - } - existing.Disabled = true - existing.Status = coreauth.StatusDisabled - if _, err := s.coreManager.Update(ctx, existing); err != nil { - log.Errorf("failed to disable auth %s: %v", id, err) - } - } + if existing, ok := s.coreManager.GetByID(id); ok && existing != nil { + if strings.EqualFold(existing.Provider, "gemini-web") { + // 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 err := conversation.RemoveMatchesByLabel(label); err != nil { + log.Debugf("failed to remove gemini web sticky entries for %s: %v", label, err) + } + } + } + existing.Disabled = true + existing.Status = coreauth.StatusDisabled + if _, err := s.coreManager.Update(ctx, existing); err != nil { + log.Errorf("failed to disable auth %s: %v", id, err) + } + } } func (s *Service) ensureExecutorsForAuth(a *coreauth.Auth) {