237 lines
7.2 KiB
Go
237 lines
7.2 KiB
Go
package ai
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"strings"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
const testFingerprintSecret = "model-cache-test-fingerprint-secret-32-bytes"
|
|
|
|
type fakeModelsCache struct {
|
|
mu sync.Mutex
|
|
values map[ModelsCacheKey]ModelsCacheValue
|
|
locks map[ModelsCacheKey]string
|
|
seen []ModelsCacheKey
|
|
getErr error
|
|
lockErr error
|
|
}
|
|
|
|
func newFakeModelsCache() *fakeModelsCache {
|
|
return &fakeModelsCache{values: make(map[ModelsCacheKey]ModelsCacheValue), locks: make(map[ModelsCacheKey]string)}
|
|
}
|
|
|
|
func (c *fakeModelsCache) Get(_ context.Context, key ModelsCacheKey) (ModelsCacheValue, bool, error) {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
c.seen = append(c.seen, key)
|
|
if c.getErr != nil {
|
|
return ModelsCacheValue{}, false, c.getErr
|
|
}
|
|
value, ok := c.values[key]
|
|
value.Models = cloneModels(value.Models)
|
|
return value, ok, nil
|
|
}
|
|
|
|
func (c *fakeModelsCache) Set(_ context.Context, key ModelsCacheKey, value ModelsCacheValue, _ time.Duration) error {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
value.Models = cloneModels(value.Models)
|
|
c.values[key] = value
|
|
return nil
|
|
}
|
|
|
|
func (c *fakeModelsCache) AcquireRefresh(_ context.Context, key ModelsCacheKey, token string, _ time.Duration) (bool, error) {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
if c.lockErr != nil {
|
|
return false, c.lockErr
|
|
}
|
|
if _, exists := c.locks[key]; exists {
|
|
return false, nil
|
|
}
|
|
c.locks[key] = token
|
|
return true, nil
|
|
}
|
|
|
|
func (c *fakeModelsCache) ReleaseRefresh(_ context.Context, key ModelsCacheKey, token string) error {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
if c.locks[key] == token {
|
|
delete(c.locks, key)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
type modelClient struct {
|
|
calls atomic.Int32
|
|
models []string
|
|
err error
|
|
started chan struct{}
|
|
release chan struct{}
|
|
}
|
|
|
|
func (c *modelClient) Complete(context.Context, string, string, string) (string, error) {
|
|
return "", nil
|
|
}
|
|
|
|
func (c *modelClient) CompleteStream(context.Context, string, string, string, func(string) error) (string, error) {
|
|
return "", nil
|
|
}
|
|
|
|
func (c *modelClient) ListModels(context.Context, string) ([]string, error) {
|
|
c.calls.Add(1)
|
|
if c.started != nil {
|
|
select {
|
|
case c.started <- struct{}{}:
|
|
default:
|
|
}
|
|
}
|
|
if c.release != nil {
|
|
<-c.release
|
|
}
|
|
return cloneModels(c.models), c.err
|
|
}
|
|
|
|
func TestRegistryListModelsCacheHit(t *testing.T) {
|
|
cache := newFakeModelsCache()
|
|
client := &modelClient{models: []string{"provider-model"}}
|
|
registry := NewRegistryWithClients(cache, testFingerprintSecret, map[string]Client{ProviderXAI: client})
|
|
cache.values[registry.cacheKey(ProviderXAI, "api-key")] = ModelsCacheValue{
|
|
Models: []string{"cached-model"}, FetchedAt: time.Now().UTC().UnixNano(),
|
|
}
|
|
|
|
models, fromCache, err := registry.ListModelsCached(context.Background(), ProviderXAI, "api-key")
|
|
if err != nil || !fromCache || fmt.Sprint(models) != "[cached-model]" {
|
|
t.Fatalf("unexpected result models=%v fromCache=%v err=%v", models, fromCache, err)
|
|
}
|
|
if client.calls.Load() != 0 {
|
|
t.Fatalf("provider called on cache hit")
|
|
}
|
|
}
|
|
|
|
func TestRegistrySeparatesProvidersAndCredentials(t *testing.T) {
|
|
cache := newFakeModelsCache()
|
|
xai := &modelClient{models: []string{"xai-model"}}
|
|
opencode := &modelClient{models: []string{"opencode-model"}}
|
|
registry := NewRegistryWithClients(cache, testFingerprintSecret, map[string]Client{
|
|
ProviderXAI: xai, ProviderOpenCodeGo: opencode,
|
|
})
|
|
|
|
for _, call := range []struct{ provider, key string }{
|
|
{ProviderXAI, "key-a"}, {ProviderXAI, "key-b"}, {ProviderOpenCodeGo, "key-a"},
|
|
} {
|
|
if _, _, err := registry.ListModelsCached(context.Background(), call.provider, call.key); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
|
|
cache.mu.Lock()
|
|
entries := len(cache.values)
|
|
cache.mu.Unlock()
|
|
if entries != 3 || xai.calls.Load() != 2 || opencode.calls.Load() != 1 {
|
|
t.Fatalf("separation failed entries=%d xai=%d opencode=%d", entries, xai.calls.Load(), opencode.calls.Load())
|
|
}
|
|
}
|
|
|
|
func TestRegistryReturnsStaleOnProviderFailure(t *testing.T) {
|
|
cache := newFakeModelsCache()
|
|
client := &modelClient{err: errors.New("provider unavailable")}
|
|
registry := NewRegistryWithClients(cache, testFingerprintSecret, map[string]Client{ProviderXAI: client})
|
|
cache.values[registry.cacheKey(ProviderXAI, "api-key")] = ModelsCacheValue{
|
|
Models: []string{"stale-model"}, FetchedAt: time.Now().UTC().Add(-10 * time.Minute).UnixNano(),
|
|
}
|
|
|
|
models, fromCache, err := registry.ListModelsCached(context.Background(), ProviderXAI, "api-key")
|
|
if err != nil || !fromCache || fmt.Sprint(models) != "[stale-model]" {
|
|
t.Fatalf("unexpected stale result models=%v fromCache=%v err=%v", models, fromCache, err)
|
|
}
|
|
}
|
|
|
|
func TestRegistryRedisFailureFallsOpenToProvider(t *testing.T) {
|
|
cache := newFakeModelsCache()
|
|
cache.getErr = errors.New("redis read failed")
|
|
cache.lockErr = errors.New("redis lock failed")
|
|
client := &modelClient{models: []string{"provider-model"}}
|
|
registry := NewRegistryWithClients(cache, testFingerprintSecret, map[string]Client{ProviderXAI: client})
|
|
|
|
models, fromCache, err := registry.ListModelsCached(context.Background(), ProviderXAI, "api-key")
|
|
if err != nil || fromCache || fmt.Sprint(models) != "[provider-model]" {
|
|
t.Fatalf("unexpected fail-open result models=%v fromCache=%v err=%v", models, fromCache, err)
|
|
}
|
|
if client.calls.Load() != 1 {
|
|
t.Fatalf("provider calls=%d, want 1", client.calls.Load())
|
|
}
|
|
}
|
|
|
|
func TestRegistryCoalescesConcurrentRefresh(t *testing.T) {
|
|
cache := newFakeModelsCache()
|
|
client := &modelClient{
|
|
models: []string{"fresh-model"}, started: make(chan struct{}, 1), release: make(chan struct{}),
|
|
}
|
|
registry := NewRegistryWithClients(cache, testFingerprintSecret, map[string]Client{ProviderXAI: client})
|
|
type result struct {
|
|
models []string
|
|
err error
|
|
}
|
|
results := make(chan result, 2)
|
|
call := func() {
|
|
models, _, err := registry.ListModelsCached(context.Background(), ProviderXAI, "api-key")
|
|
results <- result{models: models, err: err}
|
|
}
|
|
go call()
|
|
select {
|
|
case <-client.started:
|
|
case <-time.After(time.Second):
|
|
t.Fatal("provider did not start")
|
|
}
|
|
go call()
|
|
time.Sleep(50 * time.Millisecond)
|
|
close(client.release)
|
|
|
|
for range 2 {
|
|
select {
|
|
case got := <-results:
|
|
if got.err != nil || fmt.Sprint(got.models) != "[fresh-model]" {
|
|
t.Fatalf("unexpected result models=%v err=%v", got.models, got.err)
|
|
}
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("concurrent refresh timed out")
|
|
}
|
|
}
|
|
if client.calls.Load() != 1 {
|
|
t.Fatalf("provider calls=%d, want 1", client.calls.Load())
|
|
}
|
|
}
|
|
|
|
func TestRegistryCacheKeyDoesNotExposeAPIKey(t *testing.T) {
|
|
cache := newFakeModelsCache()
|
|
client := &modelClient{models: []string{"model"}}
|
|
registry := NewRegistryWithClients(cache, testFingerprintSecret, map[string]Client{ProviderXAI: client})
|
|
apiKey := "sk-plain-secret-must-never-appear"
|
|
if _, _, err := registry.ListModelsCached(context.Background(), ProviderXAI, apiKey); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
cache.mu.Lock()
|
|
seen := fmt.Sprintf("%+v %+v", cache.seen, cache.values)
|
|
var fingerprint string
|
|
for key := range cache.values {
|
|
fingerprint = key.CredentialFingerprint
|
|
}
|
|
cache.mu.Unlock()
|
|
if strings.Contains(seen, apiKey) {
|
|
t.Fatalf("cache metadata exposed API key: %s", seen)
|
|
}
|
|
if len(fingerprint) != sha256HexLength || strings.Contains(fingerprint, apiKey) {
|
|
t.Fatalf("invalid credential fingerprint %q", fingerprint)
|
|
}
|
|
}
|
|
|
|
const sha256HexLength = 64
|