chuan/CLIProxyAPI
1
0
Fork 0
mirror of https://github.com/router-for-me/CLIProxyAPI.git synced 2026-02-19 04:40:52 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(gemini-web): Correct ambiguity check in conversation lookup

Browse Source
This commit is contained in:
hkfires
2025-09-29 21:40:25 +08:00
parent 6080527e9e
commit 1d70336a91
3 changed files with 185 additions and 174 deletions
Download Patch File Download Diff File Expand all files Collapse all files

283
internal/provider/gemini-web/conversation/index.go
View File

@@ -1,16 +1,16 @@
package conversation package conversation
import ( import (
"bytes" "bytes"
"encoding/json" "encoding/json"
"errors" "errors"
"os" "os"
"path/filepath" "path/filepath"
"strings" "strings"
"sync" "sync"
"time" "time"
bolt "go.etcd.io/bbolt" bolt "go.etcd.io/bbolt"
) )
const ( const (
@@ -66,36 +66,36 @@ func indexPath() string {
// StoreMatch persists or updates a conversation hash mapping. // StoreMatch persists or updates a conversation hash mapping.
func StoreMatch(hash string, record MatchRecord) error { func StoreMatch(hash string, record MatchRecord) error {
if strings.TrimSpace(hash) == "" { if strings.TrimSpace(hash) == "" {
return errors.New("gemini-web conversation: empty hash") return errors.New("gemini-web conversation: empty hash")
} }
db, err := openIndex() db, err := openIndex()
if err != nil { if err != nil {
return err return err
} }
record.UpdatedAt = time.Now().UTC().Unix() record.UpdatedAt = time.Now().UTC().Unix()
payload, err := json.Marshal(record) payload, err := json.Marshal(record)
if err != nil { if err != nil {
return err return err
} }
return db.Update(func(tx *bolt.Tx) error { return db.Update(func(tx *bolt.Tx) error {
bucket, err := tx.CreateBucketIfNotExists([]byte(bucketMatches)) bucket, err := tx.CreateBucketIfNotExists([]byte(bucketMatches))
if err != nil { if err != nil {
return err return err
} }
// Namespace by account label to avoid cross-account collisions. // Namespace by account label to avoid cross-account collisions.
label := strings.ToLower(strings.TrimSpace(record.AccountLabel)) label := strings.ToLower(strings.TrimSpace(record.AccountLabel))
if label == "" { if label == "" {
return errors.New("gemini-web conversation: empty account label") return errors.New("gemini-web conversation: empty account label")
} }
key := []byte(hash + ":" + label) key := []byte(hash + ":" + label)
if err := bucket.Put(key, payload); err != nil { if err := bucket.Put(key, payload); err != nil {
return err return err
} }
// Best-effort cleanup of legacy single-key format (hash -> MatchRecord). // Best-effort cleanup of legacy single-key format (hash -> MatchRecord).
// We do not know its label; leave it for lookup fallback/cleanup elsewhere. // We do not know its label; leave it for lookup fallback/cleanup elsewhere.
return nil return nil
}) })
} }
// LookupMatch retrieves a stored mapping. // 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. // hash, it returns not found to avoid redirecting to the wrong credential.
// Falls back to legacy single-key entries if present. // 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 foundOne bool var foundOne bool
var single MatchRecord var ambiguous bool
err = db.View(func(tx *bolt.Tx) error { var firstLabel string
bucket := tx.Bucket([]byte(bucketMatches)) var single MatchRecord
if bucket == nil { err = db.View(func(tx *bolt.Tx) error {
return nil bucket := tx.Bucket([]byte(bucketMatches))
} if bucket == nil {
// Scan namespaced keys with prefix "hash:" return nil
prefix := []byte(hash + ":") }
c := bucket.Cursor() // Scan namespaced keys with prefix "hash:"
for k, v := c.Seek(prefix); k != nil && bytes.HasPrefix(k, prefix); k, v = c.Next() { prefix := []byte(hash + ":")
if len(v) == 0 { c := bucket.Cursor()
continue for k, v := c.Seek(prefix); k != nil && bytes.HasPrefix(k, prefix); k, v = c.Next() {
} if len(v) == 0 {
var rec MatchRecord continue
if err := json.Unmarshal(v, &rec); err != nil { }
// Ignore malformed; removal is handled elsewhere. var rec MatchRecord
continue if err := json.Unmarshal(v, &rec); err != nil {
} // Ignore malformed; removal is handled elsewhere.
if strings.TrimSpace(rec.AccountLabel) == "" || rec.PrefixLen <= 0 { continue
continue }
} if strings.TrimSpace(rec.AccountLabel) == "" || rec.PrefixLen <= 0 {
if foundOne { continue
// More than one distinct label exists for this hash; ambiguous. }
return nil label := strings.ToLower(strings.TrimSpace(rec.AccountLabel))
} if !foundOne {
single = rec firstLabel = label
foundOne = true single = rec
} foundOne = true
if foundOne { continue
return nil }
} if label != firstLabel {
// Fallback to legacy single-key format ambiguous = true
raw := bucket.Get([]byte(hash)) // Early exit scan; ambiguity detected.
if len(raw) == 0 { return nil
return nil }
} }
return json.Unmarshal(raw, &single) if foundOne {
}) return nil
if err != nil { }
return MatchRecord{}, false, err // Fallback to legacy single-key format
} raw := bucket.Get([]byte(hash))
if strings.TrimSpace(single.AccountLabel) == "" || single.PrefixLen <= 0 { if len(raw) == 0 {
return MatchRecord{}, false, nil return nil
} }
return single, true, 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). // 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 {
return err return err
} }
return db.Update(func(tx *bolt.Tx) error { return db.Update(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
} }
// Delete namespaced entries // Delete namespaced entries
prefix := []byte(hash + ":") prefix := []byte(hash + ":")
c := bucket.Cursor() c := bucket.Cursor()
for k, _ := c.Seek(prefix); k != nil && bytes.HasPrefix(k, prefix); k, _ = c.Next() { for k, _ := c.Seek(prefix); k != nil && bytes.HasPrefix(k, prefix); k, _ = c.Next() {
if err := bucket.Delete(k); err != nil { if err := bucket.Delete(k); err != nil {
return err return err
} }
} }
// Delete legacy entry // Delete legacy entry
_ = bucket.Delete([]byte(hash)) _ = bucket.Delete([]byte(hash))
return nil return nil
}) })
} }
// RemoveMatchForLabel deletes the mapping for the given hash and label only. // RemoveMatchForLabel deletes the mapping for the given hash and label only.
func RemoveMatchForLabel(hash, label string) error { func RemoveMatchForLabel(hash, label string) error {
label = strings.ToLower(strings.TrimSpace(label)) label = strings.ToLower(strings.TrimSpace(label))
if strings.TrimSpace(hash) == "" || label == "" { if strings.TrimSpace(hash) == "" || label == "" {
return nil return nil
} }
db, err := openIndex() db, err := openIndex()
if err != nil { if err != nil {
return err return err
} }
return db.Update(func(tx *bolt.Tx) error { return db.Update(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
} }
// Remove namespaced key // Remove namespaced key
_ = bucket.Delete([]byte(hash + ":" + label)) _ = bucket.Delete([]byte(hash + ":" + label))
// If legacy single-key exists and matches label, remove it as well. // If legacy single-key exists and matches label, remove it as well.
if raw := bucket.Get([]byte(hash)); len(raw) > 0 { if raw := bucket.Get([]byte(hash)); len(raw) > 0 {
var rec MatchRecord var rec MatchRecord
if err := json.Unmarshal(raw, &rec); err == nil { if err := json.Unmarshal(raw, &rec); err == nil {
if strings.EqualFold(strings.TrimSpace(rec.AccountLabel), label) { if strings.EqualFold(strings.TrimSpace(rec.AccountLabel), label) {
_ = bucket.Delete([]byte(hash)) _ = bucket.Delete([]byte(hash))
} }
} }
} }
return nil return nil
}) })
} }
// RemoveMatchesByLabel removes all entries associated with the specified label. // RemoveMatchesByLabel removes all entries associated with the specified label.

28
sdk/cliproxy/auth/selector_rr.go
View File

@@ -64,20 +64,20 @@ func (s *geminiWebStickySelector) Pick(ctx context.Context, provider, model stri
if label == "" { if label == "" {
continue continue
} }
auth := findAuthByLabel(auths, label) auth := findAuthByLabel(auths, label)
if auth != nil { if auth != nil {
if opts.Metadata != nil { if opts.Metadata != nil {
opts.Metadata[conversation.MetadataMatchKey] = &conversation.MatchResult{ opts.Metadata[conversation.MetadataMatchKey] = &conversation.MatchResult{
Hash: candidate.Hash, Hash: candidate.Hash,
Record: record, Record: record,
Model: normalizedModel, Model: normalizedModel,
} }
} }
return auth, nil return auth, nil
} }
_ = conversation.RemoveMatchForLabel(candidate.Hash, label) _ = conversation.RemoveMatchForLabel(candidate.Hash, label)
} }
} }
return s.base.Pick(ctx, provider, model, opts, auths) return s.base.Pick(ctx, provider, model, opts, auths)
} }

48
sdk/cliproxy/service.go
View File

@@ -206,30 +206,30 @@ func (s *Service) applyCoreAuthRemoval(ctx context.Context, id string) {
return return
} }
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") {
// Prefer the stable cookie label stored in metadata when available. // Prefer the stable cookie label stored in metadata when available.
var label string var label string
if existing.Metadata != nil { if existing.Metadata != nil {
if v, ok := existing.Metadata["label"].(string); ok { if v, ok := existing.Metadata["label"].(string); ok {
label = strings.TrimSpace(v) label = strings.TrimSpace(v)
} }
} }
if label == "" { if label == "" {
label = strings.TrimSpace(existing.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)
} }
} }
} }
existing.Disabled = true existing.Disabled = true
existing.Status = coreauth.StatusDisabled existing.Status = coreauth.StatusDisabled
if _, err := s.coreManager.Update(ctx, existing); err != nil { if _, err := s.coreManager.Update(ctx, existing); err != nil {
log.Errorf("failed to disable auth %s: %v", id, err) log.Errorf("failed to disable auth %s: %v", id, err)
} }
} }
} }
func (s *Service) ensureExecutorsForAuth(a *coreauth.Auth) { func (s *Service) ensureExecutorsForAuth(a *coreauth.Auth) {