From 8853dbd6eb61d4c07eb8bfae989c0bd097029a64 Mon Sep 17 00:00:00 2001 From: tbphp Date: Fri, 24 Jul 2026 11:55:09 +0800 Subject: [PATCH] =?UTF-8?q?feat(runtime):=20=E5=A2=9E=E5=8A=A0=E5=81=A5?= =?UTF-8?q?=E5=BA=B7=E7=8A=B6=E6=80=81=E4=B8=8E=E8=B7=AF=E7=94=B1=E6=A3=80?= =?UTF-8?q?=E6=9F=A5?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/container/container.go | 3 + internal/container/container_test.go | 116 +++ internal/control/health.go | 245 ++++++ internal/control/health_test.go | 336 ++++++++ internal/control/route_inspect.go | 178 +++++ internal/control/route_inspect_test.go | 769 +++++++++++++++++++ internal/control/runtime_observation.go | 53 ++ internal/control/runtime_observation_test.go | 114 +++ internal/control/server.go | 2 + internal/control/server_test.go | 5 + internal/control/service.go | 12 +- internal/control/test_helpers_test.go | 40 +- internal/gateway/dependency_test.go | 1 + internal/health/stats.go | 1 + internal/health/stats_test.go | 13 + internal/scheduler/dependency_test.go | 37 + internal/scheduler/inspect.go | 290 +++++++ internal/scheduler/inspect_test.go | 393 ++++++++++ internal/scheduler/scheduler.go | 76 +- internal/scheduler/scheduler_test.go | 10 +- internal/state/loader/loader_test.go | 19 + internal/state/registry.go | 7 +- internal/state/registry_test.go | 59 ++ internal/state/runtime_view.go | 74 ++ internal/state/snapshot.go | 132 ++-- internal/state/snapshot_test.go | 202 ++++- internal/state/weights_test.go | 82 ++ 27 files changed, 3152 insertions(+), 117 deletions(-) create mode 100644 internal/control/health.go create mode 100644 internal/control/health_test.go create mode 100644 internal/control/route_inspect.go create mode 100644 internal/control/route_inspect_test.go create mode 100644 internal/control/runtime_observation.go create mode 100644 internal/control/runtime_observation_test.go create mode 100644 internal/scheduler/dependency_test.go create mode 100644 internal/scheduler/inspect.go create mode 100644 internal/scheduler/inspect_test.go create mode 100644 internal/state/runtime_view.go diff --git a/internal/container/container.go b/internal/container/container.go index 47324a3c..8a74bf4a 100644 --- a/internal/container/container.go +++ b/internal/container/container.go @@ -56,6 +56,9 @@ func BuildContainer() (*dig.Container, error) { func(service *requestlog.Service) control.RequestLogReader { return service }, + func(service *requestlog.Service) control.RequestLogStatsReader { + return service + }, func(service *requestlog.Service) control.RequestLogCleaner { return service }, diff --git a/internal/container/container_test.go b/internal/container/container_test.go index ec99ba72..ada565de 100644 --- a/internal/container/container_test.go +++ b/internal/container/container_test.go @@ -311,6 +311,7 @@ func TestBuildContainerWiresRequestLogIntoEveryConsumer(t *testing.T) { service *requestlog.Service, sink telemetry.RequestLogSink, reader control.RequestLogReader, + statsReader control.RequestLogStatsReader, cleaner control.RequestLogCleaner, lifecycle app.RequestLogRuntime, limiter *ratelimit.AccessKeyRPM, @@ -330,6 +331,7 @@ func TestBuildContainerWiresRequestLogIntoEveryConsumer(t *testing.T) { for name, adapter := range map[string]any{ "gateway sink": sink, "control reader": reader, + "control stats": statsReader, "control cleaner": cleaner, "app lifecycle": lifecycle, } { @@ -517,6 +519,120 @@ func TestBuildContainerUsesSingletonStatsStore(t *testing.T) { } } +func TestBuildContainerWiresRuntimeReadConsumersToSingletons(t *testing.T) { + t.Setenv("AUTH_KEY", "test-auth-key") + t.Setenv("DATA_DIR", t.TempDir()) + t.Setenv("DATABASE_DSN", ":memory:") + t.Setenv("ENCRYPTION_KEY", "test-master-key-long") + + dependencyContainer, err := BuildContainer() + if err != nil { + t.Fatalf("BuildContainer() error = %v", err) + } + err = dependencyContainer.Invoke(func( + requestLogService *requestlog.Service, + requestLogStats control.RequestLogStatsReader, + statsStore *health.StatsStore, + gatewayHandler *gateway.Handler, + controlService *control.Service, + db *gorm.DB, + ) { + t.Cleanup(func() { + sqlDB, dbErr := db.DB() + if dbErr == nil { + _ = sqlDB.Close() + } + }) + if requestLogStats != requestLogService { + t.Fatalf( + "RequestLogStatsReader = %T, want singleton %p", + requestLogStats, + requestLogService, + ) + } + if statsStore == nil || gatewayHandler == nil || controlService == nil { + t.Fatalf( + "runtime read consumers unresolved: stats=%p gateway=%p control=%p", + statsStore, + gatewayHandler, + controlService, + ) + } + }) + if err != nil { + t.Fatalf("resolve runtime read consumers: %v", err) + } +} + +func TestContainerHealthEndpointReadsSharedStatsStore(t *testing.T) { + t.Setenv("AUTH_KEY", "test-auth-key") + t.Setenv("DATA_DIR", t.TempDir()) + t.Setenv("DATABASE_DSN", ":memory:") + t.Setenv("ENCRYPTION_KEY", "test-master-key-long") + if err := i18n.Init(); err != nil { + t.Fatalf("i18n.Init() error = %v", err) + } + + dependencyContainer, err := BuildContainer() + if err != nil { + t.Fatalf("BuildContainer() error = %v", err) + } + err = dependencyContainer.Invoke(func( + engine *gin.Engine, + manager *state.Manager, + registry *state.KeyRegistry, + stats *health.StatsStore, + db *gorm.DB, + ) { + t.Cleanup(func() { + sqlDB, dbErr := db.DB() + if dbErr == nil { + _ = sqlDB.Close() + } + }) + if _, publishErr := manager.Publish(state.CompileInput{Groups: []state.GroupConfig{{ + ID: 1, Name: "shared", Protocols: []protocol.Protocol{protocol.OpenAI}, + Models: []state.ModelConfig{{ID: "model"}}, Enabled: true, + }}}); publishErr != nil { + t.Fatalf("Publish() error = %v", publishErr) + } + if replaceErr := registry.Replace([]state.KeyEntry{{ + ID: 1, GroupID: 1, Status: state.KeyStatusActive, + Blacklisted: true, EncryptedValue: "cipher", + }}); replaceErr != nil { + t.Fatalf("Replace() error = %v", replaceErr) + } + stats.Record(1, false, time.Now()) + + request := httptest.NewRequest(http.MethodGet, "/api/health", nil) + request.Header.Set("Authorization", "Bearer test-auth-key") + recorder := httptest.NewRecorder() + engine.ServeHTTP(recorder, request) + if recorder.Code != http.StatusOK { + t.Fatalf("response = %d %s", recorder.Code, recorder.Body.String()) + } + var envelope struct { + Data struct { + BlacklistedKeys []struct { + KeyID uint `json:"key_id"` + RecentFailureCount uint64 `json:"recent_failure_count"` + } `json:"blacklisted_keys"` + } `json:"data"` + } + if decodeErr := json.Unmarshal(recorder.Body.Bytes(), &envelope); decodeErr != nil { + t.Fatalf("decode response: %v", decodeErr) + } + if len(envelope.Data.BlacklistedKeys) != 1 || + envelope.Data.BlacklistedKeys[0].KeyID != 1 || + envelope.Data.BlacklistedKeys[0].RecentFailureCount != 1 { + t.Fatalf("blacklisted stats = %#v", envelope.Data.BlacklistedKeys) + } + }) + if err != nil { + t.Fatalf("invoke container health endpoint: %v", err) + } +} + func TestBuildContainerRegistersWebUIControlAndGatewayRoutes(t *testing.T) { t.Setenv("AUTH_KEY", "test-auth-key") t.Setenv("DATA_DIR", t.TempDir()) diff --git a/internal/control/health.go b/internal/control/health.go new file mode 100644 index 00000000..19ad13b8 --- /dev/null +++ b/internal/control/health.go @@ -0,0 +1,245 @@ +package control + +import ( + "fmt" + "sort" + "time" + + "github.com/gin-gonic/gin" + + "gpt-load/internal/health" + app_errors "gpt-load/internal/platform/errors" + "gpt-load/internal/platform/response" + "gpt-load/internal/requestlog" + "gpt-load/internal/state" +) + +type RequestLogStatsReader interface { + Stats() requestlog.Stats +} + +type healthCountsResponse struct { + Total int `json:"total"` + Available int `json:"available"` + Cooldown int `json:"cooldown"` + Blacklisted int `json:"blacklisted"` + Disabled int `json:"disabled"` +} + +type healthGroupResponse struct { + ID uint `json:"id"` + Name string `json:"name"` + Enabled bool `json:"enabled"` + Counts healthCountsResponse `json:"counts"` +} + +type healthRecoveryResponse struct { + Automatic bool `json:"automatic"` + Mode string `json:"mode"` + At *time.Time `json:"at"` +} + +type healthProblemKeyResponse struct { + KeyID uint `json:"key_id"` + GroupID uint `json:"group_id"` + GroupName string `json:"group_name"` + CooldownUntil *time.Time `json:"cooldown_until,omitempty"` + FailureCount int `json:"failure_count"` + RecentSuccessCount uint64 `json:"recent_success_count"` + RecentFailureCount uint64 `json:"recent_failure_count"` + ConsecutiveFailureCount uint64 `json:"consecutive_failure_count"` + WeightManual *int `json:"weight_manual"` + WeightAuto int `json:"weight_auto"` + Recovery healthRecoveryResponse `json:"recovery"` +} + +type requestLogHealthResponse struct { + EnqueuedTotal uint64 `json:"enqueued_total"` + PersistedTotal uint64 `json:"persisted_total"` + DroppedNotRunningTotal uint64 `json:"dropped_not_running_total"` + DroppedQueueFullTotal uint64 `json:"dropped_queue_full_total"` + DroppedStoppingTotal uint64 `json:"dropped_stopping_total"` + DroppedPersistFailedTotal uint64 `json:"dropped_persist_failed_total"` + DroppedShutdownTotal uint64 `json:"dropped_shutdown_total"` + DroppedTotal uint64 `json:"dropped_total"` + WriteFailureTotal uint64 `json:"write_failure_total"` + RetentionInvalidSettingTotal uint64 `json:"retention_invalid_setting_total"` + RetentionDeleteFailureTotal uint64 `json:"retention_delete_failure_total"` + QueueDepth int `json:"queue_depth"` + QueueCapacity int `json:"queue_capacity"` + LastWriteFailureAt *time.Time `json:"last_write_failure_at"` + LastRetentionFailureAt *time.Time `json:"last_retention_failure_at"` +} + +type runtimeHealthResponse struct { + ObservedAt time.Time `json:"observed_at"` + SnapshotRevision uint64 `json:"snapshot_revision"` + StatsWindowSeconds int64 `json:"stats_window_seconds"` + Counts healthCountsResponse `json:"counts"` + Groups []healthGroupResponse `json:"groups"` + CooldownKeys []healthProblemKeyResponse `json:"cooldown_keys"` + BlacklistedKeys []healthProblemKeyResponse `json:"blacklisted_keys"` + RequestLog requestLogHealthResponse `json:"request_log"` +} + +type healthBucket string + +const ( + healthBucketAvailable healthBucket = "available" + healthBucketCooldown healthBucket = "cooldown" + healthBucketBlacklisted healthBucket = "blacklisted" + healthBucketDisabled healthBucket = "disabled" +) + +func classifyHealthKey( + group state.GroupCatalogView, + key state.KeyRuntimeView, + now time.Time, +) healthBucket { + if !group.Enabled || + (group.WeightManual != nil && *group.WeightManual == 0) || + key.Status != state.KeyStatusActive || + (key.WeightManual != nil && *key.WeightManual == 0) { + return healthBucketDisabled + } + switch key.RuntimeState(now) { + case state.KeyRuntimeBlacklisted: + return healthBucketBlacklisted + case state.KeyRuntimeCooldown: + return healthBucketCooldown + default: + return healthBucketAvailable + } +} + +func addHealthCount(counts *healthCountsResponse, bucket healthBucket) { + counts.Total++ + switch bucket { + case healthBucketAvailable: + counts.Available++ + case healthBucketCooldown: + counts.Cooldown++ + case healthBucketBlacklisted: + counts.Blacklisted++ + case healthBucketDisabled: + counts.Disabled++ + } +} + +func optionalUTC(value time.Time) *time.Time { + if value.IsZero() { + return nil + } + value = value.UTC() + return &value +} + +func cloneInt(value *int) *int { + if value == nil { + return nil + } + cloned := *value + return &cloned +} + +func (service *Service) RuntimeHealth() (runtimeHealthResponse, error) { + observation, err := service.captureRuntimeObservation() + if err != nil { + return runtimeHealthResponse{}, err + } + if service.stats == nil || service.requestLogStats == nil { + return runtimeHealthResponse{}, fmt.Errorf( + "runtime health dependencies unavailable: %w", + app_errors.ErrInternalServer, + ) + } + result := runtimeHealthResponse{ + ObservedAt: observation.observedAt, + SnapshotRevision: observation.snapshot.Revision, + StatsWindowSeconds: int64(health.StatsWindow / time.Second), + Groups: []healthGroupResponse{}, + CooldownKeys: []healthProblemKeyResponse{}, + BlacklistedKeys: []healthProblemKeyResponse{}, + } + groupIDs := make([]uint, 0, len(observation.snapshot.GroupCatalog)) + for groupID := range observation.snapshot.GroupCatalog { + groupIDs = append(groupIDs, groupID) + } + sort.Slice(groupIDs, func(i, j int) bool { return groupIDs[i] < groupIDs[j] }) + groupIndexes := make(map[uint]int, len(groupIDs)) + for _, groupID := range groupIDs { + group := observation.snapshot.GroupCatalog[groupID] + groupIndexes[groupID] = len(result.Groups) + result.Groups = append(result.Groups, healthGroupResponse{ + ID: group.ID, Name: group.Name, Enabled: group.Enabled, + }) + } + for _, key := range observation.keys { + group := observation.snapshot.GroupCatalog[key.GroupID] + index := groupIndexes[key.GroupID] + bucket := classifyHealthKey(group, key, observation.observedAt) + addHealthCount(&result.Counts, bucket) + addHealthCount(&result.Groups[index].Counts, bucket) + if bucket != healthBucketCooldown && bucket != healthBucketBlacklisted { + continue + } + stats := service.stats.Snapshot(key.ID, observation.observedAt) + detail := healthProblemKeyResponse{ + KeyID: key.ID, + GroupID: key.GroupID, + GroupName: group.Name, + FailureCount: key.FailureCount, + RecentSuccessCount: stats.Success, + RecentFailureCount: stats.Failure, + ConsecutiveFailureCount: stats.ConsecutiveFailure, + WeightManual: cloneInt(key.WeightManual), + WeightAuto: key.WeightAuto, + } + if bucket == healthBucketCooldown { + detail.CooldownUntil = optionalUTC(key.CooldownUntil) + detail.Recovery = healthRecoveryResponse{ + Automatic: true, + Mode: "cooldown_expiry", + At: optionalUTC(key.CooldownUntil), + } + result.CooldownKeys = append(result.CooldownKeys, detail) + } else { + detail.Recovery = healthRecoveryResponse{ + Automatic: true, + Mode: "validation_probe", + } + result.BlacklistedKeys = append(result.BlacklistedKeys, detail) + } + } + result.RequestLog = mapRequestLogHealth(service.requestLogStats.Stats()) + return result, nil +} + +func mapRequestLogHealth(stats requestlog.Stats) requestLogHealthResponse { + return requestLogHealthResponse{ + EnqueuedTotal: stats.EnqueuedTotal, + PersistedTotal: stats.PersistedTotal, + DroppedNotRunningTotal: stats.DroppedNotRunningTotal, + DroppedQueueFullTotal: stats.DroppedQueueFullTotal, + DroppedStoppingTotal: stats.DroppedStoppingTotal, + DroppedPersistFailedTotal: stats.DroppedPersistFailedTotal, + DroppedShutdownTotal: stats.DroppedShutdownTotal, + DroppedTotal: stats.DroppedTotal, + WriteFailureTotal: stats.WriteFailureTotal, + RetentionInvalidSettingTotal: stats.RetentionInvalidSettingTotal, + RetentionDeleteFailureTotal: stats.RetentionDeleteFailureTotal, + QueueDepth: stats.QueueDepth, + QueueCapacity: stats.QueueCapacity, + LastWriteFailureAt: optionalUTC(stats.LastWriteFailureAt), + LastRetentionFailureAt: optionalUTC(stats.LastRetentionFailureAt), + } +} + +func (server *Server) handleRuntimeHealth(c *gin.Context) { + result, err := server.service.RuntimeHealth() + if err != nil { + writeServiceError(c, "runtime_health", err) + return + } + response.SuccessI18n(c, "common.success", result) +} diff --git a/internal/control/health_test.go b/internal/control/health_test.go new file mode 100644 index 00000000..fa395df3 --- /dev/null +++ b/internal/control/health_test.go @@ -0,0 +1,336 @@ +package control + +import ( + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "reflect" + "strings" + "testing" + "time" + + "github.com/gin-gonic/gin" + + "gpt-load/internal/platform/config" + app_errors "gpt-load/internal/platform/errors" + "gpt-load/internal/protocol" + "gpt-load/internal/requestlog" + "gpt-load/internal/state" +) + +func healthNow() time.Time { + return time.Date(2026, time.July, 24, 10, 0, 0, 0, time.UTC) +} + +func TestRuntimeHealthReturnsMutuallyExclusiveCurrentState(t *testing.T) { + fixture := newServiceFixture(t) + now := healthNow() + fixture.service.now = func() time.Time { return now } + zero := 0 + if _, err := fixture.manager.Publish(state.CompileInput{ + Groups: []state.GroupConfig{ + { + ID: 1, Name: "active", Protocols: []protocol.Protocol{protocol.OpenAI}, + Models: []state.ModelConfig{{ID: "model"}}, Enabled: true, + }, + { + ID: 2, Name: "disabled", Protocols: []protocol.Protocol{protocol.OpenAI}, + Models: []state.ModelConfig{{ID: "model"}}, Enabled: false, + }, + { + ID: 3, Name: "zero-weight", Protocols: []protocol.Protocol{protocol.OpenAI}, + Models: []state.ModelConfig{{ID: "model"}}, + WeightManual: &zero, Enabled: true, + }, + { + ID: 4, Name: "empty", Protocols: []protocol.Protocol{protocol.OpenAI}, + Models: []state.ModelConfig{{ID: "model"}}, Enabled: true, + }, + }, + }); err != nil { + t.Fatalf("Publish() error = %v", err) + } + keyWeightZero := 0 + if err := fixture.registry.Replace([]state.KeyEntry{ + {ID: 11, GroupID: 1, Status: state.KeyStatusActive, EncryptedValue: "available"}, + { + ID: 12, GroupID: 1, Status: state.KeyStatusActive, + CooldownUntil: now.Add(time.Minute), FailureCount: 1, + EncryptedValue: "cooldown", + }, + { + ID: 13, GroupID: 1, Status: state.KeyStatusActive, + Blacklisted: true, CooldownUntil: now.Add(time.Hour), + FailureCount: 3, EncryptedValue: "blacklisted", + }, + { + ID: 14, GroupID: 1, Status: state.KeyStatusDisabled, + Blacklisted: true, EncryptedValue: "disabled", + }, + { + ID: 15, GroupID: 1, Status: state.KeyStatusActive, + WeightManual: &keyWeightZero, Blacklisted: true, + EncryptedValue: "weight-zero", + }, + {ID: 21, GroupID: 2, Status: state.KeyStatusActive, EncryptedValue: "disabled-group"}, + {ID: 31, GroupID: 3, Status: state.KeyStatusActive, EncryptedValue: "zero-group"}, + }); err != nil { + t.Fatalf("Replace() error = %v", err) + } + fixture.stats.Record(12, true, now) + fixture.stats.Record(12, false, now) + fixture.stats.Record(13, false, now) + fixture.requestLogStats.value = requestlog.Stats{ + EnqueuedTotal: 100, PersistedTotal: 98, + DroppedQueueFullTotal: 1, DroppedPersistFailedTotal: 1, + DroppedTotal: 2, WriteFailureTotal: 1, + QueueDepth: 2, QueueCapacity: 4096, + LastWriteFailureAt: now.Add(-time.Minute), + } + + got, err := fixture.service.RuntimeHealth() + if err != nil { + t.Fatalf("RuntimeHealth() error = %v", err) + } + if !got.ObservedAt.Equal(now) || got.SnapshotRevision != fixture.manager.Current().Revision || + got.StatsWindowSeconds != 300 { + t.Fatalf("observation metadata = %#v", got) + } + wantCounts := healthCountsResponse{ + Total: 7, Available: 1, Cooldown: 1, Blacklisted: 1, Disabled: 4, + } + if got.Counts != wantCounts { + t.Fatalf("global counts = %#v, want %#v", got.Counts, wantCounts) + } + if len(got.Groups) != 4 || got.Groups[0].ID != 1 || + got.Groups[1].ID != 2 || got.Groups[2].ID != 3 || got.Groups[3].ID != 4 { + t.Fatalf("group order = %#v", got.Groups) + } + if got.Groups[0].Counts != (healthCountsResponse{ + Total: 5, Available: 1, Cooldown: 1, Blacklisted: 1, Disabled: 2, + }) { + t.Fatalf("active group counts = %#v", got.Groups[0].Counts) + } + if got.Groups[1].Counts.Disabled != 1 || got.Groups[2].Counts.Disabled != 1 || + got.Groups[3].Counts.Total != 0 { + t.Fatalf("disabled/zero/empty group counts = %#v", got.Groups) + } + if len(got.CooldownKeys) != 1 || got.CooldownKeys[0].KeyID != 12 || + got.CooldownKeys[0].RecentSuccessCount != 1 || + got.CooldownKeys[0].RecentFailureCount != 1 || + got.CooldownKeys[0].Recovery.Mode != "cooldown_expiry" { + t.Fatalf("cooldown details = %#v", got.CooldownKeys) + } + if len(got.BlacklistedKeys) != 1 || got.BlacklistedKeys[0].KeyID != 13 || + got.BlacklistedKeys[0].Recovery.Mode != "validation_probe" || + got.BlacklistedKeys[0].Recovery.At != nil { + t.Fatalf("blacklisted details = %#v", got.BlacklistedKeys) + } + if got.RequestLog.DroppedTotal != 2 || + got.RequestLog.LastWriteFailureAt == nil || + got.RequestLog.LastRetentionFailureAt != nil { + t.Fatalf("request log stats = %#v", got.RequestLog) + } +} + +func TestRuntimeHealthSortsProblemKeysByGroupAndKey(t *testing.T) { + fixture := newServiceFixture(t) + now := healthNow() + fixture.service.now = func() time.Time { return now } + if _, err := fixture.manager.Publish(state.CompileInput{ + Groups: []state.GroupConfig{ + { + ID: 2, Name: "two", Protocols: []protocol.Protocol{protocol.OpenAI}, + Models: []state.ModelConfig{{ID: "model"}}, Enabled: true, + }, + { + ID: 1, Name: "one", Protocols: []protocol.Protocol{protocol.OpenAI}, + Models: []state.ModelConfig{{ID: "model"}}, Enabled: true, + }, + }, + }); err != nil { + t.Fatalf("Publish() error = %v", err) + } + if err := fixture.registry.Replace([]state.KeyEntry{ + { + ID: 22, GroupID: 2, Status: state.KeyStatusActive, + CooldownUntil: now.Add(time.Minute), EncryptedValue: "cooldown-22", + }, + { + ID: 13, GroupID: 1, Status: state.KeyStatusActive, + Blacklisted: true, EncryptedValue: "blacklisted-13", + }, + { + ID: 12, GroupID: 1, Status: state.KeyStatusActive, + CooldownUntil: now.Add(time.Minute), EncryptedValue: "cooldown-12", + }, + { + ID: 21, GroupID: 2, Status: state.KeyStatusActive, + Blacklisted: true, EncryptedValue: "blacklisted-21", + }, + { + ID: 11, GroupID: 1, Status: state.KeyStatusActive, + CooldownUntil: now.Add(time.Minute), EncryptedValue: "cooldown-11", + }, + }); err != nil { + t.Fatalf("Replace() error = %v", err) + } + got, err := fixture.service.RuntimeHealth() + if err != nil { + t.Fatalf("RuntimeHealth() error = %v", err) + } + pairs := func(items []healthProblemKeyResponse) [][2]uint { + result := make([][2]uint, 0, len(items)) + for _, item := range items { + result = append(result, [2]uint{item.GroupID, item.KeyID}) + } + return result + } + if gotPairs, want := pairs(got.CooldownKeys), [][2]uint{ + {1, 11}, {1, 12}, {2, 22}, + }; !reflect.DeepEqual(gotPairs, want) { + t.Fatalf("cooldown order = %v, want %v", gotPairs, want) + } + if gotPairs, want := pairs(got.BlacklistedKeys), [][2]uint{ + {1, 13}, {2, 21}, + }; !reflect.DeepEqual(gotPairs, want) { + t.Fatalf("blacklisted order = %v, want %v", gotPairs, want) + } +} + +func TestRuntimeHealthJSONOmitsScoresCredentialsAndZeroTimes(t *testing.T) { + fixture := newServiceFixture(t) + fixture.service.now = healthNow + if _, err := fixture.manager.Publish(state.CompileInput{Groups: []state.GroupConfig{{ + ID: 1, Name: "safe", Protocols: []protocol.Protocol{protocol.OpenAI}, + Models: []state.ModelConfig{{ID: "model"}}, Enabled: true, + }}}); err != nil { + t.Fatalf("Publish() error = %v", err) + } + if err := fixture.registry.Replace([]state.KeyEntry{{ + ID: 1, GroupID: 1, Status: state.KeyStatusActive, + Blacklisted: true, EncryptedValue: "cipher-must-not-appear", + }}); err != nil { + t.Fatalf("Replace() error = %v", err) + } + beforeSnapshot := fixture.manager.Current() + beforeKeys := fixture.registry.Snapshot() + beforeStats := fixture.stats.Snapshot(1, healthNow()) + + result, err := fixture.service.RuntimeHealth() + if err != nil { + t.Fatalf("RuntimeHealth() error = %v", err) + } + body, err := json.Marshal(result) + if err != nil { + t.Fatalf("Marshal() error = %v", err) + } + lower := strings.ToLower(string(body)) + for _, forbidden := range []string{ + "cipher-must-not-appear", "encrypted", "hash", "header_rules", + "percentage", "success_rate", "score", "average_latency", + "0001-01-01", + } { + if strings.Contains(lower, forbidden) { + t.Fatalf("health JSON exposes %q: %s", forbidden, body) + } + } + if fixture.manager.Current() != beforeSnapshot || + !reflect.DeepEqual(fixture.registry.Snapshot(), beforeKeys) || + fixture.stats.Snapshot(1, healthNow()) != beforeStats { + t.Fatal("RuntimeHealth() mutated runtime state") + } +} + +func TestRuntimeHealthDTOHasNoCredentialOrScoreFields(t *testing.T) { + forbidden := map[string]struct{}{ + "EncryptedValue": {}, "KeyHash": {}, "AccessKey": {}, + "HeaderRules": {}, "Percentage": {}, "SuccessRate": {}, "Score": {}, + } + types := []reflect.Type{ + reflect.TypeOf(runtimeHealthResponse{}), + reflect.TypeOf(healthCountsResponse{}), + reflect.TypeOf(healthGroupResponse{}), + reflect.TypeOf(healthProblemKeyResponse{}), + reflect.TypeOf(healthRecoveryResponse{}), + reflect.TypeOf(requestLogHealthResponse{}), + } + for _, typ := range types { + for index := 0; index < typ.NumField(); index++ { + name := typ.Field(index).Name + if _, denied := forbidden[name]; denied { + t.Fatalf("%s exposes forbidden field %s", typ.Name(), name) + } + } + } +} + +func TestRuntimeHealthFailsLoudForRegistryCatalogMismatch(t *testing.T) { + fixture := newServiceFixture(t) + if err := fixture.registry.Replace([]state.KeyEntry{{ + ID: 1, GroupID: 999, Status: state.KeyStatusActive, + EncryptedValue: "cipher", + }}); err != nil { + t.Fatalf("Replace() error = %v", err) + } + if _, err := fixture.service.RuntimeHealth(); !errors.Is( + err, + app_errors.ErrInternalServer, + ) { + t.Fatalf("RuntimeHealth() error = %v, want INTERNAL_SERVER_ERROR", err) + } +} + +func TestRuntimeHealthFailsWhenSnapshotIsUninitialized(t *testing.T) { + fixture := newServiceFixture(t) + fixture.service.manager = state.NewManager() + if _, err := fixture.service.RuntimeHealth(); !errors.Is( + err, + app_errors.ErrInternalServer, + ) { + t.Fatalf("RuntimeHealth() error = %v, want INTERNAL_SERVER_ERROR", err) + } +} + +func TestRuntimeHealthEndpointRequiresManagementAuthentication(t *testing.T) { + initControlI18n(t) + fixture := newServiceFixture(t) + fixture.service.now = healthNow + engine := gin.New() + NewServer(&config.Config{AuthKey: "test-auth-key"}, fixture.service).RegisterRoutes(engine) + + request := httptest.NewRequest(http.MethodGet, "/api/health", nil) + recorder := httptest.NewRecorder() + engine.ServeHTTP(recorder, request) + if recorder.Code != http.StatusUnauthorized { + t.Fatalf("response = %d %s, want 401", recorder.Code, recorder.Body.String()) + } + + authenticated := httptest.NewRequest(http.MethodGet, "/api/health", nil) + authenticated.Header.Set("Authorization", "Bearer test-auth-key") + success := httptest.NewRecorder() + engine.ServeHTTP(success, authenticated) + if success.Code != http.StatusOK { + t.Fatalf("authenticated response = %d %s, want 200", success.Code, success.Body.String()) + } + var envelope struct { + Code int `json:"code"` + Data runtimeHealthResponse `json:"data"` + } + if err := json.Unmarshal(success.Body.Bytes(), &envelope); err != nil { + t.Fatalf("decode authenticated response: %v", err) + } + if envelope.Code != 0 || !envelope.Data.ObservedAt.Equal(healthNow()) || + envelope.Data.StatsWindowSeconds != 300 { + t.Fatalf("authenticated envelope = %#v", envelope) + } + body := success.Body.String() + for _, emptyArray := range []string{ + `"groups":[]`, `"cooldown_keys":[]`, `"blacklisted_keys":[]`, + } { + if !strings.Contains(body, emptyArray) { + t.Fatalf("response must contain %s: %s", emptyArray, body) + } + } +} diff --git a/internal/control/route_inspect.go b/internal/control/route_inspect.go new file mode 100644 index 00000000..1b20fb76 --- /dev/null +++ b/internal/control/route_inspect.go @@ -0,0 +1,178 @@ +package control + +import ( + "errors" + "fmt" + "strings" + "time" + + "github.com/gin-gonic/gin" + + app_errors "gpt-load/internal/platform/errors" + "gpt-load/internal/platform/response" + "gpt-load/internal/protocol" + "gpt-load/internal/scheduler" + "gpt-load/internal/state" +) + +type routeInspectRequest struct { + Protocol protocol.Protocol `json:"protocol"` + ExternalModel string `json:"external_model"` + AccessKeyID uint `json:"access_key_id"` +} + +type routeInspectAccessKeyResponse struct { + ID uint `json:"id"` + Name string `json:"name"` + Status state.AccessKeyStatus `json:"status"` +} + +type routeInspectKeyResponse struct { + KeyID uint `json:"key_id"` + Available bool `json:"available"` + ReasonCode *scheduler.ReasonCode `json:"reason_code"` + WeightManual *int `json:"weight_manual"` + WeightAuto int `json:"weight_auto"` + EffectiveWeight int64 `json:"effective_weight"` + CooldownUntil *time.Time `json:"cooldown_until"` +} + +type routeInspectGroupResponse struct { + GroupID uint `json:"group_id"` + GroupName string `json:"group_name"` + UpstreamModel string `json:"upstream_model"` + WeightManual *int `json:"weight_manual"` + Included bool `json:"included"` + Routable bool `json:"routable"` + ReasonCode *scheduler.ReasonCode `json:"reason_code"` + Keys []routeInspectKeyResponse `json:"keys"` +} + +type routeInspectResponse struct { + ObservedAt time.Time `json:"observed_at"` + SnapshotRevision uint64 `json:"snapshot_revision"` + Protocol protocol.Protocol `json:"protocol"` + ExternalModel string `json:"external_model"` + AccessKey routeInspectAccessKeyResponse `json:"access_key"` + Routable bool `json:"routable"` + ReasonCode *scheduler.ReasonCode `json:"reason_code"` + Groups []routeInspectGroupResponse `json:"groups"` +} + +func optionalReason(value scheduler.ReasonCode) *scheduler.ReasonCode { + if value == "" { + return nil + } + cloned := value + return &cloned +} + +func validateRouteInspectRequest(request routeInspectRequest) error { + if !request.Protocol.Valid() || + request.AccessKeyID == 0 || + request.ExternalModel == "" || + strings.TrimSpace(request.ExternalModel) != request.ExternalModel { + return app_errors.ErrValidation + } + return nil +} + +func (service *Service) InspectRoute( + request routeInspectRequest, +) (routeInspectResponse, error) { + if err := validateRouteInspectRequest(request); err != nil { + return routeInspectResponse{}, err + } + observation, err := service.captureRuntimeObservation() + if err != nil { + return routeInspectResponse{}, err + } + accessKey, exists := observation.snapshot.AccessKeysByID[request.AccessKeyID] + if !exists { + return routeInspectResponse{}, app_errors.ErrResourceNotFound + } + explanation, err := scheduler.Inspect( + observation.snapshot, + observation.keys, + scheduler.Query{ + Protocol: request.Protocol, + ExternalModel: request.ExternalModel, + AccessKey: accessKey, + }, + observation.observedAt, + ) + if err != nil { + if errors.Is(err, scheduler.ErrInconsistentSnapshot) { + return routeInspectResponse{}, fmt.Errorf( + "inspect current route: %w", + app_errors.ErrInternalServer, + ) + } + return routeInspectResponse{}, err + } + return mapRouteInspectResponse( + observation, + request, + accessKey, + explanation, + ), nil +} + +func mapRouteInspectResponse( + observation runtimeObservation, + request routeInspectRequest, + accessKey state.AccessKeyView, + explanation scheduler.Inspection, +) routeInspectResponse { + result := routeInspectResponse{ + ObservedAt: observation.observedAt, + SnapshotRevision: observation.snapshot.Revision, + Protocol: request.Protocol, + ExternalModel: request.ExternalModel, + AccessKey: routeInspectAccessKeyResponse{ + ID: accessKey.ID, Name: accessKey.Name, Status: accessKey.Status, + }, + Routable: explanation.Routable, + ReasonCode: optionalReason(explanation.Reason), + Groups: []routeInspectGroupResponse{}, + } + for _, group := range explanation.Groups { + groupResponse := routeInspectGroupResponse{ + GroupID: group.GroupID, + GroupName: group.GroupName, + UpstreamModel: group.UpstreamModelID, + WeightManual: cloneInt(group.WeightManual), + Included: group.Included, + Routable: group.Routable, + ReasonCode: optionalReason(group.Reason), + Keys: []routeInspectKeyResponse{}, + } + for _, key := range group.Keys { + groupResponse.Keys = append(groupResponse.Keys, routeInspectKeyResponse{ + KeyID: key.KeyID, + Available: key.Available, + ReasonCode: optionalReason(key.Reason), + WeightManual: cloneInt(key.WeightManual), + WeightAuto: key.WeightAuto, + EffectiveWeight: key.EffectiveWeight, + CooldownUntil: optionalUTC(key.CooldownUntil), + }) + } + result.Groups = append(result.Groups, groupResponse) + } + return result +} + +func (server *Server) handleRouteInspect(c *gin.Context) { + var request routeInspectRequest + if err := bindStrictJSON(c, &request); err != nil { + writeServiceError(c, "inspect_route", mapControlJSONError(err)) + return + } + result, err := server.service.InspectRoute(request) + if err != nil { + writeServiceError(c, "inspect_route", err) + return + } + response.SuccessI18n(c, "common.success", result) +} diff --git a/internal/control/route_inspect_test.go b/internal/control/route_inspect_test.go new file mode 100644 index 00000000..04edc2dc --- /dev/null +++ b/internal/control/route_inspect_test.go @@ -0,0 +1,769 @@ +package control + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "reflect" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/gin-gonic/gin" + + "gpt-load/internal/dialect" + "gpt-load/internal/platform/config" + app_errors "gpt-load/internal/platform/errors" + "gpt-load/internal/protocol" + "gpt-load/internal/requestlog" + "gpt-load/internal/scheduler" + "gpt-load/internal/state" +) + +func performRouteInspectRequest( + engine *gin.Engine, + authKey string, + body string, +) *httptest.ResponseRecorder { + request := httptest.NewRequest( + http.MethodPost, + "/api/route/inspect", + strings.NewReader(body), + ) + request.Header.Set("Content-Type", "application/json") + if authKey != "" { + request.Header.Set("Authorization", "Bearer "+authKey) + } + recorder := httptest.NewRecorder() + engine.ServeHTTP(recorder, request) + return recorder +} + +func decodeRouteInspectSuccess( + t *testing.T, + recorder *httptest.ResponseRecorder, +) routeInspectResponse { + t.Helper() + if recorder.Code != http.StatusOK { + t.Fatalf("response = %d %s, want 200", recorder.Code, recorder.Body.String()) + } + var envelope struct { + Code int `json:"code"` + Data routeInspectResponse `json:"data"` + } + if err := json.Unmarshal(recorder.Body.Bytes(), &envelope); err != nil { + t.Fatalf("decode response: %v", err) + } + if envelope.Code != 0 { + t.Fatalf("code = %d, want 0: %s", envelope.Code, recorder.Body.String()) + } + return envelope.Data +} + +func assertRouteReason( + t *testing.T, + got *scheduler.ReasonCode, + want scheduler.ReasonCode, +) { + t.Helper() + if got == nil || *got != want { + t.Fatalf("reason = %v, want %q", got, want) + } +} + +func TestRouteInspectEndpointRejectsMalformedAndInvalidRequests(t *testing.T) { + initControlI18n(t) + fixture := newServiceFixture(t) + engine := gin.New() + NewServer(&config.Config{AuthKey: "test-auth-key"}, fixture.service).RegisterRoutes(engine) + tests := []struct { + name string + body string + wantCode string + }{ + { + name: "unknown field", + body: `{"protocol":"openai","external_model":"model","access_key_id":1,"extra":true}`, + wantCode: app_errors.ErrInvalidJSON.Code, + }, + { + name: "plaintext access key rejected", + body: `{"protocol":"openai","external_model":"model","access_key_id":1,"access_key":"secret"}`, + wantCode: app_errors.ErrInvalidJSON.Code, + }, + { + name: "multiple values", + body: `{"protocol":"openai","external_model":"model","access_key_id":1}{}`, + wantCode: app_errors.ErrInvalidJSON.Code, + }, + { + name: "trailing malformed value", + body: `{"protocol":"openai","external_model":"model","access_key_id":1}x`, + wantCode: app_errors.ErrInvalidJSON.Code, + }, + { + name: "malformed JSON", + body: `{"protocol":"openai"`, + wantCode: app_errors.ErrInvalidJSON.Code, + }, + { + name: "invalid protocol", + body: `{"protocol":"invalid","external_model":"model","access_key_id":1}`, + wantCode: app_errors.ErrValidation.Code, + }, + { + name: "empty model", + body: `{"protocol":"openai","external_model":"","access_key_id":1}`, + wantCode: app_errors.ErrValidation.Code, + }, + { + name: "trimmed model required", + body: `{"protocol":"openai","external_model":" model ","access_key_id":1}`, + wantCode: app_errors.ErrValidation.Code, + }, + { + name: "zero access key id", + body: `{"protocol":"openai","external_model":"model","access_key_id":0}`, + wantCode: app_errors.ErrValidation.Code, + }, + { + name: "protocol has wrong JSON type", + body: `{"protocol":1,"external_model":"model","access_key_id":1}`, + wantCode: app_errors.ErrInvalidJSON.Code, + }, + { + name: "external model has wrong JSON type", + body: `{"protocol":"openai","external_model":false,"access_key_id":1}`, + wantCode: app_errors.ErrInvalidJSON.Code, + }, + { + name: "access key id has wrong JSON type", + body: `{"protocol":"openai","external_model":"model","access_key_id":"1"}`, + wantCode: app_errors.ErrInvalidJSON.Code, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + recorder := performRouteInspectRequest(engine, "test-auth-key", test.body) + if recorder.Code != http.StatusBadRequest { + t.Fatalf("response = %d %s, want 400", recorder.Code, recorder.Body.String()) + } + var envelope struct { + Code string `json:"code"` + } + if err := json.Unmarshal(recorder.Body.Bytes(), &envelope); err != nil { + t.Fatalf("decode response: %v", err) + } + if envelope.Code != test.wantCode { + t.Fatalf("code = %q, want %q", envelope.Code, test.wantCode) + } + }) + } +} + +func TestRouteInspectEndpointReturnsCurrentSafeExplanation(t *testing.T) { + initControlI18n(t) + fixture := newServiceFixture(t) + now := healthNow() + fixture.service.now = func() time.Time { return now } + groupWeight := 20 + if _, err := fixture.manager.Publish(state.CompileInput{ + Groups: []state.GroupConfig{ + { + ID: 2, Name: "backup", UpstreamURL: "https://backup.invalid", + Protocols: []protocol.Protocol{protocol.OpenAI}, + Models: []state.ModelConfig{{ID: "provider-backup", Alias: "public-model"}}, + WeightManual: &groupWeight, Enabled: true, + }, + { + ID: 1, Name: "primary", UpstreamURL: "https://not-called.invalid", + Protocols: []protocol.Protocol{protocol.OpenAI}, + Models: []state.ModelConfig{{ID: "provider-model", Alias: "public-model"}}, + Enabled: true, + }, + }, + AccessKeys: []state.AccessKeyConfig{{ + ID: 10, Name: "production", KeyHash: "active-hash", + Status: state.AccessKeyStatusActive, + }}, + }); err != nil { + t.Fatalf("Publish() error = %v", err) + } + keyWeight := 25 + if err := fixture.registry.Replace([]state.KeyEntry{ + { + ID: 31, GroupID: 2, Status: state.KeyStatusActive, + WeightAuto: 30, EncryptedValue: "cipher-three", + }, + { + ID: 22, GroupID: 1, Status: state.KeyStatusActive, + CooldownUntil: now.Add(time.Minute), WeightAuto: 40, + EncryptedValue: "cipher-two", + }, + { + ID: 21, GroupID: 1, Status: state.KeyStatusActive, + WeightManual: &keyWeight, WeightAuto: 90, + EncryptedValue: "cipher-one", + }, + }); err != nil { + t.Fatalf("Replace() error = %v", err) + } + engine := gin.New() + NewServer(&config.Config{AuthKey: "test-auth-key"}, fixture.service).RegisterRoutes(engine) + recorder := performRouteInspectRequest( + engine, + "test-auth-key", + `{"protocol":"openai","external_model":"public-model","access_key_id":10}`, + ) + got := decodeRouteInspectSuccess(t, recorder) + if !got.ObservedAt.Equal(now) || + got.SnapshotRevision != fixture.manager.Current().Revision || + got.Protocol != protocol.OpenAI || + got.ExternalModel != "public-model" || + got.AccessKey != (routeInspectAccessKeyResponse{ + ID: 10, Name: "production", Status: state.AccessKeyStatusActive, + }) || + !got.Routable || got.ReasonCode != nil { + t.Fatalf("route response = %#v", got) + } + if len(got.Groups) != 2 || got.Groups[0].GroupID != 1 || + got.Groups[1].GroupID != 2 { + t.Fatalf("group order = %#v", got.Groups) + } + primary := got.Groups[0] + if primary.GroupName != "primary" || + primary.UpstreamModel != "provider-model" || + primary.WeightManual != nil || !primary.Included || + !primary.Routable || primary.ReasonCode != nil || + len(primary.Keys) != 2 || + primary.Keys[0].KeyID != 21 || primary.Keys[1].KeyID != 22 { + t.Fatalf("primary group = %#v", primary) + } + available := primary.Keys[0] + if !available.Available || available.ReasonCode != nil || + available.WeightManual == nil || *available.WeightManual != 25 || + available.WeightAuto != 90 || available.EffectiveWeight != 50*25 || + available.CooldownUntil != nil { + t.Fatalf("available key = %#v", available) + } + cooldown := primary.Keys[1] + if cooldown.Available || cooldown.WeightManual != nil || + cooldown.WeightAuto != 40 || cooldown.EffectiveWeight != 0 || + cooldown.CooldownUntil == nil || + !cooldown.CooldownUntil.Equal(now.Add(time.Minute).UTC()) { + t.Fatalf("cooldown key = %#v", cooldown) + } + assertRouteReason(t, cooldown.ReasonCode, scheduler.ReasonKeyCooldown) + backup := got.Groups[1] + if backup.GroupName != "backup" || + backup.UpstreamModel != "provider-backup" || + backup.WeightManual == nil || *backup.WeightManual != 20 || + !backup.Included || !backup.Routable || backup.ReasonCode != nil || + len(backup.Keys) != 1 || backup.Keys[0].KeyID != 31 || + !backup.Keys[0].Available || backup.Keys[0].ReasonCode != nil || + backup.Keys[0].WeightManual != nil || + backup.Keys[0].WeightAuto != 30 || + backup.Keys[0].EffectiveWeight != 20*30 || + backup.Keys[0].CooldownUntil != nil { + t.Fatalf("backup group = %#v", backup) + } + body := recorder.Body.String() + if strings.Count(body, `"reason_code":null`) != 5 || + strings.Count(body, `"cooldown_until":null`) != 2 { + t.Fatalf("success response must preserve explicit nulls: %s", body) + } + lower := strings.ToLower(body) + for _, forbidden := range []string{ + "cipher-one", "cipher-two", "cipher-three", "active-hash", "upstream_url", + "header_rules", "filters", + } { + if strings.Contains(lower, forbidden) { + t.Fatalf("response exposes %q: %s", forbidden, body) + } + } +} + +func TestRouteInspectEndpointReturnsFilterExplanations(t *testing.T) { + initControlI18n(t) + tests := []struct { + name string + filters state.FilterSet + wantReason scheduler.ReasonCode + assertGroup bool + }{ + { + name: "protocol filter", + filters: state.FilterSet{ + Protocols: map[protocol.Protocol]struct{}{protocol.Anthropic: {}}, + }, + wantReason: scheduler.ReasonProtocolFiltered, + }, + { + name: "model filter", + filters: state.FilterSet{ + Models: map[string]struct{}{"other-model": {}}, + }, + wantReason: scheduler.ReasonModelFiltered, + }, + { + name: "group filter", + filters: state.FilterSet{ + Groups: map[uint]struct{}{99: {}}, + }, + wantReason: scheduler.ReasonGroupFiltered, assertGroup: true, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + fixture := newServiceFixture(t) + now := healthNow() + fixture.service.now = func() time.Time { return now } + manual := 15 + if _, err := fixture.manager.Publish(state.CompileInput{ + Groups: []state.GroupConfig{ + { + ID: 2, Name: "second", Protocols: []protocol.Protocol{protocol.OpenAI}, + Models: []state.ModelConfig{{ID: "provider-two", Alias: "public-model"}}, + WeightManual: &manual, Enabled: true, + }, + { + ID: 1, Name: "first", Protocols: []protocol.Protocol{protocol.OpenAI}, + Models: []state.ModelConfig{{ID: "provider-one", Alias: "public-model"}}, + Enabled: true, + }, + }, + AccessKeys: []state.AccessKeyConfig{{ + ID: 10, Name: "filtered", KeyHash: "filtered-hash", + Status: state.AccessKeyStatusActive, Filters: test.filters, + }}, + }); err != nil { + t.Fatalf("Publish() error = %v", err) + } + if err := fixture.registry.Replace([]state.KeyEntry{ + {ID: 22, GroupID: 2, Status: state.KeyStatusActive, EncryptedValue: "two"}, + {ID: 11, GroupID: 1, Status: state.KeyStatusActive, EncryptedValue: "one"}, + }); err != nil { + t.Fatalf("Replace() error = %v", err) + } + engine := gin.New() + NewServer( + &config.Config{AuthKey: "test-auth-key"}, + fixture.service, + ).RegisterRoutes(engine) + recorder := performRouteInspectRequest( + engine, + "test-auth-key", + `{"protocol":"openai","external_model":"public-model","access_key_id":10}`, + ) + got := decodeRouteInspectSuccess(t, recorder) + if !got.ObservedAt.Equal(now) || + got.SnapshotRevision != fixture.manager.Current().Revision || + got.Routable { + t.Fatalf("filter response = %#v", got) + } + assertRouteReason(t, got.ReasonCode, test.wantReason) + if !test.assertGroup { + if got.Groups == nil || len(got.Groups) != 0 || + !strings.Contains(recorder.Body.String(), `"groups":[]`) { + t.Fatalf("top-level filter groups = %#v", got.Groups) + } + return + } + if len(got.Groups) != 2 || + got.Groups[0].GroupID != 1 || got.Groups[1].GroupID != 2 { + t.Fatalf("group order = %#v", got.Groups) + } + for index, group := range got.Groups { + if group.Included || group.Routable || + group.Keys == nil || len(group.Keys) != 0 { + t.Fatalf("filtered group %d = %#v", index, group) + } + assertRouteReason(t, group.ReasonCode, scheduler.ReasonGroupFiltered) + } + if got.Groups[0].WeightManual != nil || + got.Groups[1].WeightManual == nil || + *got.Groups[1].WeightManual != 15 || + got.Groups[0].UpstreamModel != "provider-one" || + got.Groups[1].UpstreamModel != "provider-two" { + t.Fatalf("filtered group mapping = %#v", got.Groups) + } + }) + } +} + +func TestRouteInspectEndpointReturnsNoRouteTargetExplanation(t *testing.T) { + initControlI18n(t) + fixture := newServiceFixture(t) + now := healthNow() + fixture.service.now = func() time.Time { return now } + if _, err := fixture.manager.Publish(state.CompileInput{ + Groups: []state.GroupConfig{{ + ID: 1, Name: "primary", Protocols: []protocol.Protocol{protocol.OpenAI}, + Models: []state.ModelConfig{{ID: "configured-model"}}, Enabled: true, + }}, + AccessKeys: []state.AccessKeyConfig{{ + ID: 10, Name: "production", KeyHash: "active-hash", + Status: state.AccessKeyStatusActive, + }}, + }); err != nil { + t.Fatalf("Publish() error = %v", err) + } + engine := gin.New() + NewServer(&config.Config{AuthKey: "test-auth-key"}, fixture.service).RegisterRoutes(engine) + recorder := performRouteInspectRequest( + engine, + "test-auth-key", + `{"protocol":"openai","external_model":"missing-model","access_key_id":10}`, + ) + got := decodeRouteInspectSuccess(t, recorder) + if !got.ObservedAt.Equal(now) || + got.SnapshotRevision != fixture.manager.Current().Revision || + got.Routable || got.Groups == nil || len(got.Groups) != 0 { + t.Fatalf("no-target response = %#v", got) + } + assertRouteReason(t, got.ReasonCode, scheduler.ReasonNoRouteTarget) + if !strings.Contains(recorder.Body.String(), `"groups":[]`) { + t.Fatalf("no-target groups must be []: %s", recorder.Body.String()) + } +} + +func TestRouteInspectEndpointReturnsNoAvailableKeyExplanation(t *testing.T) { + initControlI18n(t) + fixture := newServiceFixture(t) + now := healthNow() + fixture.service.now = func() time.Time { return now } + groupWeight := 25 + if _, err := fixture.manager.Publish(state.CompileInput{ + Groups: []state.GroupConfig{{ + ID: 1, Name: "primary", Protocols: []protocol.Protocol{protocol.OpenAI}, + Models: []state.ModelConfig{{ID: "provider-model", Alias: "public-model"}}, + WeightManual: &groupWeight, Enabled: true, + }}, + AccessKeys: []state.AccessKeyConfig{{ + ID: 10, Name: "production", KeyHash: "active-hash", + Status: state.AccessKeyStatusActive, + }}, + }); err != nil { + t.Fatalf("Publish() error = %v", err) + } + zero := 0 + disabledManual := 7 + sourceZone := time.FixedZone("source-offset", 8*60*60) + cooldownAt := now.In(sourceZone).Add(90 * time.Second) + if cooldownAt.Location() == time.UTC { + t.Fatal("cooldown fixture must use a non-UTC location") + } + if err := fixture.registry.Replace([]state.KeyEntry{ + { + ID: 14, GroupID: 1, Status: state.KeyStatusActive, + WeightAuto: 70, CooldownUntil: cooldownAt, EncryptedValue: "cooldown", + }, + { + ID: 12, GroupID: 1, Status: state.KeyStatusActive, + WeightManual: &zero, WeightAuto: 45, EncryptedValue: "zero", + }, + { + ID: 13, GroupID: 1, Status: state.KeyStatusActive, + WeightAuto: 60, Blacklisted: true, EncryptedValue: "blacklisted", + }, + { + ID: 11, GroupID: 1, Status: state.KeyStatusDisabled, + WeightManual: &disabledManual, WeightAuto: 30, EncryptedValue: "disabled", + }, + }); err != nil { + t.Fatalf("Replace() error = %v", err) + } + engine := gin.New() + NewServer(&config.Config{AuthKey: "test-auth-key"}, fixture.service).RegisterRoutes(engine) + recorder := performRouteInspectRequest( + engine, + "test-auth-key", + `{"protocol":"openai","external_model":"public-model","access_key_id":10}`, + ) + got := decodeRouteInspectSuccess(t, recorder) + if !got.ObservedAt.Equal(now) || + got.SnapshotRevision != fixture.manager.Current().Revision || + got.Routable { + t.Fatalf("unavailable response = %#v", got) + } + assertRouteReason(t, got.ReasonCode, scheduler.ReasonNoAvailableKey) + if len(got.Groups) != 1 { + t.Fatalf("groups = %#v", got.Groups) + } + group := got.Groups[0] + if group.GroupID != 1 || group.GroupName != "primary" || + group.UpstreamModel != "provider-model" || + group.WeightManual == nil || *group.WeightManual != 25 || + !group.Included || group.Routable || len(group.Keys) != 4 { + t.Fatalf("unavailable group = %#v", group) + } + assertRouteReason(t, group.ReasonCode, scheduler.ReasonNoAvailableKey) + wantReasons := []scheduler.ReasonCode{ + scheduler.ReasonKeyDisabled, + scheduler.ReasonKeyWeightZero, + scheduler.ReasonKeyBlacklisted, + scheduler.ReasonKeyCooldown, + } + wantManual := []*int{&disabledManual, &zero, nil, nil} + wantAuto := []int{30, 45, 60, 70} + for index, key := range group.Keys { + if key.KeyID != uint(11+index) || key.Available || + key.EffectiveWeight != 0 || + key.WeightAuto != wantAuto[index] { + t.Fatalf("unavailable key %d = %#v", index, key) + } + assertRouteReason(t, key.ReasonCode, wantReasons[index]) + if wantManual[index] == nil { + if key.WeightManual != nil { + t.Fatalf("key %d manual weight = %v, want nil", index, key.WeightManual) + } + } else if key.WeightManual == nil || + *key.WeightManual != *wantManual[index] { + t.Fatalf("key %d manual weight = %v, want %d", index, key.WeightManual, *wantManual[index]) + } + if index == 3 { + if key.CooldownUntil == nil || + !key.CooldownUntil.Equal(cooldownAt.UTC()) { + t.Fatalf("cooldown = %v, want %v", key.CooldownUntil, cooldownAt.UTC()) + } + if key.CooldownUntil.Location() != time.UTC { + t.Fatalf("cooldown location = %v, want UTC", key.CooldownUntil.Location()) + } + } else if key.CooldownUntil != nil { + t.Fatalf("key %d cooldown = %v, want nil", index, key.CooldownUntil) + } + } + if strings.Count(recorder.Body.String(), `"cooldown_until":null`) != 3 { + t.Fatalf("non-cooldown keys must encode null cooldown: %s", recorder.Body.String()) + } + wantCooldownJSON := `"cooldown_until":"` + + cooldownAt.UTC().Format(time.RFC3339) + `"` + if !strings.Contains(recorder.Body.String(), wantCooldownJSON) { + t.Fatalf( + "cooldown must encode as UTC: got %s, want %s", + recorder.Body.String(), + wantCooldownJSON, + ) + } +} + +func TestRouteInspectReturnsDisabledAccessKeyAsExplanation(t *testing.T) { + initControlI18n(t) + fixture := newServiceFixture(t) + if _, err := fixture.manager.Publish(state.CompileInput{ + AccessKeys: []state.AccessKeyConfig{{ + ID: 12, Name: "disabled", KeyHash: "disabled-hash", + Status: state.AccessKeyStatusDisabled, + }}, + }); err != nil { + t.Fatalf("Publish() error = %v", err) + } + result, err := fixture.service.InspectRoute(routeInspectRequest{ + Protocol: protocol.OpenAI, ExternalModel: "model", AccessKeyID: 12, + }) + if err != nil { + t.Fatalf("InspectRoute() error = %v", err) + } + if result.Routable || result.ReasonCode == nil || + *result.ReasonCode != scheduler.ReasonAccessKeyDisabled || + result.Groups == nil || len(result.Groups) != 0 { + t.Fatalf("disabled result = %#v", result) + } + engine := gin.New() + NewServer(&config.Config{AuthKey: "test-auth-key"}, fixture.service).RegisterRoutes(engine) + recorder := performRouteInspectRequest( + engine, + "test-auth-key", + `{"protocol":"openai","external_model":"model","access_key_id":12}`, + ) + if recorder.Code != http.StatusOK || + !strings.Contains(recorder.Body.String(), `"reason_code":"access_key_disabled"`) || + !strings.Contains(recorder.Body.String(), `"groups":[]`) { + t.Fatalf("disabled HTTP response = %d %s", recorder.Code, recorder.Body.String()) + } +} + +func TestRouteInspectMissingAccessKeyReturnsNotFound(t *testing.T) { + initControlI18n(t) + fixture := newServiceFixture(t) + _, err := fixture.service.InspectRoute(routeInspectRequest{ + Protocol: protocol.OpenAI, ExternalModel: "model", AccessKeyID: 404, + }) + if !errors.Is(err, app_errors.ErrResourceNotFound) { + t.Fatalf("InspectRoute() error = %v, want NOT_FOUND", err) + } + engine := gin.New() + NewServer(&config.Config{AuthKey: "test-auth-key"}, fixture.service).RegisterRoutes(engine) + recorder := performRouteInspectRequest( + engine, + "test-auth-key", + `{"protocol":"openai","external_model":"model","access_key_id":404}`, + ) + var envelope struct { + Code string `json:"code"` + } + if decodeErr := json.Unmarshal(recorder.Body.Bytes(), &envelope); decodeErr != nil { + t.Fatalf("decode missing-key response: %v", decodeErr) + } + if recorder.Code != http.StatusNotFound || + envelope.Code != app_errors.ErrResourceNotFound.Code { + t.Fatalf("missing-key response = %d %#v", recorder.Code, envelope) + } +} + +type routeInspectEncryptionSpy struct { + calls atomic.Int64 +} + +func (spy *routeInspectEncryptionSpy) Encrypt(string) (string, error) { + spy.calls.Add(1) + return "", nil +} + +func (spy *routeInspectEncryptionSpy) Decrypt(string) (string, error) { + spy.calls.Add(1) + return "", nil +} + +func (spy *routeInspectEncryptionSpy) Hash(string) string { + spy.calls.Add(1) + return "" +} + +func TestRouteInspectNeverCallsUpstreamOrMutatesRuntime(t *testing.T) { + var upstreamCalls atomic.Int64 + upstream := httptest.NewServer(http.HandlerFunc(func( + writer http.ResponseWriter, + request *http.Request, + ) { + upstreamCalls.Add(1) + writer.WriteHeader(http.StatusNoContent) + })) + defer upstream.Close() + + fixture := newServiceFixture(t) + encryptionSpy := &routeInspectEncryptionSpy{} + fixture.service.encryption = encryptionSpy + dialectCalls := 0 + fixture.service.dialects = dialect.NewSet(&recordingDiscoveryDialect{ + value: protocol.OpenAI, + listFn: func( + context.Context, + string, + string, + state.HeaderRules, + ) ([]string, error) { + dialectCalls++ + return nil, nil + }, + }) + if _, err := fixture.manager.Publish(state.CompileInput{ + Groups: []state.GroupConfig{{ + ID: 1, Name: "upstream", UpstreamURL: upstream.URL, + Protocols: []protocol.Protocol{protocol.OpenAI}, + Models: []state.ModelConfig{{ID: "model"}}, Enabled: true, + }}, + AccessKeys: []state.AccessKeyConfig{{ + ID: 1, Name: "client", KeyHash: "hash", + Status: state.AccessKeyStatusActive, + }}, + }); err != nil { + t.Fatalf("Publish() error = %v", err) + } + if err := fixture.registry.Replace([]state.KeyEntry{{ + ID: 1, GroupID: 1, Status: state.KeyStatusActive, + EncryptedValue: "cipher", + }}); err != nil { + t.Fatalf("Replace() error = %v", err) + } + beforeSnapshot := fixture.manager.Current() + beforeKeys := fixture.registry.Snapshot() + beforeStats := fixture.stats.Snapshot(1, healthNow()) + requestLogStatsCalls := 0 + fixture.requestLogStats.fn = func() requestlog.Stats { + requestLogStatsCalls++ + return requestlog.Stats{} + } + + if _, err := fixture.service.InspectRoute(routeInspectRequest{ + Protocol: protocol.OpenAI, ExternalModel: "model", AccessKeyID: 1, + }); err != nil { + t.Fatalf("InspectRoute() error = %v", err) + } + if upstreamCalls.Load() != 0 || dialectCalls != 0 || encryptionSpy.calls.Load() != 0 { + t.Fatalf( + "upstream calls = %d, Dialect calls = %d, encryption calls = %d", + upstreamCalls.Load(), + dialectCalls, + encryptionSpy.calls.Load(), + ) + } + if requestLogStatsCalls != 0 { + t.Fatalf("RequestLog stats calls = %d, want 0", requestLogStatsCalls) + } + if fixture.manager.Current() != beforeSnapshot || + !reflect.DeepEqual(fixture.registry.Snapshot(), beforeKeys) || + fixture.stats.Snapshot(1, healthNow()) != beforeStats { + t.Fatal("InspectRoute() mutated runtime state") + } +} + +func TestRouteInspectEndpointRequiresManagementAuthentication(t *testing.T) { + initControlI18n(t) + fixture := newServiceFixture(t) + beforeSnapshot := fixture.manager.Current() + beforeKeys := fixture.registry.Snapshot() + engine := gin.New() + NewServer(&config.Config{AuthKey: "test-auth-key"}, fixture.service).RegisterRoutes(engine) + recorder := performRouteInspectRequest( + engine, + "", + `{"protocol":"openai","external_model":"model","access_key_id":1}`, + ) + if recorder.Code != http.StatusUnauthorized { + t.Fatalf("response = %d %s, want 401", recorder.Code, recorder.Body.String()) + } + if fixture.manager.Current() != beforeSnapshot || + !reflect.DeepEqual(fixture.registry.Snapshot(), beforeKeys) { + t.Fatal("unauthenticated Inspector request reached runtime state") + } +} + +func TestRouteInspectCatalogMismatchReturnsInternalServerError(t *testing.T) { + initControlI18n(t) + fixture := newServiceFixture(t) + if _, err := fixture.manager.Publish(state.CompileInput{ + Groups: []state.GroupConfig{{ + ID: 1, Name: "primary", Protocols: []protocol.Protocol{protocol.OpenAI}, + Models: []state.ModelConfig{{ID: "model"}}, Enabled: true, + }}, + AccessKeys: []state.AccessKeyConfig{{ + ID: 1, Name: "client", KeyHash: "hash", Status: state.AccessKeyStatusActive, + }}, + }); err != nil { + t.Fatalf("Publish() error = %v", err) + } + fixture.manager.Current().RouteCatalog[protocol.OpenAI]["model"] = []state.RouteTarget{{ + GroupID: 999, UpstreamModelID: "model", + }} + engine := gin.New() + NewServer(&config.Config{AuthKey: "test-auth-key"}, fixture.service).RegisterRoutes(engine) + recorder := performRouteInspectRequest( + engine, + "test-auth-key", + `{"protocol":"openai","external_model":"model","access_key_id":1}`, + ) + var envelope struct { + Code string `json:"code"` + } + if err := json.Unmarshal(recorder.Body.Bytes(), &envelope); err != nil { + t.Fatalf("decode response: %v", err) + } + if recorder.Code != http.StatusInternalServerError || + envelope.Code != app_errors.ErrInternalServer.Code { + t.Fatalf("catalog mismatch response = %d %#v", recorder.Code, envelope) + } +} diff --git a/internal/control/runtime_observation.go b/internal/control/runtime_observation.go new file mode 100644 index 00000000..11c9816e --- /dev/null +++ b/internal/control/runtime_observation.go @@ -0,0 +1,53 @@ +package control + +import ( + "fmt" + "time" + + app_errors "gpt-load/internal/platform/errors" + "gpt-load/internal/state" +) + +type runtimeObservation struct { + observedAt time.Time + snapshot *state.ConfigSnapshot + keys []state.KeyRuntimeView +} + +func (service *Service) captureRuntimeObservation() (runtimeObservation, error) { + if service == nil || service.manager == nil || service.registry == nil || + service.now == nil { + return runtimeObservation{}, fmt.Errorf( + "capture runtime observation: %w", + app_errors.ErrInternalServer, + ) + } + service.writeMu.RLock() + snapshot := service.manager.Current() + if snapshot == nil { + service.writeMu.RUnlock() + return runtimeObservation{}, fmt.Errorf( + "capture runtime observation: Snapshot is nil: %w", + app_errors.ErrInternalServer, + ) + } + keys := service.registry.Snapshot() + observedAt := service.now().UTC() + service.writeMu.RUnlock() + + for _, key := range keys { + if _, exists := snapshot.GroupCatalog[key.GroupID]; !exists { + return runtimeObservation{}, fmt.Errorf( + "capture runtime observation: key %d group %d missing from catalog: %w", + key.ID, + key.GroupID, + app_errors.ErrInternalServer, + ) + } + } + return runtimeObservation{ + observedAt: observedAt, + snapshot: snapshot, + keys: keys, + }, nil +} diff --git a/internal/control/runtime_observation_test.go b/internal/control/runtime_observation_test.go new file mode 100644 index 00000000..42524c29 --- /dev/null +++ b/internal/control/runtime_observation_test.go @@ -0,0 +1,114 @@ +package control + +import ( + "sync" + "testing" + "time" + + "gorm.io/gorm" + + "gpt-load/internal/requestlog" + "gpt-load/internal/state" + "gpt-load/internal/storage/models" +) + +func TestCaptureRuntimeObservationWaitsForPublishedConfigPair(t *testing.T) { + fixture := newServiceFixture(t) + group := validControlGroup("runtime-observation") + var key models.UpstreamKey + runtimeApplied := make(chan struct{}) + allowPublish := make(chan struct{}) + var releaseOnce sync.Once + releasePublish := func() { + releaseOnce.Do(func() { close(allowPublish) }) + } + defer releasePublish() + writeDone := make(chan error, 1) + go func() { + _, err := fixture.service.writeConfig(t.Context(), func(tx *gorm.DB) error { + if err := tx.Create(group).Error; err != nil { + return err + } + key = models.UpstreamKey{ + GroupID: group.ID, KeyValue: "cipher-runtime-observation", + KeyHash: "hash-runtime-observation", + Status: models.UpstreamKeyStatusActive, + } + return tx.Create(&key).Error + }, func() error { + if err := fixture.registry.ApplyImport(group.ID, []state.KeyEntry{{ + ID: key.ID, GroupID: group.ID, Status: state.KeyStatusActive, + EncryptedValue: key.KeyValue, + }}); err != nil { + return err + } + close(runtimeApplied) + <-allowPublish + return nil + }) + writeDone <- err + }() + + select { + case <-runtimeApplied: + case <-time.After(2 * time.Second): + t.Fatal("Registry update barrier timed out") + } + if fixture.service.writeMu.TryRLock() { + fixture.service.writeMu.RUnlock() + t.Fatal("configuration writer did not retain writeMu before Publish") + } + captureStarted := make(chan struct{}) + captureDone := make(chan struct { + value runtimeObservation + err error + }, 1) + go func() { + close(captureStarted) + value, err := fixture.service.captureRuntimeObservation() + captureDone <- struct { + value runtimeObservation + err error + }{value: value, err: err} + }() + <-captureStarted + select { + case result := <-captureDone: + t.Fatalf("capture completed inside Registry/Snapshot gap: %#v", result) + case <-time.After(50 * time.Millisecond): + } + releasePublish() + if err := <-writeDone; err != nil { + t.Fatalf("writeConfig() error = %v", err) + } + result := <-captureDone + if result.err != nil { + t.Fatalf("captureRuntimeObservation() error = %v", result.err) + } + if _, exists := result.value.snapshot.GroupCatalog[group.ID]; !exists { + t.Fatalf("captured Snapshot missing Group %d", group.ID) + } + found := false + for _, view := range result.value.keys { + if view.ID == key.ID && view.GroupID == group.ID { + found = true + } + } + if !found { + t.Fatalf("captured keys = %#v, want key %d", result.value.keys, key.ID) + } +} + +func TestRuntimeHealthReleasesReadLockBeforeRequestLogMapping(t *testing.T) { + fixture := newServiceFixture(t) + fixture.requestLogStats.fn = func() requestlog.Stats { + if !fixture.service.writeMu.TryLock() { + t.Fatal("writeMu remained read-locked during RequestLog Stats mapping") + } + fixture.service.writeMu.Unlock() + return requestlog.Stats{} + } + if _, err := fixture.service.RuntimeHealth(); err != nil { + t.Fatalf("RuntimeHealth() error = %v", err) + } +} diff --git a/internal/control/server.go b/internal/control/server.go index c326c7e3..ef71670e 100644 --- a/internal/control/server.go +++ b/internal/control/server.go @@ -41,7 +41,9 @@ func (s *Server) RegisterRoutes(engine *gin.Engine) { api := engine.Group("/api") api.Use(i18n.Middleware(), s.authenticate()) api.GET("/auth/session", s.handleAuthSession) + api.GET("/health", s.handleRuntimeHealth) api.GET("/logs", s.handleListRequestLogs) + api.POST("/route/inspect", s.handleRouteInspect) api.GET("/groups", s.handleListGroups) api.POST("/groups", s.handleCreateGroup) api.POST("/groups/:group_id/keys/import", s.handleImportGroupKeys) diff --git a/internal/control/server_test.go b/internal/control/server_test.go index 1ea35f10..21500452 100644 --- a/internal/control/server_test.go +++ b/internal/control/server_test.go @@ -186,6 +186,11 @@ func TestControlJSONBodyLimitAppliesToEveryJSONEndpoint(t *testing.T) { }, jsonPrefix: `{"name":"body-limit-updated-access-key"}`, }, + { + name: "route inspector", method: http.MethodPost, + path: func(uint, uint) string { return "/api/route/inspect" }, + jsonPrefix: `{"protocol":"openai","external_model":"gpt-4o","access_key_id":1}`, + }, } { t.Run(endpoint.method+" "+endpoint.name, func(t *testing.T) { fixture := newServiceFixture(t) diff --git a/internal/control/service.go b/internal/control/service.go index 02ef3e23..6430fe25 100644 --- a/internal/control/service.go +++ b/internal/control/service.go @@ -14,6 +14,7 @@ import ( "gorm.io/gorm" "gpt-load/internal/dialect" + "gpt-load/internal/health" "gpt-load/internal/platform/encryption" app_errors "gpt-load/internal/platform/errors" "gpt-load/internal/state" @@ -32,9 +33,12 @@ type Service struct { encryption encryption.Service dialects dialect.Set requestLogs RequestLogReader + stats *health.StatsStore + requestLogStats RequestLogStatsReader modelDiscoveryTimeout time.Duration random io.Reader - writeMu sync.Mutex + now func() time.Time + writeMu sync.RWMutex } func NewService( @@ -44,11 +48,15 @@ func NewService( encryptionService encryption.Service, dialects dialect.Set, requestLogs RequestLogReader, + stats *health.StatsStore, + requestLogStats RequestLogStatsReader, ) *Service { return &Service{ db: db, manager: manager, registry: registry, encryption: encryptionService, dialects: dialects, requestLogs: requestLogs, - modelDiscoveryTimeout: defaultModelDiscoveryTimeout, random: rand.Reader, + stats: stats, requestLogStats: requestLogStats, + modelDiscoveryTimeout: defaultModelDiscoveryTimeout, + random: rand.Reader, now: time.Now, } } diff --git a/internal/control/test_helpers_test.go b/internal/control/test_helpers_test.go index efba2e9a..ffcd4ca6 100644 --- a/internal/control/test_helpers_test.go +++ b/internal/control/test_helpers_test.go @@ -12,8 +12,10 @@ import ( "gorm.io/gorm/logger" "gpt-load/internal/dialect" + "gpt-load/internal/health" "gpt-load/internal/platform/encryption" "gpt-load/internal/platform/i18n" + "gpt-load/internal/requestlog" "gpt-load/internal/state" "gpt-load/internal/storage" ) @@ -24,11 +26,25 @@ var ( ) type serviceFixture struct { - db *gorm.DB - manager *state.Manager - registry *state.KeyRegistry - encryption encryption.Service - service *Service + db *gorm.DB + manager *state.Manager + registry *state.KeyRegistry + encryption encryption.Service + stats *health.StatsStore + requestLogStats *staticRequestLogStatsReader + service *Service +} + +type staticRequestLogStatsReader struct { + value requestlog.Stats + fn func() requestlog.Stats +} + +func (reader *staticRequestLogStatsReader) Stats() requestlog.Stats { + if reader.fn != nil { + return reader.fn() + } + return reader.value } func initControlI18n(t *testing.T) { @@ -65,9 +81,21 @@ func newServiceFixtureWithDSN(t *testing.T, dsn string) serviceFixture { if _, err := manager.Publish(state.CompileInput{}); err != nil { t.Fatalf("manager.Publish(empty) error = %v", err) } + stats := health.NewStatsStore() + requestLogStats := &staticRequestLogStatsReader{} return serviceFixture{ db: db, manager: manager, registry: registry, encryption: keyService, - service: NewService(db, manager, registry, keyService, dialect.NewSet(), nil), + stats: stats, requestLogStats: requestLogStats, + service: NewService( + db, + manager, + registry, + keyService, + dialect.NewSet(), + nil, + stats, + requestLogStats, + ), } } diff --git a/internal/gateway/dependency_test.go b/internal/gateway/dependency_test.go index bc63ace7..05a2108a 100644 --- a/internal/gateway/dependency_test.go +++ b/internal/gateway/dependency_test.go @@ -10,6 +10,7 @@ func TestDataPlaneDependencyGraph(t *testing.T) { roots := []string{ "gpt-load/internal/gateway", "gpt-load/internal/scheduler", + "gpt-load/internal/state", "gpt-load/internal/health", "gpt-load/internal/telemetry", } diff --git a/internal/health/stats.go b/internal/health/stats.go index 0ee92c0b..c76f5d04 100644 --- a/internal/health/stats.go +++ b/internal/health/stats.go @@ -8,6 +8,7 @@ import ( const ( statsBucketCount = 5 statsBucketWidth = time.Minute + StatsWindow = time.Duration(statsBucketCount) * statsBucketWidth ) type KeyStats struct { diff --git a/internal/health/stats_test.go b/internal/health/stats_test.go index 5bdd11b2..377e2f0e 100644 --- a/internal/health/stats_test.go +++ b/internal/health/stats_test.go @@ -10,6 +10,19 @@ func statsBase() time.Time { return time.Date(2026, time.July, 22, 10, 0, 0, 0, time.UTC) } +func TestStatsWindowMatchesBucketConfiguration(t *testing.T) { + if StatsWindow != 5*time.Minute { + t.Fatalf("StatsWindow = %s, want 5m", StatsWindow) + } + if StatsWindow != time.Duration(statsBucketCount)*statsBucketWidth { + t.Fatalf( + "StatsWindow = %s, bucket configuration = %s", + StatsWindow, + time.Duration(statsBucketCount)*statsBucketWidth, + ) + } +} + func TestStatsStoreSnapshotUnknownKeyReturnsZero(t *testing.T) { store := NewStatsStore() now := statsBase() diff --git a/internal/scheduler/dependency_test.go b/internal/scheduler/dependency_test.go new file mode 100644 index 00000000..c5be086d --- /dev/null +++ b/internal/scheduler/dependency_test.go @@ -0,0 +1,37 @@ +package scheduler + +import ( + "os/exec" + "strings" + "testing" +) + +func TestRuntimeReadDomainDependencyGraph(t *testing.T) { + roots := []string{ + "gpt-load/internal/scheduler", + "gpt-load/internal/state", + "gpt-load/internal/health", + } + forbidden := []string{ + "gpt-load/internal/control", + "gpt-load/internal/storage", + "gpt-load/internal/gateway", + "gpt-load/internal/requestlog", + "gpt-load/internal/platform/encryption", + "gorm.io/gorm", + "github.com/gin-gonic/gin", + } + for _, root := range roots { + output, err := exec.Command("go", "list", "-deps", root).CombinedOutput() + if err != nil { + t.Fatalf("go list -deps %s: %v\n%s", root, err, output) + } + for _, dependency := range strings.Fields(string(output)) { + for _, denied := range forbidden { + if dependency == denied || strings.HasPrefix(dependency, denied+"/") { + t.Errorf("%s transitively depends on forbidden package %s", root, dependency) + } + } + } + } +} diff --git a/internal/scheduler/inspect.go b/internal/scheduler/inspect.go new file mode 100644 index 00000000..9ce9db63 --- /dev/null +++ b/internal/scheduler/inspect.go @@ -0,0 +1,290 @@ +package scheduler + +import ( + "errors" + "fmt" + "sort" + "time" + + "gpt-load/internal/protocol" + "gpt-load/internal/state" +) + +var ErrInconsistentSnapshot = errors.New("inconsistent scheduler snapshot") + +type ReasonCode string + +const ( + ReasonAccessKeyDisabled ReasonCode = "access_key_disabled" + ReasonProtocolFiltered ReasonCode = "protocol_filtered" + ReasonModelFiltered ReasonCode = "model_filtered" + ReasonNoRouteTarget ReasonCode = "no_route_target" + ReasonGroupDisabled ReasonCode = "group_disabled" + ReasonGroupFiltered ReasonCode = "group_filtered" + ReasonNoAvailableGroup ReasonCode = "no_available_group" + ReasonNoKeys ReasonCode = "no_keys" + ReasonGroupWeightZero ReasonCode = "group_weight_zero" + ReasonKeyDisabled ReasonCode = "key_disabled" + ReasonKeyBlacklisted ReasonCode = "key_blacklisted" + ReasonKeyCooldown ReasonCode = "key_cooldown" + ReasonKeyWeightZero ReasonCode = "key_weight_zero" + ReasonNoAvailableKey ReasonCode = "no_available_key" +) + +type Inspection struct { + Routable bool + Reason ReasonCode + Groups []GroupInspection +} + +type GroupInspection struct { + GroupID uint + GroupName string + UpstreamModelID string + WeightManual *int + Included bool + Routable bool + Reason ReasonCode + Keys []KeyInspection +} + +type KeyInspection struct { + KeyID uint + Available bool + Reason ReasonCode + WeightManual *int + WeightAuto int + EffectiveWeight int64 + CooldownUntil time.Time +} + +type targetDecision struct { + target state.RouteTarget + group state.GroupCatalogView + included bool + reason ReasonCode +} + +func cloneWeight(weight *int) *int { + if weight == nil { + return nil + } + value := *weight + return &value +} + +func evaluateTargets( + snapshot *state.ConfigSnapshot, + index map[protocol.Protocol]map[string][]state.RouteTarget, + query Query, +) ([]targetDecision, ReasonCode, error) { + if snapshot == nil { + return nil, "", fmt.Errorf("%w: nil ConfigSnapshot", ErrInconsistentSnapshot) + } + if len(query.AccessKey.Filters.Protocols) > 0 { + if _, allowed := query.AccessKey.Filters.Protocols[query.Protocol]; !allowed { + return []targetDecision{}, ReasonProtocolFiltered, nil + } + } + if len(query.AccessKey.Filters.Models) > 0 { + if _, allowed := query.AccessKey.Filters.Models[query.ExternalModel]; !allowed { + return []targetDecision{}, ReasonModelFiltered, nil + } + } + routes := index[query.Protocol][query.ExternalModel] + if len(routes) == 0 { + return []targetDecision{}, ReasonNoRouteTarget, nil + } + + decisions := make([]targetDecision, 0, len(routes)) + seenGroups := make(map[uint]struct{}, len(routes)) + included := 0 + for _, target := range routes { + if _, duplicate := seenGroups[target.GroupID]; duplicate { + continue + } + seenGroups[target.GroupID] = struct{}{} + group, exists := snapshot.GroupCatalog[target.GroupID] + if !exists { + return nil, "", fmt.Errorf( + "%w: route target group %d missing from catalog", + ErrInconsistentSnapshot, + target.GroupID, + ) + } + decision := targetDecision{target: target, group: group, included: true} + switch { + case !group.Enabled: + decision.included = false + decision.reason = ReasonGroupDisabled + case len(query.AccessKey.Filters.Groups) > 0: + if _, allowed := query.AccessKey.Filters.Groups[target.GroupID]; !allowed { + decision.included = false + decision.reason = ReasonGroupFiltered + } + } + if decision.included { + included++ + } + decisions = append(decisions, decision) + } + if included > 0 { + return decisions, "", nil + } + reason := decisions[0].reason + for _, decision := range decisions[1:] { + if decision.reason != reason { + return decisions, ReasonNoAvailableGroup, nil + } + } + return decisions, reason, nil +} + +func normalizedAutoWeight(weight int) int { + if weight == 0 { + return state.DefaultWeight + } + return weight +} + +func effectiveWeight(groupManual, keyManual *int, keyAuto int) int64 { + groupWeight := state.DefaultWeight + if groupManual != nil { + groupWeight = *groupManual + } + keyWeight := normalizedAutoWeight(keyAuto) + if keyManual != nil { + keyWeight = *keyManual + } + if groupWeight <= 0 || keyWeight <= 0 { + return 0 + } + return int64(groupWeight) * int64(keyWeight) +} + +func inspectKey( + group state.GroupCatalogView, + key state.KeyRuntimeView, + now time.Time, +) KeyInspection { + result := KeyInspection{ + KeyID: key.ID, WeightManual: cloneWeight(key.WeightManual), + WeightAuto: normalizedAutoWeight(key.WeightAuto), + } + if group.WeightManual != nil && *group.WeightManual == 0 { + result.Reason = ReasonGroupWeightZero + return result + } + if key.Status != state.KeyStatusActive { + result.Reason = ReasonKeyDisabled + return result + } + if key.WeightManual != nil && *key.WeightManual == 0 { + result.Reason = ReasonKeyWeightZero + return result + } + switch key.RuntimeState(now) { + case state.KeyRuntimeBlacklisted: + result.Reason = ReasonKeyBlacklisted + case state.KeyRuntimeCooldown: + result.Reason = ReasonKeyCooldown + result.CooldownUntil = key.CooldownUntil + default: + result.Available = true + result.EffectiveWeight = effectiveWeight( + group.WeightManual, + key.WeightManual, + key.WeightAuto, + ) + } + return result +} + +func Inspect( + snapshot *state.ConfigSnapshot, + keys []state.KeyRuntimeView, + query Query, + now time.Time, +) (Inspection, error) { + result := Inspection{Groups: []GroupInspection{}} + if snapshot == nil { + return Inspection{}, fmt.Errorf("%w: nil ConfigSnapshot", ErrInconsistentSnapshot) + } + if query.AccessKey.Status == state.AccessKeyStatusDisabled { + result.Reason = ReasonAccessKeyDisabled + return result, nil + } + + keysByGroup := make(map[uint][]state.KeyRuntimeView) + for _, key := range keys { + if _, exists := snapshot.GroupCatalog[key.GroupID]; !exists { + return Inspection{}, fmt.Errorf( + "%w: Registry key %d group %d missing from catalog", + ErrInconsistentSnapshot, + key.ID, + key.GroupID, + ) + } + cloned := key + cloned.WeightManual = cloneWeight(key.WeightManual) + keysByGroup[key.GroupID] = append(keysByGroup[key.GroupID], cloned) + } + for groupID := range keysByGroup { + sort.Slice(keysByGroup[groupID], func(i, j int) bool { + return keysByGroup[groupID][i].ID < keysByGroup[groupID][j].ID + }) + } + + decisions, staticReason, err := evaluateTargets( + snapshot, + snapshot.RouteCatalog, + query, + ) + if err != nil { + return Inspection{}, err + } + for _, decision := range decisions { + groupResult := GroupInspection{ + GroupID: decision.group.ID, GroupName: decision.group.Name, + UpstreamModelID: decision.target.UpstreamModelID, + WeightManual: cloneWeight(decision.group.WeightManual), + Included: decision.included, Reason: decision.reason, + Keys: []KeyInspection{}, + } + if !decision.included { + result.Groups = append(result.Groups, groupResult) + continue + } + groupKeys := keysByGroup[decision.group.ID] + groupWeightZero := decision.group.WeightManual != nil && + *decision.group.WeightManual == 0 + for _, key := range groupKeys { + keyResult := inspectKey(decision.group, key, now) + if keyResult.Available && keyResult.EffectiveWeight > 0 { + groupResult.Routable = true + } + groupResult.Keys = append(groupResult.Keys, keyResult) + } + switch { + case groupWeightZero: + groupResult.Reason = ReasonGroupWeightZero + case len(groupKeys) == 0: + groupResult.Reason = ReasonNoKeys + case !groupResult.Routable: + groupResult.Reason = ReasonNoAvailableKey + } + if groupResult.Routable { + result.Routable = true + } + result.Groups = append(result.Groups, groupResult) + } + if result.Routable { + return result, nil + } + if staticReason != "" { + result.Reason = staticReason + } else { + result.Reason = ReasonNoAvailableKey + } + return result, nil +} diff --git a/internal/scheduler/inspect_test.go b/internal/scheduler/inspect_test.go new file mode 100644 index 00000000..e35ffdd4 --- /dev/null +++ b/internal/scheduler/inspect_test.go @@ -0,0 +1,393 @@ +package scheduler + +import ( + "errors" + "fmt" + "math/rand" + "reflect" + "testing" + "time" + + "gpt-load/internal/protocol" + "gpt-load/internal/state" +) + +func inspectNow() time.Time { + return time.Date(2026, time.July, 24, 10, 0, 0, 0, time.UTC) +} + +func inspectSnapshot(t *testing.T) *state.ConfigSnapshot { + t.Helper() + zero := 0 + snapshot, err := state.Compile(state.CompileInput{ + Groups: []state.GroupConfig{ + { + ID: 2, Name: "disabled", Protocols: []protocol.Protocol{protocol.OpenAI}, + Models: []state.ModelConfig{{ID: "provider-disabled", Alias: "public"}}, + Enabled: false, + }, + { + ID: 1, Name: "active", Protocols: []protocol.Protocol{protocol.OpenAI}, + Models: []state.ModelConfig{{ID: "provider-active", Alias: "public"}}, + Enabled: true, + }, + { + ID: 3, Name: "weight-zero", Protocols: []protocol.Protocol{protocol.OpenAI}, + Models: []state.ModelConfig{{ID: "provider-zero", Alias: "public"}}, + WeightManual: &zero, Enabled: true, + }, + }, + }) + if err != nil { + t.Fatalf("Compile() error = %v", err) + } + return snapshot +} + +func TestInspectAppliesTopLevelReasonPriority(t *testing.T) { + snapshot := inspectSnapshot(t) + tests := []struct { + name string + query Query + reason ReasonCode + }{ + { + name: "access key disabled", + query: Query{ + Protocol: protocol.OpenAI, ExternalModel: "public", + AccessKey: state.AccessKeyView{Status: state.AccessKeyStatusDisabled}, + }, + reason: ReasonAccessKeyDisabled, + }, + { + name: "protocol filtered before model", + query: Query{ + Protocol: protocol.OpenAI, ExternalModel: "public", + AccessKey: state.AccessKeyView{ + Status: state.AccessKeyStatusActive, + Filters: state.FilterSet{ + Protocols: map[protocol.Protocol]struct{}{protocol.Anthropic: {}}, + Models: map[string]struct{}{"other": {}}, + }, + }, + }, + reason: ReasonProtocolFiltered, + }, + { + name: "model filtered", + query: Query{ + Protocol: protocol.OpenAI, ExternalModel: "public", + AccessKey: state.AccessKeyView{ + Status: state.AccessKeyStatusActive, + Filters: state.FilterSet{Models: map[string]struct{}{"other": {}}}, + }, + }, + reason: ReasonModelFiltered, + }, + { + name: "no route target", + query: Query{ + Protocol: protocol.OpenAI, ExternalModel: "missing", + AccessKey: state.AccessKeyView{Status: state.AccessKeyStatusActive}, + }, + reason: ReasonNoRouteTarget, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + got, err := Inspect(snapshot, nil, test.query, inspectNow()) + if err != nil { + t.Fatalf("Inspect() error = %v", err) + } + if got.Routable || got.Reason != test.reason || len(got.Groups) != 0 { + t.Fatalf("Inspection = %#v, want reason %q and empty groups", got, test.reason) + } + }) + } +} + +func TestInspectExplainsGroupsAndKeysInStableOrder(t *testing.T) { + now := inspectNow() + snapshot := inspectSnapshot(t) + zero := 0 + keys := []state.KeyRuntimeView{ + { + ID: 32, GroupID: 3, Status: state.KeyStatusActive, + WeightAuto: state.DefaultWeight, + }, + { + ID: 12, GroupID: 1, Status: state.KeyStatusActive, + WeightAuto: state.DefaultWeight, CooldownUntil: now.Add(time.Minute), + }, + { + ID: 11, GroupID: 1, Status: state.KeyStatusActive, + WeightManual: &zero, WeightAuto: state.DefaultWeight, + }, + { + ID: 13, GroupID: 1, Status: state.KeyStatusActive, + WeightAuto: 40, + }, + } + before := append([]state.KeyRuntimeView(nil), keys...) + got, err := Inspect(snapshot, keys, Query{ + Protocol: protocol.OpenAI, ExternalModel: "public", + AccessKey: state.AccessKeyView{ + Status: state.AccessKeyStatusActive, + Filters: state.FilterSet{Groups: map[uint]struct{}{1: {}, 3: {}}}, + }, + }, now) + if err != nil { + t.Fatalf("Inspect() error = %v", err) + } + if !reflect.DeepEqual(keys, before) { + t.Fatalf("Inspect() mutated input: got=%#v want=%#v", keys, before) + } + if !got.Routable || got.Reason != "" || len(got.Groups) != 3 { + t.Fatalf("Inspection = %#v", got) + } + if got.Groups[0].GroupID != 1 || got.Groups[1].GroupID != 2 || + got.Groups[2].GroupID != 3 { + t.Fatalf("group order = %#v", got.Groups) + } + if got.Groups[1].Included || got.Groups[1].Reason != ReasonGroupDisabled || + len(got.Groups[1].Keys) != 0 { + t.Fatalf("disabled group = %#v", got.Groups[1]) + } + active := got.Groups[0] + if !active.Included || !active.Routable || len(active.Keys) != 3 || + active.Keys[0].KeyID != 11 || active.Keys[1].KeyID != 12 || + active.Keys[2].KeyID != 13 { + t.Fatalf("active group = %#v", active) + } + if active.Keys[0].Reason != ReasonKeyWeightZero || + active.Keys[1].Reason != ReasonKeyCooldown || + !active.Keys[2].Available || active.Keys[2].EffectiveWeight != 50*40 { + t.Fatalf("key explanations = %#v", active.Keys) + } + if got.Groups[2].Reason != ReasonGroupWeightZero || + got.Groups[2].Keys[0].Reason != ReasonGroupWeightZero { + t.Fatalf("weight-zero group = %#v", got.Groups[2]) + } +} + +func TestInspectEligiblePoolMatchesIteratorInitialWeightedPool(t *testing.T) { + now := inspectNow() + snapshot := inspectSnapshot(t) + manual := 25 + zero := 0 + registry := state.NewKeyRegistry() + if err := registry.Replace([]state.KeyEntry{ + {ID: 11, GroupID: 1, Status: state.KeyStatusActive, WeightAuto: 40, EncryptedValue: "one"}, + { + ID: 12, GroupID: 1, Status: state.KeyStatusActive, + WeightManual: &manual, WeightAuto: 80, EncryptedValue: "two", + }, + { + ID: 13, GroupID: 1, Status: state.KeyStatusActive, + CooldownUntil: now.Add(time.Minute), EncryptedValue: "cooldown", + }, + { + ID: 14, GroupID: 1, Status: state.KeyStatusActive, + Blacklisted: true, EncryptedValue: "blacklisted", + }, + { + ID: 15, GroupID: 1, Status: state.KeyStatusDisabled, + EncryptedValue: "disabled", + }, + { + ID: 16, GroupID: 1, Status: state.KeyStatusActive, + WeightManual: &zero, EncryptedValue: "key-zero", + }, + {ID: 21, GroupID: 2, Status: state.KeyStatusActive, EncryptedValue: "disabled-group"}, + {ID: 31, GroupID: 3, Status: state.KeyStatusActive, EncryptedValue: "zero-group"}, + }); err != nil { + t.Fatalf("Replace() error = %v", err) + } + query := Query{ + Protocol: protocol.OpenAI, ExternalModel: "public", + AccessKey: state.AccessKeyView{Status: state.AccessKeyStatusActive}, + } + inspection, err := Inspect(snapshot, registry.Snapshot(), query, now) + if err != nil { + t.Fatalf("Inspect() error = %v", err) + } + inspectPool := make(map[uint]int64) + for _, group := range inspection.Groups { + for _, key := range group.Keys { + if key.Available && key.EffectiveWeight > 0 { + inspectPool[key.KeyID] = key.EffectiveWeight + } + } + } + wantPool := map[uint]int64{11: 50 * 40, 12: 50 * 25} + if !reflect.DeepEqual(inspectPool, wantPool) { + t.Fatalf("Inspector pool = %#v, want %#v", inspectPool, wantPool) + } + + iterator := newWithClock( + snapshot, + registry, + query, + rand.New(rand.NewSource(1)), + func() time.Time { return now }, + ) + weighted, _ := iterator.weightedPool(now) + iteratorPool := make(map[uint]int64, len(weighted)) + for _, key := range weighted { + iteratorPool[key.meta.ID] = key.weight + } + if !reflect.DeepEqual(inspectPool, iteratorPool) { + t.Fatalf("Inspector pool = %#v, Iterator pool = %#v", inspectPool, iteratorPool) + } + for _, group := range inspection.Groups { + for _, key := range group.Keys { + if key.Reason != "" { + if _, selected := iteratorPool[key.KeyID]; selected { + t.Fatalf("excluded key %d entered Iterator pool", key.KeyID) + } + } + } + } +} + +func TestInspectRejectsCatalogRegistryMismatch(t *testing.T) { + _, err := Inspect(inspectSnapshot(t), []state.KeyRuntimeView{{ + ID: 99, GroupID: 999, Status: state.KeyStatusActive, + }}, Query{ + Protocol: protocol.OpenAI, ExternalModel: "public", + AccessKey: state.AccessKeyView{Status: state.AccessKeyStatusActive}, + }, inspectNow()) + if !errors.Is(err, ErrInconsistentSnapshot) { + t.Fatalf("Inspect() error = %v, want ErrInconsistentSnapshot", err) + } +} + +func TestInspectReportsNoKeysForIncludedGroup(t *testing.T) { + got, err := Inspect(inspectSnapshot(t), nil, Query{ + Protocol: protocol.OpenAI, + ExternalModel: "public", + AccessKey: state.AccessKeyView{ + Status: state.AccessKeyStatusActive, + Filters: state.FilterSet{Groups: map[uint]struct{}{1: {}}}, + }, + }, inspectNow()) + if err != nil { + t.Fatalf("Inspect() error = %v", err) + } + if got.Routable || got.Reason != ReasonNoAvailableKey || + len(got.Groups) != 3 || + !got.Groups[0].Included || + got.Groups[0].Reason != ReasonNoKeys || + got.Groups[0].Keys == nil || + len(got.Groups[0].Keys) != 0 { + t.Fatalf("no-key Inspection = %#v", got) + } +} + +func TestInspectUsesExactKeyReasonPriority(t *testing.T) { + now := inspectNow() + zero := 0 + snapshot := inspectSnapshot(t) + keys := []state.KeyRuntimeView{ + { + ID: 11, GroupID: 1, Status: state.KeyStatusDisabled, + WeightManual: &zero, Blacklisted: true, + CooldownUntil: now.Add(time.Minute), + }, + { + ID: 12, GroupID: 1, Status: state.KeyStatusActive, + WeightManual: &zero, Blacklisted: true, + }, + { + ID: 13, GroupID: 1, Status: state.KeyStatusActive, + Blacklisted: true, CooldownUntil: now.Add(time.Minute), + }, + { + ID: 14, GroupID: 1, Status: state.KeyStatusActive, + CooldownUntil: now.Add(time.Minute), + }, + } + got, err := Inspect(snapshot, keys, Query{ + Protocol: protocol.OpenAI, ExternalModel: "public", + AccessKey: state.AccessKeyView{ + Status: state.AccessKeyStatusActive, + Filters: state.FilterSet{Groups: map[uint]struct{}{1: {}}}, + }, + }, now) + if err != nil { + t.Fatalf("Inspect() error = %v", err) + } + want := []ReasonCode{ + ReasonKeyDisabled, + ReasonKeyWeightZero, + ReasonKeyBlacklisted, + ReasonKeyCooldown, + } + for index, reason := range want { + if got.Groups[0].Keys[index].Reason != reason { + t.Fatalf("key %d reason = %q, want %q", index, got.Groups[0].Keys[index].Reason, reason) + } + } + if got.Routable || got.Reason != ReasonNoAvailableKey || + got.Groups[0].Reason != ReasonNoAvailableKey { + t.Fatalf("unavailable Inspection = %#v", got) + } +} + +func TestInspectSummarizesStaticGroupExclusions(t *testing.T) { + group := func(id uint, enabled bool) state.GroupConfig { + return state.GroupConfig{ + ID: id, Name: fmt.Sprintf("group-%d", id), + Protocols: []protocol.Protocol{protocol.OpenAI}, + Models: []state.ModelConfig{{ID: fmt.Sprintf("provider-%d", id), Alias: "public"}}, + Enabled: enabled, + } + } + tests := []struct { + name string + groups []state.GroupConfig + allowedGroups map[uint]struct{} + want ReasonCode + }{ + { + name: "all disabled", + groups: []state.GroupConfig{group(1, false), group(2, false)}, + want: ReasonGroupDisabled, + }, + { + name: "all filtered", + groups: []state.GroupConfig{group(1, true), group(2, true)}, + allowedGroups: map[uint]struct{}{999: {}}, + want: ReasonGroupFiltered, + }, + { + name: "mixed disabled and filtered", + groups: []state.GroupConfig{group(1, false), group(2, true)}, + allowedGroups: map[uint]struct{}{999: {}}, + want: ReasonNoAvailableGroup, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + snapshot, err := state.Compile(state.CompileInput{Groups: test.groups}) + if err != nil { + t.Fatalf("Compile() error = %v", err) + } + result, err := Inspect(snapshot, nil, Query{ + Protocol: protocol.OpenAI, + ExternalModel: "public", + AccessKey: state.AccessKeyView{ + Status: state.AccessKeyStatusActive, + Filters: state.FilterSet{Groups: test.allowedGroups}, + }, + }, inspectNow()) + if err != nil { + t.Fatalf("Inspect() error = %v", err) + } + if result.Routable || result.Reason != test.want || + len(result.Groups) != 2 { + t.Fatalf("Inspection = %#v, want reason %q", result, test.want) + } + }) + } +} diff --git a/internal/scheduler/scheduler.go b/internal/scheduler/scheduler.go index 4bce8629..31ac5955 100644 --- a/internal/scheduler/scheduler.go +++ b/internal/scheduler/scheduler.go @@ -85,14 +85,14 @@ func (iterator *Iterator) SkipGroup(groupID uint) { iterator.skippedGroups[groupID] = struct{}{} } -func (iterator *Iterator) Next() (Selection, error) { - if iterator == nil || iterator.keys == nil || iterator.random == nil || iterator.now == nil || len(iterator.groupIDs) == 0 { - return Selection{}, ErrExhausted +func (iterator *Iterator) weightedPool(now time.Time) ([]weightedKey, int64) { + if iterator == nil || iterator.keys == nil || len(iterator.groupIDs) == 0 { + return nil, 0 } pool := iterator.keys.CollectCandidates(iterator.groupIDs, func(keyID uint) bool { _, tried := iterator.tried[keyID] return tried - }, iterator.now()) + }, now) weighted := make([]weightedKey, 0, len(pool)) var total int64 for _, key := range pool { @@ -103,13 +103,25 @@ func (iterator *Iterator) Next() (Selection, error) { if !ok { continue } - weight := effectiveWeight(key, target.group) + weight := effectiveWeight( + target.group.WeightManual, + key.WeightManual, + key.WeightAuto, + ) if weight <= 0 { continue } weighted = append(weighted, weightedKey{meta: key, weight: weight}) total += weight } + return weighted, total +} + +func (iterator *Iterator) Next() (Selection, error) { + if iterator == nil || iterator.random == nil || iterator.now == nil { + return Selection{}, ErrExhausted + } + weighted, total := iterator.weightedPool(iterator.now()) if total <= 0 { return Selection{}, ErrExhausted } @@ -133,57 +145,27 @@ func (iterator *Iterator) Next() (Selection, error) { }, nil } -func effectiveWeight(key state.KeyMeta, group state.GroupView) int64 { - groupWeight := state.DefaultWeight - if group.WeightManual != nil { - groupWeight = *group.WeightManual - } - keyWeight := key.WeightAuto - if keyWeight == 0 { - keyWeight = state.DefaultWeight - } - if key.WeightManual != nil { - keyWeight = *key.WeightManual - } - if groupWeight <= 0 || keyWeight <= 0 { - return 0 - } - return int64(groupWeight) * int64(keyWeight) -} - func filterTargets(snapshot *state.ConfigSnapshot, query Query) []candidateTarget { if snapshot == nil { return nil } - if len(query.AccessKey.Filters.Protocols) > 0 { - if _, ok := query.AccessKey.Filters.Protocols[query.Protocol]; !ok { - return nil - } + decisions, _, err := evaluateTargets(snapshot, snapshot.Candidates, query) + if err != nil { + return nil } - if len(query.AccessKey.Filters.Models) > 0 { - if _, ok := query.AccessKey.Filters.Models[query.ExternalModel]; !ok { - return nil - } - } - - byModel := snapshot.Candidates[query.Protocol] - targets := make([]candidateTarget, 0, len(byModel[query.ExternalModel])) - seenGroups := make(map[uint]struct{}) - for _, target := range byModel[query.ExternalModel] { - if len(query.AccessKey.Filters.Groups) > 0 { - if _, ok := query.AccessKey.Filters.Groups[target.GroupID]; !ok { - continue - } - } - group, ok := snapshot.Groups[target.GroupID] - if !ok { + targets := make([]candidateTarget, 0, len(decisions)) + for _, decision := range decisions { + if !decision.included { continue } - if _, duplicate := seenGroups[target.GroupID]; duplicate { + group, exists := snapshot.Groups[decision.target.GroupID] + if !exists { continue } - seenGroups[target.GroupID] = struct{}{} - targets = append(targets, candidateTarget{target: target, group: group}) + targets = append(targets, candidateTarget{ + target: decision.target, + group: group, + }) } return targets } diff --git a/internal/scheduler/scheduler_test.go b/internal/scheduler/scheduler_test.go index 0b7ab3f8..6ce6eae1 100644 --- a/internal/scheduler/scheduler_test.go +++ b/internal/scheduler/scheduler_test.go @@ -326,7 +326,11 @@ func TestEffectiveWeightCombinesGroupAndKeyWeights(t *testing.T) { } for _, test := range tests { t.Run(test.name, func(t *testing.T) { - if got := effectiveWeight(test.key, test.group); got != test.want { + if got := effectiveWeight( + test.group.WeightManual, + test.key.WeightManual, + test.key.WeightAuto, + ); got != test.want { t.Fatalf("effectiveWeight() = %d, want %d", got, test.want) } }) @@ -546,5 +550,9 @@ func schedulerSnapshot() *state.ConfigSnapshot { 1: {ID: 1, Name: "one", UpstreamURL: "https://one.example"}, 2: {ID: 2, Name: "two", UpstreamURL: "https://two.example"}, }, + GroupCatalog: map[uint]state.GroupCatalogView{ + 1: {ID: 1, Name: "one", Enabled: true}, + 2: {ID: 2, Name: "two", Enabled: true}, + }, } } diff --git a/internal/state/loader/loader_test.go b/internal/state/loader/loader_test.go index b1001d19..23710ee0 100644 --- a/internal/state/loader/loader_test.go +++ b/internal/state/loader/loader_test.go @@ -329,6 +329,10 @@ func TestLoaderMapsSystemAndGroupRows(t *testing.T) { if _, ok := snapshot.Groups[disabled.ID]; ok { t.Fatalf("disabled group %d is present in snapshot", disabled.ID) } + if got := snapshot.GroupCatalog[disabled.ID]; got.ID != disabled.ID || + got.Name != disabled.Name || got.Enabled { + t.Fatalf("disabled GroupCatalog entry = %#v", got) + } if len(view.Models) != 3 || view.Models[0].Alias != "Primary" || view.Models[1].Alias != "Secondary" { t.Errorf("group models = %#v, want all aliases retained", view.Models) } @@ -358,6 +362,10 @@ func TestLoaderMapsSystemAndGroupRows(t *testing.T) { if _, ok := openAICandidates["hidden"]; ok { t.Fatal("disabled group model hidden is present in candidates") } + if got := snapshot.RouteCatalog[protocol.OpenAI]["Hidden"]; len(got) != 1 || + got[0].GroupID != disabled.ID || got[0].UpstreamModelID != "hidden" { + t.Fatalf("disabled RouteCatalog entry = %#v", got) + } } func TestLoaderMapsValidationModelIntoRuntimeSnapshot(t *testing.T) { @@ -500,6 +508,17 @@ func TestLoaderMapsAccessAndUpstreamKeys(t *testing.T) { if _, ok := snapshot.AccessKeysByHash[disabledAccess.KeyHash]; ok { t.Fatalf("disabled access hash %q is present in snapshot", disabledAccess.KeyHash) } + if got := snapshot.AccessKeysByID[disabledAccess.ID]; got.ID != disabledAccess.ID || + got.Status != state.AccessKeyStatusDisabled || got.Name != disabledAccess.Name { + t.Fatalf("disabled AccessKeysByID entry = %#v", got) + } + if got := snapshot.GroupCatalog[secondGroup.ID]; got.ID != secondGroup.ID || + got.Name != secondGroup.Name || !got.Enabled { + t.Fatalf("second GroupCatalog entry = %#v", got) + } + if len(snapshot.RouteCatalog) == 0 { + t.Fatal("RouteCatalog was not compiled from persisted groups") + } if _, ok := access.Filters.Groups[firstGroup.ID]; !ok { t.Errorf("access filters groups = %#v, want group %d", access.Filters.Groups, firstGroup.ID) } diff --git a/internal/state/registry.go b/internal/state/registry.go index 259e55bc..5f4e3b4f 100644 --- a/internal/state/registry.go +++ b/internal/state/registry.go @@ -208,12 +208,13 @@ func (r *KeyRegistry) CollectCandidates(groupIDs []uint, excluded func(uint) boo metas := make([]KeyMeta, 0) for _, groupID := range groupIDs { for _, entry := range r.buckets[groupID] { - if entry.Status != KeyStatusActive || entry.Blacklisted || entry.CooldownUntil.After(now) { + view := runtimeView(entry) + if view.RuntimeState(now) != KeyRuntimeAvailable { continue } meta := KeyMeta{ - ID: entry.ID, GroupID: entry.GroupID, - WeightManual: cloneWeight(entry.WeightManual), WeightAuto: entry.WeightAuto, + ID: view.ID, GroupID: view.GroupID, + WeightManual: cloneWeight(view.WeightManual), WeightAuto: view.WeightAuto, } metas = append(metas, meta) } diff --git a/internal/state/registry_test.go b/internal/state/registry_test.go index 803163fe..edcebae8 100644 --- a/internal/state/registry_test.go +++ b/internal/state/registry_test.go @@ -129,6 +129,56 @@ func TestKeyRegistryActiveKeyIDs(t *testing.T) { } } +func TestKeyRegistrySnapshotIsSortedDetachedAndCredentialFree(t *testing.T) { + weight := 17 + now := time.Date(2026, time.July, 24, 10, 0, 0, 0, time.UTC) + registry := NewKeyRegistry() + mustReplaceKeyEntries(t, registry, []KeyEntry{ + { + ID: 22, GroupID: 2, Status: KeyStatusDisabled, + EncryptedValue: "cipher-disabled", + }, + { + ID: 11, GroupID: 1, WeightManual: &weight, WeightAuto: 23, + Status: KeyStatusActive, CooldownUntil: now.Add(time.Minute), + Blacklisted: true, FailureCount: 3, EncryptedValue: "cipher-secret", + }, + { + ID: 12, GroupID: 1, Status: KeyStatusActive, + EncryptedValue: "cipher-default-weight", + }, + }) + + got := registry.Snapshot() + if len(got) != 3 || got[0].ID != 11 || got[1].ID != 12 || got[2].ID != 22 { + t.Fatalf("Snapshot order = %#v", got) + } + if got[0].WeightManual == nil || *got[0].WeightManual != 17 || + got[0].WeightAuto != 23 || !got[0].Blacklisted || + got[0].FailureCount != 3 || !got[0].CooldownUntil.Equal(now.Add(time.Minute)) { + t.Fatalf("Snapshot runtime values = %#v", got[0]) + } + if got[1].WeightAuto != DefaultWeight { + t.Fatalf("default WeightAuto = %d, want %d", got[1].WeightAuto, DefaultWeight) + } + *got[0].WeightManual = 99 + got[0].FailureCount = 99 + again := registry.Snapshot() + if again[0].WeightManual == nil || *again[0].WeightManual != 17 || + again[0].FailureCount != 3 { + t.Fatalf("Snapshot aliases Registry: %#v", again[0]) + } + + typ := reflect.TypeOf(KeyRuntimeView{}) + for _, forbidden := range []string{ + "EncryptedValue", "KeyHash", "Hash", "Mask", "HeaderRules", + } { + if _, exists := typ.FieldByName(forbidden); exists { + t.Fatalf("KeyRuntimeView exposes forbidden field %s", forbidden) + } + } +} + func TestKeyRegistryReplaceFailurePreservesRegistry(t *testing.T) { invalidBatches := map[string][]KeyEntry{ "duplicate ids": { @@ -423,6 +473,15 @@ func TestKeyRegistryConcurrentMutationsAndCollection(t *testing.T) { break } } + views := registry.Snapshot() + for index := 1; index < len(views); index++ { + previous, current := views[index-1], views[index] + if previous.GroupID > current.GroupID || + (previous.GroupID == current.GroupID && previous.ID > current.ID) { + errors <- fmt.Errorf("Snapshot() returned unstable order: %#v then %#v", previous, current) + break + } + } } }(reader) } diff --git a/internal/state/runtime_view.go b/internal/state/runtime_view.go new file mode 100644 index 00000000..c22723d1 --- /dev/null +++ b/internal/state/runtime_view.go @@ -0,0 +1,74 @@ +package state + +import ( + "sort" + "time" +) + +type KeyRuntimeView struct { + ID uint + GroupID uint + WeightManual *int + WeightAuto int + Status KeyStatus + CooldownUntil time.Time + Blacklisted bool + FailureCount int +} + +type KeyRuntimeState string + +const ( + KeyRuntimeAvailable KeyRuntimeState = "available" + KeyRuntimeDisabled KeyRuntimeState = "disabled" + KeyRuntimeBlacklisted KeyRuntimeState = "blacklisted" + KeyRuntimeCooldown KeyRuntimeState = "cooldown" +) + +func (view KeyRuntimeView) RuntimeState(now time.Time) KeyRuntimeState { + if view.Status != KeyStatusActive { + return KeyRuntimeDisabled + } + if view.Blacklisted { + return KeyRuntimeBlacklisted + } + if view.CooldownUntil.After(now) { + return KeyRuntimeCooldown + } + return KeyRuntimeAvailable +} + +func runtimeView(entry *KeyEntry) KeyRuntimeView { + return KeyRuntimeView{ + ID: entry.ID, + GroupID: entry.GroupID, + WeightManual: cloneWeight(entry.WeightManual), + WeightAuto: entry.WeightAuto, + Status: entry.Status, + CooldownUntil: entry.CooldownUntil, + Blacklisted: entry.Blacklisted, + FailureCount: entry.FailureCount, + } +} + +func sortRuntimeViews(views []KeyRuntimeView) { + sort.Slice(views, func(i, j int) bool { + if views[i].GroupID != views[j].GroupID { + return views[i].GroupID < views[j].GroupID + } + return views[i].ID < views[j].ID + }) +} + +func (r *KeyRegistry) Snapshot() []KeyRuntimeView { + r.mu.RLock() + views := make([]KeyRuntimeView, 0, len(r.keyGroups)) + for _, bucket := range r.buckets { + for _, entry := range bucket { + views = append(views, runtimeView(entry)) + } + } + r.mu.RUnlock() + sortRuntimeViews(views) + return views +} diff --git a/internal/state/snapshot.go b/internal/state/snapshot.go index 073443e7..a260e5fe 100644 --- a/internal/state/snapshot.go +++ b/internal/state/snapshot.go @@ -103,9 +103,17 @@ type GroupView struct { WeightManual *int } +type GroupCatalogView struct { + ID uint + Name string + Enabled bool + WeightManual *int +} + type AccessKeyView struct { ID uint Name string + Status AccessKeyStatus Filters FilterSet RPMLimit int64 } @@ -115,6 +123,9 @@ type ConfigSnapshot struct { Candidates map[protocol.Protocol]map[string][]RouteTarget Groups map[uint]GroupView AccessKeysByHash map[string]AccessKeyView + RouteCatalog map[protocol.Protocol]map[string][]RouteTarget + GroupCatalog map[uint]GroupCatalogView + AccessKeysByID map[uint]AccessKeyView } func Compile(input CompileInput) (*ConfigSnapshot, error) { @@ -126,9 +137,17 @@ func Compile(input CompileInput) (*ConfigSnapshot, error) { Candidates: make(map[protocol.Protocol]map[string][]RouteTarget), Groups: make(map[uint]GroupView), AccessKeysByHash: make(map[string]AccessKeyView), + RouteCatalog: make(map[protocol.Protocol]map[string][]RouteTarget), + GroupCatalog: make(map[uint]GroupCatalogView), + AccessKeysByID: make(map[uint]AccessKeyView), } for _, group := range input.Groups { + snapshot.GroupCatalog[group.ID] = GroupCatalogView{ + ID: group.ID, Name: group.Name, Enabled: group.Enabled, + WeightManual: cloneWeight(group.WeightManual), + } + appendRouteTarget(snapshot.RouteCatalog, group) if !group.Enabled { continue } @@ -148,36 +167,49 @@ func Compile(input CompileInput) (*ConfigSnapshot, error) { WeightManual: cloneWeight(group.WeightManual), } snapshot.Groups[group.ID] = view - - for _, model := range group.Models { - modelID := strings.TrimSpace(model.ID) - external := externalModelName(model) - for _, p := range group.Protocols { - if snapshot.Candidates[p] == nil { - snapshot.Candidates[p] = make(map[string][]RouteTarget) - } - snapshot.Candidates[p][external] = append( - snapshot.Candidates[p][external], - RouteTarget{GroupID: group.ID, UpstreamModelID: modelID}, - ) - } - } + appendRouteTarget(snapshot.Candidates, group) } for _, accessKey := range input.AccessKeys { - if accessKey.Status != AccessKeyStatusActive { - continue - } - snapshot.AccessKeysByHash[accessKey.KeyHash] = AccessKeyView{ - ID: accessKey.ID, Name: accessKey.Name, Filters: cloneFilterSet(accessKey.Filters), - RPMLimit: accessKey.RPMLimit, + snapshot.AccessKeysByID[accessKey.ID] = newAccessKeyView(accessKey) + if accessKey.Status == AccessKeyStatusActive { + snapshot.AccessKeysByHash[accessKey.KeyHash] = newAccessKeyView(accessKey) } } - for _, byModel := range snapshot.Candidates { - for modelID := range byModel { - sort.Slice(byModel[modelID], func(i, j int) bool { - left, right := byModel[modelID][i], byModel[modelID][j] + sortRouteIndex(snapshot.Candidates) + sortRouteIndex(snapshot.RouteCatalog) + return snapshot, nil +} + +func newAccessKeyView(input AccessKeyConfig) AccessKeyView { + return AccessKeyView{ + ID: input.ID, Name: input.Name, Status: input.Status, + Filters: cloneFilterSet(input.Filters), RPMLimit: input.RPMLimit, + } +} + +func appendRouteTarget(index map[protocol.Protocol]map[string][]RouteTarget, group GroupConfig) { + for _, model := range group.Models { + target := RouteTarget{ + GroupID: group.ID, + UpstreamModelID: strings.TrimSpace(model.ID), + } + external := externalModelName(model) + for _, value := range group.Protocols { + if index[value] == nil { + index[value] = make(map[string][]RouteTarget) + } + index[value][external] = append(index[value][external], target) + } + } +} + +func sortRouteIndex(index map[protocol.Protocol]map[string][]RouteTarget) { + for _, byModel := range index { + for model := range byModel { + sort.Slice(byModel[model], func(i, j int) bool { + left, right := byModel[model][i], byModel[model][j] if left.GroupID != right.GroupID { return left.GroupID < right.GroupID } @@ -185,7 +217,6 @@ func Compile(input CompileInput) (*ConfigSnapshot, error) { }) } } - return snapshot, nil } func compileRuntimeSettings(system, group config.Settings) (TimeoutConfig, HeaderRules, error) { @@ -362,66 +393,69 @@ func validHTTPHeaderValue(value string) bool { } func validateCompileInput(input CompileInput) error { - enabledGroupIDs := make(map[uint]struct{}) - activeHashes := make(map[string]struct{}) + groupIDs := make(map[uint]struct{}, len(input.Groups)) for _, group := range input.Groups { - if !group.Enabled { - continue - } if group.ID == 0 { return fmt.Errorf("group id is required") } - if _, duplicate := enabledGroupIDs[group.ID]; duplicate { + if _, duplicate := groupIDs[group.ID]; duplicate { return fmt.Errorf("duplicate group id %d", group.ID) } + groupIDs[group.ID] = struct{}{} if err := validateManualWeight(fmt.Sprintf("group %d", group.ID), group.WeightManual); err != nil { return err } - enabledGroupIDs[group.ID] = struct{}{} if len(group.Protocols) == 0 { return fmt.Errorf("group %d protocols are required", group.ID) } - protocols := make(map[protocol.Protocol]struct{}, len(group.Protocols)) - for _, p := range group.Protocols { - if !p.Valid() { - return fmt.Errorf("group %d has invalid protocol %q", group.ID, p) + seenProtocols := make(map[protocol.Protocol]struct{}, len(group.Protocols)) + for _, value := range group.Protocols { + if !value.Valid() { + return fmt.Errorf("group %d has invalid protocol %q", group.ID, value) } - if _, duplicate := protocols[p]; duplicate { - return fmt.Errorf("group %d has duplicate protocol %q", group.ID, p) + if _, duplicate := seenProtocols[value]; duplicate { + return fmt.Errorf("group %d has duplicate protocol %q", group.ID, value) } - protocols[p] = struct{}{} + seenProtocols[value] = struct{}{} } - externalModels := make(map[string]struct{}, len(group.Models)) + seenModels := make(map[string]struct{}, len(group.Models)) for _, model := range group.Models { - modelID := strings.TrimSpace(model.ID) - if modelID == "" { + if strings.TrimSpace(model.ID) == "" { return fmt.Errorf("group %d model id is required", group.ID) } external := externalModelName(model) - if _, duplicate := externalModels[external]; duplicate { + if _, duplicate := seenModels[external]; duplicate { return fmt.Errorf("group %d has duplicate external model %q", group.ID, external) } - externalModels[external] = struct{}{} + seenModels[external] = struct{}{} } } + + accessKeyIDs := make(map[uint]struct{}, len(input.AccessKeys)) + hashes := make(map[string]struct{}, len(input.AccessKeys)) for _, accessKey := range input.AccessKeys { + if accessKey.ID == 0 { + return fmt.Errorf("access key id is required") + } + if _, duplicate := accessKeyIDs[accessKey.ID]; duplicate { + return fmt.Errorf("duplicate access key id %d", accessKey.ID) + } + accessKeyIDs[accessKey.ID] = struct{}{} if accessKey.RPMLimit < 0 { return fmt.Errorf("access key %d rpm limit must not be negative", accessKey.ID) } switch accessKey.Status { - case AccessKeyStatusDisabled: - continue - case AccessKeyStatusActive: + case AccessKeyStatusActive, AccessKeyStatusDisabled: default: return fmt.Errorf("access key %d has invalid status %q", accessKey.ID, accessKey.Status) } if strings.TrimSpace(accessKey.KeyHash) == "" { return fmt.Errorf("access key %d key hash is required", accessKey.ID) } - if _, duplicate := activeHashes[accessKey.KeyHash]; duplicate { + if _, duplicate := hashes[accessKey.KeyHash]; duplicate { return fmt.Errorf("duplicate access key hash %q", accessKey.KeyHash) } - activeHashes[accessKey.KeyHash] = struct{}{} + hashes[accessKey.KeyHash] = struct{}{} if err := validateFilterSet(accessKey.ID, accessKey.Filters); err != nil { return err } diff --git a/internal/state/snapshot_test.go b/internal/state/snapshot_test.go index c7245b0f..917cf43a 100644 --- a/internal/state/snapshot_test.go +++ b/internal/state/snapshot_test.go @@ -52,6 +52,75 @@ func TestCompileIndexesExternalModelsAndPreservesUpstreamIDs(t *testing.T) { } } +func TestCompileBuildsManagementCatalogsWithoutChangingActiveIndexes(t *testing.T) { + disabledWeight := 20 + input := CompileInput{ + Groups: []GroupConfig{ + { + ID: 2, Name: "disabled", UpstreamURL: "https://disabled.example", + Protocols: []protocol.Protocol{protocol.OpenAI}, + Models: []ModelConfig{{ID: "provider-disabled", Alias: "public"}}, + WeightManual: &disabledWeight, Enabled: false, + }, + { + ID: 1, Name: "active", UpstreamURL: "https://active.example", + Protocols: []protocol.Protocol{protocol.OpenAI}, + Models: []ModelConfig{{ID: "provider-active", Alias: "public"}}, + Enabled: true, + }, + }, + AccessKeys: []AccessKeyConfig{ + { + ID: 11, Name: "active-client", KeyHash: "active-hash", + Status: AccessKeyStatusActive, + Filters: FilterSet{Groups: map[uint]struct{}{1: {}}}, + RPMLimit: 10, + }, + { + ID: 12, Name: "disabled-client", KeyHash: "disabled-hash", + Status: AccessKeyStatusDisabled, + Filters: FilterSet{Models: map[string]struct{}{"public": {}}}, + RPMLimit: 20, + }, + }, + } + + snapshot, err := Compile(input) + if err != nil { + t.Fatalf("Compile() error = %v", err) + } + if len(snapshot.Groups) != 1 || snapshot.Groups[1].Name != "active" { + t.Fatalf("active Groups = %#v", snapshot.Groups) + } + if got := snapshot.Candidates[protocol.OpenAI]["public"]; !reflect.DeepEqual(got, []RouteTarget{{ + GroupID: 1, UpstreamModelID: "provider-active", + }}) { + t.Fatalf("active Candidates = %#v", got) + } + wantRoutes := []RouteTarget{ + {GroupID: 1, UpstreamModelID: "provider-active"}, + {GroupID: 2, UpstreamModelID: "provider-disabled"}, + } + if got := snapshot.RouteCatalog[protocol.OpenAI]["public"]; !reflect.DeepEqual(got, wantRoutes) { + t.Fatalf("RouteCatalog = %#v, want %#v", got, wantRoutes) + } + if got := snapshot.GroupCatalog[2]; got.ID != 2 || got.Name != "disabled" || + got.Enabled || got.WeightManual == nil || *got.WeightManual != 20 { + t.Fatalf("disabled GroupCatalog entry = %#v", got) + } + if len(snapshot.AccessKeysByHash) != 1 || + snapshot.AccessKeysByHash["active-hash"].Status != AccessKeyStatusActive { + t.Fatalf("active AccessKeysByHash = %#v", snapshot.AccessKeysByHash) + } + if _, exists := snapshot.AccessKeysByHash["disabled-hash"]; exists { + t.Fatal("disabled key entered AccessKeysByHash") + } + if got := snapshot.AccessKeysByID[12]; got.ID != 12 || + got.Status != AccessKeyStatusDisabled || got.RPMLimit != 20 { + t.Fatalf("disabled AccessKeysByID entry = %#v", got) + } +} + func TestCompileCarriesValidationModelWithoutChangingCandidates(t *testing.T) { snapshot, err := Compile(CompileInput{Groups: []GroupConfig{{ ID: 1, @@ -95,16 +164,131 @@ func TestCompileRejectsDuplicateExternalModelWithinGroup(t *testing.T) { } } -func TestCompileSkipsDisabledGroupValidation(t *testing.T) { - snapshot, err := Compile(CompileInput{Groups: []GroupConfig{{ - ID: 1, Name: "disabled", Enabled: false, - Models: []ModelConfig{{ID: " ", Alias: "public"}, {ID: "other", Alias: "public"}}, - }}}) - if err != nil { - t.Fatalf("Compile() error = %v, want disabled group skipped", err) +func TestCompileValidatesDisabledCatalogEntries(t *testing.T) { + invalidWeight := -1 + tests := []struct { + name string + input CompileInput + wantErr string + }{ + { + name: "disabled group zero id", + input: CompileInput{Groups: []GroupConfig{{ + Protocols: []protocol.Protocol{protocol.OpenAI}, + Models: []ModelConfig{{ID: "model"}}, Enabled: false, + }}}, + wantErr: "group id is required", + }, + { + name: "disabled group invalid protocol", + input: CompileInput{Groups: []GroupConfig{{ + ID: 1, Protocols: []protocol.Protocol{"invalid"}, + Models: []ModelConfig{{ID: "model"}}, Enabled: false, + }}}, + wantErr: "invalid protocol", + }, + { + name: "disabled group invalid weight", + input: CompileInput{Groups: []GroupConfig{{ + ID: 1, Protocols: []protocol.Protocol{protocol.OpenAI}, + Models: []ModelConfig{{ID: "model"}}, + WeightManual: &invalidWeight, Enabled: false, + }}}, + wantErr: "manual weight", + }, + { + name: "disabled access key zero id", + input: CompileInput{AccessKeys: []AccessKeyConfig{{ + KeyHash: "disabled-hash", Status: AccessKeyStatusDisabled, + }}}, + wantErr: "access key id is required", + }, + { + name: "disabled access key empty hash", + input: CompileInput{AccessKeys: []AccessKeyConfig{{ + ID: 1, Status: AccessKeyStatusDisabled, + }}}, + wantErr: "key hash is required", + }, + { + name: "disabled access key invalid filter", + input: CompileInput{AccessKeys: []AccessKeyConfig{{ + ID: 1, KeyHash: "disabled-hash", Status: AccessKeyStatusDisabled, + Filters: FilterSet{Protocols: map[protocol.Protocol]struct{}{"invalid": {}}}, + }}}, + wantErr: "invalid protocol", + }, } - if len(snapshot.Groups) != 0 || len(snapshot.Candidates) != 0 { - t.Fatalf("snapshot includes disabled group: %#v", snapshot) + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + _, err := Compile(test.input) + if err == nil || !strings.Contains(err.Error(), test.wantErr) { + t.Fatalf("Compile() error = %v, want substring %q", err, test.wantErr) + } + }) + } +} + +func TestCompileRejectsDuplicateCatalogIDs(t *testing.T) { + validGroup := func(id uint, enabled bool) GroupConfig { + return GroupConfig{ + ID: id, Protocols: []protocol.Protocol{protocol.OpenAI}, + Models: []ModelConfig{{ID: "model"}}, Enabled: enabled, + } + } + if _, err := Compile(CompileInput{Groups: []GroupConfig{ + validGroup(1, true), validGroup(1, false), + }}); err == nil || !strings.Contains(err.Error(), "duplicate group id") { + t.Fatalf("duplicate Group catalog error = %v", err) + } + if _, err := Compile(CompileInput{AccessKeys: []AccessKeyConfig{ + {ID: 1, KeyHash: "one", Status: AccessKeyStatusActive}, + {ID: 1, KeyHash: "two", Status: AccessKeyStatusDisabled}, + }}); err == nil || !strings.Contains(err.Error(), "duplicate access key id") { + t.Fatalf("duplicate AccessKey catalog error = %v", err) + } +} + +func TestCompileRejectsDuplicateCatalogHashes(t *testing.T) { + _, err := Compile(CompileInput{AccessKeys: []AccessKeyConfig{ + {ID: 1, KeyHash: "duplicate", Status: AccessKeyStatusActive}, + {ID: 2, KeyHash: "duplicate", Status: AccessKeyStatusDisabled}, + }}) + if err == nil || !strings.Contains(err.Error(), "duplicate access key hash") { + t.Fatalf("duplicate AccessKey catalog hash error = %v", err) + } +} + +func TestCompileManagementCatalogsOwnCopies(t *testing.T) { + groupWeight := 20 + filterGroups := map[uint]struct{}{1: {}} + input := CompileInput{ + Groups: []GroupConfig{{ + ID: 1, Name: "group", Protocols: []protocol.Protocol{protocol.OpenAI}, + Models: []ModelConfig{{ID: "provider", Alias: "public"}}, + WeightManual: &groupWeight, Enabled: false, + }}, + AccessKeys: []AccessKeyConfig{{ + ID: 2, Name: "client", KeyHash: "hash", Status: AccessKeyStatusDisabled, + Filters: FilterSet{Groups: filterGroups}, + }}, + } + snapshot, err := Compile(input) + if err != nil { + t.Fatalf("Compile() error = %v", err) + } + groupWeight = 99 + delete(filterGroups, 1) + input.Groups[0].Models[0] = ModelConfig{ID: "changed", Alias: "changed"} + if got := snapshot.GroupCatalog[1].WeightManual; got == nil || *got != 20 { + t.Fatalf("GroupCatalog weight = %v, want 20", got) + } + if got := snapshot.RouteCatalog[protocol.OpenAI]["public"]; len(got) != 1 || + got[0].UpstreamModelID != "provider" { + t.Fatalf("RouteCatalog after caller mutation = %#v", got) + } + if _, ok := snapshot.AccessKeysByID[2].Filters.Groups[1]; !ok { + t.Fatal("AccessKeysByID filters changed with caller map") } } diff --git a/internal/state/weights_test.go b/internal/state/weights_test.go index 09e109bc..839ab2f8 100644 --- a/internal/state/weights_test.go +++ b/internal/state/weights_test.go @@ -58,6 +58,88 @@ func TestKeyRegistryCollectCandidatesExcludesRuntimeUnavailableKeys(t *testing.T } } +func TestKeyRuntimeViewClassifiesAvailability(t *testing.T) { + now := time.Date(2026, time.July, 24, 10, 0, 0, 0, time.UTC) + tests := []struct { + name string + view KeyRuntimeView + want KeyRuntimeState + }{ + { + name: "disabled wins", + view: KeyRuntimeView{ + Status: KeyStatusDisabled, Blacklisted: true, + CooldownUntil: now.Add(time.Minute), + }, + want: KeyRuntimeDisabled, + }, + { + name: "blacklist wins cooldown", + view: KeyRuntimeView{ + Status: KeyStatusActive, Blacklisted: true, + CooldownUntil: now.Add(time.Minute), + }, + want: KeyRuntimeBlacklisted, + }, + { + name: "future cooldown", + view: KeyRuntimeView{ + Status: KeyStatusActive, CooldownUntil: now.Add(time.Nanosecond), + }, + want: KeyRuntimeCooldown, + }, + { + name: "cooldown equality is available", + view: KeyRuntimeView{ + Status: KeyStatusActive, CooldownUntil: now, + }, + want: KeyRuntimeAvailable, + }, + { + name: "active", + view: KeyRuntimeView{Status: KeyStatusActive}, + want: KeyRuntimeAvailable, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + if got := test.view.RuntimeState(now); got != test.want { + t.Fatalf("RuntimeState() = %q, want %q", got, test.want) + } + }) + } +} + +func TestCollectCandidatesUsesRuntimeViewBoundary(t *testing.T) { + now := time.Date(2026, time.July, 24, 10, 0, 0, 0, time.UTC) + registry := NewKeyRegistry() + mustReplaceKeyEntries(t, registry, []KeyEntry{ + {ID: 1, GroupID: 10, Status: KeyStatusActive, EncryptedValue: "active"}, + { + ID: 2, GroupID: 10, Status: KeyStatusActive, + CooldownUntil: now.Add(time.Nanosecond), EncryptedValue: "cooling", + }, + {ID: 3, GroupID: 10, Status: KeyStatusActive, Blacklisted: true, EncryptedValue: "blacklisted"}, + {ID: 4, GroupID: 10, Status: KeyStatusActive, CooldownUntil: now, EncryptedValue: "boundary"}, + {ID: 5, GroupID: 10, Status: KeyStatusDisabled, EncryptedValue: "disabled"}, + }) + views := registry.Snapshot() + availableIDs := make([]uint, 0) + for _, view := range views { + if view.RuntimeState(now) == KeyRuntimeAvailable { + availableIDs = append(availableIDs, view.ID) + } + } + candidates := registry.CollectCandidates([]uint{10}, nil, now) + candidateIDs := make([]uint, 0, len(candidates)) + for _, candidate := range candidates { + candidateIDs = append(candidateIDs, candidate.ID) + } + if !reflect.DeepEqual(candidateIDs, availableIDs) { + t.Fatalf("CollectCandidates IDs = %v, RuntimeView IDs = %v", candidateIDs, availableIDs) + } +} + func TestKeyRegistrySetCooldownNeverShortensDeadline(t *testing.T) { now := time.Date(2026, time.July, 22, 12, 0, 0, 0, time.UTC) registry := NewKeyRegistry()