BCAI/easyai-ai-gateway
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: improve model rate limit tracking

Browse Source
This commit is contained in:
wangbo
2026-05-12 03:22:29 +08:00
parent 05632172d0
commit ba850a06c6
25 changed files with 1223 additions and 96 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/api/internal/httpapi/core_flow_integration_test.go
+35 -7
View File
@@ -665,6 +665,23 @@ WHERE reference_type = 'gateway_task'
"runtimePolicySetId": rateLimitPolicySet.ID,
"runtimePolicyOverride": map[string]any{},
}, http.StatusCreated, &rateLimitPlatformModel)
var rateLimitFailedTask struct {
Task struct {
Status string `json:"status"`
ErrorCode string `json:"errorCode"`
} `json:"task"`
}
doJSON(t, server.URL, http.MethodPost, "/api/v1/chat/completions", apiKeyResponse.Secret, map[string]any{
"model": rateLimitedModel,
"runMode": "simulation",
"simulation": true,
"simulationDurationMs": 5,
"simulationProfile": "non_retryable_failure",
"messages": []map[string]any{{"role": "user", "content": "failed first"}},
}, http.StatusAccepted, &rateLimitFailedTask)
if rateLimitFailedTask.Task.Status != "failed" || rateLimitFailedTask.Task.ErrorCode != "bad_request" {
t.Fatalf("failed rate-limited task should fail before consuming rpm: %+v", rateLimitFailedTask.Task)
}
var rateLimitTaskOne struct {
Task struct {
Status string `json:"status"`
@@ -863,7 +880,9 @@ WHERE reference_type = 'gateway_task'
{platformID: degradedPlatform.ID, runtimePolicySetID: degradePolicySet.ID},
{platformID: degradeSuccessPlatform.ID},
} {
var platformModel map[string]any
var platformModel struct {
ID string `json:"id"`
}
doJSON(t, server.URL, http.MethodPost, "/api/admin/platforms/"+item.platformID+"/models", loginResponse.AccessToken, map[string]any{
"canonicalModelKey": "openai:gpt-4o-mini",
"modelName": degradeModel,
@@ -887,14 +906,23 @@ WHERE reference_type = 'gateway_task'
"messages": []map[string]any{{"role": "user", "content": "degrade please"}},
}, http.StatusAccepted, &degradeTask)
if degradeTask.Task.Status != "succeeded" {
t.Fatalf("degrade task should fail over after cooling down failed platform: %+v", degradeTask.Task)
t.Fatalf("degrade task should fail over after cooling down failed model: %+v", degradeTask.Task)
}
var cooledDown bool
if err := testPool.QueryRow(ctx, `SELECT COALESCE(cooldown_until > now(), false) FROM integration_platforms WHERE id = $1::uuid`, degradedPlatform.ID).Scan(&cooledDown); err != nil {
t.Fatalf("read degraded platform cooldown: %v", err)
var cooledDownModel bool
var cooledDownPlatform bool
if err := testPool.QueryRow(ctx, `
SELECT COALESCE(m.cooldown_until > now(), false), COALESCE(p.cooldown_until > now(), false)
FROM platform_models m
JOIN integration_platforms p ON p.id = m.platform_id
WHERE m.platform_id = $1::uuid
AND m.model_name = $2`, degradedPlatform.ID, degradeModel).Scan(&cooledDownModel, &cooledDownPlatform); err != nil {
t.Fatalf("read degraded model cooldown: %v", err)
}
if !cooledDown {
t.Fatal("degrade policy should set platform cooldown_until")
if !cooledDownModel {
t.Fatal("degrade policy should set platform model cooldown_until")
}
if cooledDownPlatform {
t.Fatal("degrade policy should not cool down the entire platform")
}
var autoDisablePolicySet struct {
apps/api/internal/httpapi/handlers.go
+24 -3
View File
@@ -490,7 +490,7 @@ func (s *Server) estimatePricing(w http.ResponseWriter, r *http.Request) {
estimate, err := s.runner.Estimate(r.Context(), kind, model, body, user)
if err != nil {
if errors.Is(err, store.ErrNoModelCandidate) {
writeError(w, http.StatusNotFound, "no enabled platform model matches request")
writeError(w, statusFromRunError(err), err.Error(), store.ModelCandidateErrorCode(err))
return
}
s.logger.Error("estimate pricing failed", "error", err)
@@ -510,6 +510,16 @@ func (s *Server) listRateLimitWindows(w http.ResponseWriter, r *http.Request) {
writeJSON(w, http.StatusOK, map[string]any{"items": items})
}
func (s *Server) listModelRateLimitStatuses(w http.ResponseWriter, r *http.Request) {
items, err := s.store.ListModelRateLimitStatuses(r.Context())
if err != nil {
s.logger.Error("list model rate limit statuses failed", "error", err)
writeError(w, http.StatusInternalServerError, "list model rate limit statuses failed")
return
}
writeJSON(w, http.StatusOK, map[string]any{"items": items})
}
func (s *Server) createTask(kind string, compatible bool) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
user, ok := auth.UserFromContext(r.Context())
@@ -557,7 +567,7 @@ func (s *Server) createTask(kind string, compatible bool) http.Handler {
if runErr != nil {
status := statusFromRunError(runErr)
errorPayload := map[string]any{
"code": clients.ErrorCode(runErr),
"code": runErrorCode(runErr),
"message": runErr.Error(),
"status": status,
}
@@ -581,7 +591,7 @@ func (s *Server) createTask(kind string, compatible bool) http.Handler {
}
result, runErr := s.runner.Execute(r.Context(), task, user)
if runErr != nil {
writeError(w, statusFromRunError(runErr), runErr.Error(), clients.ErrorCode(runErr))
writeError(w, statusFromRunError(runErr), runErr.Error(), runErrorCode(runErr))
return
}
writeJSON(w, http.StatusOK, result.Output)
@@ -634,10 +644,14 @@ func scopeForTaskKind(kind string) string {
func statusFromRunError(err error) int {
switch {
case store.ModelCandidateErrorCode(err) == "platform_cooling_down" || store.ModelCandidateErrorCode(err) == "model_cooling_down":
return http.StatusTooManyRequests
case errors.Is(err, store.ErrNoModelCandidate):
return http.StatusNotFound
case errors.Is(err, store.ErrRateLimited):
return http.StatusTooManyRequests
case clients.ErrorCode(err) == "rate_limit":
return http.StatusTooManyRequests
case errors.Is(err, store.ErrInsufficientWalletBalance):
return http.StatusPaymentRequired
default:
@@ -645,6 +659,13 @@ func statusFromRunError(err error) int {
}
}
func runErrorCode(err error) string {
if errors.Is(err, store.ErrNoModelCandidate) {
return store.ModelCandidateErrorCode(err)
}
return clients.ErrorCode(err)
}
func (s *Server) listTasks(w http.ResponseWriter, r *http.Request) {
user, ok := auth.UserFromContext(r.Context())
if !ok {
apps/api/internal/httpapi/server.go
+1
View File
@@ -107,6 +107,7 @@ func NewServer(cfg config.Config, db *store.Store, logger *slog.Logger) http.Han
mux.Handle("GET /api/v1/models", server.auth.Require(auth.PermissionBasic, http.HandlerFunc(server.listPlayableModels)))
mux.Handle("GET /api/v1/playground/models", server.auth.Require(auth.PermissionBasic, http.HandlerFunc(server.listPlayableModels)))
mux.Handle("GET /api/admin/runtime/rate-limit-windows", server.requireAdmin(auth.PermissionPower, http.HandlerFunc(server.listRateLimitWindows)))
mux.Handle("GET /api/admin/runtime/model-rate-limits", server.requireAdmin(auth.PermissionPower, http.HandlerFunc(server.listModelRateLimitStatuses)))
mux.Handle("POST /api/v1/chat/completions", server.auth.Require(auth.PermissionBasic, server.createTask("chat.completions", false)))
mux.Handle("POST /api/v1/responses", server.auth.Require(auth.PermissionBasic, server.createTask("responses", false)))
mux.Handle("POST /api/v1/images/generations", server.auth.Require(auth.PermissionBasic, server.createTask("images.generations", false)))