fix(api): address copilot review feedback on probe cache key and test stability

This commit is contained in:
lc6464 2026-04-01 00:34:18 +08:00
parent 48c0bb1e97
commit f4c5ed12f5
No known key found for this signature in database
GPG key ID: 53C61B42FEC71D6D
2 changed files with 71 additions and 10 deletions

View file

@ -4,9 +4,11 @@ import (
"context" "context"
"encoding/json" "encoding/json"
"fmt" "fmt"
"hash/fnv"
"net" "net"
"net/http" "net/http"
"net/url" "net/url"
"strconv"
"strings" "strings"
"sync" "sync"
"time" "time"
@ -192,13 +194,17 @@ func modelProbeCacheKey(m *config.ModelConfig) string {
protocol, modelID := splitModel(m.Model) protocol, modelID := splitModel(m.Model)
modelName := strings.ToLower(strings.TrimSpace(m.ModelName)) modelName := strings.ToLower(strings.TrimSpace(m.ModelName))
apiBase := strings.ToLower(strings.TrimSpace(modelProbeAPIBase(m))) apiBaseRaw := modelProbeAPIBase(m)
apiBase := strings.ToLower(strings.TrimRight(strings.TrimSpace(apiBaseRaw), "/"))
authMethod := strings.ToLower(strings.TrimSpace(m.AuthMethod)) authMethod := strings.ToLower(strings.TrimSpace(m.AuthMethod))
connectMode := strings.ToLower(strings.TrimSpace(m.ConnectMode)) connectMode := strings.ToLower(strings.TrimSpace(m.ConnectMode))
hasAPIKey := strings.TrimSpace(m.APIKey()) != "" apiKeyFingerprint := modelProbeAPIKeyFingerprint(m.APIKey())
var b strings.Builder var b strings.Builder
b.Grow(len(modelName) + len(protocol) + len(modelID) + len(apiBase) + len(authMethod) + len(connectMode) + 8) b.Grow(
len(modelName) + len(protocol) + len(modelID) + len(apiBase) + len(authMethod) +
len(connectMode) + len(apiKeyFingerprint) + 8,
)
b.WriteString(modelName) b.WriteString(modelName)
b.WriteByte('|') b.WriteByte('|')
b.WriteString(protocol) b.WriteString(protocol)
@ -211,15 +217,22 @@ func modelProbeCacheKey(m *config.ModelConfig) string {
b.WriteByte('|') b.WriteByte('|')
b.WriteString(connectMode) b.WriteString(connectMode)
b.WriteByte('|') b.WriteByte('|')
if hasAPIKey { b.WriteString(apiKeyFingerprint)
b.WriteByte('1')
} else {
b.WriteByte('0')
}
return b.String() return b.String()
} }
func modelProbeAPIKeyFingerprint(raw string) string {
apiKey := strings.TrimSpace(raw)
if apiKey == "" {
return "none"
}
h := fnv.New64a()
_, _ = h.Write([]byte(apiKey))
return strconv.FormatUint(h.Sum64(), 36)
}
func (s *modelProbeCacheState) getCachedResult(cacheKey string, now time.Time) (bool, bool) { func (s *modelProbeCacheState) getCachedResult(cacheKey string, now time.Time) (bool, bool) {
s.mu.RLock() s.mu.RLock()
defer s.mu.RUnlock() defer s.mu.RUnlock()

View file

@ -89,6 +89,46 @@ func TestProbeLocalModelAvailability_LMStudioUsesOpenAICompatibleProbe(t *testin
} }
} }
func TestModelProbeCacheKey_DifferentAPIKeysProduceDifferentKeys(t *testing.T) {
base := &config.ModelConfig{
ModelName: "local-vllm",
Model: "vllm/custom-model",
APIBase: "http://127.0.0.1:8000/v1",
AuthMethod: "local",
ConnectMode: "",
}
m1 := *base
m1.SetAPIKey("key-a")
m2 := *base
m2.SetAPIKey("key-b")
k1 := modelProbeCacheKey(&m1)
k2 := modelProbeCacheKey(&m2)
if k1 == k2 {
t.Fatal("modelProbeCacheKey() should differ when api key changes")
}
}
func TestModelProbeCacheKey_NormalizesTrailingSlashInAPIBase(t *testing.T) {
m1 := &config.ModelConfig{
ModelName: "local-vllm",
Model: "vllm/custom-model",
APIBase: "http://127.0.0.1:8000/v1",
}
m2 := &config.ModelConfig{
ModelName: "local-vllm",
Model: "vllm/custom-model",
APIBase: "http://127.0.0.1:8000/v1/",
}
k1 := modelProbeCacheKey(m1)
k2 := modelProbeCacheKey(m2)
if k1 != k2 {
t.Fatalf("modelProbeCacheKey() mismatch for equivalent api_base values: %q vs %q", k1, k2)
}
}
func TestProbeLocalModelAvailability_SuccessBackoff(t *testing.T) { func TestProbeLocalModelAvailability_SuccessBackoff(t *testing.T) {
resetModelProbeHooks(t) resetModelProbeHooks(t)
@ -272,23 +312,27 @@ func TestProbeLocalModelAvailability_DeduplicatesInflightProbe(t *testing.T) {
const workers = 8 const workers = 8
var wg sync.WaitGroup var wg sync.WaitGroup
results := make(chan bool, workers) results := make(chan bool, workers)
workerStarted := make(chan struct{}, workers)
for range workers { for range workers {
wg.Add(1) wg.Add(1)
go func() { go func() {
defer wg.Done() defer wg.Done()
workerStarted <- struct{}{}
results <- probeLocalModelAvailability(model) results <- probeLocalModelAvailability(model)
}() }()
} }
for range workers {
<-workerStarted
}
select { select {
case <-probeStarted: case <-probeStarted:
case <-time.After(200 * time.Millisecond): case <-time.After(200 * time.Millisecond):
t.Fatal("probe did not start in time") t.Fatal("probe did not start in time")
} }
// Give waiting goroutines a short window to join the same inflight call.
time.Sleep(50 * time.Millisecond)
if got := atomic.LoadInt32(&calls); got != 1 { if got := atomic.LoadInt32(&calls); got != 1 {
t.Fatalf("concurrent probe calls = %d, want 1", got) t.Fatalf("concurrent probe calls = %d, want 1", got)
} }
@ -302,6 +346,10 @@ func TestProbeLocalModelAvailability_DeduplicatesInflightProbe(t *testing.T) {
t.Fatal("deduplicated probe result = false, want true") t.Fatal("deduplicated probe result = false, want true")
} }
} }
if got := atomic.LoadInt32(&calls); got != 1 {
t.Fatalf("final probe calls = %d, want 1", got)
}
} }
func TestOllamaModelMatches_WithTagRequiresExactTag(t *testing.T) { func TestOllamaModelMatches_WithTagRequiresExactTag(t *testing.T) {