feat(runtime): 增加健康状态与路由检查
This commit is contained in:
27 files changed
+3152
-117
No files matched your search
@@ -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
|
||||
},
|
||||
|
||||
@@ -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())
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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,
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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",
|
||||
}
|
||||
|
||||
@@ -8,6 +8,7 @@ import (
|
||||
const (
|
||||
statsBucketCount = 5
|
||||
statsBucketWidth = time.Minute
|
||||
StatsWindow = time.Duration(statsBucketCount) * statsBucketWidth
|
||||
)
|
||||
|
||||
type KeyStats struct {
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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},
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
+83
-49
@@ -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
|
||||
}
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in new issue
Block a user