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

feat: add river-backed async task queue

Browse Source
This commit is contained in:
wangbo
2026-05-12 10:11:54 +08:00
parent d69aaed444
commit 7e220b7477
30 changed files with 1342 additions and 200 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/api/internal/httpapi/core_flow_integration_test.go
+234 -13
View File
@@ -18,6 +18,7 @@ import (
"testing"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
"github.com/jackc/pgx/v5/pgxpool"
@@ -37,7 +38,11 @@ func TestCoreLocalFlow(t *testing.T) {
}
defer db.Close()
handler := NewServer(config.Config{
assertRuntimeRecoveryReleasesPendingRateReservations(t, ctx, db)
serverCtx, cancelServer := context.WithCancel(ctx)
defer cancelServer()
handler := NewServerWithContext(serverCtx, config.Config{
AppEnv: "test",
HTTPAddr: ":0",
DatabaseURL: databaseURL,
@@ -169,13 +174,14 @@ func TestCoreLocalFlow(t *testing.T) {
if err := testPool.QueryRow(ctx, `SELECT count(*) FROM gateway_tasks`).Scan(&taskCountBefore); err != nil {
t.Fatalf("count tasks before scoped request: %v", err)
}
defaultImageModel := "openai:gpt-image-1"
doJSON(t, server.URL, http.MethodPost, "/api/v1/images/generations", chatOnlyAPIKeyResponse.Secret, map[string]any{
"model": "gpt-image-1",
"model": defaultImageModel,
"prompt": "scope should block this",
}, http.StatusForbidden, nil)
doJSON(t, server.URL, http.MethodPost, "/api/v1/pricing/estimate", chatOnlyAPIKeyResponse.Secret, map[string]any{
"kind": "images.generations",
"model": "gpt-image-1",
"model": defaultImageModel,
"prompt": "scope should block this estimate",
}, http.StatusForbidden, nil)
var taskCountAfter int
@@ -309,8 +315,9 @@ VALUES ($1, 5, '{"purpose":"core-flow"}'::jsonb)`, inviteCode); err != nil {
Result map[string]any `json:"result"`
} `json:"task"`
}
defaultTextModel := "openai:gpt-4o-mini"
doJSON(t, server.URL, http.MethodPost, "/api/v1/chat/completions", apiKeyResponse.Secret, map[string]any{
"model": "gpt-4o-mini",
"model": defaultTextModel,
"runMode": "simulation",
"simulation": true,
"simulationDurationMs": 5,
@@ -334,7 +341,7 @@ VALUES ($1, 5, '{"purpose":"core-flow"}'::jsonb)`, inviteCode); err != nil {
var compatChat map[string]any
doJSON(t, server.URL, http.MethodPost, "/v1/chat/completions", apiKeyResponse.Secret, map[string]any{
"model": "gpt-4o-mini",
"model": defaultTextModel,
"runMode": "simulation",
"messages": []map[string]any{{"role": "user", "content": "ping"}},
"simulation": true,
@@ -352,7 +359,7 @@ VALUES ($1, 5, '{"purpose":"core-flow"}'::jsonb)`, inviteCode); err != nil {
} `json:"task"`
}
doJSON(t, server.URL, http.MethodPost, "/api/v1/images/generations", apiKeyResponse.Secret, map[string]any{
"model": "gpt-image-1",
"model": defaultImageModel,
"runMode": "simulation",
"prompt": "a tiny gateway console",
"size": "1024x1024",
@@ -372,7 +379,7 @@ VALUES ($1, 5, '{"purpose":"core-flow"}'::jsonb)`, inviteCode); err != nil {
} `json:"task"`
}
doJSON(t, server.URL, http.MethodPost, "/api/v1/images/edits", apiKeyResponse.Secret, map[string]any{
"model": "gpt-image-1",
"model": defaultImageModel,
"runMode": "simulation",
"prompt": "replace background with clean studio light",
"image": "https://example.com/source.png",
@@ -384,6 +391,42 @@ VALUES ($1, 5, '{"purpose":"core-flow"}'::jsonb)`, inviteCode); err != nil {
t.Fatalf("unexpected image edit task: %+v", imageEditResponse.Task)
}
doubaoLiteImageEditModel := "doubao-5.0-lite图像编辑"
var doubaoLitePlatformModel struct {
ID string `json:"id"`
}
doJSON(t, server.URL, http.MethodPost, "/api/admin/platforms/"+platform.ID+"/models", loginResponse.AccessToken, map[string]any{
"canonicalModelKey": "easyai:doubao-5.0-lite图像编辑",
"modelName": doubaoLiteImageEditModel,
"modelAlias": doubaoLiteImageEditModel,
"modelType": []string{"image_edit", "image_generate"},
"displayName": doubaoLiteImageEditModel,
}, http.StatusCreated, &doubaoLitePlatformModel)
var doubaoLiteAsyncEdit struct {
TaskID string `json:"taskId"`
Task struct {
ID string `json:"id"`
Status string `json:"status"`
AsyncMode bool `json:"asyncMode"`
} `json:"task"`
}
doJSONWithHeaders(t, server.URL, http.MethodPost, "/api/v1/images/edits", apiKeyResponse.Secret, map[string]any{
"model": doubaoLiteImageEditModel,
"runMode": "simulation",
"prompt": "turn the attached bright desktop object into a clean product-style render",
"image": "https://example.com/doubao-lite-source.png",
"size": "2048x2048",
"simulation": true,
"simulationDurationMs": 5,
}, map[string]string{"X-Async": "true"}, http.StatusAccepted, &doubaoLiteAsyncEdit)
if doubaoLiteAsyncEdit.TaskID == "" || !doubaoLiteAsyncEdit.Task.AsyncMode {
t.Fatalf("doubao-5.0-lite image edit async task should be accepted: %+v", doubaoLiteAsyncEdit)
}
doubaoLiteAsyncEditDone := waitForTaskStatus(t, server.URL, apiKeyResponse.Secret, doubaoLiteAsyncEdit.TaskID, []string{"succeeded"}, 5*time.Second)
if doubaoLiteAsyncEditDone.Status != "succeeded" {
t.Fatalf("doubao-5.0-lite image edit async task should succeed through river queue, got %+v", doubaoLiteAsyncEditDone)
}
var gptImageModelTypesRaw []byte
if err := testPool.QueryRow(ctx, `
SELECT model_type
@@ -641,6 +684,7 @@ WHERE reference_type = 'gateway_task'
}
rateLimitedModel := "rate-limit-smoke-" + suffixText
rateLimitWindowSeconds := 3
var rateLimitPolicySet struct {
ID string `json:"id"`
}
@@ -652,7 +696,7 @@ WHERE reference_type = 'gateway_task'
"maxAttempts": 1,
},
"rateLimitPolicy": map[string]any{
"rules": []map[string]any{{"metric": "rpm", "limit": 1, "windowSeconds": 60}},
"rules": []map[string]any{{"metric": "rpm", "limit": 1, "windowSeconds": rateLimitWindowSeconds}},
},
}, http.StatusCreated, &rateLimitPolicySet)
var rateLimitPlatformModel map[string]any
@@ -682,6 +726,7 @@ WHERE reference_type = 'gateway_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)
}
waitForRateLimitWindowHead(t, rateLimitWindowSeconds)
var rateLimitTaskOne struct {
Task struct {
Status string `json:"status"`
@@ -713,6 +758,38 @@ WHERE reference_type = 'gateway_task'
if rateLimitTaskTwo.Task.Status != "failed" || rateLimitTaskTwo.Task.ErrorCode != "rate_limit" {
t.Fatalf("runtime policy rate limit should fail second task with rate_limit: %+v", rateLimitTaskTwo.Task)
}
var asyncRateLimitTask struct {
TaskID string `json:"taskId"`
Task struct {
ID string `json:"id"`
Status string `json:"status"`
AsyncMode bool `json:"asyncMode"`
} `json:"task"`
}
doJSONWithHeaders(t, server.URL, http.MethodPost, "/api/v1/chat/completions", apiKeyResponse.Secret, map[string]any{
"model": rateLimitedModel,
"runMode": "simulation",
"simulation": true,
"simulationDurationMs": 5,
"messages": []map[string]any{{"role": "user", "content": "async queued"}},
}, map[string]string{"X-Async": "true"}, http.StatusAccepted, &asyncRateLimitTask)
if asyncRateLimitTask.TaskID == "" || asyncRateLimitTask.Task.ID != asyncRateLimitTask.TaskID || !asyncRateLimitTask.Task.AsyncMode {
t.Fatalf("async task response should expose task id and async mode: %+v", asyncRateLimitTask)
}
asyncRateLimitDetail := waitForTaskStatus(t, server.URL, apiKeyResponse.Secret, asyncRateLimitTask.TaskID, []string{"queued"}, 2*time.Second)
if asyncRateLimitDetail.Status != "queued" {
t.Fatalf("async rate-limited task should return to queued state, got %+v", asyncRateLimitDetail)
}
if len(asyncRateLimitDetail.Attempts) == 0 || asyncRateLimitDetail.Attempts[0].ErrorCode != "rate_limit" {
t.Fatalf("async rate-limited task should record a rate_limit attempt before requeue: %+v", asyncRateLimitDetail)
}
asyncRateLimitCompleted := waitForTaskStatus(t, server.URL, apiKeyResponse.Secret, asyncRateLimitTask.TaskID, []string{"succeeded"}, time.Duration(rateLimitWindowSeconds+3)*time.Second)
if asyncRateLimitCompleted.Status != "succeeded" {
t.Fatalf("async rate-limited task should be pulled from queue after the limit window resets, got %+v", asyncRateLimitCompleted)
}
if len(asyncRateLimitCompleted.Attempts) < 2 || asyncRateLimitCompleted.Attempts[len(asyncRateLimitCompleted.Attempts)-1].Status != "succeeded" {
t.Fatalf("async rate-limited task should create a new successful attempt after requeue: %+v", asyncRateLimitCompleted)
}
videoRouteModel := "video-route-smoke-" + suffixText
var videoRoutePlatformModel map[string]any
@@ -823,16 +900,19 @@ WHERE reference_type = 'gateway_task'
Metrics map[string]any `json:"metrics"`
}
doJSON(t, server.URL, http.MethodGet, "/api/v1/tasks/"+failoverTask.Task.ID, apiKeyResponse.Secret, nil, http.StatusOK, &failoverDetail)
if len(failoverDetail.Attempts) != 2 {
t.Fatalf("failover task history should include two attempts, got %+v", failoverDetail.Attempts)
if len(failoverDetail.Attempts) != 3 {
t.Fatalf("failover task history should include two failed retries plus one successful failover attempt, got %+v", failoverDetail.Attempts)
}
if failoverDetail.Attempts[0].PlatformName != "OpenAI Retryable Failure" || failoverDetail.Attempts[0].Status != "failed" || !failoverDetail.Attempts[0].Retryable || failoverDetail.Attempts[0].ErrorCode == "" {
t.Fatalf("first failover attempt should preserve failed platform and reason: %+v", failoverDetail.Attempts[0])
}
if failoverDetail.Attempts[1].PlatformName != "OpenAI Retry Success" || failoverDetail.Attempts[1].Status != "succeeded" || failoverDetail.Attempts[1].ResponseMS <= 0 {
t.Fatalf("second failover attempt should preserve successful platform: %+v", failoverDetail.Attempts[1])
if failoverDetail.Attempts[1].PlatformName != "OpenAI Retryable Failure" || failoverDetail.Attempts[1].Status != "failed" || !failoverDetail.Attempts[1].Retryable {
t.Fatalf("second failover attempt should preserve the same-client retry failure: %+v", failoverDetail.Attempts[1])
}
if summary, ok := failoverDetail.Metrics["attempts"].([]any); !ok || len(summary) != 2 {
if failoverDetail.Attempts[2].PlatformName != "OpenAI Retry Success" || failoverDetail.Attempts[2].Status != "succeeded" || failoverDetail.Attempts[2].ResponseMS <= 0 {
t.Fatalf("third failover attempt should preserve successful platform: %+v", failoverDetail.Attempts[2])
}
if summary, ok := failoverDetail.Metrics["attempts"].([]any); !ok || len(summary) != 3 {
t.Fatalf("task metrics should keep attempt-chain summary, got %+v", failoverDetail.Metrics)
}
@@ -1045,6 +1125,9 @@ WHERE m.platform_id = $1::uuid
if resp.StatusCode != http.StatusOK || !bytes.Contains(body, []byte("task.completed")) {
t.Fatalf("unexpected events response status=%d body=%s", resp.StatusCode, string(body))
}
if !bytes.Contains(body, []byte("task.running")) {
t.Fatalf("events response should include running transition event body=%s", string(body))
}
if !bytes.Contains(body, []byte("task.progress")) {
t.Fatalf("events response should include progress events body=%s", string(body))
}
@@ -1074,6 +1157,47 @@ WHERE m.platform_id = $1::uuid
if callbackRows == 0 {
t.Fatal("task progress callback outbox should receive events")
}
var restartAsyncTask struct {
TaskID string `json:"taskId"`
Task struct {
ID string `json:"id"`
Status string `json:"status"`
AsyncMode bool `json:"asyncMode"`
} `json:"task"`
}
doJSONWithHeaders(t, server.URL, http.MethodPost, "/api/v1/chat/completions", apiKeyResponse.Secret, map[string]any{
"model": defaultTextModel,
"runMode": "simulation",
"simulation": true,
"simulationDurationMs": 2000,
"messages": []map[string]any{{"role": "user", "content": "river worker restart"}},
}, map[string]string{"X-Async": "true"}, http.StatusAccepted, &restartAsyncTask)
if restartAsyncTask.TaskID == "" || !restartAsyncTask.Task.AsyncMode {
t.Fatalf("restart async task should be accepted as async: %+v", restartAsyncTask)
}
restartRunning := waitForTaskStatus(t, server.URL, apiKeyResponse.Secret, restartAsyncTask.TaskID, []string{"running"}, 3*time.Second)
if restartRunning.Status != "running" {
t.Fatalf("restart async task should be running before worker restart, got %+v", restartRunning)
}
cancelServer()
serverCtx2, cancelServer2 := context.WithCancel(ctx)
defer cancelServer2()
server2 := httptest.NewServer(NewServerWithContext(serverCtx2, config.Config{
AppEnv: "test",
HTTPAddr: ":0",
DatabaseURL: databaseURL,
IdentityMode: "hybrid",
JWTSecret: "test-secret",
TaskProgressCallbackEnabled: true,
TaskProgressCallbackURL: "http://callback.local/task-progress",
CORSAllowedOrigin: "*",
}, db, slog.New(slog.NewTextHandler(io.Discard, nil))))
defer server2.Close()
restartRecovered := waitForTaskStatus(t, server2.URL, apiKeyResponse.Secret, restartAsyncTask.TaskID, []string{"succeeded"}, 8*time.Second)
if restartRecovered.Status != "succeeded" {
t.Fatalf("river worker restart should recover and finish async task, got %+v", restartRecovered)
}
}
func TestOriginAllowedSupportsCommaSeparatedOrigins(t *testing.T) {
@@ -1089,6 +1213,51 @@ func TestOriginAllowedSupportsCommaSeparatedOrigins(t *testing.T) {
}
}
func assertRuntimeRecoveryReleasesPendingRateReservations(t *testing.T, ctx context.Context, db *store.Store) {
t.Helper()
suffix := strconv.FormatInt(time.Now().UnixNano(), 10)
task, err := db.CreateTask(ctx, store.CreateTaskInput{
Kind: "chat.completions",
Model: "recovery-reservation-smoke-" + suffix,
RunMode: "test",
Async: true,
Request: map[string]any{"messages": []any{}},
}, &auth.User{ID: "recovery-user-" + suffix})
if err != nil {
t.Fatalf("create recovery reservation task: %v", err)
}
scopeKey := "recovery-client-" + suffix
if _, err := db.ReserveRateLimits(ctx, task.ID, "", []store.RateLimitReservation{{
ScopeType: "client",
ScopeKey: scopeKey,
Metric: "rpm",
Limit: 10,
Amount: 3,
WindowSeconds: 60,
}}); err != nil {
t.Fatalf("reserve recovery rate limit: %v", err)
}
recovery, err := db.RecoverInterruptedRuntimeState(ctx)
if err != nil {
t.Fatalf("recover interrupted runtime state with pending reservation: %v", err)
}
if recovery.ReleasedRateReservations == 0 {
t.Fatalf("recovery should release pending rate reservation, got %+v", recovery)
}
var reservedValue float64
if err := db.Pool().QueryRow(ctx, `
SELECT COALESCE(MAX(reserved_value), 0)::float8
FROM gateway_rate_limit_counters
WHERE scope_type = 'client'
AND scope_key = $1
AND metric = 'rpm'`, scopeKey).Scan(&reservedValue); err != nil {
t.Fatalf("query recovered rate limit counter: %v", err)
}
if reservedValue != 0 {
t.Fatalf("recovery should release reserved counter value, got %f", reservedValue)
}
}
func applyMigration(t *testing.T, ctx context.Context, databaseURL string) {
t.Helper()
_, filename, _, _ := runtime.Caller(0)
@@ -1114,6 +1283,11 @@ func applyMigration(t *testing.T, ctx context.Context, databaseURL string) {
}
func doJSON(t *testing.T, baseURL string, method string, path string, token string, payload any, expectedStatus int, out any) {
t.Helper()
doJSONWithHeaders(t, baseURL, method, path, token, payload, nil, expectedStatus, out)
}
func doJSONWithHeaders(t *testing.T, baseURL string, method string, path string, token string, payload any, headers map[string]string, expectedStatus int, out any) {
t.Helper()
var body io.Reader
if payload != nil {
@@ -1133,6 +1307,9 @@ func doJSON(t *testing.T, baseURL string, method string, path string, token stri
if token != "" {
req.Header.Set("Authorization", "Bearer "+token)
}
for key, value := range headers {
req.Header.Set(key, value)
}
resp, err := http.DefaultClient.Do(req)
if err != nil {
t.Fatalf("%s %s: %v", method, path, err)
@@ -1149,6 +1326,50 @@ func doJSON(t *testing.T, baseURL string, method string, path string, token stri
}
}
type taskWaitDetail struct {
ID string `json:"id"`
Status string `json:"status"`
Attempts []struct {
Status string `json:"status"`
ErrorCode string `json:"errorCode"`
} `json:"attempts"`
}
func waitForTaskStatus(t *testing.T, baseURL string, token string, taskID string, statuses []string, timeout time.Duration) taskWaitDetail {
t.Helper()
wanted := map[string]bool{}
for _, status := range statuses {
wanted[status] = true
}
var detail taskWaitDetail
deadline := time.Now().Add(timeout)
for time.Now().Before(deadline) {
doJSON(t, baseURL, http.MethodGet, "/api/v1/tasks/"+taskID, token, nil, http.StatusOK, &detail)
if wanted[detail.Status] {
return detail
}
time.Sleep(100 * time.Millisecond)
}
return detail
}
func waitForRateLimitWindowHead(t *testing.T, windowSeconds int) {
t.Helper()
if windowSeconds <= 0 {
return
}
window := time.Duration(windowSeconds) * time.Second
deadline := time.Now().Add(window + time.Second)
for time.Now().Before(deadline) {
elapsed := time.Duration(time.Now().UnixNano() % int64(window))
if elapsed < 300*time.Millisecond {
return
}
time.Sleep(20 * time.Millisecond)
}
t.Fatalf("timed out waiting for %ds rate limit window head", windowSeconds)
}
func stringSliceContains(values []string, target string) bool {
for _, value := range values {
if value == target {
apps/api/internal/httpapi/handlers.go
+27 -7
View File
@@ -542,11 +542,13 @@ func (s *Server) createTask(kind string, compatible bool) http.Handler {
writeError(w, http.StatusForbidden, "api key scope does not allow this capability")
return
}
asyncMode := asyncRequest(r)
task, err := s.store.CreateTask(r.Context(), store.CreateTaskInput{
Kind: kind,
Model: model,
RunMode: runModeFromRequest(body),
Async: asyncMode,
Request: body,
}, user)
if err != nil {
@@ -554,6 +556,14 @@ func (s *Server) createTask(kind string, compatible bool) http.Handler {
writeError(w, http.StatusInternalServerError, "create task failed")
return
}
if asyncMode {
if err := s.runner.EnqueueAsyncTask(r.Context(), task); err != nil {
writeError(w, http.StatusInternalServerError, err.Error(), "enqueue_failed")
return
}
writeTaskAccepted(w, task)
return
}
if compatible {
if boolValue(body, "stream") {
flusher := prepareCompatibleStream(w)
@@ -602,13 +612,23 @@ func (s *Server) createTask(kind string, compatible bool) http.Handler {
s.logger.Warn("task completed with failure", "kind", kind, "taskId", task.ID, "error", runErr)
}
writeJSON(w, http.StatusAccepted, map[string]any{
"task": result.Task,
"next": map[string]string{
"events": fmt.Sprintf("/api/v1/tasks/%s/events", task.ID),
"detail": fmt.Sprintf("/api/v1/tasks/%s", task.ID),
},
})
writeTaskAccepted(w, result.Task)
})
}
func asyncRequest(r *http.Request) bool {
value := strings.TrimSpace(strings.ToLower(r.Header.Get("x-async")))
return value == "1" || value == "true" || value == "yes" || value == "on"
}
func writeTaskAccepted(w http.ResponseWriter, task store.GatewayTask) {
writeJSON(w, http.StatusAccepted, map[string]any{
"taskId": task.ID,
"task": task,
"next": map[string]string{
"events": fmt.Sprintf("/api/v1/tasks/%s/events", task.ID),
"detail": fmt.Sprintf("/api/v1/tasks/%s", task.ID),
},
})
}
apps/api/internal/httpapi/server.go
+9 -1
View File
@@ -1,6 +1,7 @@
package httpapi
import (
"context"
"log/slog"
"net/http"
"strings"
@@ -12,6 +13,7 @@ import (
)
type Server struct {
ctx context.Context
cfg config.Config
store *store.Store
auth *auth.Authenticator
@@ -20,7 +22,12 @@ type Server struct {
}
func NewServer(cfg config.Config, db *store.Store, logger *slog.Logger) http.Handler {
return NewServerWithContext(context.Background(), cfg, db, logger)
}
func NewServerWithContext(ctx context.Context, cfg config.Config, db *store.Store, logger *slog.Logger) http.Handler {
server := &Server{
ctx: ctx,
cfg: cfg,
store: db,
auth: auth.New(cfg.JWTSecret, cfg.ServerMainBaseURL, cfg.ServerMainInternalToken),
@@ -28,6 +35,7 @@ func NewServer(cfg config.Config, db *store.Store, logger *slog.Logger) http.Han
logger: logger,
}
server.auth.LocalAPIKeyVerifier = db.VerifyLocalAPIKey
server.runner.StartAsyncQueueWorker(ctx)
mux := http.NewServeMux()
mux.HandleFunc("GET /healthz", server.health)
@@ -146,7 +154,7 @@ func (s *Server) cors(next http.Handler) http.Handler {
w.Header().Set("Access-Control-Allow-Origin", origin)
w.Header().Set("Vary", "Origin")
w.Header().Set("Access-Control-Allow-Credentials", "true")
w.Header().Set("Access-Control-Allow-Headers", "Authorization, Content-Type, X-Comfy-Api-Key")
w.Header().Set("Access-Control-Allow-Headers", "Authorization, Content-Type, X-Comfy-Api-Key, X-Async")
w.Header().Set("Access-Control-Allow-Methods", "GET, POST, PUT, PATCH, DELETE, OPTIONS")
}
if r.Method == http.MethodOptions {