feat(api): add load-aware client fallback

2026-05-12 16:59:51 +08:00
parent c5cede2359
commit c2696e7bbe
7 changed files with 558 additions and 5 deletions
@@ -1193,6 +1193,8 @@ WHERE m.platform_id = $1::uuid
 		t.Fatalf("failover events should include retrying event status=%d body=%s", resp.StatusCode, string(body))
 	}

+	assertLoadAvoidanceSimulatedRetryChain(t, ctx, testPool, server.URL, loginResponse.AccessToken, apiKeyResponse.Secret, suffixText)
+
 	var callbackRows int
 	if err := testPool.QueryRow(ctx, `SELECT count(*) FROM gateway_task_callback_outbox WHERE task_id = $1::uuid`, taskResponse.Task.ID).Scan(&callbackRows); err != nil {
 		t.Fatalf("read callback outbox: %v", err)
@@ -1396,6 +1398,212 @@ func waitForTaskStatus(t *testing.T, baseURL string, token string, taskID string
 	return detail
 }

+func assertLoadAvoidanceSimulatedRetryChain(t *testing.T, ctx context.Context, testPool *pgxpool.Pool, baseURL string, adminToken string, runtimeToken string, suffixText string) {
+	t.Helper()
+	model := "load-avoidance-smoke-" + suffixText
+	type scenario struct {
+		keySuffix string
+		name      string
+		failure   string
+		priority  int
+		full      bool
+	}
+	scenarios := []scenario{
+		{keySuffix: "hard-stop", name: "Load Avoidance Hard Stop", failure: "fatal_failure", priority: 20},
+		{keySuffix: "retryable", name: "Load Avoidance Retryable", failure: "retryable_failure", priority: 30},
+		{keySuffix: "full-rate-limit", name: "Load Avoidance Full Rate Limit", failure: "rate_limit", priority: 1, full: true},
+		{keySuffix: "full-overloaded", name: "Load Avoidance Full Overloaded", failure: "overloaded", priority: 2, full: true},
+		{keySuffix: "full-fatal", name: "Load Avoidance Full Fatal", failure: "fatal_failure", priority: 3, full: true},
+	}
+	for _, item := range scenarios {
+		var platform struct {
+			ID          string `json:"id"`
+			PlatformKey string `json:"platformKey"`
+			Name        string `json:"name"`
+		}
+		credentials := map[string]any{"mode": "simulation"}
+		if item.failure != "" {
+			credentials["simulationFailure"] = item.failure
+		}
+		doJSON(t, baseURL, http.MethodPost, "/api/admin/platforms", adminToken, map[string]any{
+			"provider":    "openai",
+			"platformKey": "openai-load-" + item.keySuffix + "-" + suffixText,
+			"name":        item.name,
+			"baseUrl":     "https://api.openai.com/v1",
+			"authType":    "bearer",
+			"credentials": credentials,
+			"priority":    item.priority,
+		}, http.StatusCreated, &platform)
+		if platform.ID == "" || platform.PlatformKey == "" {
+			t.Fatalf("load avoidance platform was not created: %+v", platform)
+		}
+
+		var platformModel struct {
+			ID string `json:"id"`
+		}
+		payload := map[string]any{
+			"canonicalModelKey": "openai:gpt-4o-mini",
+			"modelName":         model,
+			"providerModelName": model,
+			"modelAlias":        model,
+			"modelType":         []string{"text_generate"},
+			"displayName":       item.name,
+			"retryPolicy":       map[string]any{"enabled": true, "maxAttempts": 1},
+		}
+		if item.full {
+			payload["rateLimitPolicy"] = map[string]any{
+				"rules": []map[string]any{
+					{"metric": "rpm", "limit": 10, "windowSeconds": 60},
+					{"metric": "tpm_total", "limit": 200, "windowSeconds": 60},
+					{"metric": "concurrent", "limit": 1, "leaseTtlSeconds": 120},
+				},
+			}
+		}
+		doJSON(t, baseURL, http.MethodPost, "/api/admin/platforms/"+platform.ID+"/models", adminToken, payload, http.StatusCreated, &platformModel)
+		if platformModel.ID == "" {
+			t.Fatalf("load avoidance platform model was not created for %s: %+v", item.name, platformModel)
+		}
+		if item.full {
+			seedQueuedConcurrencyLoad(t, ctx, testPool, platform.PlatformKey, model, platformModel.ID)
+		}
+	}
+
+	var taskResponse struct {
+		Task struct {
+			ID        string `json:"id"`
+			Status    string `json:"status"`
+			ErrorCode string `json:"errorCode"`
+		} `json:"task"`
+	}
+	doJSON(t, baseURL, http.MethodPost, "/api/v1/chat/completions", runtimeToken, map[string]any{
+		"model":                model,
+		"runMode":              "simulation",
+		"simulation":           true,
+		"simulationDurationMs": 5,
+		"messages":             []map[string]any{{"role": "user", "content": "load avoidance retry chain"}},
+	}, http.StatusAccepted, &taskResponse)
+	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)
+	}
+
+	var detail struct {
+		ID        string `json:"id"`
+		Status    string `json:"status"`
+		ErrorCode string `json:"errorCode"`
+		Attempts  []struct {
+			AttemptNo    int            `json:"attemptNo"`
+			PlatformName string         `json:"platformName"`
+			Status       string         `json:"status"`
+			Retryable    bool           `json:"retryable"`
+			ErrorCode    string         `json:"errorCode"`
+			Metrics      map[string]any `json:"metrics"`
+		} `json:"attempts"`
+		Metrics map[string]any `json:"metrics"`
+	}
+	doJSON(t, baseURL, http.MethodGet, "/api/v1/tasks/"+taskResponse.Task.ID, runtimeToken, nil, http.StatusOK, &detail)
+	if detail.Status != "failed" || len(detail.Attempts) != len(scenarios) {
+		t.Fatalf("load avoidance detail should expose every attempted client, got status=%s attempts=%+v", detail.Status, detail.Attempts)
+	}
+	expected := []struct {
+		name      string
+		code      string
+		retryable bool
+		avoided   bool
+	}{
+		{name: "Load Avoidance Hard Stop", code: "bad_request"},
+		{name: "Load Avoidance Retryable", code: "server_error", retryable: true},
+		{name: "Load Avoidance Full Rate Limit", code: "rate_limit", retryable: true, avoided: true},
+		{name: "Load Avoidance Full Overloaded", code: "overloaded", retryable: true, avoided: true},
+		{name: "Load Avoidance Full Fatal", code: "bad_request", avoided: true},
+	}
+	attemptSummary := make([]string, 0, len(detail.Attempts))
+	for index, want := range expected {
+		got := detail.Attempts[index]
+		attemptSummary = append(attemptSummary, got.PlatformName+":"+got.Status+":"+got.ErrorCode)
+		if got.AttemptNo != index+1 || got.PlatformName != want.name || got.Status != "failed" || got.ErrorCode != want.code || got.Retryable != want.retryable {
+			t.Fatalf("unexpected load avoidance attempt %d: got %+v want %+v", index+1, got, want)
+		}
+		if boolFromTestMap(got.Metrics, "loadAvoided") != want.avoided {
+			t.Fatalf("loadAvoided mismatch for %s metrics=%+v", got.PlatformName, got.Metrics)
+		}
+		if !want.avoided && floatFromTestAny(got.Metrics["loadRatio"]) != 0 {
+			t.Fatalf("non-full candidate should not carry load pressure, got %s metrics=%+v", got.PlatformName, got.Metrics)
+		}
+		if want.avoided {
+			if ratio := floatFromTestAny(got.Metrics["loadRatio"]); ratio < 1 {
+				t.Fatalf("avoided candidate should expose full load ratio, got %s ratio=%v metrics=%+v", got.PlatformName, ratio, got.Metrics)
+			}
+			loadMetrics, _ := got.Metrics["loadMetrics"].(map[string]any)
+			concurrent, _ := loadMetrics["concurrent"].(map[string]any)
+			if ratio := floatFromTestAny(concurrent["ratio"]); ratio < 1 {
+				t.Fatalf("avoided candidate should expose concurrent saturation, got %s loadMetrics=%+v", got.PlatformName, loadMetrics)
+			}
+			if queued := floatFromTestAny(loadMetrics["queued"]); queued < 1 {
+				t.Fatalf("avoided candidate should expose queued waiting load, got %s loadMetrics=%+v", got.PlatformName, loadMetrics)
+			}
+		}
+	}
+	if !attemptTraceHasReason(detail.Attempts[0].Metrics, "load_avoidance_fallback") {
+		t.Fatalf("first hard-stop attempt should continue because avoided candidates remain, metrics=%+v", detail.Attempts[0].Metrics)
+	}
+	if summary, ok := detail.Metrics["attempts"].([]any); !ok || len(summary) != len(scenarios) {
+		t.Fatalf("final task metrics should preserve load-avoidance attempt summary, got %+v", detail.Metrics)
+	}
+	t.Logf("load avoidance retry chain: %s", strings.Join(attemptSummary, " -> "))
+}
+
+func seedQueuedConcurrencyLoad(t *testing.T, ctx context.Context, testPool *pgxpool.Pool, platformKey string, model string, platformModelID string) {
+	t.Helper()
+	queueKey := platformKey + ":text_generate:" + model
+	if _, err := testPool.Exec(ctx, `
+INSERT INTO gateway_tasks (
+  kind, run_mode, user_id, model, requested_model, model_type,
+  request, normalized_request, status, queue_key, priority, async_mode,
+  next_run_at, result, billings
+)
+VALUES (
+  'chat.completions', 'simulation', $1, $2, $2, 'text_generate',
+  '{}'::jsonb, '{}'::jsonb, 'queued', $3, 999, true,
+  now() + interval '1 hour', '{}'::jsonb, '[]'::jsonb
+)`, "load-avoidance-seed-"+platformModelID, model, queueKey); err != nil {
+		t.Fatalf("seed queued load for %s: %v", queueKey, err)
+	}
+}
+
+func boolFromTestMap(values map[string]any, key string) bool {
+	value, _ := values[key].(bool)
+	return value
+}
+
+func floatFromTestAny(value any) float64 {
+	switch typed := value.(type) {
+	case float64:
+		return typed
+	case float32:
+		return float64(typed)
+	case int:
+		return float64(typed)
+	case int64:
+		return float64(typed)
+	case json.Number:
+		out, _ := typed.Float64()
+		return out
+	default:
+		return 0
+	}
+}
+
+func attemptTraceHasReason(metrics map[string]any, reason string) bool {
+	trace, _ := metrics["trace"].([]any)
+	for _, raw := range trace {
+		item, _ := raw.(map[string]any)
+		if item["reason"] == reason {
+			return true
+		}
+	}
+	return false
+}
+
 func waitForRateLimitWindowHead(t *testing.T, windowSeconds int) {
 	t.Helper()
 	if windowSeconds <= 0 {