feat: add runner failover rules and cache affinity

This commit is contained in:
2026-06-28 20:50:23 +08:00
parent 41bb3a6525
commit 0e675e259c
31 changed files with 2939 additions and 470 deletions
@@ -275,12 +275,61 @@ VALUES ($1, 5, '{"purpose":"core-flow"}'::jsonb)`, inviteCode); err != nil {
CanonicalModelKey string `json:"canonicalModelKey"`
ProviderModelName string `json:"providerModelName"`
ModelType []string `json:"modelType"`
ModelAlias string `json:"modelAlias"`
} `json:"items"`
}
doJSON(t, server.URL, http.MethodGet, "/api/admin/catalog/base-models", loginResponse.AccessToken, nil, http.StatusOK, &baseModels)
if len(baseModels.Items) < 300 {
t.Fatalf("server-main seed should include the migrated base model catalog: got %d", len(baseModels.Items))
}
requiredMiniMaxModels := map[string]struct {
providerModelName string
modelType string
alias string
}{
"minimax:speech-2.8-hd": {providerModelName: "speech-2.8-hd", modelType: "text_to_speech", alias: "MiniMax-Speech-2.8-HD"},
"minimax:speech-2.8-turbo": {providerModelName: "speech-2.8-turbo", modelType: "text_to_speech", alias: "MiniMax-Speech-2.8-Turbo"},
"minimax:voice-clone": {providerModelName: "voice_clone", modelType: "voice_clone", alias: "MiniMax-Voice-Clone"},
}
seenMiniMaxModels := map[string]bool{}
for _, model := range baseModels.Items {
required, ok := requiredMiniMaxModels[model.CanonicalModelKey]
if !ok {
continue
}
seenMiniMaxModels[model.CanonicalModelKey] = true
if model.ProviderModelName != required.providerModelName || model.ModelAlias != required.alias || !stringSliceContains(model.ModelType, required.modelType) {
t.Fatalf("MiniMax base model is not aligned with server-main: %+v", model)
}
if strings.Contains(model.ModelAlias, " ") {
t.Fatalf("MiniMax Agent-facing alias should not contain spaces: %+v", model)
}
}
for canonicalKey := range requiredMiniMaxModels {
if !seenMiniMaxModels[canonicalKey] {
t.Fatalf("missing MiniMax base model %s in catalog", canonicalKey)
}
}
var modelCatalog struct {
Items []miniMaxModelCatalogItem `json:"items"`
}
doJSON(t, server.URL, http.MethodGet, "/api/v1/model-catalog", apiKeyResponse.Secret, nil, http.StatusOK, &modelCatalog)
for _, required := range requiredMiniMaxModels {
item := findModelCatalogItem(modelCatalog.Items, required.alias)
if item == nil {
t.Fatalf("missing MiniMax model %s in development model catalog", required.alias)
}
if item.ModelName != required.providerModelName || !stringSliceContains(item.ModelType, required.modelType) {
t.Fatalf("MiniMax model catalog item is not aligned with base model: %+v", *item)
}
if strings.Contains(item.Alias, " ") {
t.Fatalf("MiniMax development catalog alias should not contain spaces: %+v", *item)
}
if !modelCatalogItemHasEnabledSource(item.Sources, required.alias, required.providerModelName, required.modelType) {
t.Fatalf("MiniMax development catalog item has no enabled source: %+v", *item)
}
}
baseModelInput := map[string]any{
"providerKey": "openai",
@@ -486,7 +535,7 @@ VALUES ($1, 5, '{"purpose":"core-flow"}'::jsonb)`, inviteCode); err != nil {
speechMarker := "speech-simulation-" + suffixText
var speechResult map[string]any
doJSON(t, server.URL, http.MethodPost, "/api/v1/speech/generations", apiKeyResponse.Secret, map[string]any{
"model": "speech-2.6-turbo",
"model": "MiniMax-Speech-2.8-HD",
"runMode": "simulation",
"text": "hello gateway speech",
"voice_id": "female-shaonv",
@@ -516,6 +565,46 @@ VALUES ($1, 5, '{"purpose":"core-flow"}'::jsonb)`, inviteCode); err != nil {
t.Fatalf("speech simulation task should succeed with text_to_speech billing: %+v", speechTaskDetail)
}
voiceCloneMarker := "voice-clone-simulation-" + suffixText
voiceID := "voice_clone_" + suffixText
var voiceCloneResult map[string]any
doJSON(t, server.URL, http.MethodPost, "/api/v1/voice_clone", apiKeyResponse.Secret, map[string]any{
"model": "MiniMax-Voice-Clone",
"runMode": "simulation",
"voice_id": voiceID,
"audio_url": "/static/simulation/audio.wav",
"prompt_audio_url": "/static/simulation/audio.wav",
"prompt_text": "hello prompt audio",
"text": "hello voice clone preview",
"preview_model": "speech-2.8-hd",
"simulation": true,
"simulationDurationMs": 5,
"integrationTestMarker": voiceCloneMarker,
}, http.StatusOK, &voiceCloneResult)
if voiceCloneResult["status"] != "success" || voiceCloneResult["voice_id"] != voiceID {
t.Fatalf("unexpected voice clone compatible result: %+v", voiceCloneResult)
}
clonedVoice, _ := voiceCloneResult["cloned_voice"].(map[string]any)
if clonedVoice == nil || clonedVoice["voiceId"] != voiceID || clonedVoice["previewModel"] != "speech-2.8-hd" || clonedVoice["platformId"] == "" {
t.Fatalf("voice clone should persist a platform-bound cloned voice: %+v", voiceCloneResult)
}
var voiceCloneTaskDetail struct {
Status string `json:"status"`
ModelType string `json:"modelType"`
Model string `json:"model"`
Attempts []struct {
PlatformID string `json:"platformId"`
} `json:"attempts"`
}
voiceCloneTaskID := waitForTaskIDByRequestField(t, ctx, testPool, "integrationTestMarker", voiceCloneMarker, 2*time.Second)
doJSON(t, server.URL, http.MethodGet, "/api/v1/tasks/"+voiceCloneTaskID, apiKeyResponse.Secret, nil, http.StatusOK, &voiceCloneTaskDetail)
if voiceCloneTaskDetail.Status != "succeeded" || voiceCloneTaskDetail.ModelType != "voice_clone" || voiceCloneTaskDetail.Model != "MiniMax-Voice-Clone" || len(voiceCloneTaskDetail.Attempts) == 0 {
t.Fatalf("voice clone simulation task should succeed through the development interface: %+v", voiceCloneTaskDetail)
}
if voiceCloneTaskDetail.Attempts[0].PlatformID != clonedVoice["platformId"] {
t.Fatalf("voice clone persisted platform should match task attempt: task=%+v voice=%+v", voiceCloneTaskDetail, clonedVoice)
}
doubaoLiteImageEditModel := "doubao-5.0-lite图像编辑"
var doubaoLitePlatformModel struct {
ID string `json:"id"`
@@ -857,7 +946,7 @@ WHERE reference_type = 'gateway_task'
"simulationDurationMs": 5,
"simulationProfile": "non_retryable_failure",
"messages": []map[string]any{{"role": "user", "content": "failed first"}},
}, "rate-limit-failed-first-"+suffixText, http.StatusBadGateway, nil, &rateLimitFailedTask.Task)
}, "rate-limit-failed-first-"+suffixText, http.StatusBadRequest, nil, &rateLimitFailedTask.Task)
if rateLimitFailedTask.Task.Status != "failed" || rateLimitFailedTask.Task.ErrorCode != "bad_request" {
t.Fatalf("failed rate-limited task should fail before consuming rpm: %+v", rateLimitFailedTask.Task)
}
@@ -1460,6 +1549,142 @@ WHERE m.platform_id = $1::uuid
}
}
func TestCacheAffinityRoutingFlow(t *testing.T) {
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
if databaseURL == "" {
t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run the PostgreSQL integration flow")
}
ctx := context.Background()
applyMigration(t, ctx, databaseURL)
db, err := store.Connect(ctx, databaseURL)
if err != nil {
t.Fatalf("connect store: %v", err)
}
defer db.Close()
serverCtx, cancelServer := context.WithCancel(ctx)
defer cancelServer()
server := httptest.NewServer(NewServerWithContext(serverCtx, config.Config{
AppEnv: "test",
HTTPAddr: ":0",
DatabaseURL: databaseURL,
IdentityMode: "hybrid",
JWTSecret: "test-secret",
CORSAllowedOrigin: "*",
}, db, slog.New(slog.NewTextHandler(io.Discard, nil))))
defer server.Close()
suffixText := strconv.FormatInt(time.Now().UnixNano(), 10)
username := "cache_affinity_admin_" + suffixText
password := "password123"
var registerResponse struct {
AccessToken string `json:"accessToken"`
}
doJSON(t, server.URL, http.MethodPost, "/api/v1/auth/register", "", map[string]any{
"username": username,
"email": username + "@example.com",
"password": password,
}, http.StatusCreated, &registerResponse)
testPool, err := pgxpool.New(ctx, databaseURL)
if err != nil {
t.Fatalf("connect test pool: %v", err)
}
defer testPool.Close()
if _, err := testPool.Exec(ctx, `UPDATE gateway_users SET roles = '["admin"]'::jsonb WHERE username = $1`, username); err != nil {
t.Fatalf("promote cache affinity user: %v", err)
}
var loginResponse struct {
AccessToken string `json:"accessToken"`
}
doJSON(t, server.URL, http.MethodPost, "/api/v1/auth/login", "", map[string]any{
"account": username,
"password": password,
}, http.StatusOK, &loginResponse)
var apiKeyResponse struct {
Secret string `json:"secret"`
APIKey struct {
ID string `json:"id"`
} `json:"apiKey"`
}
doJSON(t, server.URL, http.MethodPost, "/api/v1/api-keys", loginResponse.AccessToken, map[string]any{
"name": "cache affinity key",
}, http.StatusCreated, &apiKeyResponse)
model := "cache-affinity-smoke-" + suffixText
cacheKey := "cache-affinity-key-" + suffixText
lowPlatform, lowPlatformModelID := createSimulationTextPlatformModel(t, server.URL, loginResponse.AccessToken, "cache-low-"+suffixText, "Cache Affinity Low Priority", model, 20, map[string]any{
"rules": []map[string]any{
{"metric": "concurrent", "limit": 1, "leaseTtlSeconds": 120},
},
})
for index := 0; index < 3; index++ {
detail := runCacheAffinityChatTask(t, ctx, testPool, server.URL, apiKeyResponse.Secret, model, cacheKey, "prime-"+strconv.Itoa(index)+"-"+suffixText, map[string]any{
"inputTokens": 1000,
"cachedInputTokens": 800,
"outputTokens": 20,
})
if detail.Status != "succeeded" || len(detail.Attempts) != 1 || detail.Attempts[0].PlatformName != "Cache Affinity Low Priority" {
t.Fatalf("priming request should use the only low-priority platform, got %+v", detail)
}
}
highPlatform, _ := createSimulationTextPlatformModel(t, server.URL, loginResponse.AccessToken, "cache-high-"+suffixText, "Cache Affinity High Priority", model, 1, nil)
_ = highPlatform
sameKeyDetail := runCacheAffinityChatTask(t, ctx, testPool, server.URL, apiKeyResponse.Secret, model, cacheKey, "same-key-"+suffixText, map[string]any{
"inputTokens": 1000,
"cachedInputTokens": 700,
"outputTokens": 20,
})
if sameKeyDetail.Status != "succeeded" || len(sameKeyDetail.Attempts) != 1 || sameKeyDetail.Attempts[0].PlatformName != "Cache Affinity Low Priority" {
t.Fatalf("same cache key should prefer cached low-priority platform, got %+v", sameKeyDetail)
}
if !boolFromTestMap(sameKeyDetail.Attempts[0].Metrics, "cacheAffinityApplied") || floatFromTestAny(sameKeyDetail.Attempts[0].Metrics["cacheAffinityBoost"]) <= 0 {
t.Fatalf("same cache key attempt should expose cache affinity metrics: %+v", sameKeyDetail.Attempts[0].Metrics)
}
if intFromTestAny(sameKeyDetail.Usage["inputTokens"]) != 1000 ||
intFromTestAny(sameKeyDetail.Usage["cachedInputTokens"]) != 700 ||
!boolFromTestMap(sameKeyDetail.Usage, "cachedInputTokensKnown") {
t.Fatalf("same cache key task detail should expose simulated usage, got %+v", sameKeyDetail.Usage)
}
newKeyDetail := runCacheAffinityChatTask(t, ctx, testPool, server.URL, apiKeyResponse.Secret, model, "cache-affinity-new-"+suffixText, "new-key-"+suffixText, map[string]any{
"inputTokens": 1000,
"cachedInputTokens": 0,
"outputTokens": 20,
})
if newKeyDetail.Status != "succeeded" || len(newKeyDetail.Attempts) != 1 || newKeyDetail.Attempts[0].PlatformName != "Cache Affinity High Priority" {
t.Fatalf("new cache key should fall back to base priority, got %+v", newKeyDetail)
}
seedQueuedConcurrencyLoad(t, ctx, testPool, lowPlatform.PlatformKey, model, lowPlatformModelID)
fullAvoidedDetail := runCacheAffinityChatTask(t, ctx, testPool, server.URL, apiKeyResponse.Secret, model, cacheKey, "full-avoid-"+suffixText, map[string]any{
"inputTokens": 1000,
"cachedInputTokens": 0,
"outputTokens": 20,
})
if fullAvoidedDetail.Status != "succeeded" || len(fullAvoidedDetail.Attempts) != 1 || fullAvoidedDetail.Attempts[0].PlatformName != "Cache Affinity High Priority" {
t.Fatalf("full cached platform should be avoided before cache affinity boost, got %+v", fullAvoidedDetail)
}
var requestCount int
var inputTokens int
var cachedTokens int
var emaHitRatio float64
if err := testPool.QueryRow(ctx, `
SELECT request_count::int, input_tokens::int, cached_input_tokens::int, ema_hit_ratio::float8
FROM gateway_cache_affinity_stats
WHERE client_id = $1
ORDER BY updated_at DESC
LIMIT 1`, lowPlatform.PlatformKey+":text_generate:"+model).Scan(&requestCount, &inputTokens, &cachedTokens, &emaHitRatio); err != nil {
t.Fatalf("read cache affinity stats: %v", err)
}
if requestCount < 4 || inputTokens < 4000 || cachedTokens < 3100 || emaHitRatio <= 0 {
t.Fatalf("cache affinity stats should aggregate low-priority platform observations, count=%d input=%d cached=%d ema=%v", requestCount, inputTokens, cachedTokens, emaHitRatio)
}
}
func TestOriginAllowedSupportsCommaSeparatedOrigins(t *testing.T) {
allowed := "http://localhost:5178, http://127.0.0.1:5178"
if !originAllowed("http://localhost:5178", allowed) {
@@ -1636,6 +1861,108 @@ type taskWaitDetail struct {
} `json:"attempts"`
}
type cacheAffinityPlatformFixture struct {
ID string `json:"id"`
PlatformKey string `json:"platformKey"`
Name string `json:"name"`
}
type cacheAffinityTaskDetail struct {
ID string `json:"id"`
Status string `json:"status"`
Usage map[string]any `json:"usage"`
Metrics map[string]any `json:"metrics"`
Attempts []struct {
AttemptNo int `json:"attemptNo"`
PlatformName string `json:"platformName"`
Status string `json:"status"`
Metrics map[string]any `json:"metrics"`
} `json:"attempts"`
}
func createSimulationTextPlatformModel(t *testing.T, baseURL string, adminToken string, platformKey string, platformName string, model string, priority int, rateLimitPolicy map[string]any) (cacheAffinityPlatformFixture, string) {
t.Helper()
var platform cacheAffinityPlatformFixture
doJSON(t, baseURL, http.MethodPost, "/api/admin/platforms", adminToken, map[string]any{
"provider": "openai",
"platformKey": platformKey,
"name": platformName,
"baseUrl": "https://api.openai.com/v1",
"authType": "bearer",
"credentials": map[string]any{"mode": "simulation"},
"priority": priority,
}, http.StatusCreated, &platform)
if platform.ID == "" || platform.PlatformKey == "" {
t.Fatalf("cache affinity platform was not created: %+v", platform)
}
payload := map[string]any{
"canonicalModelKey": "openai:gpt-4o-mini",
"modelName": model,
"providerModelName": model,
"modelAlias": model,
"modelType": []string{"text_generate"},
"displayName": platformName,
}
if len(rateLimitPolicy) > 0 {
payload["rateLimitPolicy"] = rateLimitPolicy
}
var platformModel struct {
ID string `json:"id"`
}
doJSON(t, baseURL, http.MethodPost, "/api/admin/platforms/"+platform.ID+"/models", adminToken, payload, http.StatusCreated, &platformModel)
if platformModel.ID == "" {
t.Fatalf("cache affinity platform model was not created: %+v", platformModel)
}
return platform, platformModel.ID
}
func runCacheAffinityChatTask(t *testing.T, ctx context.Context, pool *pgxpool.Pool, baseURL string, token string, model string, cacheAffinityKey string, marker string, simulationUsage map[string]any) cacheAffinityTaskDetail {
t.Helper()
var detail cacheAffinityTaskDetail
doAPIV1ChatCompletionAndLoadTask(t, ctx, pool, baseURL, token, map[string]any{
"model": model,
"runMode": "simulation",
"simulation": true,
"simulationDurationMs": 5,
"cacheAffinityKey": cacheAffinityKey,
"simulationUsage": simulationUsage,
"messages": []map[string]any{{"role": "user", "content": "cache affinity route"}},
}, marker, http.StatusOK, nil, &detail)
return detail
}
type miniMaxModelCatalogItem struct {
Alias string `json:"alias"`
ModelName string `json:"modelName"`
ModelType []string `json:"modelType"`
Sources []miniMaxModelCatalogSource `json:"sources"`
}
type miniMaxModelCatalogSource struct {
ModelAlias string `json:"modelAlias"`
ModelName string `json:"modelName"`
ModelType []string `json:"modelType"`
Enabled bool `json:"enabled"`
}
func findModelCatalogItem(items []miniMaxModelCatalogItem, alias string) *miniMaxModelCatalogItem {
for index := range items {
if items[index].Alias == alias {
return &items[index]
}
}
return nil
}
func modelCatalogItemHasEnabledSource(sources []miniMaxModelCatalogSource, alias string, modelName string, modelType string) bool {
for _, source := range sources {
if source.Enabled && source.ModelAlias == alias && source.ModelName == modelName && stringSliceContains(source.ModelType, modelType) {
return true
}
}
return false
}
func waitForTaskStatus(t *testing.T, baseURL string, token string, taskID string, statuses []string, timeout time.Duration) taskWaitDetail {
t.Helper()
wanted := map[string]bool{}
@@ -1762,7 +2089,7 @@ func assertLoadAvoidanceSimulatedRetryChain(t *testing.T, ctx context.Context, t
"simulation": true,
"simulationDurationMs": 5,
"messages": []map[string]any{{"role": "user", "content": "load avoidance retry chain"}},
}, "load-avoidance-"+suffixText, http.StatusBadGateway, nil, &taskResponse.Task)
}, "load-avoidance-"+suffixText, http.StatusBadRequest, nil, &taskResponse.Task)
if taskResponse.Task.ID == "" || taskResponse.Task.Status != "failed" || taskResponse.Task.ErrorCode != "bad_request" {
t.Fatalf("load avoidance task should only fail after avoided clients are retried, got %+v", taskResponse.Task)
}
@@ -1856,6 +2183,22 @@ func boolFromTestMap(values map[string]any, key string) bool {
return value
}
func intFromTestAny(value any) int {
switch typed := value.(type) {
case int:
return typed
case int64:
return int(typed)
case float64:
return int(typed)
case json.Number:
out, _ := typed.Int64()
return int(out)
default:
return 0
}
}
func floatFromTestAny(value any) float64 {
switch typed := value.(type) {
case float64:
+1
View File
@@ -80,6 +80,7 @@ func (s *compatibleStreamWriter) writeDone(w http.ResponseWriter, output map[str
}
if s.includeUsage && !s.sentUsage {
if usage, ok := output["usage"].(map[string]any); ok && len(usage) > 0 {
usage = clients.NormalizeChatCompletionUsage(usage)
s.writeChatData(w, s.chatChunk([]any{}, usage))
s.sentUsage = true
}