internal/registry/web_search_index_test.go
Rama main · réplica de lectura del repositorio original.
package registry
import (
"fmt"
"sync"
"testing"
)
// referenceHasClientSupportingWebSearch recomputes the web-search aggregate with
// the original full scan. The index must agree with it in every state the
// registry can reach, so tests compare the fast path against this oracle
// instead of only asserting hand-picked expectations.
func referenceHasClientSupportingWebSearch(r *ModelRegistry, modelID, provider, excludeClientID string) bool {
for clientID, infos := range r.clientModelInfos {
if clientID == excludeClientID {
continue
}
if provider != "" && r.clientProviders[clientID] != provider {
continue
}
info := infos[modelID]
if info != nil && info.SupportsWebSearch {
return true
}
}
return false
}
// assertWebSearchIndexMatchesScan verifies the incremental index against the
// full scan for every client/model/provider combination currently registered.
func assertWebSearchIndexMatchesScan(t *testing.T, r *ModelRegistry, context string) {
t.Helper()
r.mutex.RLock()
defer r.mutex.RUnlock()
providers := make(map[string]struct{})
modelIDs := make(map[string]struct{})
for clientID, infos := range r.clientModelInfos {
providers[r.clientProviders[clientID]] = struct{}{}
for modelID := range infos {
modelIDs[modelID] = struct{}{}
}
}
for modelID := range modelIDs {
for provider := range providers {
for _, exclude := range []string{"", "__none__"} {
want := referenceHasClientSupportingWebSearch(r, modelID, provider, exclude)
got := r.hasClientSupportingWebSearchLocked(modelID, provider, exclude)
if got != want {
t.Fatalf("%s: model=%q provider=%q exclude=%q: index=%v scan=%v", context, modelID, provider, exclude, got, want)
}
}
}
// Cross-check per-client exclusion, which the index answers by
// subtracting the excluded client instead of skipping it.
for clientID := range r.clientModelInfos {
want := referenceHasClientSupportingWebSearch(r, modelID, "", clientID)
got := r.hasClientSupportingWebSearchLocked(modelID, "", clientID)
if got != want {
t.Fatalf("%s: model=%q exclude=%q: index=%v scan=%v", context, modelID, clientID, got, want)
}
}
// The published Info flags are what callers actually read, so verify
// them against the scan too: a self-consistent index would still be a
// bug if the stored aggregate had been computed from a stale one.
registration := r.models[modelID]
if registration == nil || registration.Info == nil {
continue
}
if want, got := referenceHasClientSupportingWebSearch(r, modelID, "", ""), registration.Info.SupportsWebSearch; got != want {
t.Fatalf("%s: model=%q Info.SupportsWebSearch=%v scan=%v", context, modelID, got, want)
}
for provider, info := range registration.InfoByProvider {
if info == nil {
continue
}
if want, got := referenceHasClientSupportingWebSearch(r, modelID, provider, ""), info.SupportsWebSearch; got != want {
t.Fatalf("%s: model=%q provider=%q InfoByProvider.SupportsWebSearch=%v scan=%v", context, modelID, provider, got, want)
}
}
}
}
func newWebSearchIndexRegistry() *ModelRegistry {
reg := newTestModelRegistry()
reg.RegisterClient("a", "codex", []*ModelInfo{
{ID: "solo", SupportsWebSearch: true},
{ID: "shared", SupportsWebSearch: true},
})
reg.RegisterClient("b", "codex", []*ModelInfo{
{ID: "shared", SupportsWebSearch: true},
})
reg.RegisterClient("c", "gemini", []*ModelInfo{
{ID: "shared", SupportsWebSearch: true},
{ID: "off"},
})
return reg
}
func TestWebSearchIndexMatchesScanAcrossMutations(t *testing.T) {
reg := newWebSearchIndexRegistry()
assertWebSearchIndexMatchesScan(t, reg, "initial")
// Re-registration may add or drop contributions; both paths must retract
// the old counts before writing the new ones.
reg.RegisterClient("b", "codex", []*ModelInfo{{ID: "shared"}})
assertWebSearchIndexMatchesScan(t, reg, "re-register drops web search")
reg.RegisterClient("b", "codex", []*ModelInfo{{ID: "shared", SupportsWebSearch: true}})
assertWebSearchIndexMatchesScan(t, reg, "re-register restores web search")
reg.RegisterClient("d", "codex", []*ModelInfo{{ID: "solo", SupportsWebSearch: true}})
assertWebSearchIndexMatchesScan(t, reg, "second supporter of solo")
reg.ApplyClientModelCapabilities("d", reg.ClientRegistrationEpoch("d"), func(_ string, info *ModelInfo) {
info.SupportsWebSearch = false
})
assertWebSearchIndexMatchesScan(t, reg, "capabilities clear web search")
reg.ApplyClientModelCapabilities("a", reg.ClientRegistrationEpoch("a"), func(_ string, info *ModelInfo) {
info.SupportsWebSearch = true
})
assertWebSearchIndexMatchesScan(t, reg, "capabilities set web search")
reg.UnregisterClient("a")
assertWebSearchIndexMatchesScan(t, reg, "unregister sole supporter")
reg.UnregisterClient("d")
reg.UnregisterClient("b")
reg.UnregisterClient("c")
assertWebSearchIndexMatchesScan(t, reg, "unregister all")
}
func TestWebSearchIndexDropsOrphanCountsOnPureAddition(t *testing.T) {
// The pure-addition path assumes the client has no prior bookkeeping, but
// clientModelInfos can outlive clientModels (any code path that rewrites one
// without the other). A stale entry still contributes to the index, so
// registering the client again must retract it instead of counting it twice.
reg := newTestModelRegistry()
reg.mutex.Lock()
reg.clientModelInfos["ghost"] = map[string]*ModelInfo{
"orphaned": {ID: "orphaned", SupportsWebSearch: true},
}
reg.clientProviders["ghost"] = "codex"
// The orphan is seeded in the shape any buggy path would leave behind: the
// index still agrees with clientModelInfos, but clientModels no longer
// carries the client, so the next registration takes the pure-addition path.
reg.incrementWebSearchCountLocked("orphaned", "codex")
reg.mutex.Unlock()
assertWebSearchIndexMatchesScan(t, reg, "orphan seeded")
// Registering the same client with a catalog that advertises no web search
// must leave a clean index, not the orphaned contribution plus nothing.
reg.RegisterClient("ghost", "codex", []*ModelInfo{{ID: "other"}})
assertWebSearchIndexMatchesScan(t, reg, "orphan replaced by pure addition")
if reg.hasClientSupportingWebSearchLocked("orphaned", "", "") {
t.Fatalf("index kept the orphaned web-search contribution after re-registration")
}
if reg.hasClientSupportingWebSearchLocked("orphaned", "codex", "") {
t.Fatalf("provider index kept the orphaned web-search contribution after re-registration")
}
}
func assertPublishedWebSearchSupport(t *testing.T, reg *ModelRegistry, modelID, provider string, want bool) {
t.Helper()
for _, view := range []string{"", provider} {
info := reg.GetModelInfo(modelID, view)
if info == nil {
t.Errorf("GetModelInfo(%q, %q) returned nil", modelID, view)
continue
}
if info.SupportsWebSearch != want {
t.Errorf("GetModelInfo(%q, %q).SupportsWebSearch = %v, want %v", modelID, view, info.SupportsWebSearch, want)
}
}
}
func TestWebSearchIndexEmptyClientExclusion(t *testing.T) {
for _, clientID := range []string{"", "supporter"} {
t.Run(fmt.Sprintf("client=%q", clientID), func(t *testing.T) {
reg := newTestModelRegistry()
reg.RegisterClient(clientID, "antigravity", []*ModelInfo{{ID: "m", SupportsWebSearch: true}})
reg.RegisterClient("observer", "antigravity", []*ModelInfo{{ID: "m"}})
reg.mutex.RLock()
defer reg.mutex.RUnlock()
for _, provider := range []string{"", "antigravity", "gemini"} {
for _, exclude := range []string{"", "supporter", "observer", "missing"} {
want := provider != "gemini" && clientID != exclude
if got := referenceHasClientSupportingWebSearch(reg, "m", provider, exclude); got != want {
t.Errorf("provider=%q exclude=%q: scan=%v, want %v", provider, exclude, got, want)
}
if got := reg.hasClientSupportingWebSearchLocked("m", provider, exclude); got != want {
t.Errorf("provider=%q exclude=%q: index=%v, want %v", provider, exclude, got, want)
}
}
}
})
}
}
func TestWebSearchIndexPublishedViewsAfterSupporterUnregister(t *testing.T) {
for _, clientID := range []string{"", "supporter"} {
t.Run(fmt.Sprintf("client=%q", clientID), func(t *testing.T) {
reg := newTestModelRegistry()
reg.RegisterClient(clientID, "antigravity", []*ModelInfo{{ID: "m", SupportsWebSearch: true}})
reg.RegisterClient("observer", "antigravity", []*ModelInfo{{ID: "m"}})
assertPublishedWebSearchSupport(t, reg, "m", "antigravity", true)
reg.UnregisterClient(clientID)
assertPublishedWebSearchSupport(t, reg, "m", "antigravity", false)
assertWebSearchIndexMatchesScan(t, reg, "supporter unregistered")
})
}
}
func TestWebSearchIndexPreservesRealSupporterOnOrphanAddition(t *testing.T) {
reg := newTestModelRegistry()
reg.RegisterClient("real", "antigravity", []*ModelInfo{{ID: "m", SupportsWebSearch: true}})
reg.RegisterClient("ghost", "antigravity", []*ModelInfo{{ID: "m", SupportsWebSearch: true}})
// Artificially leave orphaned metadata and index entries. Normal registry
// operations keep these maps together; this exercises defensive cleanup.
reg.mutex.Lock()
delete(reg.clientModels, "ghost")
reg.mutex.Unlock()
assertWebSearchIndexMatchesScan(t, reg, "orphan seeded beside real supporter")
reg.RegisterClient("ghost", "antigravity", []*ModelInfo{{ID: "m"}})
assertPublishedWebSearchSupport(t, reg, "m", "antigravity", true)
assertWebSearchIndexMatchesScan(t, reg, "orphan replaced beside real supporter")
}
func TestWebSearchIndexRefreshesOrphanedModelsOnPureAddition(t *testing.T) {
for _, realSupportsWebSearch := range []bool{false, true} {
t.Run(fmt.Sprintf("real_supports_web_search=%v", realSupportsWebSearch), func(t *testing.T) {
reg := newTestModelRegistry()
reg.RegisterClient("real", "antigravity", []*ModelInfo{{ID: "m", SupportsWebSearch: realSupportsWebSearch}})
reg.RegisterClient("ghost", "antigravity", []*ModelInfo{{ID: "m", SupportsWebSearch: true}})
// Artificially orphan the old catalog to exercise defensive cleanup.
reg.mutex.Lock()
delete(reg.clientModels, "ghost")
reg.mutex.Unlock()
assertPublishedWebSearchSupport(t, reg, "m", "antigravity", true)
assertWebSearchIndexMatchesScan(t, reg, "orphan seeded before catalog replacement")
// The old model is absent from the new catalog, so additions cannot
// refresh its published capability as a side effect.
reg.RegisterClient("ghost", "antigravity", []*ModelInfo{{ID: "n"}})
assertPublishedWebSearchSupport(t, reg, "m", "antigravity", realSupportsWebSearch)
assertPublishedWebSearchSupport(t, reg, "n", "antigravity", false)
assertWebSearchIndexMatchesScan(t, reg, "orphan replaced with a different catalog")
})
}
}
func TestWebSearchIndexCleansOrphanOnUnregister(t *testing.T) {
for _, realSupportsWebSearch := range []bool{false, true} {
t.Run(fmt.Sprintf("real_supports_web_search=%v", realSupportsWebSearch), func(t *testing.T) {
reg := newTestModelRegistry()
reg.RegisterClient("real", "antigravity", []*ModelInfo{{ID: "m", SupportsWebSearch: realSupportsWebSearch}})
reg.RegisterClient("ghost", "antigravity", []*ModelInfo{{ID: "m", SupportsWebSearch: true}})
// This inconsistent state is constructed only to test orphan cleanup.
reg.mutex.Lock()
delete(reg.clientModels, "ghost")
reg.mutex.Unlock()
reg.UnregisterClient("ghost")
reg.mutex.RLock()
if _, exists := reg.clientModelInfos["ghost"]; exists {
t.Error("unregister left orphaned client model metadata")
}
reg.mutex.RUnlock()
assertPublishedWebSearchSupport(t, reg, "m", "antigravity", realSupportsWebSearch)
// Repeated cleanup must not retract the real client's contribution.
reg.UnregisterClient("ghost")
reg.mutex.RLock()
wantCount := 0
if realSupportsWebSearch {
wantCount = 1
}
if got := reg.webSearchClientCounts["m"]; got != wantCount {
t.Errorf("global count after repeated orphan cleanup = %d, want %d", got, wantCount)
}
if got := reg.webSearchProviderCounts["m"]["antigravity"]; got != wantCount {
t.Errorf("provider count after repeated orphan cleanup = %d, want %d", got, wantCount)
}
reg.mutex.RUnlock()
assertWebSearchIndexMatchesScan(t, reg, "orphan unregistered twice")
reg.RegisterClient("ghost", "antigravity", []*ModelInfo{{ID: "m"}})
assertPublishedWebSearchSupport(t, reg, "m", "antigravity", realSupportsWebSearch)
assertWebSearchIndexMatchesScan(t, reg, "orphan re-registered after cleanup")
})
}
}
func TestWebSearchIndexTracksSoleSupporterAfterUnregister(t *testing.T) {
reg := newWebSearchIndexRegistry()
if !reg.hasClientSupportingWebSearchLocked("solo", "", "") {
t.Fatalf("expected solo to be supported before its only supporter left")
}
reg.UnregisterClient("a")
if reg.hasClientSupportingWebSearchLocked("solo", "", "") {
t.Fatalf("sole supporter left, but the aggregate still reports web search")
}
// The published flag must follow the index, otherwise callers keep
// advertising web search for a model no client can serve.
if reg.models["solo"] != nil && reg.models["solo"].Info != nil && reg.models["solo"].Info.SupportsWebSearch {
t.Fatalf("published Info.SupportsWebSearch stayed true after the sole supporter left")
}
}
func TestWebSearchIndexToleratesNilIndexMaps(t *testing.T) {
// The index maps are nil in a partially constructed registry (tests and any
// future literal). The helpers must lazily initialize instead of panicking
// on assignment to a nil map.
reg := &ModelRegistry{mutex: &sync.RWMutex{}}
if reg.hasClientSupportingWebSearchLocked("m", "", "") {
t.Fatalf("registry without index must not report web-search support")
}
reg.incrementWebSearchCountLocked("m", "codex")
if !reg.hasClientSupportingWebSearchLocked("m", "codex", "") {
t.Fatalf("expected the increment to be visible after lazy initialization")
}
reg.decrementWebSearchCountLocked("m", "codex")
if reg.hasClientSupportingWebSearchLocked("m", "codex", "") {
t.Fatalf("expected the decrement to clear the count")
}
}
func TestWebSearchIndexProviderChangeRetractsOldProvider(t *testing.T) {
// Re-registering the same client under a different provider must retract
// the old provider's contribution. The old-provider aggregate is recomputed
// while this client's counts are still indexed, so a wrong drop/add ordering
// leaves a phantom count for the provider the client no longer serves.
reg := newTestModelRegistry()
reg.RegisterClient("x", "codex", []*ModelInfo{
{ID: "m", SupportsWebSearch: true},
})
if !reg.hasClientSupportingWebSearchLocked("m", "codex", "") {
t.Fatalf("expected codex to support web search before the provider change")
}
reg.RegisterClient("x", "gemini", []*ModelInfo{
{ID: "m", SupportsWebSearch: true},
})
if reg.hasClientSupportingWebSearchLocked("m", "codex", "") {
t.Fatalf("codex must no longer report web search after the client moved to gemini")
}
if !reg.hasClientSupportingWebSearchLocked("m", "gemini", "") {
t.Fatalf("gemini must report web search after the provider change")
}
if got := reg.webSearchProviderCounts["m"]["codex"]; got != 0 {
t.Fatalf("stale codex provider count = %d, want 0", got)
}
if got := reg.webSearchProviderCounts["m"]["gemini"]; got != 1 {
t.Fatalf("gemini provider count = %d, want 1", got)
}
assertWebSearchIndexMatchesScan(t, reg, "provider change")
}
// BenchmarkRegisterClientWebSearchAggregate measures whole-registry
// registration, the hot path that blocked startup: every registered model
// recomputes the web-search aggregate while the registry write lock is held.
// Sizes are reported separately because the pre-fix implementation is quadratic
// in the client count and only the growth curve exposes that.
func BenchmarkRegisterClientWebSearchAggregate(b *testing.B) {
const modelsPerClient = 20
infos := make([]*ModelInfo, 0, modelsPerClient)
for i := 0; i < modelsPerClient; i++ {
// Only the first model advertises web search, and it must be set on
// SupportsWebSearch: the aggregate is keyed off that field, not off
// NativeCapabilities.WebSearch. Leaving it false makes every lookup miss
// the client set, which silently benchmarks a different code path.
infos = append(infos, &ModelInfo{
ID: fmt.Sprintf("model-%d", i),
SupportsWebSearch: i == 0,
})
}
for _, clients := range []int{250, 500, 1000} {
b.Run(fmt.Sprintf("clients=%d", clients), func(b *testing.B) {
b.ResetTimer()
for i := 0; i < b.N; i++ {
reg := newTestModelRegistry()
for c := 0; c < clients; c++ {
reg.RegisterClient(fmt.Sprintf("client-%d", c), "codex", infos)
}
}
})
}
}