internal/registry/model_registry_safety_test.go
Rama main · réplica de lectura del repositorio original.
package registry
import (
"testing"
"time"
)
func TestGetModelInfoReturnsClone(t *testing.T) {
r := newTestModelRegistry()
r.RegisterClient("client-1", "gemini", []*ModelInfo{{
ID: "m1",
DisplayName: "Model One",
Thinking: &ThinkingSupport{Min: 1, Max: 2, Levels: []string{"low", "high"}},
}})
first := r.GetModelInfo("m1", "gemini")
if first == nil {
t.Fatal("expected model info")
}
first.DisplayName = "mutated"
first.Thinking.Levels[0] = "mutated"
second := r.GetModelInfo("m1", "gemini")
if second.DisplayName != "Model One" {
t.Fatalf("expected cloned display name, got %q", second.DisplayName)
}
if second.Thinking == nil || len(second.Thinking.Levels) == 0 || second.Thinking.Levels[0] != "low" {
t.Fatalf("expected cloned thinking levels, got %+v", second.Thinking)
}
}
func TestGetModelsForClientReturnsClones(t *testing.T) {
r := newTestModelRegistry()
r.RegisterClient("client-1", "gemini", []*ModelInfo{{
ID: "m1",
DisplayName: "Model One",
Thinking: &ThinkingSupport{Levels: []string{"low", "high"}},
}})
first := r.GetModelsForClient("client-1")
if len(first) != 1 || first[0] == nil {
t.Fatalf("expected one model, got %+v", first)
}
first[0].DisplayName = "mutated"
first[0].Thinking.Levels[0] = "mutated"
second := r.GetModelsForClient("client-1")
if len(second) != 1 || second[0] == nil {
t.Fatalf("expected one model on second fetch, got %+v", second)
}
if second[0].DisplayName != "Model One" {
t.Fatalf("expected cloned display name, got %q", second[0].DisplayName)
}
if second[0].Thinking == nil || len(second[0].Thinking.Levels) == 0 || second[0].Thinking.Levels[0] != "low" {
t.Fatalf("expected cloned thinking levels, got %+v", second[0].Thinking)
}
}
func TestGetAvailableModelsByProviderReturnsClones(t *testing.T) {
r := newTestModelRegistry()
r.RegisterClient("client-1", "gemini", []*ModelInfo{{
ID: "m1",
DisplayName: "Model One",
Thinking: &ThinkingSupport{Levels: []string{"low", "high"}},
}})
first := r.GetAvailableModelsByProvider("gemini")
if len(first) != 1 || first[0] == nil {
t.Fatalf("expected one model, got %+v", first)
}
first[0].DisplayName = "mutated"
first[0].Thinking.Levels[0] = "mutated"
second := r.GetAvailableModelsByProvider("gemini")
if len(second) != 1 || second[0] == nil {
t.Fatalf("expected one model on second fetch, got %+v", second)
}
if second[0].DisplayName != "Model One" {
t.Fatalf("expected cloned display name, got %q", second[0].DisplayName)
}
if second[0].Thinking == nil || len(second[0].Thinking.Levels) == 0 || second[0].Thinking.Levels[0] != "low" {
t.Fatalf("expected cloned thinking levels, got %+v", second[0].Thinking)
}
}
func TestCleanupExpiredQuotasInvalidatesAvailableModelsCache(t *testing.T) {
r := newTestModelRegistry()
r.RegisterClient("client-1", "openai", []*ModelInfo{{ID: "m1", Created: 1}})
r.SetModelQuotaExceeded("client-1", "m1")
if models := r.GetAvailableModels("openai"); len(models) != 1 {
t.Fatalf("expected cooldown model to remain listed before cleanup, got %d", len(models))
}
r.mutex.Lock()
quotaTime := time.Now().Add(-6 * time.Minute)
r.models["m1"].QuotaExceededClients["client-1"] = "aTime
r.mutex.Unlock()
r.CleanupExpiredQuotas()
if count := r.GetModelCount("m1"); count != 1 {
t.Fatalf("expected model count 1 after cleanup, got %d", count)
}
models := r.GetAvailableModels("openai")
if len(models) != 1 {
t.Fatalf("expected model to stay available after cleanup, got %d", len(models))
}
if got := models[0]["id"]; got != "m1" {
t.Fatalf("expected model id m1, got %v", got)
}
}
func TestGetAvailableModelsReturnsClonedSupportedParameters(t *testing.T) {
r := newTestModelRegistry()
r.RegisterClient("client-1", "openai", []*ModelInfo{{
ID: "m1",
DisplayName: "Model One",
SupportedParameters: []string{"temperature", "top_p"},
}})
first := r.GetAvailableModels("openai")
if len(first) != 1 {
t.Fatalf("expected one model, got %d", len(first))
}
params, ok := first[0]["supported_parameters"].([]string)
if !ok || len(params) != 2 {
t.Fatalf("expected supported_parameters slice, got %#v", first[0]["supported_parameters"])
}
params[0] = "mutated"
second := r.GetAvailableModels("openai")
params, ok = second[0]["supported_parameters"].([]string)
if !ok || len(params) != 2 || params[0] != "temperature" {
t.Fatalf("expected cloned supported_parameters, got %#v", second[0]["supported_parameters"])
}
}
func TestGetAvailableModelsIncludesMaxContextLengthOverride(t *testing.T) {
r := newTestModelRegistry()
const want = 1048576
r.RegisterClient("client-1", "openai", []*ModelInfo{{
ID: "deepseek-v4-flash",
ContextLength: want,
MaxContextLength: want,
}})
models := r.GetAvailableModels("openai")
if len(models) != 1 {
t.Fatalf("models length = %d, want 1", len(models))
}
if got := models[0]["context_length"]; got != want {
t.Fatalf("context_length = %#v, want %d", got, want)
}
if got := models[0]["max_context_length"]; got != want {
t.Fatalf("max_context_length = %#v, want %d", got, want)
}
}
func TestLookupModelInfoReturnsCloneForStaticDefinitions(t *testing.T) {
first := LookupModelInfo("claude-sonnet-4-6")
if first == nil || first.Thinking == nil || len(first.Thinking.Levels) == 0 {
t.Fatalf("expected static model with thinking levels, got %+v", first)
}
first.Thinking.Levels[0] = "mutated"
second := LookupModelInfo("claude-sonnet-4-6")
if second == nil || second.Thinking == nil || len(second.Thinking.Levels) == 0 || second.Thinking.Levels[0] == "mutated" {
t.Fatalf("expected static lookup clone, got %+v", second)
}
}
func TestLookupModelInfoIncludesClaudeSonnet5(t *testing.T) {
model := LookupModelInfo("claude-sonnet-5")
if model == nil {
t.Fatal("expected Claude Sonnet 5 static model")
}
if model.Type != "claude" {
t.Fatalf("Claude Sonnet 5 type = %q, want claude", model.Type)
}
if model.ContextLength != 1000000 {
t.Fatalf("Claude Sonnet 5 context length = %d, want 1000000", model.ContextLength)
}
if model.MaxCompletionTokens != 128000 {
t.Fatalf("Claude Sonnet 5 max completion tokens = %d, want 128000", model.MaxCompletionTokens)
}
if model.Thinking == nil || !model.Thinking.ZeroAllowed || !model.Thinking.DynamicAllowed || model.Thinking.Min != 0 || model.Thinking.Max != 0 {
t.Fatalf("expected Claude Sonnet 5 dynamic level-only thinking with zero allowed, got %+v", model.Thinking)
}
expectedLevels := []string{"low", "medium", "high", "xhigh", "max"}
if len(model.Thinking.Levels) != len(expectedLevels) {
t.Fatalf("Claude Sonnet 5 thinking levels = %+v, want %+v", model.Thinking.Levels, expectedLevels)
}
for i, level := range expectedLevels {
if model.Thinking.Levels[i] != level {
t.Fatalf("Claude Sonnet 5 thinking levels = %+v, want %+v", model.Thinking.Levels, expectedLevels)
}
}
}
func TestApplyClientModelCapabilities_UpdatesAllViews(t *testing.T) {
r := newTestModelRegistry()
r.RegisterClient("ag-client-1", "antigravity", []*ModelInfo{{
ID: "gemini-3.1-flash-lite",
SupportsWebSearch: false,
}})
epoch := r.ClientRegistrationEpoch("ag-client-1")
applied := r.ApplyClientModelCapabilities("ag-client-1", epoch, func(modelID string, info *ModelInfo) {
if modelID == "gemini-3.1-flash-lite" {
info.SupportsWebSearch = true
}
})
if !applied {
t.Fatal("expected capabilities to be applied")
}
// 1. Check GetModelsForClient
clientModels := r.GetModelsForClient("ag-client-1")
if len(clientModels) != 1 || !clientModels[0].SupportsWebSearch {
t.Fatalf("GetModelsForClient: expected SupportsWebSearch=true, got %+v", clientModels)
}
// 2. Check GetModelInfo (both with provider and global)
infoByProv := r.GetModelInfo("gemini-3.1-flash-lite", "antigravity")
if infoByProv == nil || !infoByProv.SupportsWebSearch {
t.Fatalf("GetModelInfo with provider: expected SupportsWebSearch=true, got %+v", infoByProv)
}
infoGlobal := r.GetModelInfo("gemini-3.1-flash-lite", "")
if infoGlobal == nil || !infoGlobal.SupportsWebSearch {
t.Fatalf("GetModelInfo global: expected SupportsWebSearch=true, got %+v", infoGlobal)
}
// 3. Check GetAvailableModelInfos
availInfos := r.GetAvailableModelInfos()
foundAvail := false
for _, m := range availInfos {
if m.ID == "gemini-3.1-flash-lite" && m.SupportsWebSearch {
foundAvail = true
break
}
}
if !foundAvail {
t.Fatal("GetAvailableModelInfos: expected gemini-3.1-flash-lite with SupportsWebSearch=true")
}
}
func TestMultiClientRegistration_PreservesProbedCapabilities(t *testing.T) {
r := newTestModelRegistry()
// 1. Client A registers model
r.RegisterClient("client-A", "antigravity", []*ModelInfo{{
ID: "gemini-3.1-flash-lite",
SupportsWebSearch: false,
}})
// A probes successfully -> SupportsWebSearch = true
epochA := r.ClientRegistrationEpoch("client-A")
r.ApplyClientModelCapabilities("client-A", epochA, func(modelID string, info *ModelInfo) {
if modelID == "gemini-3.1-flash-lite" {
info.SupportsWebSearch = true
}
})
// 2. Client B registers the same model with static info (SupportsWebSearch = false)
r.RegisterClient("client-B", "antigravity", []*ModelInfo{{
ID: "gemini-3.1-flash-lite",
SupportsWebSearch: false,
}})
// Check that global views preserve SupportsWebSearch=true because Client A is still registered and has the capability
info := r.GetModelInfo("gemini-3.1-flash-lite", "antigravity")
if info == nil || !info.SupportsWebSearch {
t.Fatalf("GetModelInfo: expected SupportsWebSearch=true to be preserved after client B registered, got %+v", info)
}
avail := r.GetAvailableModelInfos()
foundSearch := false
for _, m := range avail {
if m.ID == "gemini-3.1-flash-lite" && m.SupportsWebSearch {
foundSearch = true
break
}
}
if !foundSearch {
t.Fatal("GetAvailableModelInfos: expected SupportsWebSearch=true to be preserved after client B registered")
}
// 3. Client A unregisters -> Now neither client supports web search, so global view should drop it
r.UnregisterClient("client-A")
infoAfterA := r.GetModelInfo("gemini-3.1-flash-lite", "antigravity")
if infoAfterA == nil || infoAfterA.SupportsWebSearch {
t.Fatalf("GetModelInfo: expected SupportsWebSearch=false after client A unregistered, got %+v", infoAfterA)
}
}
func TestReRegisterClient_ClearsStaleProbedCapabilitiesWhenNoOtherClientSupports(t *testing.T) {
r := newTestModelRegistry()
// 1. Client A registers model and probes successfully (SupportsWebSearch=true)
r.RegisterClient("client-A", "antigravity", []*ModelInfo{{
ID: "gemini-3.1-flash-lite",
SupportsWebSearch: false,
}})
epochA := r.ClientRegistrationEpoch("client-A")
r.ApplyClientModelCapabilities("client-A", epochA, func(modelID string, info *ModelInfo) {
if modelID == "gemini-3.1-flash-lite" {
info.SupportsWebSearch = true
}
})
// 2. Client B registers the same model without web search support
r.RegisterClient("client-B", "antigravity", []*ModelInfo{{
ID: "gemini-3.1-flash-lite",
SupportsWebSearch: false,
}})
// 3. Client A now re-registers with new model definitions where SupportsWebSearch=false
// (e.g. config reload or model catalog update)
r.RegisterClient("client-A", "antigravity", []*ModelInfo{{
ID: "gemini-3.1-flash-lite",
SupportsWebSearch: false,
}})
// Global view must NOT keep client A's stale pre-re-registration capability snapshot!
// Since neither A nor B now has SupportsWebSearch=true, the global and provider view must be false.
info := r.GetModelInfo("gemini-3.1-flash-lite", "antigravity")
if info == nil || info.SupportsWebSearch {
t.Fatalf("GetModelInfo: expected SupportsWebSearch=false after client A re-registered with false, got %+v", info)
}
avail := r.GetAvailableModelInfos()
for _, m := range avail {
if m.ID == "gemini-3.1-flash-lite" && m.SupportsWebSearch {
t.Fatalf("GetAvailableModelInfos: expected SupportsWebSearch=false on gemini-3.1-flash-lite, got %+v", m)
}
}
clientAModels := r.GetModelsForClient("client-A")
if len(clientAModels) != 1 || clientAModels[0].SupportsWebSearch {
t.Fatalf("GetModelsForClient(A): expected SupportsWebSearch=false, got %+v", clientAModels)
}
}