diff --git a/internal/app/app.go b/internal/app/app.go index dffdfa66..648c004a 100644 --- a/internal/app/app.go +++ b/internal/app/app.go @@ -24,13 +24,14 @@ import ( "gorm.io/gorm" ) -// App owns the M0 process lifecycle and infrastructure resources. +// App owns the process lifecycle, infrastructure resources, and runtime state. type App struct { - engine *gin.Engine - config *config.Config - encryption encryption.Service - store store.Store - db *gorm.DB + engine *gin.Engine + config *config.Config + encryption encryption.Service + store store.Store + db *gorm.DB + runtimeState RuntimeStateLoader mu sync.Mutex httpServer *http.Server @@ -38,18 +39,24 @@ type App struct { serveErrors chan error } +// RuntimeStateLoader initializes the in-memory runtime state from persistence. +type RuntimeStateLoader interface { + Load(context.Context) error +} + // AppParams defines dependencies injected into App. type AppParams struct { dig.In - Engine *gin.Engine - Config *config.Config - Encryption encryption.Service - Store store.Store - DB *gorm.DB + Engine *gin.Engine + Config *config.Config + Encryption encryption.Service + Store store.Store + DB *gorm.DB + RuntimeState RuntimeStateLoader } -// NewEngine creates the M0 HTTP engine and health endpoint. +// NewEngine creates the HTTP engine and health endpoint. func NewEngine() *gin.Engine { gin.SetMode(gin.ReleaseMode) engine := gin.New() @@ -66,12 +73,13 @@ func NewEngine() *gin.Engine { // NewApp creates the application lifecycle manager. func NewApp(params AppParams) *App { return &App{ - engine: params.Engine, - config: params.Config, - encryption: params.Encryption, - store: params.Store, - db: params.DB, - serveErrors: make(chan error, 1), + engine: params.Engine, + config: params.Config, + encryption: params.Encryption, + store: params.Store, + db: params.DB, + runtimeState: params.RuntimeState, + serveErrors: make(chan error, 1), } } @@ -89,6 +97,9 @@ func (a *App) Start() error { if err := storage.AutoMigrate(a.db); err != nil { return err } + if err := a.runtimeState.Load(context.Background()); err != nil { + return fmt.Errorf("load runtime state: %w", err) + } address := net.JoinHostPort(a.config.Server.Host, strconv.Itoa(a.config.Server.Port)) listener, err := net.Listen("tcp", address) diff --git a/internal/app/app_test.go b/internal/app/app_test.go index 7c0bd34b..5a6981d5 100644 --- a/internal/app/app_test.go +++ b/internal/app/app_test.go @@ -4,6 +4,7 @@ import ( "bytes" "context" "encoding/json" + "errors" "net/http" "net/http/httptest" "strings" @@ -14,13 +15,22 @@ import ( "gpt-load/internal/platform/encryption" app_errors "gpt-load/internal/platform/errors" "gpt-load/internal/platform/version" + "gpt-load/internal/state" + "gpt-load/internal/state/loader" "gpt-load/internal/storage" "gpt-load/internal/storage/store" "github.com/gin-gonic/gin" "github.com/sirupsen/logrus" + "gorm.io/gorm" ) +type runtimeStateLoaderFunc func(context.Context) error + +func (f runtimeStateLoaderFunc) Load(ctx context.Context) error { + return f(ctx) +} + func TestNewEngineRecoversWithoutLoggingCredentials(t *testing.T) { previousMode := gin.Mode() gin.SetMode(gin.DebugMode) @@ -109,17 +119,24 @@ func TestAppStartMigratesDatabaseAndServesHTTP(t *testing.T) { if err != nil { t.Fatalf("encryption.NewService() error = %v", err) } + manager, runtimeState := newTestRuntimeState(db) application := NewApp(AppParams{ - Engine: NewEngine(), - Config: testConfig(t), - Encryption: keyService, - Store: store.NewMemoryStore(), - DB: db, + Engine: NewEngine(), + Config: testConfig(t), + Encryption: keyService, + Store: store.NewMemoryStore(), + DB: db, + RuntimeState: runtimeState, }) + cleanupApp(t, application) if err := application.Start(); err != nil { t.Fatalf("Start() error = %v", err) } + snapshot := manager.Current() + if snapshot == nil || snapshot.Revision != 1 { + t.Fatalf("runtime snapshot = %#v, want revision 1", snapshot) + } if application.httpServer.ReadTimeout != 2*time.Second { t.Fatalf("ReadTimeout = %s, want 2s", application.httpServer.ReadTimeout) } @@ -132,14 +149,6 @@ func TestAppStartMigratesDatabaseAndServesHTTP(t *testing.T) { if application.httpServer.WriteTimeout != 0 { t.Fatalf("WriteTimeout = %s, want 0 for streaming responses", application.httpServer.WriteTimeout) } - t.Cleanup(func() { - ctx, cancel := context.WithTimeout(context.Background(), time.Second) - defer cancel() - if err := application.Stop(ctx); err != nil { - t.Errorf("Stop() error = %v", err) - } - }) - for _, table := range []string{ "groups", "upstream_keys", "access_keys", "request_logs", "usage_stats", "model_prices", "system_settings", "jobs", "schema_info", @@ -159,7 +168,7 @@ func TestAppStartMigratesDatabaseAndServesHTTP(t *testing.T) { } } -func TestAppReportsUnexpectedHTTPServeFailure(t *testing.T) { +func TestAppStartMigratesBeforeLoadingRuntimeState(t *testing.T) { db, err := storage.Open(":memory:") if err != nil { t.Fatalf("storage.Open() error = %v", err) @@ -169,23 +178,95 @@ func TestAppReportsUnexpectedHTTPServeFailure(t *testing.T) { t.Fatalf("encryption.NewService() error = %v", err) } + loadCalled := false application := NewApp(AppParams{ Engine: NewEngine(), Config: testConfig(t), Encryption: keyService, Store: store.NewMemoryStore(), DB: db, + RuntimeState: runtimeStateLoaderFunc(func(context.Context) error { + loadCalled = true + if !db.Migrator().HasTable("groups") { + return errors.New("groups table does not exist") + } + return nil + }), }) + cleanupApp(t, application) + if err := application.Start(); err != nil { t.Fatalf("Start() error = %v", err) } - t.Cleanup(func() { - ctx, cancel := context.WithTimeout(context.Background(), time.Second) - defer cancel() - if err := application.Stop(ctx); err != nil { - t.Errorf("Stop() error = %v", err) - } + if !loadCalled { + t.Fatal("runtime state loader was not called") + } +} + +func TestAppStartRejectsRuntimeStateLoadFailureBeforeListen(t *testing.T) { + db, err := storage.Open(":memory:") + if err != nil { + t.Fatalf("storage.Open() error = %v", err) + } + keyService, err := encryption.NewService("test-master-key") + if err != nil { + t.Fatalf("encryption.NewService() error = %v", err) + } + + loadErr := errors.New("corrupt runtime config") + application := NewApp(AppParams{ + Engine: NewEngine(), + Config: testConfig(t), + Encryption: keyService, + Store: store.NewMemoryStore(), + DB: db, + RuntimeState: runtimeStateLoaderFunc(func(context.Context) error { + return loadErr + }), }) + cleanupApp(t, application) + + err = application.Start() + if !errors.Is(err, loadErr) { + t.Fatalf("Start() error = %v, want wrapped runtime state error", err) + } + if !strings.Contains(err.Error(), "load runtime state") { + t.Fatalf("Start() error = %q, want runtime state context", err) + } + if got := application.Address(); got != "" { + t.Fatalf("Address() = %q after failed load, want empty", got) + } + if application.httpServer != nil { + t.Fatalf("httpServer = %#v after failed load, want nil", application.httpServer) + } + if application.listener != nil { + t.Fatalf("listener = %#v after failed load, want nil", application.listener) + } +} + +func TestAppReportsUnexpectedHTTPServeFailure(t *testing.T) { + db, err := storage.Open(":memory:") + if err != nil { + t.Fatalf("storage.Open() error = %v", err) + } + keyService, err := encryption.NewService("test-master-key") + if err != nil { + t.Fatalf("encryption.NewService() error = %v", err) + } + _, runtimeState := newTestRuntimeState(db) + + application := NewApp(AppParams{ + Engine: NewEngine(), + Config: testConfig(t), + Encryption: keyService, + Store: store.NewMemoryStore(), + DB: db, + RuntimeState: runtimeState, + }) + cleanupApp(t, application) + if err := application.Start(); err != nil { + t.Fatalf("Start() error = %v", err) + } if err := application.listener.Close(); err != nil { t.Fatalf("close live listener: %v", err) @@ -200,6 +281,23 @@ func TestAppReportsUnexpectedHTTPServeFailure(t *testing.T) { } } +func newTestRuntimeState(db *gorm.DB) (*state.Manager, *loader.Loader) { + manager := state.NewManager() + registry := state.NewKeyRegistry() + return manager, loader.New(db, manager, registry) +} + +func cleanupApp(t *testing.T, application *App) { + t.Helper() + t.Cleanup(func() { + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + if err := application.Stop(ctx); err != nil { + t.Errorf("Stop() error = %v", err) + } + }) +} + func testConfig(t *testing.T) *config.Config { t.Helper() return &config.Config{ diff --git a/internal/container/container.go b/internal/container/container.go index def82d14..0cf1e3b0 100644 --- a/internal/container/container.go +++ b/internal/container/container.go @@ -1,18 +1,20 @@ -// Package container assembles the M0 dependency graph with dig. +// Package container assembles the 2.0 dependency graph with dig. package container import ( + "go.uber.org/dig" + "gorm.io/gorm" + "gpt-load/internal/app" "gpt-load/internal/platform/config" "gpt-load/internal/platform/encryption" + "gpt-load/internal/state" + stateloader "gpt-load/internal/state/loader" "gpt-load/internal/storage" "gpt-load/internal/storage/store" - - "go.uber.org/dig" - "gorm.io/gorm" ) -// BuildContainer creates the 2.0 M0 dependency graph. +// BuildContainer creates the 2.0 runtime foundation dependency graph. func BuildContainer() (*dig.Container, error) { dependencyContainer := dig.New() @@ -28,6 +30,11 @@ func BuildContainer() (*dig.Container, error) { return store.NewStore(cfg.RedisDSN) }, app.NewEngine, + state.NewManager, + state.NewKeyRegistry, + func(db *gorm.DB, manager *state.Manager, registry *state.KeyRegistry) app.RuntimeStateLoader { + return stateloader.New(db, manager, registry) + }, app.NewApp, } diff --git a/internal/container/container_test.go b/internal/container/container_test.go index 28c4155d..b2e47de9 100644 --- a/internal/container/container_test.go +++ b/internal/container/container_test.go @@ -1,20 +1,23 @@ package container import ( + "context" "os" "path/filepath" "testing" + "github.com/gin-gonic/gin" + "gorm.io/gorm" + "gpt-load/internal/app" "gpt-load/internal/platform/config" "gpt-load/internal/platform/encryption" + "gpt-load/internal/state" + "gpt-load/internal/storage" "gpt-load/internal/storage/store" - - "github.com/gin-gonic/gin" - "gorm.io/gorm" ) -func TestBuildContainerResolvesM0DependencyGraph(t *testing.T) { +func TestBuildContainerResolvesS2DependencyGraph(t *testing.T) { dataDir := t.TempDir() t.Setenv("AUTH_KEY", "test-auth-key") t.Setenv("DATA_DIR", dataDir) @@ -35,8 +38,10 @@ func TestBuildContainerResolvesM0DependencyGraph(t *testing.T) { db *gorm.DB, storageStore store.Store, _ *gin.Engine, + manager *state.Manager, + registry *state.KeyRegistry, + runtimeState app.RuntimeStateLoader, ) { - resolved = true t.Cleanup(func() { _ = storageStore.Close() sqlDB, dbErr := db.DB() @@ -44,12 +49,26 @@ func TestBuildContainerResolvesM0DependencyGraph(t *testing.T) { _ = sqlDB.Close() } }) + if err := storage.AutoMigrate(db); err != nil { + t.Fatalf("AutoMigrate() error = %v", err) + } + if err := runtimeState.Load(context.Background()); err != nil { + t.Fatalf("runtimeState.Load() error = %v", err) + } + snapshot := manager.Current() + if snapshot == nil || snapshot.Revision != 1 { + t.Fatalf("current snapshot = %#v, want revision 1", snapshot) + } + if got := registry.CollectCandidates(nil, nil); len(got) != 0 { + t.Fatalf("empty registry candidates = %#v", got) + } + resolved = true }) if err != nil { - t.Fatalf("resolve M0 dependency graph: %v", err) + t.Fatalf("resolve S2 dependency graph: %v", err) } if !resolved { - t.Fatal("M0 dependency graph was not invoked") + t.Fatal("S2 dependency graph was not invoked") } if _, err := os.Stat(filepath.Join(dataDir, encryption.KeyFileName)); err != nil { t.Fatalf("container did not initialize encryption keyfile: %v", err) diff --git a/internal/protocol/protocol.go b/internal/protocol/protocol.go new file mode 100644 index 00000000..3fb1bc4b --- /dev/null +++ b/internal/protocol/protocol.go @@ -0,0 +1,20 @@ +// Package protocol defines protocol identifiers shared by runtime domains. +package protocol + +type Protocol string + +const ( + OpenAI Protocol = "openai" + Anthropic Protocol = "anthropic" + Gemini Protocol = "gemini" + OpenAIResponse Protocol = "openai-response" +) + +func (p Protocol) Valid() bool { + switch p { + case OpenAI, Anthropic, Gemini, OpenAIResponse: + return true + default: + return false + } +} diff --git a/internal/protocol/protocol_test.go b/internal/protocol/protocol_test.go new file mode 100644 index 00000000..76ae056a --- /dev/null +++ b/internal/protocol/protocol_test.go @@ -0,0 +1,23 @@ +package protocol + +import "testing" + +func TestProtocolValid(t *testing.T) { + t.Parallel() + tests := []struct { + protocol Protocol + want bool + }{ + {protocol: OpenAI, want: true}, + {protocol: Anthropic, want: true}, + {protocol: Gemini, want: true}, + {protocol: OpenAIResponse, want: true}, + {protocol: Protocol("unknown"), want: false}, + {protocol: "", want: false}, + } + for _, tt := range tests { + if got := tt.protocol.Valid(); got != tt.want { + t.Errorf("Protocol(%q).Valid() = %t, want %t", tt.protocol, got, tt.want) + } + } +} diff --git a/internal/state/doc.go b/internal/state/doc.go deleted file mode 100644 index ec01f750..00000000 --- a/internal/state/doc.go +++ /dev/null @@ -1,2 +0,0 @@ -// Package state owns ConfigSnapshot and KeyRegistry runtime views starting in M1. -package state diff --git a/internal/state/loader/loader.go b/internal/state/loader/loader.go new file mode 100644 index 00000000..df344210 --- /dev/null +++ b/internal/state/loader/loader.go @@ -0,0 +1,230 @@ +// Package loader adapts persisted storage rows into runtime state. +package loader + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "strings" + + "gorm.io/gorm" + + "gpt-load/internal/platform/config" + "gpt-load/internal/protocol" + "gpt-load/internal/state" + "gpt-load/internal/storage/models" +) + +type Loader struct { + db *gorm.DB + manager *state.Manager + registry *state.KeyRegistry +} + +type storedRows struct { + settings []models.SystemSetting + groups []models.Group + accessKeys []models.AccessKey + upstreamKeys []models.UpstreamKey +} + +type modelDTO struct { + ID string `json:"id"` + Alias string `json:"alias"` +} + +type filterDTO struct { + Groups []uint `json:"groups"` + Protocols []protocol.Protocol `json:"protocols"` + Models []string `json:"models"` +} + +func (f filterDTO) toState() state.FilterSet { + filters := state.FilterSet{} + if len(f.Groups) > 0 { + filters.Groups = make(map[uint]struct{}, len(f.Groups)) + for _, id := range f.Groups { + filters.Groups[id] = struct{}{} + } + } + if len(f.Protocols) > 0 { + filters.Protocols = make(map[protocol.Protocol]struct{}, len(f.Protocols)) + for _, p := range f.Protocols { + filters.Protocols[p] = struct{}{} + } + } + if len(f.Models) > 0 { + filters.Models = make(map[string]struct{}, len(f.Models)) + for _, model := range f.Models { + filters.Models[model] = struct{}{} + } + } + return filters +} + +func New(db *gorm.DB, manager *state.Manager, registry *state.KeyRegistry) *Loader { + return &Loader{db: db, manager: manager, registry: registry} +} + +func (l *Loader) Load(ctx context.Context) error { + input, entries, err := l.read(ctx) + if err != nil { + return fmt.Errorf("read runtime state: %w", err) + } + if err := state.ValidateKeyEntries(entries); err != nil { + return fmt.Errorf("validate upstream keys: %w", err) + } + if _, err := l.manager.Publish(input); err != nil { + return fmt.Errorf("publish config snapshot: %w", err) + } + if err := l.registry.Replace(entries); err != nil { + return fmt.Errorf("replace key registry: %w", err) + } + return nil +} + +func (l *Loader) queryRows(ctx context.Context) (storedRows, error) { + db := l.db.WithContext(ctx) + var rows storedRows + if err := db.Order("key ASC").Find(&rows.settings).Error; err != nil { + return storedRows{}, fmt.Errorf("query system settings: %w", err) + } + if err := db.Order("id ASC").Find(&rows.groups).Error; err != nil { + return storedRows{}, fmt.Errorf("query groups: %w", err) + } + if err := db.Order("id ASC").Find(&rows.accessKeys).Error; err != nil { + return storedRows{}, fmt.Errorf("query access keys: %w", err) + } + if err := db.Order("id ASC").Find(&rows.upstreamKeys).Error; err != nil { + return storedRows{}, fmt.Errorf("query upstream keys: %w", err) + } + return rows, nil +} + +func (l *Loader) read(ctx context.Context) (state.CompileInput, []state.KeyEntry, error) { + rows, err := l.queryRows(ctx) + if err != nil { + return state.CompileInput{}, nil, err + } + input, err := mapSystemAndGroups(rows) + if err != nil { + return state.CompileInput{}, nil, err + } + accessKeys, err := mapAccessKeys(rows.accessKeys) + if err != nil { + return state.CompileInput{}, nil, err + } + input.AccessKeys = accessKeys + return input, mapUpstreamKeys(rows.upstreamKeys), nil +} + +func decodeJSON(raw models.JSON, target any) error { + return decodeJSONDocument(raw, target, false) +} + +func decodeFilterJSON(raw models.JSON, target *filterDTO) error { + return decodeJSONDocument(raw, target, true) +} + +func decodeJSONDocument(raw models.JSON, target any, disallowUnknownFields bool) error { + if len(raw) == 0 { + return nil + } + decoder := json.NewDecoder(bytes.NewReader(raw)) + decoder.UseNumber() + if disallowUnknownFields { + decoder.DisallowUnknownFields() + } + if err := decoder.Decode(target); err != nil { + return err + } + var extra any + if err := decoder.Decode(&extra); !errors.Is(err, io.EOF) { + if err == nil { + return fmt.Errorf("multiple JSON values") + } + return err + } + return nil +} + +func decodeSettingValue(raw string) (any, error) { + if !json.Valid([]byte(raw)) { + return raw, nil + } + var value any + decoder := json.NewDecoder(strings.NewReader(raw)) + decoder.UseNumber() + if err := decoder.Decode(&value); err != nil { + return nil, err + } + return value, nil +} + +func mapSystemAndGroups(rows storedRows) (state.CompileInput, error) { + input := state.CompileInput{ + SystemSettings: make(config.Settings, len(rows.settings)), + Groups: make([]state.GroupConfig, 0, len(rows.groups)), + } + for _, row := range rows.settings { + value, err := decodeSettingValue(row.Value) + if err != nil { + return state.CompileInput{}, fmt.Errorf("decode system setting %q: %w", row.Key, err) + } + input.SystemSettings[row.Key] = value + } + for _, row := range rows.groups { + var protocols []protocol.Protocol + if err := decodeJSON(row.Protocols, &protocols); err != nil { + return state.CompileInput{}, fmt.Errorf("decode group %d protocols: %w", row.ID, err) + } + var storedModels []modelDTO + if err := decodeJSON(row.Models, &storedModels); err != nil { + return state.CompileInput{}, fmt.Errorf("decode group %d models: %w", row.ID, err) + } + settings := make(config.Settings) + if err := decodeJSON(row.Config, &settings); err != nil { + return state.CompileInput{}, fmt.Errorf("decode group %d config: %w", row.ID, err) + } + + runtimeModels := make([]state.ModelConfig, 0, len(storedModels)) + for _, model := range storedModels { + runtimeModels = append(runtimeModels, state.ModelConfig{ID: model.ID, Alias: model.Alias}) + } + input.Groups = append(input.Groups, state.GroupConfig{ + ID: row.ID, Name: row.Name, UpstreamURL: row.UpstreamURL, + Protocols: append([]protocol.Protocol(nil), protocols...), + Models: runtimeModels, Settings: settings, Enabled: row.Enabled, + }) + } + return input, nil +} + +func mapAccessKeys(rows []models.AccessKey) ([]state.AccessKeyConfig, error) { + result := make([]state.AccessKeyConfig, 0, len(rows)) + for _, row := range rows { + var filters filterDTO + if err := decodeFilterJSON(row.Filters, &filters); err != nil { + return nil, fmt.Errorf("decode access key %d filters: %w", row.ID, err) + } + result = append(result, state.AccessKeyConfig{ + ID: row.ID, Name: row.Name, KeyHash: row.KeyHash, + Status: state.AccessKeyStatus(row.Status), Filters: filters.toState(), + }) + } + return result, nil +} + +func mapUpstreamKeys(rows []models.UpstreamKey) []state.KeyEntry { + result := make([]state.KeyEntry, 0, len(rows)) + for _, row := range rows { + result = append(result, state.KeyEntry{ + ID: row.ID, GroupID: row.GroupID, WeightManual: row.WeightManual, + Status: state.KeyStatus(row.Status), EncryptedValue: row.KeyValue, + }) + } + return result +} diff --git a/internal/state/loader/loader_test.go b/internal/state/loader/loader_test.go new file mode 100644 index 00000000..9bb2cd7e --- /dev/null +++ b/internal/state/loader/loader_test.go @@ -0,0 +1,391 @@ +package loader_test + +import ( + "context" + "fmt" + "strings" + "testing" + "time" + + "gorm.io/gorm" + + "gpt-load/internal/protocol" + "gpt-load/internal/state" + "gpt-load/internal/state/loader" + "gpt-load/internal/storage" + "gpt-load/internal/storage/models" +) + +func TestLoaderLoadsEmptyMigratedDatabase(t *testing.T) { + db := openMigratedDatabase(t) + manager := state.NewManager() + registry := state.NewKeyRegistry() + + if err := loader.New(db, manager, registry).Load(context.Background()); err != nil { + t.Fatalf("Load() error = %v", err) + } + + snapshot := manager.Current() + if snapshot == nil { + t.Fatal("Current() = nil, want an initialized snapshot") + } + if snapshot.Revision != 1 { + t.Errorf("snapshot revision = %d, want 1", snapshot.Revision) + } + if snapshot.Candidates == nil || len(snapshot.Candidates) != 0 { + t.Errorf("snapshot candidates = %#v, want initialized empty map", snapshot.Candidates) + } + if snapshot.Groups == nil || len(snapshot.Groups) != 0 { + t.Errorf("snapshot groups = %#v, want initialized empty map", snapshot.Groups) + } + if snapshot.AccessKeysByHash == nil || len(snapshot.AccessKeysByHash) != 0 { + t.Errorf("snapshot access keys = %#v, want initialized empty map", snapshot.AccessKeysByHash) + } + if got := registry.CollectCandidates([]uint{1}, nil); len(got) != 0 { + t.Errorf("registry candidates = %#v, want empty", got) + } +} + +func TestLoaderMapsSystemAndGroupRows(t *testing.T) { + db := openMigratedDatabase(t) + mustCreate(t, db, &models.SystemSetting{Key: "connect_timeout", Value: "20"}) + mustCreate(t, db, &models.SystemSetting{ + Key: "header_rules", + Value: `{"set":{"X-System":"system"},"remove":["X-System-Remove"]}`, + }) + + enabled := models.Group{ + Name: "enabled", + UpstreamURL: "https://enabled.example.com/v1", + Signature: "enabled-signature", + Protocols: models.JSON(`["openai"]`), + Models: models.JSON(`[ + {"id":"gpt-4o","alias":"Primary"}, + {"id":"gpt-4o","alias":"Secondary"}, + {"id":"gpt-4.1","alias":"Other"} + ]`), + Config: models.JSON(`{ + "request_timeout":30, + "header_rules":{"set":{"X-Group":"group"},"remove":["X-Group-Remove"]} + }`), + Enabled: true, + } + mustCreate(t, db, &enabled) + disabled := models.Group{ + Name: "disabled", + UpstreamURL: "https://disabled.example.com/v1", + Signature: "disabled-signature", + Protocols: models.JSON(`["openai"]`), + Models: models.JSON(`[{"id":"hidden","alias":"Hidden"}]`), + Config: models.JSON(`{}`), + Enabled: true, + } + mustCreate(t, db, &disabled) + if err := db.Model(&disabled).Update("enabled", false).Error; err != nil { + t.Fatalf("disable group: %v", err) + } + + manager := state.NewManager() + registry := state.NewKeyRegistry() + if err := loader.New(db, manager, registry).Load(context.Background()); err != nil { + t.Fatalf("Load() error = %v", err) + } + + snapshot := manager.Current() + if snapshot == nil { + t.Fatal("Current() = nil, want snapshot") + } + if len(snapshot.Groups) != 1 { + t.Fatalf("snapshot groups = %#v, want enabled group only", snapshot.Groups) + } + view, ok := snapshot.Groups[enabled.ID] + if !ok { + t.Fatalf("snapshot groups = %#v, want group %d", snapshot.Groups, enabled.ID) + } + if _, ok := snapshot.Groups[disabled.ID]; ok { + t.Fatalf("disabled group %d is present in snapshot", disabled.ID) + } + 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) + } + if view.Timeouts.Connect != 20*time.Second || view.Timeouts.Request != 30*time.Second { + t.Errorf("group timeouts = %#v, want connect 20s and request 30s", view.Timeouts) + } + if len(view.HeaderRules.Set) != 1 || view.HeaderRules.Set["X-Group"] != "group" { + t.Errorf("group header set rules = %#v, want whole group override", view.HeaderRules.Set) + } + if len(view.HeaderRules.Remove) != 1 || view.HeaderRules.Remove[0] != "X-Group-Remove" { + t.Errorf("group header remove rules = %#v, want group override", view.HeaderRules.Remove) + } + + openAICandidates := snapshot.Candidates[protocol.OpenAI] + if len(openAICandidates) != 2 { + t.Fatalf("OpenAI candidates = %#v, want two model IDs", openAICandidates) + } + if got := openAICandidates["gpt-4o"]; len(got) != 1 || got[0].GroupID != enabled.ID || got[0].UpstreamModelID != "gpt-4o" { + t.Errorf("gpt-4o candidates = %#v, want one id-only route for group %d", got, enabled.ID) + } + if _, ok := openAICandidates["hidden"]; ok { + t.Fatal("disabled group model hidden is present in candidates") + } +} + +func TestLoaderRejectsInvalidGroupRowsWithoutPublishing(t *testing.T) { + tests := []struct { + name string + protocols models.JSON + models models.JSON + config models.JSON + wantError string + }{ + {name: "protocols object", protocols: models.JSON(`{}`), models: models.JSON(`[]`), config: models.JSON(`{}`), wantError: "protocols"}, + {name: "models object", protocols: models.JSON(`["openai"]`), models: models.JSON(`{}`), config: models.JSON(`{}`), wantError: "models"}, + {name: "config array", protocols: models.JSON(`["openai"]`), models: models.JSON(`[]`), config: models.JSON(`[]`), wantError: "config"}, + {name: "unknown group setting", protocols: models.JSON(`["openai"]`), models: models.JSON(`[{"id":"gpt-4o"}]`), config: models.JSON(`{"unknown":true}`), wantError: "unknown group setting"}, + {name: "duplicate protocol", protocols: models.JSON(`["openai","openai"]`), models: models.JSON(`[{"id":"gpt-4o"}]`), config: models.JSON(`{}`), wantError: "duplicate protocol"}, + {name: "blank model id", protocols: models.JSON(`["openai"]`), models: models.JSON(`[{"id":" "}]`), config: models.JSON(`{}`), wantError: "model id is required"}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + db := openMigratedDatabase(t) + group := models.Group{ + Name: "invalid", UpstreamURL: "https://invalid.example.com/v1", + Signature: "invalid-signature", Protocols: test.protocols, + Models: test.models, Config: test.config, Enabled: true, + } + mustCreate(t, db, &group) + manager := state.NewManager() + registry := state.NewKeyRegistry() + + err := loader.New(db, manager, registry).Load(context.Background()) + if err == nil { + t.Fatal("Load() error = nil, want invalid group rejection") + } + if !strings.Contains(err.Error(), test.wantError) { + t.Errorf("Load() error = %q, want context containing %q", err, test.wantError) + } + if manager.Current() != nil { + t.Fatalf("Current() = %#v after failed load, want nil", manager.Current()) + } + if got := registry.CollectCandidates([]uint{group.ID}, nil); len(got) != 0 { + t.Fatalf("registry candidates after failed load = %#v, want empty", got) + } + }) + } +} + +func TestLoaderMapsAccessAndUpstreamKeys(t *testing.T) { + db := openMigratedDatabase(t) + firstGroup := createRuntimeGroup(t, db, "first", protocol.OpenAI, "gpt-4o") + secondGroup := createRuntimeGroup(t, db, "second", protocol.Anthropic, "claude-3-5-sonnet") + + activeAccess := models.AccessKey{ + Name: "active access", KeyValue: "access-cipher-active", KeyHash: "active-hash", + Status: "active", + Filters: models.JSON(fmt.Sprintf( + `{"groups":[%d,9999],"protocols":["openai"],"models":["gpt-4o"]}`, + firstGroup.ID, + )), + } + disabledAccess := models.AccessKey{ + Name: "disabled access", KeyValue: "access-cipher-disabled", KeyHash: "disabled-hash", + Status: "disabled", Filters: models.JSON(`{}`), + } + mustCreate(t, db, &activeAccess) + mustCreate(t, db, &disabledAccess) + + firstWeight := 7 + keys := []models.UpstreamKey{ + {GroupID: firstGroup.ID, KeyValue: "upstream-cipher-one", KeyHash: "upstream-hash-one", Status: models.UpstreamKeyStatusActive, WeightManual: &firstWeight}, + {GroupID: firstGroup.ID, KeyValue: "upstream-cipher-two", KeyHash: "upstream-hash-two", Status: models.UpstreamKeyStatusDisabled}, + {GroupID: secondGroup.ID, KeyValue: "upstream-cipher-three", KeyHash: "upstream-hash-three", Status: models.UpstreamKeyStatusActive}, + {GroupID: secondGroup.ID, KeyValue: "upstream-cipher-four", KeyHash: "upstream-hash-four", Status: models.UpstreamKeyStatusDisabled}, + } + for index := range keys { + mustCreate(t, db, &keys[index]) + } + + manager := state.NewManager() + registry := state.NewKeyRegistry() + if err := loader.New(db, manager, registry).Load(context.Background()); err != nil { + t.Fatalf("Load() error = %v", err) + } + + snapshot := manager.Current() + if snapshot == nil { + t.Fatal("Current() = nil, want snapshot") + } + if len(snapshot.AccessKeysByHash) != 1 { + t.Fatalf("snapshot access keys = %#v, want active key only", snapshot.AccessKeysByHash) + } + access, ok := snapshot.AccessKeysByHash[activeAccess.KeyHash] + if !ok { + t.Fatalf("snapshot access keys = %#v, want hash %q", snapshot.AccessKeysByHash, activeAccess.KeyHash) + } + if _, ok := snapshot.AccessKeysByHash[disabledAccess.KeyHash]; ok { + t.Fatalf("disabled access hash %q is present in snapshot", disabledAccess.KeyHash) + } + if _, ok := access.Filters.Groups[firstGroup.ID]; !ok { + t.Errorf("access filters groups = %#v, want group %d", access.Filters.Groups, firstGroup.ID) + } + if _, ok := access.Filters.Groups[9999]; !ok { + t.Errorf("access filters groups = %#v, want dangling group 9999 retained", access.Filters.Groups) + } + if _, ok := access.Filters.Protocols[protocol.OpenAI]; !ok { + t.Errorf("access filters protocols = %#v, want OpenAI", access.Filters.Protocols) + } + if _, ok := access.Filters.Models["gpt-4o"]; !ok { + t.Errorf("access filters models = %#v, want gpt-4o", access.Filters.Models) + } + + candidates := registry.CollectCandidates([]uint{firstGroup.ID, secondGroup.ID}, nil) + if len(candidates) != 2 { + t.Fatalf("registry candidates = %#v, want two active keys", candidates) + } + if candidates[0].ID != keys[0].ID || candidates[0].GroupID != firstGroup.ID || candidates[0].WeightManual == nil || *candidates[0].WeightManual != firstWeight { + t.Errorf("first candidate = %#v, want active weighted key %d", candidates[0], keys[0].ID) + } + if candidates[1].ID != keys[2].ID || candidates[1].GroupID != secondGroup.ID { + t.Errorf("second candidate = %#v, want active key %d", candidates[1], keys[2].ID) + } + for _, key := range keys { + got, ok := registry.EncryptedValue(key.ID) + if !ok || got != key.KeyValue { + t.Errorf("EncryptedValue(%d) = %q, %t, want %q, true", key.ID, got, ok, key.KeyValue) + } + } + + snapshotText := fmt.Sprintf("%#v", snapshot) + for _, secret := range []string{ + activeAccess.KeyValue, disabledAccess.KeyValue, + keys[0].KeyValue, keys[1].KeyValue, keys[2].KeyValue, keys[3].KeyValue, + } { + if strings.Contains(snapshotText, secret) { + t.Errorf("snapshot exposes credential material %q", secret) + } + } +} + +func TestLoaderRejectsInvalidCredentialRowsWithoutPublishing(t *testing.T) { + tests := []struct { + name string + insert func(*testing.T, *gorm.DB, models.Group) + wantError string + }{ + { + name: "unknown access status", + insert: func(t *testing.T, db *gorm.DB, _ models.Group) { + mustCreate(t, db, &models.AccessKey{ + Name: "invalid", KeyValue: "access-cipher", KeyHash: "invalid-status-hash", + Status: "revoked", Filters: models.JSON(`{}`), + }) + }, + wantError: "invalid status", + }, + { + name: "invalid filter protocol", + insert: func(t *testing.T, db *gorm.DB, _ models.Group) { + mustCreate(t, db, &models.AccessKey{ + Name: "invalid", KeyValue: "access-cipher", KeyHash: "invalid-protocol-hash", + Status: "active", Filters: models.JSON(`{"protocols":["unknown"]}`), + }) + }, + wantError: "invalid protocol", + }, + { + name: "blank filter model", + insert: func(t *testing.T, db *gorm.DB, _ models.Group) { + mustCreate(t, db, &models.AccessKey{ + Name: "invalid", KeyValue: "access-cipher", KeyHash: "blank-model-hash", + Status: "active", Filters: models.JSON(`{"models":[" "]}`), + }) + }, + wantError: "filter model is required", + }, + { + name: "unknown filter field", + insert: func(t *testing.T, db *gorm.DB, _ models.Group) { + mustCreate(t, db, &models.AccessKey{ + Name: "invalid", KeyValue: "access-cipher", KeyHash: "unknown-filter-field-hash", + Status: "active", Filters: models.JSON(`{"protcols":["openai"]}`), + }) + }, + wantError: "unknown field", + }, + { + name: "empty upstream ciphertext", + insert: func(t *testing.T, db *gorm.DB, group models.Group) { + mustCreate(t, db, &models.UpstreamKey{ + GroupID: group.ID, KeyValue: "", KeyHash: "empty-cipher-hash", + Status: models.UpstreamKeyStatusActive, + }) + }, + wantError: "encrypted value is required", + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + db := openMigratedDatabase(t) + group := createRuntimeGroup(t, db, "valid", protocol.OpenAI, "gpt-4o") + test.insert(t, db, group) + manager := state.NewManager() + registry := state.NewKeyRegistry() + + err := loader.New(db, manager, registry).Load(context.Background()) + if err == nil { + t.Fatal("Load() error = nil, want invalid credential rejection") + } + if !strings.Contains(err.Error(), test.wantError) { + t.Errorf("Load() error = %q, want context containing %q", err, test.wantError) + } + if manager.Current() != nil { + t.Fatalf("Current() = %#v after failed load, want nil", manager.Current()) + } + if got := registry.CollectCandidates([]uint{group.ID}, nil); len(got) != 0 { + t.Fatalf("registry candidates after failed load = %#v, want empty", got) + } + }) + } +} + +func openMigratedDatabase(t *testing.T) *gorm.DB { + t.Helper() + + db, err := storage.Open(":memory:") + if err != nil { + t.Fatalf("storage.Open(:memory:) error = %v", err) + } + sqlDB, err := db.DB() + if err != nil { + t.Fatalf("db.DB() error = %v", err) + } + t.Cleanup(func() { + if err := sqlDB.Close(); err != nil { + t.Errorf("close database: %v", err) + } + }) + if err := storage.AutoMigrate(db); err != nil { + t.Fatalf("storage.AutoMigrate() error = %v", err) + } + return db +} + +func mustCreate(t *testing.T, db *gorm.DB, value any) { + t.Helper() + if err := db.Create(value).Error; err != nil { + t.Fatalf("create %T: %v", value, err) + } +} + +func createRuntimeGroup(t *testing.T, db *gorm.DB, name string, p protocol.Protocol, model string) models.Group { + t.Helper() + group := models.Group{ + Name: name, UpstreamURL: "https://" + name + ".example.com/v1", + Signature: name + "-signature", Protocols: models.JSON(fmt.Sprintf(`[%q]`, p)), + Models: models.JSON(fmt.Sprintf(`[{"id":%q}]`, model)), Config: models.JSON(`{}`), Enabled: true, + } + mustCreate(t, db, &group) + return group +} diff --git a/internal/state/manager.go b/internal/state/manager.go new file mode 100644 index 00000000..1bf8cac4 --- /dev/null +++ b/internal/state/manager.go @@ -0,0 +1,35 @@ +package state + +import ( + "sync" + "sync/atomic" +) + +type Manager struct { + publishMu sync.Mutex + current atomic.Pointer[ConfigSnapshot] +} + +func NewManager() *Manager { + return &Manager{} +} + +func (m *Manager) Current() *ConfigSnapshot { + return m.current.Load() +} + +func (m *Manager) Publish(input CompileInput) (*ConfigSnapshot, error) { + next, err := Compile(input) + if err != nil { + return nil, err + } + + m.publishMu.Lock() + defer m.publishMu.Unlock() + next.Revision = 1 + if current := m.current.Load(); current != nil { + next.Revision = current.Revision + 1 + } + m.current.Store(next) + return next, nil +} diff --git a/internal/state/manager_test.go b/internal/state/manager_test.go new file mode 100644 index 00000000..16499d18 --- /dev/null +++ b/internal/state/manager_test.go @@ -0,0 +1,114 @@ +package state + +import ( + "sync" + "testing" + + "gpt-load/internal/protocol" +) + +func TestManagerPublishStartsAtRevisionOne(t *testing.T) { + manager := NewManager() + + snapshot, err := manager.Publish(managerCompileInput(1)) + if err != nil { + t.Fatalf("Publish() error = %v", err) + } + if snapshot.Revision != 1 { + t.Fatalf("Publish().Revision = %d, want 1", snapshot.Revision) + } + if current := manager.Current(); current != snapshot { + t.Fatalf("Current() = %p, want published snapshot %p", current, snapshot) + } +} + +func TestManagerPublishPreservesOldSnapshotReference(t *testing.T) { + manager := NewManager() + first, err := manager.Publish(managerCompileInput(1)) + if err != nil { + t.Fatalf("first Publish() error = %v", err) + } + second, err := manager.Publish(managerCompileInput(2)) + if err != nil { + t.Fatalf("second Publish() error = %v", err) + } + + if second.Revision != 2 { + t.Fatalf("second Publish().Revision = %d, want 2", second.Revision) + } + if first.Revision != 1 { + t.Fatalf("first snapshot Revision = %d after second publish, want 1", first.Revision) + } + if _, ok := first.Groups[1]; !ok { + t.Fatal("first snapshot reference changed after second publish") + } + if current := manager.Current(); current != second { + t.Fatalf("Current() = %p, want second snapshot %p", current, second) + } +} + +func TestManagerPublishFailureKeepsCurrentSnapshot(t *testing.T) { + manager := NewManager() + first, err := manager.Publish(managerCompileInput(1)) + if err != nil { + t.Fatalf("first Publish() error = %v", err) + } + + invalid := managerCompileInput(2) + invalid.Groups[0].Protocols = []protocol.Protocol{protocol.Protocol("invalid")} + if _, err := manager.Publish(invalid); err == nil { + t.Fatal("Publish() error = nil, want invalid protocol error") + } + if current := manager.Current(); current != first { + t.Fatalf("Current() = %p after failed publish, want %p", current, first) + } +} + +func TestManagerConcurrentPublishAndCurrent(t *testing.T) { + const goroutines = 32 + + manager := NewManager() + start := make(chan struct{}) + var wg sync.WaitGroup + for i := 0; i < goroutines; i++ { + groupID := uint(i + 1) + wg.Add(1) + go func() { + defer wg.Done() + <-start + if _, err := manager.Publish(managerCompileInput(groupID)); err != nil { + t.Errorf("Publish(group %d) error = %v", groupID, err) + } + }() + + wg.Add(1) + go func() { + defer wg.Done() + <-start + for j := 0; j < goroutines; j++ { + _ = manager.Current() + } + }() + } + close(start) + wg.Wait() + + current := manager.Current() + if current == nil { + t.Fatal("Current() = nil after concurrent publishes") + } + if current.Revision != goroutines { + t.Fatalf("Current().Revision = %d, want %d", current.Revision, goroutines) + } +} + +func managerCompileInput(groupID uint) CompileInput { + return CompileInput{Groups: []GroupConfig{{ + ID: groupID, + Name: "group", + UpstreamURL: "https://upstream.example.com", + Protocols: []protocol.Protocol{protocol.OpenAI}, + Models: []ModelConfig{{ID: "model"}}, + Enabled: true, + }}} +} diff --git a/internal/state/registry.go b/internal/state/registry.go new file mode 100644 index 00000000..7412049a --- /dev/null +++ b/internal/state/registry.go @@ -0,0 +1,245 @@ +package state + +import ( + "fmt" + "sort" + "sync" + "time" +) + +type KeyStatus string + +const ( + KeyStatusActive KeyStatus = "active" + KeyStatusDisabled KeyStatus = "disabled" +) + +type KeyEntry struct { + ID uint + GroupID uint + WeightManual *int + Status KeyStatus + CooldownUntil time.Time + Blacklisted bool + FailureCount int + EncryptedValue string +} + +type KeyMeta struct { + ID uint + GroupID uint + WeightManual *int +} + +type KeyRegistry struct { + mu sync.RWMutex + buckets map[uint]map[uint]*KeyEntry + keyGroups map[uint]uint +} + +func NewKeyRegistry() *KeyRegistry { + return &KeyRegistry{ + buckets: make(map[uint]map[uint]*KeyEntry), + keyGroups: make(map[uint]uint), + } +} + +func ValidateKeyEntries(entries []KeyEntry) error { + seen := make(map[uint]struct{}, len(entries)) + for _, entry := range entries { + if entry.ID == 0 { + return fmt.Errorf("key id is required") + } + if entry.GroupID == 0 { + return fmt.Errorf("key %d group id is required", entry.ID) + } + if entry.Status != KeyStatusActive && entry.Status != KeyStatusDisabled { + return fmt.Errorf("key %d has invalid status %q", entry.ID, entry.Status) + } + if entry.EncryptedValue == "" { + return fmt.Errorf("key %d encrypted value is required", entry.ID) + } + if _, duplicate := seen[entry.ID]; duplicate { + return fmt.Errorf("duplicate key id %d", entry.ID) + } + seen[entry.ID] = struct{}{} + } + return nil +} + +func (r *KeyRegistry) Replace(entries []KeyEntry) error { + if err := ValidateKeyEntries(entries); err != nil { + return err + } + + buckets := make(map[uint]map[uint]*KeyEntry) + keyGroups := make(map[uint]uint, len(entries)) + for _, entry := range entries { + if buckets[entry.GroupID] == nil { + buckets[entry.GroupID] = make(map[uint]*KeyEntry) + } + cloned := cloneKeyEntry(entry) + buckets[entry.GroupID][entry.ID] = &cloned + keyGroups[entry.ID] = entry.GroupID + } + + r.mu.Lock() + r.buckets = buckets + r.keyGroups = keyGroups + r.mu.Unlock() + return nil +} + +func (r *KeyRegistry) ApplyImport(groupID uint, entries []KeyEntry) error { + if groupID == 0 { + return fmt.Errorf("group id is required") + } + if err := ValidateKeyEntries(entries); err != nil { + return err + } + for _, entry := range entries { + if entry.GroupID != groupID { + return fmt.Errorf("key %d belongs to group %d, want %d", entry.ID, entry.GroupID, groupID) + } + } + + r.mu.Lock() + defer r.mu.Unlock() + for _, entry := range entries { + if existingGroupID, exists := r.keyGroups[entry.ID]; exists && existingGroupID != groupID { + return fmt.Errorf("key %d already belongs to group %d", entry.ID, existingGroupID) + } + } + for _, entry := range entries { + if r.buckets[groupID] == nil { + r.buckets[groupID] = make(map[uint]*KeyEntry) + } + cloned := cloneKeyEntry(entry) + r.buckets[groupID][entry.ID] = &cloned + r.keyGroups[entry.ID] = groupID + } + return nil +} + +func (r *KeyRegistry) RemoveKey(keyID uint) bool { + r.mu.Lock() + defer r.mu.Unlock() + groupID, ok := r.keyGroups[keyID] + if !ok { + return false + } + delete(r.buckets[groupID], keyID) + if len(r.buckets[groupID]) == 0 { + delete(r.buckets, groupID) + } + delete(r.keyGroups, keyID) + return true +} + +func (r *KeyRegistry) SetKeyStatus(keyID uint, status KeyStatus) error { + if status != KeyStatusActive && status != KeyStatusDisabled { + return fmt.Errorf("invalid key status %q", status) + } + r.mu.Lock() + defer r.mu.Unlock() + groupID, ok := r.keyGroups[keyID] + if !ok { + return fmt.Errorf("key %d not found", keyID) + } + r.buckets[groupID][keyID].Status = status + return nil +} + +func (r *KeyRegistry) EncryptedValue(keyID uint) (string, bool) { + r.mu.RLock() + defer r.mu.RUnlock() + groupID, ok := r.keyGroups[keyID] + if !ok { + return "", false + } + return r.buckets[groupID][keyID].EncryptedValue, true +} + +func (r *KeyRegistry) CollectCandidates(groupIDs []uint, excluded func(uint) bool) []KeyMeta { + r.mu.RLock() + metas := make([]KeyMeta, 0) + for _, groupID := range groupIDs { + for _, entry := range r.buckets[groupID] { + if entry.Status != KeyStatusActive { + continue + } + meta := KeyMeta{ID: entry.ID, GroupID: entry.GroupID} + if entry.WeightManual != nil { + weight := *entry.WeightManual + meta.WeightManual = &weight + } + metas = append(metas, meta) + } + } + r.mu.RUnlock() + + filtered := metas[:0] + for _, meta := range metas { + if excluded != nil && excluded(meta.ID) { + continue + } + filtered = append(filtered, meta) + } + sort.Slice(filtered, func(i, j int) bool { + if filtered[i].GroupID != filtered[j].GroupID { + return filtered[i].GroupID < filtered[j].GroupID + } + return filtered[i].ID < filtered[j].ID + }) + return filtered +} + +func (r *KeyRegistry) SetCooldown(keyID uint, until time.Time) bool { + r.mu.Lock() + defer r.mu.Unlock() + entry, ok := r.entryLocked(keyID) + if !ok { + return false + } + entry.CooldownUntil = until + return true +} + +func (r *KeyRegistry) SetBlacklisted(keyID uint) bool { + r.mu.Lock() + defer r.mu.Unlock() + entry, ok := r.entryLocked(keyID) + if !ok { + return false + } + entry.Blacklisted = true + return true +} + +func (r *KeyRegistry) IncrFailure(keyID uint) (int, bool) { + r.mu.Lock() + defer r.mu.Unlock() + entry, ok := r.entryLocked(keyID) + if !ok { + return 0, false + } + entry.FailureCount++ + return entry.FailureCount, true +} + +func (r *KeyRegistry) entryLocked(keyID uint) (*KeyEntry, bool) { + groupID, ok := r.keyGroups[keyID] + if !ok { + return nil, false + } + entry, ok := r.buckets[groupID][keyID] + return entry, ok +} + +func cloneKeyEntry(entry KeyEntry) KeyEntry { + if entry.WeightManual != nil { + weight := *entry.WeightManual + entry.WeightManual = &weight + } + return entry +} diff --git a/internal/state/registry_test.go b/internal/state/registry_test.go new file mode 100644 index 00000000..905282bf --- /dev/null +++ b/internal/state/registry_test.go @@ -0,0 +1,547 @@ +package state + +import ( + "fmt" + "reflect" + "strings" + "sync" + "testing" + "time" +) + +func TestKeyRegistryReplaceAndEncryptedValue(t *testing.T) { + registry := NewKeyRegistry() + weight := 7 + entries := []KeyEntry{ + { + ID: 1, GroupID: 10, WeightManual: &weight, + Status: KeyStatusActive, EncryptedValue: "cipher-one", + }, + { + ID: 2, GroupID: 20, + Status: KeyStatusDisabled, EncryptedValue: "cipher-two", + }, + } + if err := registry.Replace(entries); err != nil { + t.Fatalf("Replace() error = %v", err) + } + + for _, test := range []struct { + keyID uint + want string + }{ + {keyID: 1, want: "cipher-one"}, + {keyID: 2, want: "cipher-two"}, + } { + got, ok := registry.EncryptedValue(test.keyID) + if !ok || got != test.want { + t.Errorf("EncryptedValue(%d) = %q, %t, want %q, true", test.keyID, got, ok, test.want) + } + } + + weight = 99 + registry.mu.RLock() + storedWeight := *registry.buckets[10][1].WeightManual + registry.mu.RUnlock() + if storedWeight != 7 { + t.Fatalf("stored WeightManual = %d after caller mutation, want 7", storedWeight) + } + + if err := registry.Replace([]KeyEntry{{ + ID: 3, GroupID: 20, + Status: KeyStatusActive, EncryptedValue: "cipher-three", + }}); err != nil { + t.Fatalf("second Replace() error = %v", err) + } + for _, removedID := range []uint{1, 2} { + if got, ok := registry.EncryptedValue(removedID); ok || got != "" { + t.Errorf("EncryptedValue(%d) after replacement = %q, %t, want empty, false", removedID, got, ok) + } + } + if got, ok := registry.EncryptedValue(3); !ok || got != "cipher-three" { + t.Fatalf("EncryptedValue(3) = %q, %t, want %q, true", got, ok, "cipher-three") + } +} + +func TestKeyRegistryReplaceFailurePreservesRegistry(t *testing.T) { + invalidBatches := map[string][]KeyEntry{ + "duplicate ids": { + {ID: 2, GroupID: 20, Status: KeyStatusActive, EncryptedValue: "cipher-two"}, + {ID: 2, GroupID: 20, Status: KeyStatusDisabled, EncryptedValue: "duplicate"}, + }, + "invalid status": { + {ID: 3, GroupID: 30, Status: KeyStatus("cooldown"), EncryptedValue: "cipher-three"}, + }, + } + + for name, batch := range invalidBatches { + t.Run(name, func(t *testing.T) { + registry := NewKeyRegistry() + if err := registry.Replace([]KeyEntry{{ + ID: 1, GroupID: 10, + Status: KeyStatusActive, EncryptedValue: "original-cipher", + }}); err != nil { + t.Fatalf("seed Replace() error = %v", err) + } + + if err := registry.Replace(batch); err == nil { + t.Fatal("invalid Replace() error = nil, want error") + } + if got, ok := registry.EncryptedValue(1); !ok || got != "original-cipher" { + t.Errorf("original EncryptedValue(1) = %q, %t after failed replacement, want %q, true", got, ok, "original-cipher") + } + for _, entry := range batch { + if got, ok := registry.EncryptedValue(entry.ID); ok || got != "" { + t.Errorf("invalid EncryptedValue(%d) = %q, %t after failed replacement, want empty, false", entry.ID, got, ok) + } + } + }) + } +} + +func TestKeyRegistryApplyImportUpsertsOnlyRequestedGroup(t *testing.T) { + registry := NewKeyRegistry() + mustReplaceKeyEntries(t, registry, []KeyEntry{ + {ID: 1, GroupID: 10, Status: KeyStatusActive, EncryptedValue: "old-one"}, + {ID: 2, GroupID: 20, Status: KeyStatusDisabled, EncryptedValue: "cipher-two"}, + }) + + weight := 12 + if err := registry.ApplyImport(10, []KeyEntry{ + {ID: 1, GroupID: 10, Status: KeyStatusDisabled, EncryptedValue: "new-one"}, + {ID: 3, GroupID: 10, WeightManual: &weight, Status: KeyStatusActive, EncryptedValue: "cipher-three"}, + }); err != nil { + t.Fatalf("ApplyImport() error = %v", err) + } + + assertEncryptedValue(t, registry, 1, "new-one", true) + assertEncryptedValue(t, registry, 2, "cipher-two", true) + assertEncryptedValue(t, registry, 3, "cipher-three", true) + + weight = 99 + registry.mu.RLock() + storedWeight := *registry.buckets[10][3].WeightManual + registry.mu.RUnlock() + if storedWeight != 12 { + t.Fatalf("imported WeightManual = %d after caller mutation, want 12", storedWeight) + } +} + +func TestKeyRegistryRemoveKey(t *testing.T) { + registry := NewKeyRegistry() + mustReplaceKeyEntries(t, registry, []KeyEntry{ + {ID: 1, GroupID: 10, Status: KeyStatusActive, EncryptedValue: "cipher-one"}, + {ID: 2, GroupID: 10, Status: KeyStatusActive, EncryptedValue: "cipher-two"}, + {ID: 3, GroupID: 20, Status: KeyStatusActive, EncryptedValue: "cipher-three"}, + }) + + if removed := registry.RemoveKey(1); !removed { + t.Fatal("RemoveKey(1) = false, want true") + } + assertEncryptedValue(t, registry, 1, "", false) + assertEncryptedValue(t, registry, 2, "cipher-two", true) + if removed := registry.RemoveKey(1); removed { + t.Fatal("second RemoveKey(1) = true, want false") + } + + if removed := registry.RemoveKey(2); !removed { + t.Fatal("RemoveKey(2) = false, want true") + } + registry.mu.RLock() + _, emptyBucketRetained := registry.buckets[10] + registry.mu.RUnlock() + if emptyBucketRetained { + t.Fatal("empty group 10 bucket was retained") + } + assertEncryptedValue(t, registry, 3, "cipher-three", true) +} + +func TestKeyRegistrySetKeyStatus(t *testing.T) { + registry := NewKeyRegistry() + mustReplaceKeyEntries(t, registry, []KeyEntry{{ + ID: 1, GroupID: 10, Status: KeyStatusActive, EncryptedValue: "cipher-one", + }}) + + if err := registry.SetKeyStatus(1, KeyStatusDisabled); err != nil { + t.Fatalf("SetKeyStatus(disabled) error = %v", err) + } + if got := keyStatus(t, registry, 1); got != KeyStatusDisabled { + t.Fatalf("key status = %q, want %q", got, KeyStatusDisabled) + } + + if err := registry.SetKeyStatus(1, KeyStatus("cooldown")); err == nil { + t.Fatal("SetKeyStatus(invalid) error = nil, want error") + } + if got := keyStatus(t, registry, 1); got != KeyStatusDisabled { + t.Fatalf("key status after invalid update = %q, want %q", got, KeyStatusDisabled) + } + if err := registry.SetKeyStatus(99, KeyStatusActive); err == nil { + t.Fatal("SetKeyStatus(missing key) error = nil, want error") + } +} + +func TestKeyRegistryCollectCandidatesFiltersStatusAndExcluded(t *testing.T) { + registry := NewKeyRegistry() + weight := 7 + mustReplaceKeyEntries(t, registry, []KeyEntry{ + {ID: 4, GroupID: 30, Status: KeyStatusActive, EncryptedValue: "unselected"}, + { + ID: 3, GroupID: 20, Status: KeyStatusActive, EncryptedValue: "cipher-three", + CooldownUntil: time.Date(2026, time.July, 16, 12, 0, 0, 0, time.UTC), + Blacklisted: true, FailureCount: 9, + }, + {ID: 2, GroupID: 10, Status: KeyStatusDisabled, EncryptedValue: "disabled"}, + {ID: 5, GroupID: 20, Status: KeyStatusActive, EncryptedValue: "excluded"}, + {ID: 1, GroupID: 10, WeightManual: &weight, Status: KeyStatusActive, EncryptedValue: "cipher-one"}, + }) + + excluded := map[uint]bool{5: true} + got := registry.CollectCandidates([]uint{20, 10}, func(keyID uint) bool { + return excluded[keyID] + }) + if len(got) != 2 { + t.Fatalf("CollectCandidates() length = %d, want 2: %#v", len(got), got) + } + if got[0].GroupID != 10 || got[0].ID != 1 { + t.Errorf("CollectCandidates()[0] = %#v, want group 10 key 1", got[0]) + } + if got[1].GroupID != 20 || got[1].ID != 3 { + t.Errorf("CollectCandidates()[1] = %#v, want group 20 key 3", got[1]) + } + if got[0].WeightManual == nil || *got[0].WeightManual != 7 { + t.Fatalf("CollectCandidates()[0].WeightManual = %v, want 7", got[0].WeightManual) + } + + *got[0].WeightManual = 99 + again := registry.CollectCandidates([]uint{10}, nil) + if len(again) != 1 || again[0].WeightManual == nil || *again[0].WeightManual != 7 { + t.Fatalf("CollectCandidates() after caller mutation = %#v, want isolated weight 7", again) + } + + typ := reflect.TypeOf(KeyMeta{}) + for _, field := range []string{"EncryptedValue", "CooldownUntil", "Blacklisted", "FailureCount"} { + if _, ok := typ.FieldByName(field); ok { + t.Fatalf("KeyMeta exposes private field %s", field) + } + } +} + +func TestKeyRegistryReservedRuntimeMutationsAreAtomic(t *testing.T) { + registry := NewKeyRegistry() + mustReplaceKeyEntries(t, registry, []KeyEntry{{ + ID: 1, GroupID: 10, Status: KeyStatusActive, EncryptedValue: "cipher-one", + }}) + + until := time.Date(2026, time.July, 16, 12, 0, 0, 0, time.UTC) + if ok := registry.SetCooldown(1, until); !ok { + t.Fatal("SetCooldown(1) = false, want true") + } + if ok := registry.SetBlacklisted(1); !ok { + t.Fatal("SetBlacklisted(1) = false, want true") + } + for want := 1; want <= 2; want++ { + got, ok := registry.IncrFailure(1) + if !ok || got != want { + t.Fatalf("IncrFailure(1) = %d, %t, want %d, true", got, ok, want) + } + } + + registry.mu.RLock() + entry := *registry.buckets[10][1] + registry.mu.RUnlock() + if !entry.CooldownUntil.Equal(until) { + t.Errorf("CooldownUntil = %v, want %v", entry.CooldownUntil, until) + } + if !entry.Blacklisted { + t.Error("Blacklisted = false, want true") + } + if entry.FailureCount != 2 { + t.Errorf("FailureCount = %d, want 2", entry.FailureCount) + } + + if ok := registry.SetCooldown(99, until); ok { + t.Error("SetCooldown(missing key) = true, want false") + } + if ok := registry.SetBlacklisted(99); ok { + t.Error("SetBlacklisted(missing key) = true, want false") + } + if got, ok := registry.IncrFailure(99); ok || got != 0 { + t.Errorf("IncrFailure(missing key) = %d, %t, want 0, false", got, ok) + } +} + +func TestKeyRegistryConcurrentMutationsAndCollection(t *testing.T) { + const ( + writerCount = 6 + readerCount = 6 + operations = 48 + ) + + registry := NewKeyRegistry() + groupIDs := make([]uint, writerCount) + for writer := 0; writer < writerCount; writer++ { + groupIDs[writer] = uint(writer + 1) + } + + start := make(chan struct{}) + errors := make(chan error, writerCount*operations*6+readerCount*writerCount*operations) + var wg sync.WaitGroup + for writer := 0; writer < writerCount; writer++ { + groupID := uint(writer + 1) + keyBase := uint((writer + 1) * 1000) + wg.Add(1) + go func() { + defer wg.Done() + <-start + for operation := 0; operation < operations; operation++ { + keyID := keyBase + uint(operation) + entry := KeyEntry{ + ID: keyID, GroupID: groupID, Status: KeyStatusActive, + EncryptedValue: fmt.Sprintf("cipher-%d", keyID), + } + if err := registry.ApplyImport(groupID, []KeyEntry{entry}); err != nil { + errors <- fmt.Errorf("ApplyImport(group %d, key %d): %w", groupID, keyID, err) + continue + } + if err := registry.SetKeyStatus(keyID, KeyStatusDisabled); err != nil { + errors <- fmt.Errorf("SetKeyStatus(%d, disabled): %w", keyID, err) + } + if ok := registry.SetCooldown(keyID, time.Unix(int64(operation+1), 0)); !ok { + errors <- fmt.Errorf("SetCooldown(%d) = false", keyID) + } + if ok := registry.SetBlacklisted(keyID); !ok { + errors <- fmt.Errorf("SetBlacklisted(%d) = false", keyID) + } + if _, ok := registry.IncrFailure(keyID); !ok { + errors <- fmt.Errorf("IncrFailure(%d) = false", keyID) + } + if err := registry.SetKeyStatus(keyID, KeyStatusActive); err != nil { + errors <- fmt.Errorf("SetKeyStatus(%d, active): %w", keyID, err) + } + if operation%3 == 0 { + if removed := registry.RemoveKey(keyID); !removed { + errors <- fmt.Errorf("RemoveKey(%d) = false", keyID) + } + } else if _, ok := registry.EncryptedValue(keyID); !ok { + errors <- fmt.Errorf("EncryptedValue(%d) missing after import", keyID) + } + } + }() + } + + for reader := 0; reader < readerCount; reader++ { + wg.Add(1) + go func(reader int) { + defer wg.Done() + <-start + for operation := 0; operation < writerCount*operations; operation++ { + writer := operation % writerCount + keyID := uint((writer+1)*1000 + operation/writerCount%operations) + registry.EncryptedValue(keyID) + candidates := registry.CollectCandidates(groupIDs, func(candidateID uint) bool { + return candidateID%19 == uint(reader)%19 + }) + for index := 1; index < len(candidates); index++ { + previous := candidates[index-1] + current := candidates[index] + if previous.GroupID > current.GroupID || + (previous.GroupID == current.GroupID && previous.ID > current.ID) { + errors <- fmt.Errorf("CollectCandidates() returned unsorted entries: %#v then %#v", previous, current) + break + } + } + } + }(reader) + } + + close(start) + wg.Wait() + close(errors) + for err := range errors { + t.Error(err) + } + + finalWeight := 13 + finalEntries := []KeyEntry{ + {ID: 42, GroupID: 4, Status: KeyStatusActive, EncryptedValue: "final-forty-two"}, + {ID: 43, GroupID: 4, Status: KeyStatusDisabled, EncryptedValue: "final-disabled"}, + {ID: 7, GroupID: 2, WeightManual: &finalWeight, Status: KeyStatusActive, EncryptedValue: "final-seven"}, + } + if err := registry.Replace(finalEntries); err != nil { + t.Fatalf("final Replace() error = %v", err) + } + + got := registry.CollectCandidates([]uint{4, 2}, nil) + if len(got) != 2 { + t.Fatalf("final CollectCandidates() length = %d, want 2: %#v", len(got), got) + } + if got[0].GroupID != 2 || got[0].ID != 7 || got[0].WeightManual == nil || *got[0].WeightManual != 13 { + t.Errorf("final CollectCandidates()[0] = %#v, want group 2 key 7 weight 13", got[0]) + } + if got[1].GroupID != 4 || got[1].ID != 42 || got[1].WeightManual != nil { + t.Errorf("final CollectCandidates()[1] = %#v, want group 4 key 42 without manual weight", got[1]) + } + + for keyID, want := range map[uint]string{ + 7: "final-seven", + 42: "final-forty-two", + 43: "final-disabled", + } { + assertEncryptedValue(t, registry, keyID, want, true) + } + for writer := 0; writer < writerCount; writer++ { + assertEncryptedValue(t, registry, uint((writer+1)*1000), "", false) + } +} + +func TestValidateKeyEntriesRejectsMalformedEntries(t *testing.T) { + tests := []struct { + name string + entries []KeyEntry + wantInError string + }{ + { + name: "zero id", + entries: []KeyEntry{{GroupID: 10, Status: KeyStatusActive, EncryptedValue: "cipher"}}, + wantInError: "id is required", + }, + { + name: "zero group id", + entries: []KeyEntry{{ID: 1, Status: KeyStatusActive, EncryptedValue: "cipher"}}, + wantInError: "group id is required", + }, + { + name: "invalid status", + entries: []KeyEntry{{ID: 1, GroupID: 10, Status: KeyStatus("cooldown"), EncryptedValue: "cipher"}}, + wantInError: "invalid status", + }, + { + name: "empty ciphertext", + entries: []KeyEntry{{ID: 1, GroupID: 10, Status: KeyStatusActive}}, + wantInError: "encrypted value is required", + }, + { + name: "duplicate id", + entries: []KeyEntry{ + {ID: 1, GroupID: 10, Status: KeyStatusActive, EncryptedValue: "cipher-one"}, + {ID: 1, GroupID: 20, Status: KeyStatusDisabled, EncryptedValue: "cipher-two"}, + }, + wantInError: "duplicate key id 1", + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + err := ValidateKeyEntries(test.entries) + if err == nil { + t.Fatal("ValidateKeyEntries() error = nil, want error") + } + if !strings.Contains(err.Error(), test.wantInError) { + t.Fatalf("ValidateKeyEntries() error = %q, want substring %q", err, test.wantInError) + } + }) + } + + if err := ValidateKeyEntries([]KeyEntry{ + {ID: 1, GroupID: 10, Status: KeyStatusActive, EncryptedValue: "cipher-one"}, + {ID: 2, GroupID: 10, Status: KeyStatusDisabled, EncryptedValue: "cipher-two"}, + }); err != nil { + t.Fatalf("ValidateKeyEntries(valid) error = %v", err) + } +} + +func TestKeyRegistryApplyImportFailureDoesNotMutateRegistry(t *testing.T) { + tests := []struct { + name string + groupID uint + batch []KeyEntry + }{ + { + name: "invalid entry after valid entry", + groupID: 10, + batch: []KeyEntry{ + {ID: 1, GroupID: 10, Status: KeyStatusActive, EncryptedValue: "updated-one"}, + {ID: 2, GroupID: 10, Status: KeyStatusActive}, + }, + }, + { + name: "entry group differs from requested group", + groupID: 10, + batch: []KeyEntry{ + {ID: 1, GroupID: 10, Status: KeyStatusActive, EncryptedValue: "updated-one"}, + {ID: 3, GroupID: 20, Status: KeyStatusActive, EncryptedValue: "cipher-three"}, + }, + }, + { + name: "zero requested group", + batch: []KeyEntry{ + {ID: 1, GroupID: 10, Status: KeyStatusActive, EncryptedValue: "updated-one"}, + }, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + registry := NewKeyRegistry() + mustReplaceKeyEntries(t, registry, []KeyEntry{{ + ID: 1, GroupID: 10, Status: KeyStatusActive, EncryptedValue: "original-one", + }}) + + if err := registry.ApplyImport(test.groupID, test.batch); err == nil { + t.Fatal("ApplyImport(invalid batch) error = nil, want error") + } + assertEncryptedValue(t, registry, 1, "original-one", true) + for _, entry := range test.batch { + if entry.ID != 1 { + assertEncryptedValue(t, registry, entry.ID, "", false) + } + } + }) + } +} + +func TestKeyRegistryApplyImportRejectsExistingIDFromAnotherGroup(t *testing.T) { + registry := NewKeyRegistry() + mustReplaceKeyEntries(t, registry, []KeyEntry{{ + ID: 1, GroupID: 10, Status: KeyStatusActive, EncryptedValue: "original-one", + }}) + + err := registry.ApplyImport(20, []KeyEntry{ + {ID: 2, GroupID: 20, Status: KeyStatusActive, EncryptedValue: "cipher-two"}, + {ID: 1, GroupID: 20, Status: KeyStatusActive, EncryptedValue: "moved-one"}, + }) + if err == nil { + t.Fatal("ApplyImport(cross-group key id) error = nil, want error") + } + if !strings.Contains(err.Error(), "already belongs to group 10") { + t.Fatalf("ApplyImport(cross-group key id) error = %q, want existing group context", err) + } + assertEncryptedValue(t, registry, 1, "original-one", true) + assertEncryptedValue(t, registry, 2, "", false) +} + +func mustReplaceKeyEntries(t *testing.T, registry *KeyRegistry, entries []KeyEntry) { + t.Helper() + if err := registry.Replace(entries); err != nil { + t.Fatalf("Replace() error = %v", err) + } +} + +func assertEncryptedValue(t *testing.T, registry *KeyRegistry, keyID uint, want string, wantOK bool) { + t.Helper() + got, ok := registry.EncryptedValue(keyID) + if got != want || ok != wantOK { + t.Errorf("EncryptedValue(%d) = %q, %t, want %q, %t", keyID, got, ok, want, wantOK) + } +} + +func keyStatus(t *testing.T, registry *KeyRegistry, keyID uint) KeyStatus { + t.Helper() + registry.mu.RLock() + defer registry.mu.RUnlock() + groupID, ok := registry.keyGroups[keyID] + if !ok { + t.Fatalf("key %d does not exist", keyID) + } + return registry.buckets[groupID][keyID].Status +} diff --git a/internal/state/snapshot.go b/internal/state/snapshot.go new file mode 100644 index 00000000..c76eba2d --- /dev/null +++ b/internal/state/snapshot.go @@ -0,0 +1,383 @@ +// Package state owns immutable runtime configuration snapshots. +package state + +import ( + "encoding/json" + "fmt" + "math" + "sort" + "strings" + "time" + + "gpt-load/internal/platform/config" + "gpt-load/internal/protocol" +) + +type CompileInput struct { + SystemSettings config.Settings + Groups []GroupConfig + AccessKeys []AccessKeyConfig +} + +type GroupConfig struct { + ID uint + Name string + UpstreamURL string + Protocols []protocol.Protocol + Models []ModelConfig + Settings config.Settings + Enabled bool +} + +type ModelConfig struct { + ID string + Alias string +} + +type AccessKeyConfig struct { + ID uint + Name string + KeyHash string + Status AccessKeyStatus + Filters FilterSet +} + +type AccessKeyStatus string + +const ( + AccessKeyStatusActive AccessKeyStatus = "active" + AccessKeyStatusDisabled AccessKeyStatus = "disabled" +) + +type FilterSet struct { + Groups map[uint]struct{} + Protocols map[protocol.Protocol]struct{} + Models map[string]struct{} +} + +type RouteTarget struct { + GroupID uint + UpstreamModelID string +} + +type TimeoutConfig struct { + Connect time.Duration + FirstByte time.Duration + Request time.Duration + StreamIdle time.Duration +} + +var defaultTimeouts = TimeoutConfig{ + Connect: 15 * time.Second, + FirstByte: 120 * time.Second, + Request: 600 * time.Second, + StreamIdle: 300 * time.Second, +} + +type HeaderRules struct { + Set map[string]string + Remove []string +} + +type GroupView struct { + ID uint + Name string + UpstreamURL string + Protocols []protocol.Protocol + Models []ModelConfig + Timeouts TimeoutConfig + HeaderRules HeaderRules +} + +type AccessKeyView struct { + ID uint + Name string + Filters FilterSet +} + +type ConfigSnapshot struct { + Revision uint64 + Candidates map[protocol.Protocol]map[string][]RouteTarget + Groups map[uint]GroupView + AccessKeysByHash map[string]AccessKeyView +} + +func Compile(input CompileInput) (*ConfigSnapshot, error) { + if err := validateCompileInput(input); err != nil { + return nil, err + } + + snapshot := &ConfigSnapshot{ + Candidates: make(map[protocol.Protocol]map[string][]RouteTarget), + Groups: make(map[uint]GroupView), + AccessKeysByHash: make(map[string]AccessKeyView), + } + + for _, group := range input.Groups { + if !group.Enabled { + continue + } + timeouts, headerRules, err := compileRuntimeSettings(input.SystemSettings, group.Settings) + if err != nil { + return nil, err + } + view := GroupView{ + ID: group.ID, Name: group.Name, UpstreamURL: group.UpstreamURL, + Protocols: append([]protocol.Protocol(nil), group.Protocols...), + Models: append([]ModelConfig(nil), group.Models...), + Timeouts: timeouts, + HeaderRules: headerRules, + } + snapshot.Groups[group.ID] = view + + seen := make(map[protocol.Protocol]map[string]struct{}) + for _, model := range group.Models { + modelID := strings.TrimSpace(model.ID) + for _, p := range group.Protocols { + if seen[p] == nil { + seen[p] = make(map[string]struct{}) + } + if _, duplicate := seen[p][modelID]; duplicate { + continue + } + seen[p][modelID] = struct{}{} + if snapshot.Candidates[p] == nil { + snapshot.Candidates[p] = make(map[string][]RouteTarget) + } + snapshot.Candidates[p][modelID] = append( + snapshot.Candidates[p][modelID], + RouteTarget{GroupID: group.ID, UpstreamModelID: modelID}, + ) + } + } + } + + 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), + } + } + + 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] + if left.GroupID != right.GroupID { + return left.GroupID < right.GroupID + } + return left.UpstreamModelID < right.UpstreamModelID + }) + } + } + return snapshot, nil +} + +func compileRuntimeSettings(system, group config.Settings) (TimeoutConfig, HeaderRules, error) { + allowedGroupSettings := map[string]struct{}{ + "connect_timeout": {}, "first_byte_timeout": {}, + "request_timeout": {}, "stream_idle_timeout": {}, "header_rules": {}, + } + for key := range group { + if _, ok := allowedGroupSettings[key]; !ok { + return TimeoutConfig{}, HeaderRules{}, fmt.Errorf("unknown group setting %q", key) + } + } + + merged := config.MergeSettings(system, group) + timeouts := defaultTimeouts + fields := []struct { + key string + target *time.Duration + }{ + {key: "connect_timeout", target: &timeouts.Connect}, + {key: "first_byte_timeout", target: &timeouts.FirstByte}, + {key: "request_timeout", target: &timeouts.Request}, + {key: "stream_idle_timeout", target: &timeouts.StreamIdle}, + } + for _, field := range fields { + value, ok := merged[field.key] + if !ok { + continue + } + seconds, err := positiveWholeSeconds(field.key, value) + if err != nil { + return TimeoutConfig{}, HeaderRules{}, err + } + *field.target = time.Duration(seconds) * time.Second + } + + rules, err := parseHeaderRules(merged["header_rules"]) + if err != nil { + return TimeoutConfig{}, HeaderRules{}, err + } + return timeouts, rules, nil +} + +func positiveWholeSeconds(path string, value any) (int64, error) { + const maxTimeoutSeconds = int64((1<<63)-1) / int64(time.Second) + + var seconds int64 + switch typed := value.(type) { + case int: + seconds = int64(typed) + case int64: + seconds = typed + case float64: + if math.IsNaN(typed) || math.IsInf(typed, 0) || math.Trunc(typed) != typed || typed > float64(maxTimeoutSeconds) { + return 0, fmt.Errorf("%s must be a positive whole number", path) + } + seconds = int64(typed) + case json.Number: + parsed, err := typed.Int64() + if err != nil { + return 0, fmt.Errorf("%s must be a positive whole number", path) + } + seconds = parsed + default: + return 0, fmt.Errorf("%s must be a positive whole number", path) + } + if seconds <= 0 || seconds > maxTimeoutSeconds { + return 0, fmt.Errorf("%s must be a positive whole number within duration range", path) + } + return seconds, nil +} + +func parseHeaderRules(value any) (HeaderRules, error) { + rules := HeaderRules{Set: make(map[string]string)} + if value == nil { + return rules, nil + } + object, ok := value.(map[string]any) + if !ok { + return HeaderRules{}, fmt.Errorf("header_rules must be an object") + } + for key := range object { + if key != "set" && key != "remove" { + return HeaderRules{}, fmt.Errorf("unknown header_rules field %q", key) + } + } + if rawSet, exists := object["set"]; exists { + set, ok := rawSet.(map[string]any) + if !ok { + return HeaderRules{}, fmt.Errorf("header_rules.set must be an object") + } + for name, rawValue := range set { + text, ok := rawValue.(string) + if !ok { + return HeaderRules{}, fmt.Errorf("header_rules.set.%s must be a string", name) + } + rules.Set[name] = text + } + } + if rawRemove, exists := object["remove"]; exists { + remove, ok := rawRemove.([]any) + if !ok { + return HeaderRules{}, fmt.Errorf("header_rules.remove must be an array") + } + rules.Remove = make([]string, 0, len(remove)) + for index, rawName := range remove { + name, ok := rawName.(string) + if !ok { + return HeaderRules{}, fmt.Errorf("header_rules.remove[%d] must be a string", index) + } + rules.Remove = append(rules.Remove, name) + } + } + return rules, nil +} + +func validateCompileInput(input CompileInput) error { + enabledGroupIDs := make(map[uint]struct{}) + activeHashes := make(map[string]struct{}) + 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 { + return fmt.Errorf("duplicate group id %d", group.ID) + } + 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) + } + if _, duplicate := protocols[p]; duplicate { + return fmt.Errorf("group %d has duplicate protocol %q", group.ID, p) + } + protocols[p] = struct{}{} + } + for _, model := range group.Models { + modelID := strings.TrimSpace(model.ID) + if modelID == "" { + return fmt.Errorf("group %d model id is required", group.ID) + } + } + } + for _, accessKey := range input.AccessKeys { + switch accessKey.Status { + case AccessKeyStatusDisabled: + continue + case AccessKeyStatusActive: + 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 { + return fmt.Errorf("duplicate access key hash %q", accessKey.KeyHash) + } + activeHashes[accessKey.KeyHash] = struct{}{} + if err := validateFilterSet(accessKey.ID, accessKey.Filters); err != nil { + return err + } + } + return nil +} + +func validateFilterSet(accessKeyID uint, filters FilterSet) error { + for p := range filters.Protocols { + if !p.Valid() { + return fmt.Errorf("access key %d filter has invalid protocol %q", accessKeyID, p) + } + } + for model := range filters.Models { + if strings.TrimSpace(model) == "" { + return fmt.Errorf("access key %d filter model is required", accessKeyID) + } + } + return nil +} + +func cloneFilterSet(source FilterSet) FilterSet { + cloned := FilterSet{} + if source.Groups != nil { + cloned.Groups = make(map[uint]struct{}, len(source.Groups)) + for id := range source.Groups { + cloned.Groups[id] = struct{}{} + } + } + if source.Protocols != nil { + cloned.Protocols = make(map[protocol.Protocol]struct{}, len(source.Protocols)) + for p := range source.Protocols { + cloned.Protocols[p] = struct{}{} + } + } + if source.Models != nil { + cloned.Models = make(map[string]struct{}, len(source.Models)) + for model := range source.Models { + cloned.Models[model] = struct{}{} + } + } + return cloned +} diff --git a/internal/state/snapshot_test.go b/internal/state/snapshot_test.go new file mode 100644 index 00000000..b5c2bd0a --- /dev/null +++ b/internal/state/snapshot_test.go @@ -0,0 +1,469 @@ +package state + +import ( + "encoding/json" + "reflect" + "strings" + "testing" + "time" + + "gpt-load/internal/platform/config" + "gpt-load/internal/protocol" +) + +func TestCompileBuildsIDOnlyCandidatesAndActiveAccessKeys(t *testing.T) { + input := CompileInput{ + Groups: []GroupConfig{ + { + ID: 1, Name: "enabled", UpstreamURL: "https://one.example.com", + Protocols: []protocol.Protocol{protocol.OpenAI, protocol.Anthropic}, + Models: []ModelConfig{ + {ID: "model-real", Alias: "model-alias"}, + {ID: "model-real", Alias: "second-alias"}, + }, + Enabled: true, + }, + { + ID: 2, Name: "disabled", UpstreamURL: "https://two.example.com", + Protocols: []protocol.Protocol{protocol.OpenAI}, + Models: []ModelConfig{{ID: "hidden"}}, Enabled: false, + }, + }, + AccessKeys: []AccessKeyConfig{ + {ID: 10, Name: "active", KeyHash: "active-hash", Status: AccessKeyStatusActive}, + {ID: 11, Name: "disabled", KeyHash: "disabled-hash", Status: AccessKeyStatusDisabled}, + }, + } + + snapshot, err := Compile(input) + if err != nil { + t.Fatalf("Compile() error = %v", err) + } + if _, ok := snapshot.Candidates[protocol.OpenAI]["model-alias"]; ok { + t.Fatal("alias unexpectedly entered the M1 candidate index") + } + targets := snapshot.Candidates[protocol.OpenAI]["model-real"] + if len(targets) != 1 || targets[0] != (RouteTarget{GroupID: 1, UpstreamModelID: "model-real"}) { + t.Fatalf("OpenAI targets = %#v", targets) + } + if got := len(snapshot.Groups[1].Models); got != 2 { + t.Fatalf("GroupView.Models length = %d, want both alias entries", got) + } + if _, ok := snapshot.Groups[2]; ok { + t.Fatal("disabled group entered snapshot") + } + if _, ok := snapshot.Candidates[protocol.OpenAI]["hidden"]; ok { + t.Fatal("disabled group model entered candidate index") + } + if _, ok := snapshot.AccessKeysByHash["active-hash"]; !ok { + t.Fatal("active key missing") + } + if _, ok := snapshot.AccessKeysByHash["disabled-hash"]; ok { + t.Fatal("disabled key entered snapshot") + } +} + +func TestCompileMergesTimeoutsAndHeaderRules(t *testing.T) { + systemSettings := config.Settings{ + "connect_timeout": 10.0, + "first_byte_timeout": 90.0, + "header_rules": map[string]any{ + "set": map[string]any{"X-System": "system"}, + "remove": []any{"X-System-Remove"}, + }, + } + groupSettings := config.Settings{ + "connect_timeout": 20.0, + "request_timeout": 500.0, + "stream_idle_timeout": 250.0, + "header_rules": map[string]any{ + "set": map[string]any{"X-Group": "${API_KEY}"}, + "remove": []any{"X-Group-Remove"}, + }, + } + + snapshot, err := Compile(runtimeSettingsInput(systemSettings, groupSettings)) + if err != nil { + t.Fatalf("Compile() error = %v", err) + } + + group := snapshot.Groups[1] + wantTimeouts := TimeoutConfig{ + Connect: 20 * time.Second, + FirstByte: 90 * time.Second, + Request: 500 * time.Second, + StreamIdle: 250 * time.Second, + } + if group.Timeouts != wantTimeouts { + t.Errorf("GroupView.Timeouts = %#v, want %#v", group.Timeouts, wantTimeouts) + } + wantRules := HeaderRules{ + Set: map[string]string{"X-Group": "${API_KEY}"}, + Remove: []string{"X-Group-Remove"}, + } + if !reflect.DeepEqual(group.HeaderRules, wantRules) { + t.Errorf("GroupView.HeaderRules = %#v, want %#v", group.HeaderRules, wantRules) + } +} + +func TestCompileUsesDefaultRuntimeSettings(t *testing.T) { + snapshot, err := Compile(runtimeSettingsInput(nil, nil)) + if err != nil { + t.Fatalf("Compile() error = %v", err) + } + + group := snapshot.Groups[1] + wantTimeouts := TimeoutConfig{ + Connect: 15 * time.Second, + FirstByte: 120 * time.Second, + Request: 600 * time.Second, + StreamIdle: 300 * time.Second, + } + if group.Timeouts != wantTimeouts { + t.Errorf("GroupView.Timeouts = %#v, want %#v", group.Timeouts, wantTimeouts) + } + wantRules := HeaderRules{Set: map[string]string{}} + if !reflect.DeepEqual(group.HeaderRules, wantRules) { + t.Errorf("GroupView.HeaderRules = %#v, want %#v", group.HeaderRules, wantRules) + } +} + +func TestCompileRejectsMalformedRuntimeSettings(t *testing.T) { + tests := []struct { + name string + groupSettings config.Settings + wantErr string + }{ + { + name: "zero timeout", + groupSettings: config.Settings{"connect_timeout": 0}, + wantErr: "connect_timeout must be a positive whole number", + }, + { + name: "negative timeout", + groupSettings: config.Settings{"first_byte_timeout": int64(-1)}, + wantErr: "first_byte_timeout must be a positive whole number", + }, + { + name: "fractional timeout", + groupSettings: config.Settings{"request_timeout": 1.5}, + wantErr: "request_timeout must be a positive whole number", + }, + { + name: "overflow timeout", + groupSettings: config.Settings{"stream_idle_timeout": json.Number("9223372037")}, + wantErr: "stream_idle_timeout must be a positive whole number", + }, + { + name: "non-object header rules", + groupSettings: config.Settings{"header_rules": []any{}}, + wantErr: "header_rules must be an object", + }, + { + name: "non-object header set", + groupSettings: config.Settings{"header_rules": map[string]any{ + "set": []any{}, + }}, + wantErr: "header_rules.set must be an object", + }, + { + name: "non-string header set value", + groupSettings: config.Settings{"header_rules": map[string]any{ + "set": map[string]any{"X-Invalid": 1.0}, + }}, + wantErr: "header_rules.set.X-Invalid must be a string", + }, + { + name: "non-array header remove", + groupSettings: config.Settings{"header_rules": map[string]any{ + "remove": "X-Invalid", + }}, + wantErr: "header_rules.remove must be an array", + }, + { + name: "non-string header remove entry", + groupSettings: config.Settings{"header_rules": map[string]any{ + "remove": []any{"X-Valid", 1.0}, + }}, + wantErr: "header_rules.remove[1] must be a string", + }, + { + name: "unknown header rules field", + groupSettings: config.Settings{"header_rules": map[string]any{ + "append": map[string]any{"X-Invalid": "value"}, + }}, + wantErr: "unknown header_rules field \"append\"", + }, + { + name: "unknown group setting", + groupSettings: config.Settings{"retry_count": 3.0}, + wantErr: "unknown group setting \"retry_count\"", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := Compile(runtimeSettingsInput(nil, tt.groupSettings)) + if err == nil || !strings.Contains(err.Error(), tt.wantErr) { + t.Fatalf("Compile() error = %v, want error containing %q", err, tt.wantErr) + } + }) + } +} + +func TestCompileRuntimeSettingsOwnsIndependentCopies(t *testing.T) { + set := map[string]any{"X-Original": "original"} + remove := []any{"X-Original-Remove"} + headerRules := map[string]any{"set": set, "remove": remove} + groupSettings := config.Settings{"header_rules": headerRules} + + snapshot, err := Compile(runtimeSettingsInput(nil, groupSettings)) + if err != nil { + t.Fatalf("Compile() error = %v", err) + } + + set["X-Original"] = "mutated" + set["X-Added"] = "added" + remove[0] = "X-Mutated-Remove" + headerRules["set"] = map[string]any{"X-Replaced": "replaced"} + headerRules["remove"] = []any{"X-Replaced-Remove"} + groupSettings["header_rules"] = map[string]any{} + + wantRules := HeaderRules{ + Set: map[string]string{"X-Original": "original"}, + Remove: []string{"X-Original-Remove"}, + } + if got := snapshot.Groups[1].HeaderRules; !reflect.DeepEqual(got, wantRules) { + t.Errorf("GroupView.HeaderRules = %#v after input mutation, want %#v", got, wantRules) + } +} + +func runtimeSettingsInput(systemSettings, groupSettings config.Settings) CompileInput { + return CompileInput{ + SystemSettings: systemSettings, + Groups: []GroupConfig{{ + ID: 1, + Protocols: []protocol.Protocol{protocol.OpenAI}, + Models: []ModelConfig{{ID: "model"}}, + Settings: groupSettings, + Enabled: true, + }}, + } +} + +func TestCompileRejectsInvalidCandidateConfiguration(t *testing.T) { + tests := []struct { + name string + input CompileInput + wantErr string + }{ + { + name: "zero group id", + input: CompileInput{Groups: []GroupConfig{{ + Protocols: []protocol.Protocol{protocol.OpenAI}, + Models: []ModelConfig{{ID: "model"}}, + Enabled: true, + }}}, + wantErr: "group id is required", + }, + { + name: "duplicate enabled group id", + input: CompileInput{Groups: []GroupConfig{ + { + ID: 1, + Protocols: []protocol.Protocol{protocol.OpenAI}, + Models: []ModelConfig{{ID: "model-one"}}, + Enabled: true, + }, + { + ID: 1, + Protocols: []protocol.Protocol{protocol.Anthropic}, + Models: []ModelConfig{{ID: "model-two"}}, + Enabled: true, + }, + }}, + wantErr: "duplicate group id", + }, + { + name: "empty protocol list", + input: CompileInput{Groups: []GroupConfig{{ + ID: 1, + Models: []ModelConfig{{ID: "model"}}, + Enabled: true, + }}}, + wantErr: "protocols are required", + }, + { + name: "unknown protocol", + input: CompileInput{Groups: []GroupConfig{{ + ID: 1, + Protocols: []protocol.Protocol{"unknown"}, + Models: []ModelConfig{{ID: "model"}}, + Enabled: true, + }}}, + wantErr: "invalid protocol", + }, + { + name: "duplicate protocol", + input: CompileInput{Groups: []GroupConfig{{ + ID: 1, + Protocols: []protocol.Protocol{protocol.OpenAI, protocol.OpenAI}, + Models: []ModelConfig{{ID: "model"}}, + Enabled: true, + }}}, + wantErr: "duplicate protocol", + }, + { + name: "whitespace model id", + input: CompileInput{Groups: []GroupConfig{{ + ID: 1, + Protocols: []protocol.Protocol{protocol.OpenAI}, + Models: []ModelConfig{{ID: " \t"}}, + Enabled: true, + }}}, + wantErr: "model id is required", + }, + { + name: "active access key without hash", + input: CompileInput{AccessKeys: []AccessKeyConfig{{ + ID: 10, + Status: AccessKeyStatusActive, + }}}, + wantErr: "key hash is required", + }, + { + name: "invalid access key status", + input: CompileInput{AccessKeys: []AccessKeyConfig{{ + ID: 10, + KeyHash: "hash", + Status: AccessKeyStatus("invalid"), + }}}, + wantErr: "invalid status", + }, + { + name: "duplicate active access key hash", + input: CompileInput{AccessKeys: []AccessKeyConfig{ + {ID: 10, KeyHash: "duplicate", Status: AccessKeyStatusActive}, + {ID: 11, KeyHash: "duplicate", Status: AccessKeyStatusActive}, + }}, + wantErr: "duplicate access key hash", + }, + { + name: "invalid filter protocol", + input: CompileInput{AccessKeys: []AccessKeyConfig{{ + ID: 10, + KeyHash: "hash", + Status: AccessKeyStatusActive, + Filters: FilterSet{Protocols: map[protocol.Protocol]struct{}{"invalid": {}}}, + }}}, + wantErr: "invalid protocol", + }, + { + name: "blank filter model", + input: CompileInput{AccessKeys: []AccessKeyConfig{{ + ID: 10, + KeyHash: "hash", + Status: AccessKeyStatusActive, + Filters: FilterSet{Models: map[string]struct{}{" ": {}}}, + }}}, + wantErr: "filter model is required", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := Compile(tt.input) + if err == nil || !strings.Contains(err.Error(), tt.wantErr) { + t.Fatalf("Compile() error = %v, want error containing %q", err, tt.wantErr) + } + }) + } +} + +func TestCompileAllowsDanglingFilterGroupIDs(t *testing.T) { + input := CompileInput{AccessKeys: []AccessKeyConfig{{ + ID: 10, + KeyHash: "hash", + Status: AccessKeyStatusActive, + Filters: FilterSet{Groups: map[uint]struct{}{999: {}}}, + }}} + + snapshot, err := Compile(input) + if err != nil { + t.Fatalf("Compile() error = %v", err) + } + accessKey, ok := snapshot.AccessKeysByHash["hash"] + if !ok { + t.Fatal("active access key missing") + } + if _, ok := accessKey.Filters.Groups[999]; !ok { + t.Fatal("dangling group filter id was not preserved") + } +} + +func TestCompileOwnsIndependentCopiesOfInput(t *testing.T) { + protocols := []protocol.Protocol{protocol.OpenAI} + models := []ModelConfig{{ID: "model-real", Alias: "model-alias"}} + filterGroups := map[uint]struct{}{1: {}} + filterProtocols := map[protocol.Protocol]struct{}{protocol.OpenAI: {}} + filterModels := map[string]struct{}{"model-real": {}} + input := CompileInput{ + Groups: []GroupConfig{{ + ID: 1, + Protocols: protocols, + Models: models, + Enabled: true, + }}, + AccessKeys: []AccessKeyConfig{{ + ID: 10, + KeyHash: "hash", + Status: AccessKeyStatusActive, + Filters: FilterSet{ + Groups: filterGroups, + Protocols: filterProtocols, + Models: filterModels, + }, + }}, + } + + snapshot, err := Compile(input) + if err != nil { + t.Fatalf("Compile() error = %v", err) + } + + protocols[0] = protocol.Anthropic + models[0] = ModelConfig{ID: "changed", Alias: "changed"} + delete(filterGroups, 1) + filterGroups[2] = struct{}{} + delete(filterProtocols, protocol.OpenAI) + filterProtocols[protocol.Gemini] = struct{}{} + delete(filterModels, "model-real") + filterModels["changed"] = struct{}{} + + group := snapshot.Groups[1] + if got := group.Protocols[0]; got != protocol.OpenAI { + t.Errorf("GroupView.Protocols[0] = %q, want %q", got, protocol.OpenAI) + } + if got := group.Models[0]; got != (ModelConfig{ID: "model-real", Alias: "model-alias"}) { + t.Errorf("GroupView.Models[0] = %#v, want original model", got) + } + filters := snapshot.AccessKeysByHash["hash"].Filters + if _, ok := filters.Groups[1]; !ok { + t.Error("AccessKeyView.Filters.Groups lost original group") + } + if _, ok := filters.Groups[2]; ok { + t.Error("AccessKeyView.Filters.Groups contains caller mutation") + } + if _, ok := filters.Protocols[protocol.OpenAI]; !ok { + t.Error("AccessKeyView.Filters.Protocols lost original protocol") + } + if _, ok := filters.Protocols[protocol.Gemini]; ok { + t.Error("AccessKeyView.Filters.Protocols contains caller mutation") + } + if _, ok := filters.Models["model-real"]; !ok { + t.Error("AccessKeyView.Filters.Models lost original model") + } + if _, ok := filters.Models["changed"]; ok { + t.Error("AccessKeyView.Filters.Models contains caller mutation") + } +}