diff --git a/.env.example b/.env.example index d62e114603c0..9388b904bd93 100644 --- a/.env.example +++ b/.env.example @@ -61,6 +61,10 @@ # 流模式无响应超时时间，单位秒，如果出现空补全可以尝试改为更大值 # STREAMING_TIMEOUT=300 +# 异步图片中转允许的同步上游基地址（逗号分隔；路径仅允许空或 /v1） +# ASYNC_YUNWU_ALLOWED_BASE_URLS=https://yunwu.ai +# ASYNC_GRSAI_ALLOWED_BASE_URLS=https://grsaiapi.com,https://grsai.dakka.com.cn + # TLS / HTTP 跳过验证设置 # TLS_INSECURE_SKIP_VERIFY=false diff --git a/common/async_yunwu.go b/common/async_yunwu.go new file mode 100644 index 000000000000..1fd02eb0d8f9 --- /dev/null +++ b/common/async_yunwu.go @@ -0,0 +1,86 @@ +package common + +import ( + "net/url" + "os" + "strings" +) + +type AsyncImageProvider string + +const ( + AsyncImageProviderYunwu AsyncImageProvider = "yunwu" + AsyncImageProviderGRSAI AsyncImageProvider = "grsai" +) + +// IsAllowedYunwuBaseURL restricts the async wrapper to an explicitly allowed +// Yunwu origin and the only two base paths that can produce the whitelisted +// /v1/images/generations endpoint. +func IsAllowedYunwuBaseURL(raw string) bool { + parsed, ok := normalizeAsyncImageBaseURL(raw) + if !ok { + return false + } + allowed := strings.TrimSpace(os.Getenv("ASYNC_YUNWU_ALLOWED_BASE_URLS")) + if allowed == "" { + return parsed == "https://yunwu.ai" + } + for _, item := range strings.Split(allowed, ",") { + candidate, valid := normalizeAsyncImageBaseURL(item) + if valid && parsed == candidate { + return true + } + } + return false +} + +// IsAllowedGRSAIBaseURL limits GRS AI workers to the two documented API +// origins, unless an explicit allowlist is configured for integration tests or +// private relay nodes. The dashboard origin is intentionally not accepted. +func IsAllowedGRSAIBaseURL(raw string) bool { + parsed, ok := normalizeAsyncImageBaseURL(raw) + if !ok { + return false + } + allowed := strings.TrimSpace(os.Getenv("ASYNC_GRSAI_ALLOWED_BASE_URLS")) + if allowed == "" { + return parsed == "https://grsaiapi.com" || parsed == "https://grsai.dakka.com.cn" + } + for _, item := range strings.Split(allowed, ",") { + candidate, valid := normalizeAsyncImageBaseURL(item) + if valid && parsed == candidate { + return true + } + } + return false +} + +func AsyncImageProviderForBaseURL(raw string) (AsyncImageProvider, bool) { + if IsAllowedYunwuBaseURL(raw) { + return AsyncImageProviderYunwu, true + } + if IsAllowedGRSAIBaseURL(raw) { + return AsyncImageProviderGRSAI, true + } + return "", false +} + +func IsAllowedAsyncImageBaseURL(raw string) bool { + _, ok := AsyncImageProviderForBaseURL(raw) + return ok +} + +func normalizeAsyncImageBaseURL(raw string) (string, bool) { + parsed, err := url.Parse(strings.TrimSpace(raw)) + if err != nil || parsed.User != nil || parsed.Hostname() == "" || parsed.RawQuery != "" || parsed.Fragment != "" { + return "", false + } + if parsed.Scheme != "http" && parsed.Scheme != "https" { + return "", false + } + path := strings.TrimRight(parsed.Path, "/") + if path != "" && path != "/v1" { + return "", false + } + return strings.ToLower(parsed.Scheme + "://" + parsed.Host), true +} diff --git a/constant/task.go b/constant/task.go index ecccf4dfe119..c05f2be02943 100644 --- a/constant/task.go +++ b/constant/task.go @@ -5,6 +5,7 @@ type TaskPlatform string const ( TaskPlatformSuno TaskPlatform = "suno" TaskPlatformMidjourney = "mj" + TaskPlatformAsyncImage = "async_image" ) const ( diff --git a/controller/async_job.go b/controller/async_job.go new file mode 100644 index 000000000000..3bb3d262c465 --- /dev/null +++ b/controller/async_job.go @@ -0,0 +1,410 @@ +package controller + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "strconv" + "strings" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/dto" + "github.com/QuantumNous/new-api/logger" + "github.com/QuantumNous/new-api/model" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/relay/helper" + "github.com/QuantumNous/new-api/service" + "github.com/QuantumNous/new-api/storage" + "github.com/QuantumNous/new-api/types" + "github.com/gin-gonic/gin" +) + +type asyncControllerError struct { + status int + code string + message string +} + +func (e *asyncControllerError) Error() string { return e.message } + +var newAsyncArtifactStore = func(ctx context.Context) (storage.ArtifactStore, error) { + return storage.NewS3ArtifactStore(ctx) +} + +func respondAsyncError(c *gin.Context, status int, code, message string) { + c.JSON(status, gin.H{"error": gin.H{ + "message": message, + "type": "async_task_error", + "code": code, + }}) +} + +func asyncStatusText(status model.AsyncExecutionStatus) string { + return strings.ToLower(string(status)) +} + +func asyncSubmitResponse(job *model.AsyncJob) dto.AsyncSubmitResponse { + publicID := job.Task.TaskID + return dto.AsyncSubmitResponse{ + ID: publicID, + Status: asyncStatusText(job.ExecutionStatus), + StatusURL: "/v1/async/tasks/" + publicID, + ResultURL: "/v1/async/tasks/" + publicID + "/result", + } +} + +func SubmitAsyncImageTask(c *gin.Context) { + idempotencyKey := strings.TrimSpace(c.GetHeader("Idempotency-Key")) + if idempotencyKey == "" { + respondAsyncError(c, http.StatusBadRequest, "idempotency_key_required", "Idempotency-Key header is required") + return + } + if len(idempotencyKey) > 191 { + respondAsyncError(c, http.StatusBadRequest, "idempotency_key_too_long", "Idempotency-Key must not exceed 191 bytes") + return + } + if common.BatchUpdateEnabled { + respondAsyncError(c, http.StatusServiceUnavailable, "async_immediate_billing_required", "asynchronous image tasks require BATCH_UPDATE_ENABLED=false") + return + } + + request, err := helper.GetAndValidOpenAIImageRequest(c, 0) + if err != nil { + respondAsyncError(c, http.StatusBadRequest, "invalid_request", err.Error()) + return + } + bodyStorage, err := common.GetBodyStorage(c) + if err != nil { + respondAsyncError(c, http.StatusBadRequest, "read_request_body_failed", err.Error()) + return + } + rawBody, err := bodyStorage.Bytes() + if err != nil { + respondAsyncError(c, http.StatusBadRequest, "read_request_body_failed", err.Error()) + return + } + if err := service.ValidateAsyncImageRequest(request, rawBody); err != nil { + respondAsyncError(c, http.StatusBadRequest, "invalid_request", err.Error()) + return + } + requestHash, err := service.HashAsyncRequest(rawBody) + if err != nil { + respondAsyncError(c, http.StatusBadRequest, "invalid_request", err.Error()) + return + } + encryptedPayload, err := service.EncryptAsyncPayload(rawBody) + if err != nil { + respondAsyncError(c, http.StatusServiceUnavailable, "async_encryption_unavailable", err.Error()) + return + } + + tokenID := c.GetInt("token_id") + var queuedJob *model.AsyncJob + err = model.WithAsyncIdempotencyLock(c.Request.Context(), tokenID, idempotencyKey, func() error { + existing, lookupErr := model.GetAsyncJobByTokenAndKey(c.Request.Context(), tokenID, idempotencyKey) + if lookupErr != nil { + return lookupErr + } + if existing != nil { + if existing.RequestHash != requestHash { + return &asyncControllerError{status: http.StatusConflict, code: "idempotency_key_conflict", message: "Idempotency-Key was already used with a different request"} + } + queuedJob = existing + return nil + } + + relayInfo, infoErr := relaycommon.GenRelayInfo(c, types.RelayFormatOpenAIImage, request, nil) + if infoErr != nil { + return &asyncControllerError{status: http.StatusInternalServerError, code: "relay_context_failed", message: infoErr.Error()} + } + relayInfo.InitChannelMeta(c) + provider, allowedProvider := service.AsyncImageProviderForBaseURL(relayInfo.ChannelBaseUrl) + if relayInfo.ChannelMeta == nil || !allowedProvider { + return &asyncControllerError{status: http.StatusBadRequest, code: "async_image_channel_required", message: "selected channel is not an allowed synchronous image wrapper channel"} + } + if validationErr := service.ValidateAsyncImageProviderRequest(request, provider); validationErr != nil { + return &asyncControllerError{status: http.StatusBadRequest, code: "invalid_provider_request", message: validationErr.Error()} + } + + meta := request.GetTokenCountMeta() + tokens, countErr := service.EstimateRequestToken(c, meta, relayInfo) + if countErr != nil { + return &asyncControllerError{status: http.StatusBadRequest, code: "count_token_failed", message: countErr.Error()} + } + relayInfo.SetEstimatePromptTokens(tokens) + priceData, priceErr := helper.ModelPriceHelper(c, relayInfo, tokens, meta) + if priceErr != nil { + return &asyncControllerError{status: http.StatusBadRequest, code: "model_price_error", message: priceErr.Error()} + } + relayInfo.ForcePreConsume = true + if !priceData.FreeModel { + if apiErr := service.PreConsumeBilling(c, priceData.QuotaToPreConsume, relayInfo); apiErr != nil { + return &asyncControllerError{status: apiErr.StatusCode, code: string(apiErr.GetErrorCode()), message: apiErr.Error()} + } + } + + refundOnFailure := true + defer func() { + if refundOnFailure && relayInfo.Billing != nil { + relayInfo.Billing.Refund(c) + } + }() + + task := model.InitTask(constant.TaskPlatformAsyncImage, relayInfo) + task.Status = model.TaskStatusQueued + task.Progress = "0%" + task.Action = model.AsyncEndpointImageGeneration + task.Quota = relayInfo.FinalPreConsumedQuota + task.PrivateData.BillingSource = relayInfo.BillingSource + task.PrivateData.SubscriptionId = relayInfo.SubscriptionId + task.PrivateData.TokenId = relayInfo.TokenId + task.PrivateData.NodeName = common.NodeName + task.PrivateData.BillingContext = &model.TaskBillingContext{ + ModelPrice: relayInfo.PriceData.ModelPrice, + GroupRatio: relayInfo.PriceData.GroupRatioInfo.GroupRatio, + ModelRatio: relayInfo.PriceData.ModelRatio, + OtherRatios: relayInfo.PriceData.OtherRatios(), + OriginModelName: relayInfo.OriginModelName, + PerCallBilling: relayInfo.PriceData.UsePrice, + } + task.SetData(map[string]any{"model": request.Model, "endpoint_type": model.AsyncEndpointImageGeneration}) + + job := &model.AsyncJob{ + TokenID: relayInfo.TokenId, + ChannelID: relayInfo.ChannelId, + EndpointType: model.AsyncEndpointImageGeneration, + RequestPayload: encryptedPayload, + RequestHash: requestHash, + IdempotencyKey: idempotencyKey, + ExecutionStatus: model.AsyncStatusQueued, + BillingStatus: model.AsyncBillingReserved, + BillingRequestID: relayInfo.RequestId, + } + if createErr := model.CreateAsyncTask(task, job); createErr != nil { + return createErr + } + job.Task = *task + queuedJob = job + refundOnFailure = false + return nil + }) + if err != nil { + var controllerErr *asyncControllerError + if errors.As(err, &controllerErr) { + respondAsyncError(c, controllerErr.status, controllerErr.code, controllerErr.message) + return + } + logger.LogError(c, "create async image task failed: "+err.Error()) + respondAsyncError(c, http.StatusInternalServerError, "create_task_failed", "failed to persist async image task") + return + } + c.JSON(http.StatusAccepted, asyncSubmitResponse(queuedJob)) +} + +func GetAsyncTask(c *gin.Context) { + job := ownedAsyncJob(c) + if job == nil { + return + } + c.JSON(http.StatusOK, asyncTaskStatusResponse(job)) +} + +func asyncTaskStatusResponse(job *model.AsyncJob) dto.AsyncTaskStatusResponse { + response := dto.AsyncTaskStatusResponse{ + ID: job.Task.TaskID, + Status: asyncStatusText(job.ExecutionStatus), + Progress: parseTaskProgress(job.Task.Progress), + CreatedAt: job.Task.SubmitTime, + } + if job.Task.StartTime > 0 { + started := job.Task.StartTime + response.StartedAt = &started + } + if job.Task.FinishTime > 0 { + finished := job.Task.FinishTime + response.FinishedAt = &finished + } + if job.ErrorCode != "" || job.Task.FailReason != "" { + response.Error = &dto.AsyncTaskError{Phase: job.ErrorPhase, Code: job.ErrorCode, Message: job.Task.FailReason} + } + return response +} + +func parseTaskProgress(raw string) int { + value, _ := strconv.Atoi(strings.TrimSuffix(raw, "%")) + if value < 0 { + return 0 + } + if value > 100 { + return 100 + } + return value +} + +func ownedAsyncJob(c *gin.Context) *model.AsyncJob { + publicID := c.Param("task_id") + job, err := model.GetAsyncJobByPublicTaskID(c.Request.Context(), publicID, c.GetInt("token_id")) + if err != nil { + logger.LogError(c, "query async task failed: "+err.Error()) + respondAsyncError(c, http.StatusInternalServerError, "query_task_failed", "failed to query async task") + return nil + } + if job == nil { + respondAsyncError(c, http.StatusNotFound, "task_not_found", "async task was not found") + return nil + } + return job +} + +func GetAsyncTaskResult(c *gin.Context) { + job := ownedAsyncJob(c) + if job == nil { + return + } + switch job.ExecutionStatus { + case model.AsyncStatusFailure: + respondAsyncError(c, http.StatusUnprocessableEntity, defaultString(job.ErrorCode, "task_failed"), defaultString(job.Task.FailReason, "async task failed")) + return + case model.AsyncStatusUncertain: + respondAsyncError(c, http.StatusConflict, defaultString(job.ErrorCode, "task_uncertain"), "the upstream request may have executed; automatic retry is unsafe") + return + case model.AsyncStatusCancelled: + respondAsyncError(c, http.StatusConflict, "task_cancelled", "async task was cancelled before execution") + return + case model.AsyncStatusSuccess: + // continue below + default: + respondAsyncError(c, http.StatusConflict, "task_not_ready", "async task is "+asyncStatusText(job.ExecutionStatus)) + return + } + + artifacts, err := model.ListArtifactsByTaskID(c.Request.Context(), job.TaskID) + if err != nil { + respondAsyncError(c, http.StatusInternalServerError, "artifact_query_failed", "failed to query task artifacts") + return + } + for _, artifact := range artifacts { + if artifact.ExpiresAt <= time.Now().Unix() { + respondAsyncError(c, http.StatusGone, "result_expired", "async task result has expired") + return + } + } + if len(artifacts) == 0 && len(job.ResultPayload) == 0 { + respondAsyncError(c, http.StatusGone, "result_expired", "async task result has expired") + return + } + store, err := newAsyncArtifactStore(c.Request.Context()) + if err != nil { + respondAsyncError(c, http.StatusServiceUnavailable, "artifact_store_unavailable", "artifact store is unavailable") + return + } + ttl := time.Duration(common.GetEnvOrDefault("ASYNC_SIGNED_URL_TTL_SECONDS", 900)) * time.Second + artifactResponses := make([]dto.AsyncArtifactResponse, 0, len(artifacts)) + signedURLs := make([]string, 0, len(artifacts)) + for _, artifact := range artifacts { + signedURL, signErr := store.SignedURL(c.Request.Context(), artifact.ObjectKey, ttl) + if signErr != nil { + respondAsyncError(c, http.StatusServiceUnavailable, "artifact_sign_failed", "failed to create artifact download URL") + return + } + signedURLs = append(signedURLs, signedURL) + artifactResponses = append(artifactResponses, dto.AsyncArtifactResponse{ + ContentType: artifact.ContentType, + SizeBytes: artifact.SizeBytes, + SHA256: artifact.SHA256, + ExpiresAt: artifact.ExpiresAt, + URL: signedURL, + }) + } + upstreamResponse := json.RawMessage(job.ResultPayload) + normalized := normalizedAsyncImageResponse(upstreamResponse, signedURLs) + if c.Query("include_upstream") == "false" { + upstreamResponse = nil + } + c.JSON(http.StatusOK, dto.AsyncTaskResultResponse{ + ID: job.Task.TaskID, + Status: asyncStatusText(job.ExecutionStatus), + Response: normalized, + UpstreamResponse: upstreamResponse, + Artifacts: artifactResponses, + }) +} + +func normalizedAsyncImageResponse(raw json.RawMessage, signedURLs []string) json.RawMessage { + var response map[string]any + if err := common.Unmarshal(raw, &response); err != nil { + return raw + } + data, ok := response["data"].([]any) + if ok { + for index, value := range data { + if index >= len(signedURLs) { + break + } + if item, ok := value.(map[string]any); ok { + item["url"] = signedURLs[index] + delete(item, "b64_json") + } + } + } else if len(signedURLs) > 0 { + data = make([]any, 0, len(signedURLs)) + for _, signedURL := range signedURLs { + data = append(data, map[string]any{"url": signedURL}) + } + response = map[string]any{"data": data} + } + encoded, err := common.Marshal(response) + if err != nil { + return raw + } + return encoded +} + +func CancelAsyncTask(c *gin.Context) { + job := ownedAsyncJob(c) + if job == nil { + return + } + if job.ExecutionStatus == model.AsyncStatusRunning { + respondAsyncError(c, http.StatusConflict, "upstream_cancel_unsupported", "the upstream has no cancellation API; the running request was not interrupted") + return + } + if job.ExecutionStatus != model.AsyncStatusQueued { + c.JSON(http.StatusOK, asyncTaskStatusResponse(job)) + return + } + cancelled, changed, err := model.CancelQueuedAsyncJob(c.Request.Context(), job.Task.TaskID, c.GetInt("token_id")) + if err != nil { + respondAsyncError(c, http.StatusInternalServerError, "cancel_task_failed", "failed to cancel async task") + return + } + if !changed { + latest := ownedAsyncJob(c) + if latest != nil { + c.JSON(http.StatusOK, asyncTaskStatusResponse(latest)) + } + return + } + if _, err := model.RefundAsyncJobBilling(c.Request.Context(), cancelled.ID); err != nil { + logger.LogError(c, fmt.Sprintf("refund cancelled async task %s failed: %v", cancelled.Task.TaskID, err)) + respondAsyncError(c, http.StatusInternalServerError, "refund_failed", "task was cancelled but quota refund is pending reconciliation") + return + } + refreshed, err := model.GetAsyncJobByPublicTaskID(c.Request.Context(), cancelled.Task.TaskID, c.GetInt("token_id")) + if err != nil || refreshed == nil { + respondAsyncError(c, http.StatusInternalServerError, "query_task_failed", "task was cancelled but its final state could not be loaded") + return + } + c.JSON(http.StatusOK, asyncTaskStatusResponse(refreshed)) +} + +func defaultString(value, fallback string) string { + if strings.TrimSpace(value) == "" { + return fallback + } + return value +} diff --git a/controller/async_job_test.go b/controller/async_job_test.go new file mode 100644 index 000000000000..d9bde3cae164 --- /dev/null +++ b/controller/async_job_test.go @@ -0,0 +1,186 @@ +package controller + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/storage" + "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +type controllerArtifactStore struct{} + +func (controllerArtifactStore) Put(context.Context, string, io.Reader, string) error { return nil } +func (controllerArtifactStore) Delete(context.Context, string) error { return nil } +func (controllerArtifactStore) SignedURL(_ context.Context, key string, _ time.Duration) (string, error) { + return "https://objects.example/" + key + "?signed=test", nil +} + +func setupAsyncControllerTestDB(t *testing.T) *gorm.DB { + t.Helper() + gin.SetMode(gin.TestMode) + common.SetDatabaseTypes(common.DatabaseTypeSQLite, common.DatabaseTypeSQLite) + common.RedisEnabled = false + common.BatchUpdateEnabled = false + dsn := fmt.Sprintf("file:%s?mode=memory&cache=shared", strings.ReplaceAll(t.Name(), "/", "_")) + db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{}) + require.NoError(t, err) + model.DB = db + model.LOG_DB = db + require.NoError(t, db.AutoMigrate(&model.Task{}, &model.AsyncJob{}, &model.Artifact{}, &model.TaskEvent{}, &model.User{}, &model.Token{}, &model.Channel{})) + t.Cleanup(func() { + if sqlDB, err := db.DB(); err == nil { + _ = sqlDB.Close() + } + }) + return db +} + +func createAsyncControllerFixture(t *testing.T, status model.AsyncExecutionStatus) (*model.Task, *model.AsyncJob) { + t.Helper() + var taskStatus model.TaskStatus = model.TaskStatusQueued + progress := "0%" + finish := int64(0) + if status == model.AsyncStatusRunning { + taskStatus = model.TaskStatusInProgress + progress = "50%" + } else if status == model.AsyncStatusSuccess { + taskStatus = model.TaskStatusSuccess + progress = "100%" + finish = time.Now().Unix() + } else if status == model.AsyncStatusUncertain { + taskStatus = model.TaskStatusUncertain + finish = time.Now().Unix() + } + task := &model.Task{ + TaskID: "task_controller_fixture", + Platform: constant.TaskPlatformAsyncImage, + UserId: 1, + ChannelId: 2, + Status: taskStatus, + Progress: progress, + SubmitTime: time.Now().Add(-time.Minute).Unix(), + FinishTime: finish, + Data: json.RawMessage(`{}`), + } + job := &model.AsyncJob{TokenID: 11, ChannelID: 2, EndpointType: model.AsyncEndpointImageGeneration, RequestPayload: []byte("encrypted"), RequestHash: strings.Repeat("a", 64), IdempotencyKey: "controller-fixture", ExecutionStatus: status, BillingStatus: model.AsyncBillingReserved, ResultPayload: model.JSONValue(`{"created":1,"data":[{"url":"https://temporary.example/image.png"}]}`)} + require.NoError(t, model.CreateAsyncTask(task, job)) + job.Task = *task + return task, job +} + +func asyncControllerContext(method, path, taskID string, tokenID int) (*gin.Context, *httptest.ResponseRecorder) { + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(method, path, nil) + ctx.Params = gin.Params{{Key: "task_id", Value: taskID}} + ctx.Set("token_id", tokenID) + return ctx, recorder +} + +func TestAsyncTaskOwnershipIsEnforced(t *testing.T) { + setupAsyncControllerTestDB(t) + task, _ := createAsyncControllerFixture(t, model.AsyncStatusQueued) + ctx, recorder := asyncControllerContext(http.MethodGet, "/v1/async/tasks/"+task.TaskID, task.TaskID, 999) + GetAsyncTask(ctx) + assert.Equal(t, http.StatusNotFound, recorder.Code) + assert.Contains(t, recorder.Body.String(), "task_not_found") +} + +func TestAsyncResultUsesSignedArchivedURLs(t *testing.T) { + setupAsyncControllerTestDB(t) + task, _ := createAsyncControllerFixture(t, model.AsyncStatusSuccess) + require.NoError(t, model.DB.Create(&model.Artifact{TaskID: task.ID, ObjectKey: "async/task/image.png", ContentType: "image/png", SizeBytes: 12, SHA256: strings.Repeat("b", 64), SourceURLHash: strings.Repeat("c", 64), ExpiresAt: time.Now().Add(24 * time.Hour).Unix()}).Error) + + originalFactory := newAsyncArtifactStore + newAsyncArtifactStore = func(context.Context) (storage.ArtifactStore, error) { return controllerArtifactStore{}, nil } + t.Cleanup(func() { newAsyncArtifactStore = originalFactory }) + + ctx, recorder := asyncControllerContext(http.MethodGet, "/v1/async/tasks/"+task.TaskID+"/result", task.TaskID, 11) + GetAsyncTaskResult(ctx) + require.Equal(t, http.StatusOK, recorder.Code) + assert.Contains(t, recorder.Body.String(), "https://objects.example/async/task/image.png") + assert.Contains(t, recorder.Body.String(), "https://temporary.example/image.png") +} + +func TestAsyncResultCanOmitUpstreamPayload(t *testing.T) { + setupAsyncControllerTestDB(t) + task, _ := createAsyncControllerFixture(t, model.AsyncStatusSuccess) + require.NoError(t, model.DB.Create(&model.Artifact{TaskID: task.ID, ObjectKey: "async/task/image.png", ContentType: "image/png", SizeBytes: 12, SHA256: strings.Repeat("b", 64), SourceURLHash: strings.Repeat("c", 64), ExpiresAt: time.Now().Add(24 * time.Hour).Unix()}).Error) + + originalFactory := newAsyncArtifactStore + newAsyncArtifactStore = func(context.Context) (storage.ArtifactStore, error) { return controllerArtifactStore{}, nil } + t.Cleanup(func() { newAsyncArtifactStore = originalFactory }) + + ctx, recorder := asyncControllerContext(http.MethodGet, "/v1/async/tasks/"+task.TaskID+"/result?include_upstream=false", task.TaskID, 11) + GetAsyncTaskResult(ctx) + require.Equal(t, http.StatusOK, recorder.Code) + assert.Contains(t, recorder.Body.String(), "https://objects.example/async/task/image.png") + assert.NotContains(t, recorder.Body.String(), "https://temporary.example/image.png") + assert.NotContains(t, recorder.Body.String(), "upstream_response") +} + +func TestAsyncResultReturnsGoneAfterArtifactRetentionExpires(t *testing.T) { + setupAsyncControllerTestDB(t) + task, _ := createAsyncControllerFixture(t, model.AsyncStatusSuccess) + require.NoError(t, model.DB.Create(&model.Artifact{TaskID: task.ID, ObjectKey: "async/task/expired.png", ContentType: "image/png", SizeBytes: 12, SHA256: strings.Repeat("d", 64), SourceURLHash: strings.Repeat("e", 64), ExpiresAt: time.Now().Add(-time.Second).Unix()}).Error) + + ctx, recorder := asyncControllerContext(http.MethodGet, "/v1/async/tasks/"+task.TaskID+"/result", task.TaskID, 11) + GetAsyncTaskResult(ctx) + require.Equal(t, http.StatusGone, recorder.Code) + assert.Contains(t, recorder.Body.String(), "result_expired") +} + +func TestRunningAsyncTaskCannotBeFalselyCancelled(t *testing.T) { + setupAsyncControllerTestDB(t) + task, _ := createAsyncControllerFixture(t, model.AsyncStatusRunning) + ctx, recorder := asyncControllerContext(http.MethodPost, "/v1/async/tasks/"+task.TaskID+"/cancel", task.TaskID, 11) + CancelAsyncTask(ctx) + assert.Equal(t, http.StatusConflict, recorder.Code) + assert.Contains(t, recorder.Body.String(), "upstream_cancel_unsupported") +} + +func TestAdminRetryRequiresExplicitRiskConfirmation(t *testing.T) { + setupAsyncControllerTestDB(t) + task, job := createAsyncControllerFixture(t, model.AsyncStatusUncertain) + + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodPost, "/api/task/async/"+task.TaskID+"/retry", strings.NewReader(`{"confirm_risk":false}`)) + ctx.Request.Header.Set("Content-Type", "application/json") + ctx.Params = gin.Params{{Key: "task_id", Value: task.TaskID}} + ctx.Set("id", 99) + RetryAdminAsyncTask(ctx) + assert.Equal(t, http.StatusConflict, recorder.Code) + + var unchanged model.AsyncJob + require.NoError(t, model.DB.First(&unchanged, job.ID).Error) + assert.Equal(t, model.AsyncStatusUncertain, unchanged.ExecutionStatus) + + recorder = httptest.NewRecorder() + ctx, _ = gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodPost, "/api/task/async/"+task.TaskID+"/retry", strings.NewReader(`{"confirm_risk":true}`)) + ctx.Request.Header.Set("Content-Type", "application/json") + ctx.Params = gin.Params{{Key: "task_id", Value: task.TaskID}} + ctx.Set("id", 99) + RetryAdminAsyncTask(ctx) + assert.Equal(t, http.StatusOK, recorder.Code) + + var retried model.AsyncJob + require.NoError(t, model.DB.First(&retried, job.ID).Error) + assert.Equal(t, model.AsyncStatusQueued, retried.ExecutionStatus) +} diff --git a/controller/async_task_management.go b/controller/async_task_management.go new file mode 100644 index 000000000000..760b07b867eb --- /dev/null +++ b/controller/async_task_management.go @@ -0,0 +1,166 @@ +package controller + +import ( + "errors" + "net/http" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/dto" + "github.com/QuantumNous/new-api/model" + "github.com/gin-gonic/gin" +) + +type asyncRetryRequest struct { + ConfirmRisk bool `json:"confirm_risk"` +} + +type asyncTaskManagementDetail struct { + Task *dto.TaskDto `json:"task"` + UpstreamResponse any `json:"upstream_response,omitempty"` + Artifacts []dto.AsyncArtifactResponse `json:"artifacts"` + Events []model.TaskEvent `json:"events"` +} + +func GetAdminAsyncTaskDetail(c *gin.Context) { getManagedAsyncTaskDetail(c, true) } +func GetUserAsyncTaskDetail(c *gin.Context) { getManagedAsyncTaskDetail(c, false) } +func CancelAdminAsyncTask(c *gin.Context) { cancelManagedAsyncTask(c, true) } +func CancelUserAsyncTask(c *gin.Context) { cancelManagedAsyncTask(c, false) } + +func getManagedAsyncJob(c *gin.Context, administrator bool) *model.AsyncJob { + job, err := model.GetAsyncJobForSession(c.Request.Context(), c.Param("task_id"), c.GetInt("id"), administrator) + if err != nil { + common.ApiError(c, err) + return nil + } + if job == nil { + c.JSON(http.StatusNotFound, gin.H{"success": false, "message": "async task was not found"}) + return nil + } + return job +} + +func getManagedAsyncTaskDetail(c *gin.Context, administrator bool) { + job := getManagedAsyncJob(c, administrator) + if job == nil { + return + } + items, err := tasksToDto(c.Request.Context(), []*model.Task{&job.Task}, administrator) + if err != nil { + common.ApiError(c, err) + return + } + artifacts, err := managedAsyncArtifacts(c, job.TaskID) + if err != nil { + common.ApiError(c, err) + return + } + events, err := model.ListTaskEvents(c.Request.Context(), job.TaskID) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, asyncTaskManagementDetail{ + Task: items[0], + UpstreamResponse: job.ResultPayload, + Artifacts: artifacts, + Events: events, + }) +} + +func managedAsyncArtifacts(c *gin.Context, taskID int64) ([]dto.AsyncArtifactResponse, error) { + artifacts, err := model.ListArtifactsByTaskID(c.Request.Context(), taskID) + if err != nil || len(artifacts) == 0 { + return []dto.AsyncArtifactResponse{}, err + } + store, err := newAsyncArtifactStore(c.Request.Context()) + if err != nil { + return nil, err + } + ttl := time.Duration(common.GetEnvOrDefault("ASYNC_SIGNED_URL_TTL_SECONDS", 900)) * time.Second + result := make([]dto.AsyncArtifactResponse, 0, len(artifacts)) + for _, artifact := range artifacts { + url, signErr := store.SignedURL(c.Request.Context(), artifact.ObjectKey, ttl) + if signErr != nil { + return nil, signErr + } + result = append(result, dto.AsyncArtifactResponse{ + ContentType: artifact.ContentType, + SizeBytes: artifact.SizeBytes, + SHA256: artifact.SHA256, + ExpiresAt: artifact.ExpiresAt, + URL: url, + }) + } + return result, nil +} + +func cancelManagedAsyncTask(c *gin.Context, administrator bool) { + job := getManagedAsyncJob(c, administrator) + if job == nil { + return + } + if job.ExecutionStatus == model.AsyncStatusRunning { + c.JSON(http.StatusConflict, gin.H{"success": false, "message": "the upstream has no cancellation API; the running request was not interrupted"}) + return + } + if job.ExecutionStatus != model.AsyncStatusQueued { + common.ApiSuccess(c, asyncTaskStatusResponse(job)) + return + } + actorType := "user" + if administrator { + actorType = "admin" + } + cancelled, changed, err := model.CancelQueuedAsyncJobByID(c.Request.Context(), job.ID, actorType, c.GetInt("id")) + if err != nil { + common.ApiError(c, err) + return + } + if changed { + if _, err := model.RefundAsyncJobBilling(c.Request.Context(), cancelled.ID); err != nil { + common.ApiError(c, err) + return + } + } + latest, err := model.GetAsyncJobForSession(c.Request.Context(), job.Task.TaskID, c.GetInt("id"), administrator) + if err != nil { + common.ApiError(c, err) + return + } + if latest == nil { + c.JSON(http.StatusNotFound, gin.H{"success": false, "message": "async task was not found after cancellation"}) + return + } + common.ApiSuccess(c, asyncTaskStatusResponse(latest)) +} + +func RetryAdminAsyncTask(c *gin.Context) { + job := getManagedAsyncJob(c, true) + if job == nil { + return + } + var request asyncRetryRequest + if err := c.ShouldBindJSON(&request); err != nil || !request.ConfirmRisk { + c.JSON(http.StatusConflict, gin.H{"success": false, "message": "manual retry requires explicit confirmation of duplicate generation and billing risk"}) + return + } + if job.ExecutionStatus != model.AsyncStatusFailure && job.ExecutionStatus != model.AsyncStatusUncertain { + c.JSON(http.StatusConflict, gin.H{"success": false, "message": "only FAILURE or UNCERTAIN tasks can be retried"}) + return + } + retried, changed, err := model.RetryAsyncJob(c.Request.Context(), job.ID, c.GetInt("id")) + if err != nil { + if errors.Is(err, model.ErrAsyncRetryQuotaInsufficient) { + c.JSON(http.StatusForbidden, gin.H{"success": false, "message": err.Error()}) + return + } + common.ApiError(c, err) + return + } + if !changed { + c.JSON(http.StatusConflict, gin.H{"success": false, "message": "task state changed before it could be retried"}) + return + } + common.ApiSuccess(c, asyncTaskStatusResponse(retried)) +} diff --git a/controller/log.go b/controller/log.go index 470c759fc1a1..199802cc03e7 100644 --- a/controller/log.go +++ b/controller/log.go @@ -122,6 +122,43 @@ func GetLogsStat(c *gin.Context) { return } +func GetUpstreamCostStat(c *gin.Context) { + logType, _ := strconv.Atoi(c.Query("type")) + startTimestamp, _ := strconv.ParseInt(c.Query("start_timestamp"), 10, 64) + endTimestamp, _ := strconv.ParseInt(c.Query("end_timestamp"), 10, 64) + modelName := c.Query("model_name") + username := c.Query("username") + tokenName := c.Query("token_name") + group := c.Query("group") + requestId := c.Query("request_id") + upstreamRequestId := c.Query("upstream_request_id") + channel, _ := strconv.Atoi(c.Query("channel")) + var stat model.UpstreamCostStat + if logType == model.LogTypeUnknown || logType == model.LogTypeConsume { + var err error + stat, err = model.GetUpstreamCostStat( + startTimestamp, + endTimestamp, + modelName, + username, + tokenName, + group, + requestId, + upstreamRequestId, + channel, + ) + if err != nil { + common.ApiError(c, err) + return + } + } + c.JSON(http.StatusOK, gin.H{ + "success": true, + "message": "", + "data": stat, + }) +} + func GetLogsSelfStat(c *gin.Context) { username := c.GetString("username") logType, _ := strconv.Atoi(c.Query("type")) diff --git a/controller/pricing.go b/controller/pricing.go index 8252327244c4..feffc188f66a 100644 --- a/controller/pricing.go +++ b/controller/pricing.go @@ -72,7 +72,7 @@ func GetPricing(c *gin.Context) { "usable_group": usableGroup, "supported_endpoint": model.GetSupportedEndpointMap(), "auto_groups": service.GetUserAutoGroup(group), - "pricing_version": "a42d372ccf0b5dd13ecf71203521f9d2", + "pricing_version": "c93f4990684023eef6ec35670795311442418348207af6cd284468468f812f69", }) } diff --git a/controller/task.go b/controller/task.go index a80f1a687aab..f080716ee9b2 100644 --- a/controller/task.go +++ b/controller/task.go @@ -1,6 +1,7 @@ package controller import ( + "context" "strconv" "github.com/QuantumNous/new-api/common" @@ -32,7 +33,12 @@ func GetAllTask(c *gin.Context) { items := model.TaskGetAllTasks(pageInfo.GetStartIdx(), pageInfo.GetPageSize(), queryParams) total := model.TaskCountAllTasks(queryParams) pageInfo.SetTotal(int(total)) - pageInfo.SetItems(tasksToDto(items, true)) + dtos, err := tasksToDto(c.Request.Context(), items, true) + if err != nil { + common.ApiError(c, err) + return + } + pageInfo.SetItems(dtos) common.ApiSuccess(c, pageInfo) } @@ -56,11 +62,16 @@ func GetUserTask(c *gin.Context) { items := model.TaskGetAllUserTask(userId, pageInfo.GetStartIdx(), pageInfo.GetPageSize(), queryParams) total := model.TaskCountAllUserTask(userId, queryParams) pageInfo.SetTotal(int(total)) - pageInfo.SetItems(tasksToDto(items, false)) + dtos, err := tasksToDto(c.Request.Context(), items, false) + if err != nil { + common.ApiError(c, err) + return + } + pageInfo.SetItems(dtos) common.ApiSuccess(c, pageInfo) } -func tasksToDto(tasks []*model.Task, fillUser bool) []*dto.TaskDto { +func tasksToDto(ctx context.Context, tasks []*model.Task, fillUser bool) ([]*dto.TaskDto, error) { var userIdMap map[int]*model.UserBase if fillUser { userIdMap = make(map[int]*model.UserBase) @@ -76,13 +87,35 @@ func tasksToDto(tasks []*model.Task, fillUser bool) []*dto.TaskDto { } } result := make([]*dto.TaskDto, len(tasks)) + asyncTaskIDs := make([]int64, 0, len(tasks)) + for _, task := range tasks { + if task.Platform == constant.TaskPlatformAsyncImage { + asyncTaskIDs = append(asyncTaskIDs, task.ID) + } + } + asyncJobs, err := model.ListAsyncJobsByTaskIDs(ctx, asyncTaskIDs) + if err != nil { + return nil, err + } for i, task := range tasks { if fillUser { if user, ok := userIdMap[task.UserId]; ok { task.Username = user.Username } } - result[i] = relay.TaskModel2Dto(task) + item := relay.TaskModel2Dto(task) + if job, ok := asyncJobs[task.ID]; ok { + item.Async = &dto.AsyncTaskMeta{ + ExecutionStatus: string(job.ExecutionStatus), + WorkerID: job.WorkerID, + Attempt: job.Attempt, + RequestSentAt: job.RequestSentAt, + ErrorPhase: job.ErrorPhase, + ErrorCode: job.ErrorCode, + BillingStatus: job.BillingStatus, + } + } + result[i] = item } - return result + return result, nil } diff --git a/deploy/Caddyfile b/deploy/Caddyfile new file mode 100644 index 000000000000..456c8732a3cc --- /dev/null +++ b/deploy/Caddyfile @@ -0,0 +1,20 @@ +{ + admin off +} + +http://localhost { + redir https://localhost{uri} permanent +} + +https://localhost { + tls internal + encode zstd gzip + reverse_proxy new-api-api:3000 + + header { + X-Content-Type-Options nosniff + X-Frame-Options DENY + Referrer-Policy no-referrer + -Server + } +} diff --git a/deploy/compose.staging.yml b/deploy/compose.staging.yml new file mode 100644 index 000000000000..9c5b1dd67ee3 --- /dev/null +++ b/deploy/compose.staging.yml @@ -0,0 +1,190 @@ +name: new-api-async-staging + +x-new-api-environment: &new-api-environment + SQL_DSN: postgresql://${POSTGRES_USER}:${POSTGRES_PASSWORD}@postgres:5432/${POSTGRES_DB}?sslmode=disable + REDIS_CONN_STRING: redis://:${REDIS_PASSWORD}@redis:6379/0 + SESSION_SECRET: ${SESSION_SECRET} + SESSION_COOKIE_SECURE: "true" + SESSION_COOKIE_TRUSTED_URL: https://async-api.nexaapp.cn + CRYPTO_SECRET: ${CRYPTO_SECRET} + ASYNC_REQUEST_ENCRYPTION_KEY: ${ASYNC_REQUEST_ENCRYPTION_KEY} + ASYNC_YUNWU_ALLOWED_BASE_URLS: https://yunwu.ai + ASYNC_GRSAI_ALLOWED_BASE_URLS: https://grsaiapi.com,https://grsai.dakka.com.cn + ASYNC_MAX_REQUEST_BODY_KB: 256 + ASYNC_MAX_PROMPT_CHARS: 8000 + ASYNC_MAX_INPUT_URLS: 8 + ASYNC_ARTIFACT_MAX_FILES: 8 + ASYNC_ARTIFACT_MAX_FILE_MB: 25 + ASYNC_ARTIFACT_MAX_TOTAL_MB: 100 + ASYNC_ARTIFACT_DOWNLOAD_TIMEOUT_SECONDS: 120 + ASYNC_ARTIFACT_ARCHIVE_TIMEOUT_SECONDS: 300 + ASYNC_SIGNED_URL_TTL_SECONDS: 900 + ASYNC_UPSTREAM_MAX_RESPONSE_MB: 64 + ASYNC_YUNWU_ROUTE_SUFFIX: stable + S3_ENDPOINT: http://minio:9000 + S3_PUBLIC_ENDPOINT: https://async-files.nexaapp.cn + S3_REGION: us-east-1 + S3_BUCKET: new-api-staging-artifacts + S3_ACCESS_KEY_ID: ${MINIO_ROOT_USER} + S3_SECRET_ACCESS_KEY: ${MINIO_ROOT_PASSWORD} + S3_USE_PATH_STYLE: "true" + BATCH_UPDATE_ENABLED: "false" + ERROR_LOG_ENABLED: "true" + CRITICAL_RATE_LIMIT_ENABLE: "true" + CRITICAL_RATE_LIMIT: 1000 + CRITICAL_RATE_LIMIT_DURATION: 1200 + TZ: Asia/Shanghai + +services: + new-api: + image: new-api-async-relay:staging-20260718-grsai-amd64 + restart: unless-stopped + command: ["--log-dir", "/app/logs"] + environment: + <<: *new-api-environment + APP_ROLE: api + NODE_TYPE: master + NODE_NAME: new-api-async-staging-api + ports: + - "127.0.0.1:33001:3000" + volumes: + - new_api_data:/data + - new_api_logs:/app/logs + depends_on: + postgres: + condition: service_healthy + redis: + condition: service_healthy + minio-init: + condition: service_completed_successfully + networks: [backend] + cpus: 0.75 + mem_limit: 768m + pids_limit: 256 + healthcheck: + test: ["CMD-SHELL", "wget -q -O - http://localhost:3000/api/status | grep -q '\"success\"[[:space:]]*:[[:space:]]*true'"] + interval: 10s + timeout: 5s + retries: 12 + start_period: 30s + + new-api-worker: + image: new-api-async-relay:staging-20260718-grsai-amd64 + restart: unless-stopped + command: ["--log-dir", "/app/logs"] + environment: + <<: *new-api-environment + APP_ROLE: worker + NODE_TYPE: slave + NODE_NAME: new-api-async-staging-worker + ASYNC_WORKER_ID: new-api-async-staging-worker-1 + ASYNC_WORKER_CONCURRENCY: 1 + ASYNC_CHANNEL_DEFAULT_CONCURRENCY: 1 + ASYNC_WORKER_LEASE_SECONDS: 90 + ASYNC_WORKER_POLL_MILLISECONDS: 1000 + ASYNC_JOB_TIMEOUT_SECONDS: 1800 + ASYNC_RESULT_RETENTION_MINUTES: 60 + ASYNC_ARTIFACT_CLEANUP_INTERVAL_SECONDS: 60 + volumes: + - worker_logs:/app/logs + depends_on: + new-api: + condition: service_healthy + minio-init: + condition: service_completed_successfully + networks: [backend] + cpus: 0.75 + mem_limit: 768m + pids_limit: 256 + stop_grace_period: 35m + + postgres: + image: postgres:16.10-alpine + restart: unless-stopped + environment: + POSTGRES_USER: ${POSTGRES_USER} + POSTGRES_PASSWORD: ${POSTGRES_PASSWORD} + POSTGRES_DB: ${POSTGRES_DB} + volumes: + - postgres_data:/var/lib/postgresql/data + networks: [backend] + cpus: 0.5 + mem_limit: 512m + pids_limit: 128 + healthcheck: + test: ["CMD-SHELL", "pg_isready -U $${POSTGRES_USER} -d $${POSTGRES_DB}"] + interval: 5s + timeout: 5s + retries: 20 + + redis: + image: redis:7.4.5-alpine + restart: unless-stopped + command: ["redis-server", "--appendonly", "yes", "--requirepass", "${REDIS_PASSWORD}"] + environment: + REDIS_PASSWORD: ${REDIS_PASSWORD} + volumes: + - redis_data:/data + networks: [backend] + cpus: 0.25 + mem_limit: 256m + pids_limit: 128 + healthcheck: + test: ["CMD-SHELL", "redis-cli -a $${REDIS_PASSWORD} ping 2>/dev/null | grep -q PONG"] + interval: 5s + timeout: 5s + retries: 20 + + minio: + image: minio/minio:RELEASE.2025-09-07T16-13-09Z + restart: unless-stopped + command: ["server", "/data", "--console-address", ":9001"] + environment: + MINIO_ROOT_USER: ${MINIO_ROOT_USER} + MINIO_ROOT_PASSWORD: ${MINIO_ROOT_PASSWORD} + volumes: + - minio_data:/data + ports: + - "127.0.0.1:33900:9000" + - "127.0.0.1:33901:9001" + networks: [backend] + cpus: 0.5 + mem_limit: 512m + pids_limit: 128 + healthcheck: + test: ["CMD-SHELL", "curl -fsS http://localhost:9000/minio/health/live >/dev/null"] + interval: 5s + timeout: 5s + retries: 20 + + minio-init: + image: minio/mc:RELEASE.2025-08-13T08-35-41Z + restart: "no" + entrypoint: ["/bin/sh", "-c"] + command: + - >- + until mc alias set local http://minio:9000 "$${MINIO_ROOT_USER}" "$${MINIO_ROOT_PASSWORD}"; do sleep 2; done; + mc mb --ignore-existing local/new-api-staging-artifacts; + mc anonymous set private local/new-api-staging-artifacts + environment: + MINIO_ROOT_USER: ${MINIO_ROOT_USER} + MINIO_ROOT_PASSWORD: ${MINIO_ROOT_PASSWORD} + depends_on: + minio: + condition: service_healthy + networks: [backend] + cpus: 0.25 + mem_limit: 128m + pids_limit: 64 + +volumes: + new_api_data: + new_api_logs: + worker_logs: + postgres_data: + redis_data: + minio_data: + +networks: + backend: + driver: bridge diff --git a/deploy/nginx/async-api.nexaapp.cn.bootstrap.conf b/deploy/nginx/async-api.nexaapp.cn.bootstrap.conf new file mode 100644 index 000000000000..af27212cb89e --- /dev/null +++ b/deploy/nginx/async-api.nexaapp.cn.bootstrap.conf @@ -0,0 +1,55 @@ +server { + listen 80; + listen [::]:80; + server_name async-api.nexaapp.cn; + + client_max_body_size 64m; + + location ^~ /.well-known/acme-challenge/ { + root /var/www/certbot; + } + + location / { + proxy_pass http://127.0.0.1:33001; + proxy_http_version 1.1; + proxy_set_header Host $host; + proxy_set_header X-Real-IP $remote_addr; + proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for; + proxy_set_header X-Forwarded-Proto $scheme; + proxy_set_header X-Forwarded-Host $host; + proxy_set_header Connection ""; + proxy_buffering off; + proxy_request_buffering off; + proxy_cache off; + proxy_read_timeout 1h; + proxy_send_timeout 1h; + } +} + +server { + listen 80; + listen [::]:80; + server_name async-files.nexaapp.cn; + + client_max_body_size 64m; + + location ^~ /.well-known/acme-challenge/ { + root /var/www/certbot; + } + + location / { + proxy_pass http://127.0.0.1:33900; + proxy_http_version 1.1; + proxy_set_header Host $host; + proxy_set_header X-Real-IP $remote_addr; + proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for; + proxy_set_header X-Forwarded-Proto $scheme; + proxy_set_header X-Forwarded-Host $host; + proxy_set_header Connection ""; + proxy_buffering off; + proxy_request_buffering off; + proxy_cache off; + proxy_read_timeout 5m; + proxy_send_timeout 5m; + } +} diff --git a/deploy/nginx/async-api.nexaapp.cn.conf b/deploy/nginx/async-api.nexaapp.cn.conf new file mode 100644 index 000000000000..8a8c63a4b58d --- /dev/null +++ b/deploy/nginx/async-api.nexaapp.cn.conf @@ -0,0 +1,121 @@ +server { + listen 80; + listen [::]:80; + server_name async-api.nexaapp.cn; + + location ^~ /.well-known/acme-challenge/ { + root /var/www/certbot; + } + + location / { + return 301 https://$host$request_uri; + } +} + +server { + listen 443 ssl; + listen [::]:443 ssl; + server_name async-api.nexaapp.cn; + + ssl_certificate /etc/letsencrypt/live/async-api.nexaapp.cn/fullchain.pem; + ssl_certificate_key /etc/letsencrypt/live/async-api.nexaapp.cn/privkey.pem; + include /etc/letsencrypt/options-ssl-nginx.conf; + ssl_dhparam /etc/letsencrypt/ssl-dhparams.pem; + + client_max_body_size 64m; + root /var/www/async-api-frontend/current; + + location = / { + try_files /index.html =404; + add_header Cache-Control "no-cache"; + } + + location ^~ /static/ { + try_files $uri @new_api_backend; + add_header Cache-Control "public, max-age=31536000, immutable"; + } + + location ~ ^/(favicon\.ico|logo\.png|pay-(google|apple|card)\.png|waffo-logo-(light|dark)\.svg)$ { + try_files $uri @new_api_backend; + } + + location ~ ^/(401|403|404|500|503|about|async-image-lab|channels|chat|chat2link|dashboard|docs|errors|forgot-password|keys|models|oauth|otp|playground|pricing|privacy-policy|profile|rankings|redemption-codes|register|reset|setup|sign-in|sign-up|subscriptions|system-info|system-settings|usage-logs|user-agreement|user/reset|users|wallet)(/.*)?$ { + try_files /index.html =404; + add_header Cache-Control "no-cache"; + } + + location / { + proxy_pass http://127.0.0.1:33001; + proxy_http_version 1.1; + proxy_set_header Host $host; + proxy_set_header X-Real-IP $remote_addr; + proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for; + proxy_set_header X-Forwarded-Proto $scheme; + proxy_set_header X-Forwarded-Host $host; + proxy_set_header Connection ""; + proxy_buffering off; + proxy_request_buffering off; + proxy_cache off; + proxy_read_timeout 1h; + proxy_send_timeout 1h; + } + + location @new_api_backend { + proxy_pass http://127.0.0.1:33001; + proxy_http_version 1.1; + proxy_set_header Host $host; + proxy_set_header X-Real-IP $remote_addr; + proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for; + proxy_set_header X-Forwarded-Proto $scheme; + proxy_set_header X-Forwarded-Host $host; + proxy_set_header Connection ""; + proxy_buffering off; + proxy_request_buffering off; + proxy_cache off; + proxy_read_timeout 1h; + proxy_send_timeout 1h; + } +} + +server { + listen 80; + listen [::]:80; + server_name async-files.nexaapp.cn; + + location ^~ /.well-known/acme-challenge/ { + root /var/www/certbot; + } + + location / { + return 301 https://$host$request_uri; + } +} + +server { + listen 443 ssl; + listen [::]:443 ssl; + server_name async-files.nexaapp.cn; + + ssl_certificate /etc/letsencrypt/live/async-api.nexaapp.cn/fullchain.pem; + ssl_certificate_key /etc/letsencrypt/live/async-api.nexaapp.cn/privkey.pem; + include /etc/letsencrypt/options-ssl-nginx.conf; + ssl_dhparam /etc/letsencrypt/ssl-dhparams.pem; + + client_max_body_size 64m; + + location / { + proxy_pass http://127.0.0.1:33900; + proxy_http_version 1.1; + proxy_set_header Host $host; + proxy_set_header X-Real-IP $remote_addr; + proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for; + proxy_set_header X-Forwarded-Proto $scheme; + proxy_set_header X-Forwarded-Host $host; + proxy_set_header Connection ""; + proxy_buffering off; + proxy_request_buffering off; + proxy_cache off; + proxy_read_timeout 5m; + proxy_send_timeout 5m; + } +} diff --git a/deploy/nginx/certbot-reload-nginx.sh b/deploy/nginx/certbot-reload-nginx.sh new file mode 100644 index 000000000000..a5a5e407d0f7 --- /dev/null +++ b/deploy/nginx/certbot-reload-nginx.sh @@ -0,0 +1,4 @@ +#!/bin/sh +set -eu + +/usr/bin/systemctl reload nginx diff --git a/deploy/sub2api/.env.example b/deploy/sub2api/.env.example new file mode 100644 index 000000000000..0159397bd345 --- /dev/null +++ b/deploy/sub2api/.env.example @@ -0,0 +1,19 @@ +# The default matches deploy/compose.staging.yml (name: new-api-async-staging). +NEW_API_DOCKER_NETWORK=new-api-async-staging_backend + +# Keep the management UI bound to loopback. Use an SSH tunnel for remote access. +SUB2API_BIND_HOST=127.0.0.1 +SUB2API_PORT=38080 +SUB2API_TESTBENCH_PORT=38081 + +# Pinned and verified multi-architecture image. Review release notes before changing it. +SUB2API_IMAGE=ghcr.io/wei-shaw/sub2api:0.1.160 + +SUB2API_ADMIN_EMAIL=admin@example.com +SUB2API_ADMIN_PASSWORD=replace_with_a_strong_admin_password +SUB2API_POSTGRES_PASSWORD=replace_with_a_random_database_password +SUB2API_REDIS_PASSWORD=replace_with_a_random_redis_password +SUB2API_JWT_SECRET=replace_with_at_least_32_random_bytes +SUB2API_TOTP_ENCRYPTION_KEY=replace_with_at_least_32_random_bytes + +TZ=Asia/Shanghai diff --git a/deploy/sub2api/compose.yml b/deploy/sub2api/compose.yml new file mode 100644 index 000000000000..6d6654cca400 --- /dev/null +++ b/deploy/sub2api/compose.yml @@ -0,0 +1,142 @@ +name: new-api-sub2api + +services: + sub2api: + image: ${SUB2API_IMAGE:-sub2api-local:0.1.160-model-test} + build: + context: ./custom + dockerfile: Dockerfile + args: + VERSION: 0.1.160 + restart: unless-stopped + environment: + AUTO_SETUP: "true" + SERVER_HOST: 0.0.0.0 + SERVER_PORT: 8080 + SERVER_MODE: release + RUN_MODE: simple + DATABASE_HOST: sub2api-postgres + DATABASE_PORT: 5432 + DATABASE_USER: sub2api + DATABASE_PASSWORD: ${SUB2API_POSTGRES_PASSWORD:?SUB2API_POSTGRES_PASSWORD is required} + DATABASE_DBNAME: sub2api + DATABASE_SSLMODE: disable + DATABASE_MAX_OPEN_CONNS: ${SUB2API_DATABASE_MAX_OPEN_CONNS:-50} + DATABASE_MAX_IDLE_CONNS: ${SUB2API_DATABASE_MAX_IDLE_CONNS:-10} + REDIS_HOST: sub2api-redis + REDIS_PORT: 6379 + REDIS_PASSWORD: ${SUB2API_REDIS_PASSWORD:?SUB2API_REDIS_PASSWORD is required} + REDIS_DB: 0 + ADMIN_EMAIL: ${SUB2API_ADMIN_EMAIL:?SUB2API_ADMIN_EMAIL is required} + ADMIN_PASSWORD: ${SUB2API_ADMIN_PASSWORD:?SUB2API_ADMIN_PASSWORD is required} + JWT_SECRET: ${SUB2API_JWT_SECRET:?SUB2API_JWT_SECRET is required} + JWT_EXPIRE_HOUR: ${SUB2API_JWT_EXPIRE_HOUR:-24} + TOTP_ENCRYPTION_KEY: ${SUB2API_TOTP_ENCRYPTION_KEY:?SUB2API_TOTP_ENCRYPTION_KEY is required} + SECURITY_URL_ALLOWLIST_ENABLED: ${SUB2API_URL_ALLOWLIST_ENABLED:-false} + SECURITY_URL_ALLOWLIST_ALLOW_INSECURE_HTTP: "false" + TZ: ${TZ:-Asia/Shanghai} + ports: + - "${SUB2API_BIND_HOST:-127.0.0.1}:${SUB2API_PORT:-38080}:8080" + volumes: + - sub2api_data:/app/data + depends_on: + sub2api-postgres: + condition: service_healthy + sub2api-redis: + condition: service_healthy + networks: + sub2api-internal: + new-api-backend: + aliases: + - sub2api-account-pool + healthcheck: + test: + [ + "CMD", + "wget", + "-q", + "-T", + "5", + "-O", + "/dev/null", + "http://localhost:8080/health", + ] + interval: 15s + timeout: 5s + retries: 10 + start_period: 30s + + sub2api-postgres: + image: postgres:16.10-alpine + restart: unless-stopped + environment: + POSTGRES_USER: sub2api + POSTGRES_PASSWORD: ${SUB2API_POSTGRES_PASSWORD:?SUB2API_POSTGRES_PASSWORD is required} + POSTGRES_DB: sub2api + TZ: ${TZ:-Asia/Shanghai} + volumes: + - sub2api_postgres_data:/var/lib/postgresql/data + networks: + - sub2api-internal + healthcheck: + test: ["CMD-SHELL", "pg_isready -U sub2api -d sub2api"] + interval: 5s + timeout: 5s + retries: 20 + + sub2api-redis: + image: redis:7.4.5-alpine + restart: unless-stopped + command: + [ + "redis-server", + "--appendonly", + "yes", + "--appendfsync", + "everysec", + "--requirepass", + "${SUB2API_REDIS_PASSWORD:?SUB2API_REDIS_PASSWORD is required}", + ] + environment: + REDISCLI_AUTH: ${SUB2API_REDIS_PASSWORD:?SUB2API_REDIS_PASSWORD is required} + TZ: ${TZ:-Asia/Shanghai} + volumes: + - sub2api_redis_data:/data + networks: + - sub2api-internal + healthcheck: + test: ["CMD", "redis-cli", "ping"] + interval: 5s + timeout: 5s + retries: 20 + + sub2api-testbench: + image: nginx:1.27-alpine + restart: unless-stopped + ports: + - "${SUB2API_BIND_HOST:-127.0.0.1}:${SUB2API_TESTBENCH_PORT:-38081}:8080" + volumes: + - ./testbench/index.html:/usr/share/nginx/html/index.html:ro + - ./testbench/nginx.conf:/etc/nginx/conf.d/default.conf:ro + depends_on: + sub2api: + condition: service_healthy + networks: + - new-api-backend + healthcheck: + test: ["CMD", "wget", "-q", "-T", "5", "-O", "/dev/null", "http://127.0.0.1:8080/"] + interval: 10s + timeout: 5s + retries: 10 + start_period: 5s + +volumes: + sub2api_data: + sub2api_postgres_data: + sub2api_redis_data: + +networks: + sub2api-internal: + new-api-backend: + external: true + name: ${NEW_API_DOCKER_NETWORK:-new-api-async-staging_backend} diff --git a/deploy/sub2api/custom/Dockerfile b/deploy/sub2api/custom/Dockerfile new file mode 100644 index 000000000000..64cfded88f24 --- /dev/null +++ b/deploy/sub2api/custom/Dockerfile @@ -0,0 +1,106 @@ +# syntax=docker/dockerfile:1.7 + +ARG NODE_IMAGE=node:24-alpine +ARG GOLANG_IMAGE=golang:1.26.5-alpine +ARG ALPINE_IMAGE=alpine:3.21 +ARG POSTGRES_IMAGE=postgres:16.10-alpine +ARG SUB2API_REPOSITORY=https://github.com/Wei-Shaw/sub2api.git +ARG SUB2API_REF=v0.1.160 +ARG SUB2API_COMMIT=8bfbc5ca99bf2c0ac96e0f29ffd35eb6aca27e62 +ARG GOPROXY=https://goproxy.cn,direct +ARG GOSUMDB=sum.golang.google.cn +ARG NPM_CONFIG_REGISTRY= + +FROM alpine/git:2.47.2 AS source +ARG SUB2API_REPOSITORY +ARG SUB2API_REF +ARG SUB2API_COMMIT + +RUN git clone --depth 1 --branch "${SUB2API_REF}" "${SUB2API_REPOSITORY}" /src && \ + test "$(git -C /src rev-parse HEAD)" = "${SUB2API_COMMIT}" +COPY model-test.patch /tmp/model-test.patch +RUN git -C /src apply --check /tmp/model-test.patch && \ + git -C /src apply /tmp/model-test.patch + +FROM ${NODE_IMAGE} AS frontend-builder +ARG NPM_CONFIG_REGISTRY + +WORKDIR /app/frontend +RUN corepack enable && corepack prepare pnpm@9 --activate + +COPY --from=source /src/frontend/package.json /src/frontend/pnpm-lock.yaml ./ +RUN --mount=type=cache,id=sub2api-pnpm-store,target=/root/.local/share/pnpm/store \ + if [ -n "${NPM_CONFIG_REGISTRY}" ]; then pnpm config set registry "${NPM_CONFIG_REGISTRY}"; fi && \ + pnpm install --frozen-lockfile --prefer-offline + +COPY --from=source /src/frontend/ ./ +COPY --from=source /src/docs/legal/ /app/docs/legal/ +RUN pnpm run build + +FROM ${GOLANG_IMAGE} AS backend-builder +ARG VERSION=0.1.160 +ARG COMMIT=8bfbc5ca99bf2c0ac96e0f29ffd35eb6aca27e62-model-test +ARG DATE +ARG GOPROXY +ARG GOSUMDB + +ENV GOPROXY=${GOPROXY} +ENV GOSUMDB=${GOSUMDB} + +RUN apk add --no-cache git ca-certificates tzdata +WORKDIR /app/backend + +COPY --from=source /src/backend/go.mod /src/backend/go.sum ./ +RUN go mod download +COPY --from=source /src/backend/ ./ +COPY --from=frontend-builder /app/backend/internal/web/dist ./internal/web/dist + +RUN DATE_VALUE="${DATE:-$(date -u +%Y-%m-%dT%H:%M:%SZ)}" && \ + CGO_ENABLED=0 GOOS=linux go build \ + -tags embed \ + -ldflags="-s -w -X main.Version=${VERSION} -X main.Commit=${COMMIT} -X main.Date=${DATE_VALUE} -X main.BuildType=release" \ + -trimpath \ + -o /app/sub2api \ + ./cmd/server + +FROM ${POSTGRES_IMAGE} AS pg-client + +FROM ${ALPINE_IMAGE} + +LABEL maintainer="Wei-Shaw " +LABEL description="Sub2API - AI API Gateway Platform" +LABEL org.opencontainers.image.source="https://github.com/Wei-Shaw/sub2api" + +RUN apk add --no-cache \ + ca-certificates \ + tzdata \ + su-exec \ + libpq \ + zstd-libs \ + lz4-libs \ + krb5-libs \ + libldap \ + libedit \ + && rm -rf /var/cache/apk/* + +COPY --from=pg-client /usr/local/bin/pg_dump /usr/local/bin/pg_dump +COPY --from=pg-client /usr/local/bin/psql /usr/local/bin/psql +COPY --from=pg-client /usr/local/lib/libpq.so.5* /usr/local/lib/ + +RUN addgroup -g 1000 sub2api && \ + adduser -u 1000 -G sub2api -s /bin/sh -D sub2api + +WORKDIR /app +COPY --from=backend-builder --chown=sub2api:sub2api /app/sub2api /app/sub2api +COPY --from=backend-builder --chown=sub2api:sub2api /app/backend/resources /app/resources +RUN mkdir -p /app/data && chown sub2api:sub2api /app/data + +COPY --from=source /src/deploy/docker-entrypoint.sh /app/docker-entrypoint.sh +RUN chmod +x /app/docker-entrypoint.sh + +EXPOSE 8080 +HEALTHCHECK --interval=30s --timeout=10s --start-period=10s --retries=3 \ + CMD wget -q -T 5 -O /dev/null http://localhost:${SERVER_PORT:-8080}/health || exit 1 + +ENTRYPOINT ["/app/docker-entrypoint.sh"] +CMD ["/app/sub2api"] diff --git a/deploy/sub2api/custom/model-test.patch b/deploy/sub2api/custom/model-test.patch new file mode 100644 index 000000000000..ae6844688540 --- /dev/null +++ b/deploy/sub2api/custom/model-test.patch @@ -0,0 +1,772 @@ +diff --git a/frontend/src/components/layout/AppSidebar.vue b/frontend/src/components/layout/AppSidebar.vue +index ea89b8d..1cb4149 100644 +--- a/frontend/src/components/layout/AppSidebar.vue ++++ b/frontend/src/components/layout/AppSidebar.vue +@@ -769,6 +769,7 @@ const adminNavItems = computed((): NavItem[] => { + }, + { path: '/admin/subscriptions', label: t('nav.subscriptions'), icon: CreditCardIcon, hideInSimpleMode: true }, + { path: '/admin/accounts', label: t('nav.accounts'), icon: GlobeIcon }, ++ { path: '/admin/model-test', label: t('admin.accounts.modelTest.nav'), icon: BatchImageIcon }, + { path: '/admin/announcements', label: t('nav.announcements'), icon: BellIcon }, + { path: '/admin/proxies', label: t('nav.proxies'), icon: ServerIcon }, + { +diff --git a/frontend/src/i18n/locales/en/admin/accounts.ts b/frontend/src/i18n/locales/en/admin/accounts.ts +index f60dca1..d424be8 100644 +--- a/frontend/src/i18n/locales/en/admin/accounts.ts ++++ b/frontend/src/i18n/locales/en/admin/accounts.ts +@@ -1193,6 +1193,50 @@ export default { + grokAccount: 'Grok Account', + inputMethod: 'Input Method', + reAuthorizedSuccess: 'Account re-authorized successfully', ++ modelTest: { ++ nav: 'Model Test', ++ title: 'Model Test', ++ description: 'Send real text, image-generation, or image-editing requests with an active API key', ++ actualRequestNotice: 'This page sends real requests and incurs actual usage. Test only when needed.', ++ apiKeyMemoryHint: 'The key stays in page memory only. It is not written to Local Storage or included in test output.', ++ configuration: 'Test Configuration', ++ selectKey: 'API Key', ++ selectedGroup: 'Current group: {group}', ++ noActiveKey: 'There is no active API key for this account.', ++ goCreateKey: 'Create an API key', ++ loadKeysFailed: 'Failed to load API keys. Refresh the page and try again.', ++ mode: 'Test Type', ++ textMode: 'Text', ++ imageMode: 'Image Generation', ++ imageEditMode: 'Image Editing', ++ textProtocol: 'Protocol: Anthropic Messages (/v1/messages)', ++ imageProtocol: 'Protocol: Antigravity native Gemini (/antigravity/v1beta)', ++ model: 'Model', ++ imageVerifiedHint: 'The listed models have been verified locally to return images.', ++ aspectRatio: 'Aspect Ratio', ++ imageSize: 'Image Size', ++ sourceImage: 'Source Image', ++ sourceImageHint: 'PNG, JPEG, or WebP, up to 20 MB. The image stays in page memory and is sent inline only for this request.', ++ removeSourceImage: 'Remove', ++ invalidSourceImage: 'Select a PNG, JPEG, or WebP image no larger than 20 MB.', ++ sourceImageReadFailed: 'Failed to read the source image. Select it again.', ++ prompt: 'Prompt', ++ textPromptPlaceholder: 'Enter a text test prompt', ++ imagePromptPlaceholder: 'Describe the image to generate', ++ imageEditPromptPlaceholder: 'Describe how the source image should be changed', ++ send: 'Send Real Request', ++ sending: 'Requesting...', ++ cancel: 'Cancel', ++ result: 'Test Result', ++ idle: 'Configure the request, then click “Send Real Request”.', ++ requestFailed: 'Request Failed', ++ noText: 'The request succeeded, but the response contains no displayable text.', ++ noImage: 'The request succeeded, but no image data was found in the response.', ++ textOutput: 'Model Output', ++ generatedImage: 'Generated Image', ++ download: 'Download Original', ++ responseSummary: 'View response summary (image Base64 hidden)' ++ }, + // Test Modal + testAccountConnection: 'Test Account Connection', + account: 'Account', +diff --git a/frontend/src/i18n/locales/zh/admin/accounts.ts b/frontend/src/i18n/locales/zh/admin/accounts.ts +index 0a6c739..9a530e0 100644 +--- a/frontend/src/i18n/locales/zh/admin/accounts.ts ++++ b/frontend/src/i18n/locales/zh/admin/accounts.ts +@@ -1275,6 +1275,50 @@ export default { + grokAccount: 'Grok 账号', + inputMethod: '输入方式', + reAuthorizedSuccess: '账号重新授权成功', ++ modelTest: { ++ nav: '模型测试', ++ title: '模型测试', ++ description: '通过当前 API 密钥发送真实文本、图片生成或图片改图请求', ++ actualRequestNotice: '这里会发送真实请求并产生实际用量，请按需测试。', ++ apiKeyMemoryHint: '密钥仅保留在当前页面内存中，不写入 Local Storage，也不会显示在测试结果中。', ++ configuration: '测试配置', ++ selectKey: 'API 密钥', ++ selectedGroup: '当前分组：{group}', ++ noActiveKey: '当前账号没有可用的 API 密钥。', ++ goCreateKey: '前往创建密钥', ++ loadKeysFailed: '读取 API 密钥失败，请刷新页面重试。', ++ mode: '测试类型', ++ textMode: '文本生成', ++ imageMode: '图片生成', ++ imageEditMode: '图像改图', ++ textProtocol: '协议：Anthropic Messages（/v1/messages）', ++ imageProtocol: '协议：Antigravity 原生 Gemini（/antigravity/v1beta）', ++ model: '模型', ++ imageVerifiedHint: '当前列出的是已经在本机验证可返回图片的模型。', ++ aspectRatio: '宽高比', ++ imageSize: '图片尺寸', ++ sourceImage: '原始图片', ++ sourceImageHint: '支持 PNG、JPEG、WebP，最大 20 MB。图片仅保留在当前页面内存中，并随本次请求内联发送。', ++ removeSourceImage: '移除原图', ++ invalidSourceImage: '请选择不超过 20 MB 的 PNG、JPEG 或 WebP 图片。', ++ sourceImageReadFailed: '读取原图失败，请重新选择。', ++ prompt: '提示词', ++ textPromptPlaceholder: '输入文本测试提示词', ++ imagePromptPlaceholder: '描述要生成的图片', ++ imageEditPromptPlaceholder: '描述希望如何修改原图', ++ send: '发送真实请求', ++ sending: '请求中...', ++ cancel: '取消', ++ result: '测试结果', ++ idle: '配置参数后点击“发送真实请求”开始测试。', ++ requestFailed: '请求失败', ++ noText: '请求成功，但响应中没有可显示的文本。', ++ noImage: '请求成功，但响应中没有找到图片数据。', ++ textOutput: '模型输出', ++ generatedImage: '生成图片', ++ download: '下载原图', ++ responseSummary: '查看响应摘要（图片 Base64 已隐藏）' ++ }, + // Test Modal + testAccountConnection: '测试账号连接', + account: '账号', +diff --git a/frontend/src/router/index.ts b/frontend/src/router/index.ts +index 259ea55..a6161c2 100644 +--- a/frontend/src/router/index.ts ++++ b/frontend/src/router/index.ts +@@ -514,6 +514,18 @@ const routes: RouteRecordRaw[] = [ + descriptionKey: 'admin.accounts.description' + } + }, ++ { ++ path: '/admin/model-test', ++ name: 'AdminModelTest', ++ component: () => import('@/views/admin/ModelTestView.vue'), ++ meta: { ++ requiresAuth: true, ++ requiresAdmin: true, ++ title: 'Model Test', ++ titleKey: 'admin.accounts.modelTest.title', ++ descriptionKey: 'admin.accounts.modelTest.description' ++ } ++ }, + { + path: '/admin/announcements', + name: 'AdminAnnouncements', +diff --git a/frontend/src/views/admin/ModelTestView.vue b/frontend/src/views/admin/ModelTestView.vue +new file mode 100644 +index 0000000..816e52f +--- /dev/null ++++ b/frontend/src/views/admin/ModelTestView.vue +@@ -0,0 +1,621 @@ ++ ++ ++ ++ ++ diff --git a/deploy/sub2api/testbench/index.html b/deploy/sub2api/testbench/index.html new file mode 100644 index 000000000000..18b3a0fdb5f2 --- /dev/null +++ b/deploy/sub2api/testbench/index.html @@ -0,0 +1,399 @@ + + + + + + Sub2API 账号测试台 + + + +
+
+
+ +
+

Sub2API 账号测试台

+

直接检查密钥、模型列表和真实生成请求

+
+
+
正在检查本地服务
+
+ +
+ + +
+

生成测试

非流式
+ +
+
+ HTTP -- + 延迟 -- + Request ID -- +
+
+ + + +
+
+ +
+
模型输出尚未发送
+

在左侧填入 API 密钥，读取模型后发送测试请求。

+
+ 查看原始 JSON + 
{}
+
+
+
+
+
+ + + + diff --git a/deploy/sub2api/testbench/nginx.conf b/deploy/sub2api/testbench/nginx.conf new file mode 100644 index 000000000000..5fd4f9aac381 --- /dev/null +++ b/deploy/sub2api/testbench/nginx.conf @@ -0,0 +1,27 @@ +server { + listen 8080; + server_name _; + + root /usr/share/nginx/html; + index index.html; + + add_header Cache-Control "no-store" always; + add_header X-Content-Type-Options "nosniff" always; + add_header Referrer-Policy "no-referrer" always; + add_header Content-Security-Policy "default-src 'self'; style-src 'self' 'unsafe-inline'; script-src 'self' 'unsafe-inline'; connect-src 'self'; img-src 'self' data:; base-uri 'none'; frame-ancestors 'none'; form-action 'none'" always; + + location = / { + try_files /index.html =404; + } + + location /gateway/ { + proxy_pass http://sub2api-account-pool:8080/; + proxy_http_version 1.1; + proxy_set_header Host $proxy_host; + proxy_set_header Authorization $http_authorization; + proxy_set_header X-Request-ID $request_id; + proxy_read_timeout 120s; + proxy_send_timeout 120s; + proxy_buffering off; + } +} diff --git "a/docs/New-API\345\274\202\346\255\245\344\273\273\345\212\241\344\270\255\350\275\254\347\253\231\345\274\200\345\217\221\346\226\207\346\241\243.md" "b/docs/New-API\345\274\202\346\255\245\344\273\273\345\212\241\344\270\255\350\275\254\347\253\231\345\274\200\345\217\221\346\226\207\346\241\243.md" new file mode 100644 index 000000000000..df5cce970961 --- /dev/null +++ "b/docs/New-API\345\274\202\346\255\245\344\273\273\345\212\241\344\270\255\350\275\254\347\253\231\345\274\200\345\217\221\346\226\207\346\241\243.md" @@ -0,0 +1,518 @@ +# New API 异步任务中转站开发文档 + +## 1. 文档信息 + +- 项目名称：New API 异步任务中转站 +- 上游项目：[QuantumNous/new-api](https://github.com/QuantumNous/new-api) +- 计划基线：`v1.0.0-rc.21` +- 首版上游：`https://yunwu.ai` +- 首版使用对象：仅项目所有者自己的多个项目 +- 首版范围：同步图片生成接口异步化 +- 文档状态：开发基线 + +## 2. 项目目标 + +在保留 New API 面板、用户、令牌、渠道、额度、日志和任务管理能力的基础上，增加一套持久化异步执行系统，将云雾的同步图片生成接口包装为异步任务。 + +客户端提交请求后立即取得本地 `task_id`，无需继续保持网络连接。云端 Worker 负责等待云雾返回最终结果，并将响应和媒体文件持久化。客户端之后可通过任务查询接口或回调取得结果。 + +### 2.1 成功标准 + +- 本地客户端提交任务后可以立即断开连接。 +- 客户端断线不会影响云端任务继续执行。 +- API 服务重启不会丢失尚未执行的排队任务。 +- 同一幂等键不会产生重复任务或重复扣费。 +- 云雾返回的临时媒体地址会被归档到自有对象存储。 +- 任务状态、结果、失败原因和计费信息可在 New API 面板查看。 +- 原有 `/v1/images/generations` 同步接口保持兼容。 + +### 2.2 首版不包含 + +- 不开放公众注册、在线充值和对外销售。 +- 不改造聊天流式、语音和文件接口。 +- 不承诺上游同步请求已经被接收后的严格一次执行。 +- 不对没有取消能力的上游任务伪造“已取消”。 +- 不同时维护 New API classic 前端的新增功能。 + +## 3. 总体架构 + +```mermaid +flowchart LR + C["本地项目"] -->|"提交任务，HTTP 202"| A["New API 异步接口"] + A --> K["Token 鉴权、限流、预扣费"] + A --> D[("PostgreSQL")] + D --> W["异步 Worker"] + W -->|"同步长连接"| Y["yunwu.ai"] + W -->|"下载生成结果"| O[("OSS / COS / S3")] + W --> D + C -->|"查询 task_id"| A + A -->|"任务状态和签名下载地址"| C +``` + +### 3.1 服务角色 + +同一个代码仓库和 Docker 镜像支持两种运行角色： + +```env +APP_ROLE=api +``` + +```env +APP_ROLE=worker +ASYNC_WORKER_CONCURRENCY=50 +ASYNC_JOB_TIMEOUT_SECONDS=1800 +``` + +- API 服务负责鉴权、校验、预扣费、任务入库和查询。 +- Worker 负责领取任务、调用上游、归档文件和完成结算。 +- API 请求处理器中禁止通过普通 goroutine 执行长任务。 + +## 4. 代码基线与升级策略 + +1. Fork 官方 `QuantumNous/new-api` 仓库。 +2. 从发布标签 `v1.0.0-rc.21` 建立开发分支，不直接以持续变化的 `main` 作为生产基线。 +3. 保留官方仓库为 `upstream` remote。 +4. 新功能尽量放入独立目录、独立数据表和独立路由，减少与上游核心转发代码的冲突。 +5. 定期合并上游安全修复，合并后必须运行完整测试矩阵。 +6. 生产镜像固定 commit SHA 和镜像 digest，禁止直接使用未固定的 `latest`。 + +建议新增模块： + +```text +model/async_job.go +controller/async_job.go +service/async_queue.go +service/async_worker.go +router/async-router.go +relay/asyncwrap/yunwu.go +storage/artifact_store.go +``` + +## 5. 对外 API + +### 5.1 提交异步图片任务 + +```http +POST /v1/async/images/generations +Authorization: Bearer +Idempotency-Key: +Content-Type: application/json +``` + +请求体继续使用云雾/OpenAI 图片生成格式： + +```json +{ + "model": "doubao-seedream-4-0-250828", + "prompt": "一座未来城市", + "size": "1728x2304", + "response_format": "url" +} +``` + +成功响应为 HTTP `202 Accepted`： + +```json +{ + "id": "task_xxx", + "status": "queued", + "status_url": "/v1/async/tasks/task_xxx", + "result_url": "/v1/async/tasks/task_xxx/result" +} +``` + +规则： + +- `Idempotency-Key` 为必填请求头。 +- 唯一约束为 `token_id + idempotency_key`。 +- 重复提交返回原任务，不重复预扣费。 +- 请求模型必须存在于当前 Token 和渠道允许列表中。 +- 首版只允许配置过异步包装能力的云雾渠道。 + +### 5.2 查询任务 + +```http +GET /v1/async/tasks/{task_id} +Authorization: Bearer +``` + +响应示例： + +```json +{ + "id": "task_xxx", + "status": "running", + "progress": 50, + "created_at": 1784300000, + "started_at": 1784300003, + "finished_at": null, + "error": null +} +``` + +Token 只能查询自己创建的任务，管理员可以通过面板查询全部任务。 + +### 5.3 获取结果 + +```http +GET /v1/async/tasks/{task_id}/result +Authorization: Bearer +``` + +- `SUCCESS`：返回归一化响应、原始上游响应和自有存储的签名下载地址。 +- 非终态：返回 HTTP `409` 和当前任务状态。 +- `FAILURE`：返回稳定错误码和失败阶段。 +- `UNCERTAIN`：明确提示任务可能已在上游执行，禁止客户端自动重试。 + +### 5.4 取消任务 + +```http +POST /v1/async/tasks/{task_id}/cancel +Authorization: Bearer +``` + +- `QUEUED` 任务可以取消并退款。 +- `RUNNING` 任务若上游没有取消接口，不中断连接；接口返回 HTTP `409` 并说明无法确认上游取消。 +- 已进入终态的任务保持原状态。 + +### 5.5 保留原同步接口 + +```http +POST /v1/images/generations +``` + +该接口维持 New API 原有同步语义，避免破坏现有 OpenAI SDK 和第三方客户端。 + +### 5.6 已接入的三个云雾模型 + +预发布环境已经启用并实测以下模型： + +| 模型 | 客户端请求格式 | Worker 上游协议 | 已验收产物 | +| --- | --- | --- | --- | +| `gemini-3.1-flash-image-preview` | 本站统一图片请求 | Gemini `generateContent` | JPEG | +| `gemini-3-pro-image-preview` | 本站统一图片请求 | Gemini `generateContent` | JPEG | +| `gpt-image-2` | OpenAI 图片请求 | `/v1/images/generations` | PNG | + +New API 渠道需要同时满足以下配置： + +- 类型选择 `OpenAI`，API 地址填写 `https://yunwu.ai`，不要手动附加 `/v1`。 +- 渠道模型列表加入上述三个模型。 +- 在渠道高级设置中启用“异步图片包装”，异步模型允许列表也加入上述三个模型。 +- 启用自动归档；预发布环境渠道并发暂设为 `1`，完成稳定性测试后再逐级提高。 +- 给用户 Token 开启相同的模型权限，否则请求会在入队前返回 `403`。 +- Worker 设置 `ASYNC_YUNWU_ROUTE_SUFFIX=stable`，由适配器对每次上游请求追加稳定路由，不修改渠道密钥。 + +Gemini 请求示例： + +```json +{ + "model": "gemini-3.1-flash-image-preview", + "prompt": "一只站在窗边的橘猫", + "n": 1, + "size": "1:1", + "quality": "1K" +} +``` + +Gemini 模型当前限制 `n=1`；`size` 使用宽高比，`quality` 使用 `1K`、`2K` 或 `4K`。适配器会转换为云雾的 Gemini 原生请求，并把 `inlineData` 图片归档到对象存储。 + +`gpt-image-2` 请求示例： + +```json +{ + "model": "gpt-image-2", + "prompt": "一只站在窗边的橘猫", + "n": 1, + "size": "1024x1024", + "quality": "low" +} +``` + +两类上游响应最终都由本站归一化为 `response.data[].url`，URL 指向本站对象存储的短期签名地址。 + +## 6. 任务状态机 + +```mermaid +stateDiagram-v2 + [*] --> QUEUED + QUEUED --> RUNNING: Worker 取得租约 + QUEUED --> CANCELLED: 用户取消 + RUNNING --> SUCCESS: 收到并持久化完整结果 + RUNNING --> FAILURE: 明确确认失败 + RUNNING --> UNCERTAIN: 请求已发送但结果不可确认 + FAILURE --> QUEUED: 管理员确认后手动重试 + UNCERTAIN --> QUEUED: 管理员接受重复扣费风险后重试 +``` + +状态定义： + +- `QUEUED`：任务已持久化，尚未向上游发送。 +- `RUNNING`：Worker 已领取任务并开始执行。 +- `SUCCESS`：上游响应和产物均已成功持久化。 +- `FAILURE`：可以明确确认任务失败或未被上游接受。 +- `UNCERTAIN`：请求可能已被上游接受，但最终结果无法确认。 +- `CANCELLED`：仅适用于尚未发送上游的排队任务。 + +## 7. 数据模型 + +继续复用 New API 的 `Task` 表记录用户可见状态、渠道、额度和结果摘要；新增一对一的 `async_jobs` 表记录后台执行细节。 + +### 7.1 async_jobs + +| 字段 | 用途 | +| --- | --- | +| `id` | 数据库主键 | +| `task_id` | 关联 New API Task，唯一 | +| `token_id` | 创建任务的 Token | +| `channel_id` | 入队时选定的上游渠道 | +| `endpoint_type` | 首版固定为图片生成 | +| `request_payload` | 加密或受控保存的请求内容 | +| `request_hash` | 请求一致性验证 | +| `idempotency_key` | 客户端幂等键 | +| `execution_status` | 后台执行状态 | +| `worker_id` | 当前执行节点 | +| `lease_until` | Worker 租约到期时间 | +| `attempt` | 尝试次数 | +| `request_sent_at` | 请求体开始向上游发送的时间 | +| `result_payload` | 原始上游响应或其对象存储引用 | +| `error_phase` | 失败发生阶段 | +| `error_code` | 稳定内部错误码 | +| `created_at` | 创建时间 | +| `updated_at` | 更新时间 | + +### 7.2 artifacts + +用于记录一项任务产生的一个或多个文件： + +```text +id +task_id +object_key +content_type +size_bytes +sha256 +source_url_hash +created_at +expires_at +``` + +### 7.3 task_events + +记录任务状态变化、Worker、错误阶段和管理员操作，禁止写入上游密钥和完整 Token。 + +## 8. Worker 与队列设计 + +### 8.1 领取任务 + +- 使用项目统一的 `lockForUpdate(tx)` 领取任务，兼容 SQLite、MySQL 5.7 和 PostgreSQL 9.6。 +- 领取后写入 `worker_id`、`lease_until` 和 `RUNNING`。 +- Worker 定期续租。 +- 单 Worker 内使用渠道并发信号量；预发布阶段固定一个 Worker 副本，避免横向扩容绕过进程内渠道上限。 +- 全局并发和渠道并发均可配置；增加多 Worker 副本前必须先实现数据库或 Redis 级分布式并发配额。 + +### 8.2 崩溃恢复 + +- `request_sent_at` 为空且租约过期：可以安全重新入队。 +- `request_sent_at` 非空且租约过期：标记 `UNCERTAIN`。 +- 原生异步上游若已经取得上游任务 ID，可以重新进入查询流程，不标记 `UNCERTAIN`。 +- Worker 发布时先停止领取新任务，再等待运行中任务完成。 + +### 8.3 上游调用 + +- 使用独立 HTTP Client 和连接池。 +- 设置连接、TLS、响应头和整体任务超时。 +- 禁止将客户端请求 Context 传给后台 Worker。 +- 只允许访问配置的云雾基地址和白名单路径。 +- 请求和响应日志默认脱敏。 + +## 9. 计费与重试 + +### 9.1 计费 + +- 入队成功后预扣额度。 +- 重复幂等请求不重复扣费。 +- 成功时按 New API 现有计费上下文完成结算。 +- 明确确认未被上游接受的失败任务退款。 +- 排队取消任务退款。 +- `UNCERTAIN` 默认不退款，由管理员人工核查。 +- 所有退款和结算必须具备幂等保护。 + +### 9.2 自动重试 + +可以自动重试： + +- DNS 解析失败。 +- TCP/TLS 连接建立失败。 +- 请求体尚未发送时的本地错误。 +- 明确收到可重试的 429，并遵守退避时间。 + +禁止自动重试： + +- 请求体已经发送后的读取超时。 +- Worker 在已发送请求后崩溃。 +- 上游返回的错误无法证明任务未执行。 + +管理员手动重试 `UNCERTAIN` 任务时，面板必须显示可能重复生成和重复扣费的确认提示。 + +## 10. 文件归档 + +新增统一对象存储接口： + +```go +type ArtifactStore interface { + Put(ctx context.Context, key string, body io.Reader, contentType string) error + SignedURL(ctx context.Context, key string, ttl time.Duration) (string, error) + Delete(ctx context.Context, key string) error +} +``` + +首版要求兼容 S3 协议，可对接阿里云 OSS、腾讯云 COS、AWS S3 或 MinIO。 + +归档流程： + +1. 解析上游响应中的全部媒体 URL。 +2. 校验协议、主机、DNS 解析结果和重定向目标。 +3. 阻止环回、内网、链路本地地址和云元数据地址，防止 SSRF。 +4. 限制文件数量、单文件大小、总大小、超时和 MIME 类型。 +5. 流式下载并计算 SHA-256，避免一次性加载到内存。 +6. 上传对象存储并写入 artifacts。 +7. 只有全部必需产物持久化成功后，任务才进入 `SUCCESS`。 + +默认结果保留 30 天，由清理任务删除过期对象和对应记录。 + +## 11. 面板改造 + +仅维护 New API 默认新版前端。 + +任务列表增加： + +- 异步包装任务类型。 +- 状态、进度、排队时长和执行时长。 +- Worker 节点和尝试次数。 +- 错误阶段和稳定错误码。 +- `UNCERTAIN` 高风险标识。 +- 原始响应查看和产物下载。 +- 排队任务取消。 +- 管理员手动重试。 + +渠道设置增加： + +- 是否启用同步接口异步包装。 +- 允许的图片模型。 +- 最大并发数。 +- 单任务超时时间。 +- 结果保留天数。 +- 是否自动归档文件。 + +## 12. 安全要求 + +- New API Token 只保存系统原有安全表示，不在日志输出明文。 +- 云雾密钥保存在服务端 Secret 中，不返回给客户端。 +- 任务查询必须校验 Token 所有权。 +- 管理后台不直接暴露数据库和 Redis。 +- PostgreSQL、Redis 只监听 Docker 内部网络。 +- 对象存储 Bucket 默认私有，结果使用短期签名地址。 +- Caddy 负责 HTTPS，HTTP 自动跳转 HTTPS。 +- 限制请求体大小、提示词长度、图片数量和输入 URL 数量。 +- 管理员操作写入审计事件。 +- 对外提供服务前必须另行确认 New API AGPLv3 义务和云雾转售授权；首版不对外销售。 + +## 13. 部署方案 + +### 13.1 服务器 + +- Ubuntu 24.04 LTS，x86_64。 +- 推荐 4 核 8GB、80GB SSD。 +- 独立公网 IP。 +- 开放 22、80、443 端口。 +- 服务器能够稳定访问 `yunwu.ai` 和对象存储。 + +### 13.2 Docker Compose 服务 + +```text +caddy +new-api-api +new-api-worker +postgres +redis +``` + +- 数据库和 Redis 使用强随机密码。 +- 设置 `SESSION_SECRET`、`CRYPTO_SECRET` 和各类存储 Secret。 +- 数据目录使用持久卷。 +- 每日备份 PostgreSQL，备份文件上传到独立存储位置。 +- 生产环境不直接暴露 3000、5432 和 6379 端口。 + +## 14. 测试计划 + +### 14.1 单元测试 + +- 状态迁移合法性。 +- 幂等键唯一性。 +- 任务租约领取和续租。 +- 计费预扣、结算和退款幂等性。 +- 可重试和不可重试错误分类。 +- URL、DNS、重定向和 MIME 安全校验。 + +### 14.2 集成测试 + +- 模拟云雾延迟成功、明确失败、429、5xx 和读取超时。 +- 客户端提交后立即断开连接，任务仍然成功。 +- API 容器重启不影响 Worker 任务。 +- Worker 在请求发送前崩溃，任务重新领取。 +- Worker 在请求发送后崩溃，任务进入 `UNCERTAIN`。 +- 同一幂等键并发提交只创建一个任务。 +- 超出渠道并发上限的任务保持排队。 +- 多图片响应全部归档并生成签名 URL。 +- 不同 Token 无法读取对方任务。 + +### 14.3 部署验收 + +- HTTPS、健康检查和自动重启正常。 +- 数据库和 Redis 不可从公网访问。 +- 备份和恢复演练成功。 +- 服务器重启后排队任务仍存在。 +- 日志中不存在完整 Token、云雾密钥和敏感请求体。 +- 原同步接口行为与改造前一致。 + +## 15. 实施顺序 + +1. 建立 Fork、固定基线和开发分支。 +2. 准备 PostgreSQL 开发环境并加入数据迁移。 +3. 实现 AsyncJob、Artifact 和 TaskEvent 数据模型。 +4. 实现提交、查询、结果和取消接口。 +5. 实现数据库租约队列和独立 Worker 角色。 +6. 实现云雾同步图片适配器。 +7. 接入 New API 预扣、结算和退款流程。 +8. 实现 S3 兼容对象存储和安全下载。 +9. 改造默认新版面板。 +10. 完成单元、集成、故障注入和并发测试。 +11. 在全新云服务器部署预发布环境。 +12. 完成备份恢复、断线和重启验收后切换正式使用。 + +## 16. 第一版完成定义 + +满足以下条件后，第一版才视为完成： + +- 四个异步公开接口可用并有稳定响应格式。 +- 云雾同步图片任务可以在客户端断线后继续完成。 +- 排队任务、运行任务和不确定任务均有正确状态处理。 +- 额度不会因重复提交或重复结算产生异常。 +- 所有成功产物已进入自有对象存储。 +- 面板可以查询、筛选和诊断任务。 +- 故障注入测试和部署验收全部通过。 +- 已生成部署、备份、恢复和升级操作说明。 + +## 17. 2026-07-18 服务器实测记录 + +隔离预发布环境目录为 `/opt/new-api-async-staging`，Compose 项目名为 `new-api-async-staging`。PostgreSQL、Redis、MinIO、API 和单 Worker 均只使用该 Compose 项目的网络、容器和数据卷。 + +三个模型通过本站异步 API 同时入队后，由单 Worker 依次执行，结果如下： + +| 模型 | 任务状态 | 结算状态 | 归档 | 下载验证 | +| --- | --- | --- | --- | --- | +| `gemini-3.1-flash-image-preview` | `SUCCESS` | `SETTLED` | 1 个 JPEG，326987 字节 | HTTP 200，SHA-256 已记录 | +| `gemini-3-pro-image-preview` | `SUCCESS` | `SETTLED` | 1 个 JPEG，573119 字节 | HTTP 200，SHA-256 已记录 | +| `gpt-image-2` | `SUCCESS` | `SETTLED` | 1 个 PNG，98649 字节 | HTTP 200，SHA-256 已记录 | + +真实服务器测试同时发现并修复了 PostgreSQL `TEXT` 字段返回 `string` 时无法扫描到 `json.RawMessage` 的问题。异步结果和事件详情现使用项目已有的跨驱动 `JSONValue`，可读取 `string` 与 `[]byte` 两种数据库驱动返回值。 diff --git "a/docs/Nexa-API\344\270\255\350\275\254\347\253\231\346\216\245\345\205\245\346\226\207\346\241\243.md" "b/docs/Nexa-API\344\270\255\350\275\254\347\253\231\346\216\245\345\205\245\346\226\207\346\241\243.md" new file mode 100644 index 000000000000..ebec825ee661 --- /dev/null +++ "b/docs/Nexa-API\344\270\255\350\275\254\347\253\231\346\216\245\345\205\245\346\226\207\346\241\243.md" @@ -0,0 +1,177 @@ +# Nexa 公共 API 接入文档 + +更新时间：2026-08-13 + +本文面向需要把聊天或图片生成能力接入后端服务、桌面客户端及自动化工作流的开发者。所有示例均以当前线上接口为准。 + +## 1. 开始接入 + +1. 访问 `https://async-api.nexaapp.cn/sign-up` 注册并登录。 +2. 进入「控制台 → API 密钥」创建下游 API Key。 +3. 调用付费模型前确认钱包余额充足。 +4. 使用 HTTPS、Bearer Token 和下方基础地址发起请求。 + +| 配置项 | 内容 | +| --- | --- | +| 网站 | `https://async-api.nexaapp.cn` | +| API Base URL | `https://async-api.nexaapp.cn/v1` | +| 协议 | OpenAI Compatible API | +| 鉴权 | `Authorization: Bearer sk-your-api-key` | +| 模型与价格 | `https://async-api.nexaapp.cn/pricing` | + +API Key 仅可保存在服务端、系统安全凭据存储或环境变量中。禁止把密钥写入公开网页、代码仓库、日志或截图。 + +## 2. 查询当前模型 + +模型会根据审核结果上架或下架，客户端应以 `GET /v1/models` 的实时响应为准，不要在程序中写死列表。 + +```bash +curl 'https://async-api.nexaapp.cn/v1/models' \ + -H 'Authorization: Bearer sk-your-api-key' +``` + +第三方客户端可能缓存以前获取过的模型。刷新模型列表后，仍保留的旧条目需要在客户端本地手动删除。客户端显示模型名称，不代表服务器仍会路由该模型。 + +## 3. 聊天补全 + +### 3.1 非流式请求 + +```bash +curl 'https://async-api.nexaapp.cn/v1/chat/completions' \ + -H 'Authorization: Bearer sk-your-api-key' \ + -H 'Content-Type: application/json' \ + -d '{ + "model": "gemini-3.5-flash", + "messages": [ + {"role": "user", "content": "你好，请介绍一下你自己。"} + ], + "stream": false + }' +``` + +标准响应包含 `choices` 和 `usage`。文本位于 `choices[0].message.content`，Token 用量位于 `usage`。 + +### 3.2 流式请求 + +```bash +curl -N 'https://async-api.nexaapp.cn/v1/chat/completions' \ + -H 'Authorization: Bearer sk-your-api-key' \ + -H 'Content-Type: application/json' \ + -d '{ + "model": "gpt-5.6-sol", + "messages": [ + {"role": "user", "content": "写一段简短的产品介绍。"} + ], + "stream": true, + "stream_options": {"include_usage": true} + }' +``` + +客户端应按 SSE 读取数据，依次拼接 `choices[].delta.content`，收到 `data: [DONE]` 后结束。流式响应已经输出内容后不要自动重试，否则可能产生重复内容和重复费用。 + +### 3.3 Python SDK + +```python +from openai import OpenAI + +client = OpenAI( + api_key="sk-your-api-key", + base_url="https://async-api.nexaapp.cn/v1", +) + +response = client.chat.completions.create( + model="gemini-3.5-flash", + messages=[{"role": "user", "content": "你好"}], +) + +print(response.choices[0].message.content) +``` + +## 4. 异步图片生成 + +图片模型必须使用图片接口调用，不能当作普通聊天模型使用。推荐异步接口：提交成功后任务会继续在服务器执行，客户端可以断开并稍后轮询结果。 + +完整参数、状态机与错误码见同目录的《异步图片生成 API 调用文档》。最小调用流程如下。 + +### 4.1 提交任务 + +```bash +curl 'https://async-api.nexaapp.cn/v1/async/images/generations' \ + -H 'Authorization: Bearer sk-your-api-key' \ + -H 'Content-Type: application/json' \ + -H 'Idempotency-Key: project-a-image-0001' \ + -d '{ + "model": "nano-banana-2", + "prompt": "一座被云海环绕的未来城市", + "n": 1, + "size": "16:9", + "quality": "2K", + "response_format": "url" + }' +``` + +提交成功返回 HTTP `202 Accepted`： + +```json +{ + "id": "task_xxx", + "status": "queued", + "status_url": "/v1/async/tasks/task_xxx", + "result_url": "/v1/async/tasks/task_xxx/result" +} +``` + +`Idempotency-Key` 必填。网络超时后重试完全相同的请求时复用原键；新的生成任务必须使用新键。 + +### 4.2 轮询并获取结果 + +```bash +curl 'https://async-api.nexaapp.cn/v1/async/tasks/task_xxx' \ + -H 'Authorization: Bearer sk-your-api-key' + +curl 'https://async-api.nexaapp.cn/v1/async/tasks/task_xxx/result?include_upstream=false' \ + -H 'Authorization: Bearer sk-your-api-key' +``` + +建议每 2～5 秒轮询一次。状态变为 `success`、`failure`、`uncertain` 或 `cancelled` 后停止。最终图片 URL 位于 `response.data[].url`。 + +符合退款条件的普通失败任务会退回预扣额度；`uncertain` 表示上游可能已经执行，因此可能计费，禁止自动重新提交。 + +## 5. 价格与计费 + +- 用户侧当前价格、币种和计费单位以模型广场实时展示为准。 +- 聊天模型通常按输入 Token 和输出 Token 分别计费。 +- 图片模型通常按次计费，质量、分辨率等参数可能影响价格。 +- 钱包最终变动及「控制台 → 使用日志」是实际计费的权威记录。 +- 上游价格和模型状态可能调整，客户端不得写死价格或可用性。 + +## 6. 错误处理 + +| HTTP 状态码 | 含义 | 建议处理 | +| ---: | --- | --- | +| `400` | 请求格式错误或参数不受支持 | 修正请求，不要原样重试 | +| `401` | API Key 缺失、无效或已禁用 | 检查或更换密钥 | +| `403` | 账号、分组或模型权限不足 | 检查账号和密钥权限 | +| `429` | 速率限制、配额压力或上游策略拒绝 | 先读取错误消息，仅对临时限流退避重试 | +| `500`、`502`、`503`、`504` | 网关或上游服务异常 | 设置次数上限并指数退避重试 | + +部分图片上游也使用 `429` 表示安全策略拦截。如果错误包含 `content blocked by upstream safety policy`，调整本站速率限制不会解决，应修改提示词或更换模型。 + +## 7. Cherry Studio 配置 + +1. 供应商类型选择 `OpenAI Compatible`。 +2. 填入本站创建的 API Key。 +3. 如果地址预览会自动拼接 `/v1/chat/completions`，API 地址填写 `https://async-api.nexaapp.cn`，避免出现 `/v1/v1`。 +4. 其他要求填写 Base URL 的 SDK 通常使用 `https://async-api.nexaapp.cn/v1`。 +5. 服务端模型调整后重新获取模型列表，并手动删除客户端保留的旧条目。 + +图片模型要求客户端支持图片生成接口。仅在文本聊天界面中选择图片模型，不能完成图片生成。 + +## 8. 上线检查 + +- API Key 从服务端环境变量读取，泄露后立即轮换。 +- 为连接和读取分别设置超时。 +- 只对临时 `429` 和 `5xx` 使用带抖动的指数退避重试，并设置次数上限。 +- 记录模型、HTTP 状态、请求 ID 和任务 ID，但不要记录完整 API Key。 +- 向最终用户展示模型选择器前先获取实时模型列表。 +- 异步图片任务持久化 `Idempotency-Key` 与 `task_id`，应用重启后继续查询，不要重新提交。 diff --git a/docs/migrations/001_async_image_jobs_postgresql.sql b/docs/migrations/001_async_image_jobs_postgresql.sql new file mode 100644 index 000000000000..ef48028257bb --- /dev/null +++ b/docs/migrations/001_async_image_jobs_postgresql.sql @@ -0,0 +1,89 @@ +BEGIN; + +CREATE TABLE IF NOT EXISTS async_jobs ( + id BIGSERIAL PRIMARY KEY, + task_id BIGINT NOT NULL UNIQUE, + token_id BIGINT NOT NULL, + channel_id BIGINT NOT NULL, + endpoint_type VARCHAR(40) NOT NULL, + request_payload BYTEA NOT NULL, + request_hash CHAR(64) NOT NULL, + idempotency_key VARCHAR(191) NOT NULL, + execution_status VARCHAR(20) NOT NULL, + worker_id VARCHAR(128) NOT NULL DEFAULT '', + lease_until BIGINT NOT NULL DEFAULT 0, + attempt INTEGER NOT NULL DEFAULT 0, + request_sent_at BIGINT NOT NULL DEFAULT 0, + result_payload JSONB, + error_phase VARCHAR(40) NOT NULL DEFAULT '', + error_code VARCHAR(80) NOT NULL DEFAULT '', + refund_eligible BOOLEAN NOT NULL DEFAULT FALSE, + billing_status VARCHAR(20) NOT NULL DEFAULT 'RESERVED', + billing_request_id VARCHAR(64) NOT NULL DEFAULT '', + created_at BIGINT NOT NULL, + updated_at BIGINT NOT NULL, + CONSTRAINT async_jobs_token_idempotency_unique UNIQUE (token_id, idempotency_key), + CONSTRAINT fk_async_jobs_task FOREIGN KEY (task_id) REFERENCES tasks(id) ON UPDATE CASCADE ON DELETE CASCADE +); + +CREATE INDEX IF NOT EXISTS idx_async_jobs_status_lease ON async_jobs(execution_status, lease_until); +CREATE INDEX IF NOT EXISTS idx_async_jobs_channel_id ON async_jobs(channel_id); +CREATE INDEX IF NOT EXISTS idx_async_jobs_request_sent_at ON async_jobs(request_sent_at); +CREATE INDEX IF NOT EXISTS idx_async_jobs_billing_status ON async_jobs(billing_status); + +CREATE TABLE IF NOT EXISTS artifacts ( + id BIGSERIAL PRIMARY KEY, + task_id BIGINT NOT NULL, + object_key VARCHAR(512) NOT NULL UNIQUE, + content_type VARCHAR(128) NOT NULL, + size_bytes BIGINT NOT NULL, + sha256 CHAR(64) NOT NULL, + source_url_hash CHAR(64) NOT NULL, + created_at BIGINT NOT NULL, + expires_at BIGINT NOT NULL, + CONSTRAINT fk_artifacts_task FOREIGN KEY (task_id) REFERENCES tasks(id) ON UPDATE CASCADE ON DELETE CASCADE +); + +CREATE INDEX IF NOT EXISTS idx_artifacts_task_id ON artifacts(task_id); +CREATE INDEX IF NOT EXISTS idx_artifacts_expires_at ON artifacts(expires_at); + +CREATE TABLE IF NOT EXISTS task_events ( + id BIGSERIAL PRIMARY KEY, + task_id BIGINT NOT NULL, + event_type VARCHAR(40) NOT NULL, + from_status VARCHAR(20) NOT NULL DEFAULT '', + to_status VARCHAR(20) NOT NULL DEFAULT '', + worker_id VARCHAR(128) NOT NULL DEFAULT '', + error_phase VARCHAR(40) NOT NULL DEFAULT '', + error_code VARCHAR(80) NOT NULL DEFAULT '', + actor_type VARCHAR(20) NOT NULL DEFAULT '', + actor_id BIGINT NOT NULL DEFAULT 0, + details JSONB, + created_at BIGINT NOT NULL, + CONSTRAINT fk_task_events_task FOREIGN KEY (task_id) REFERENCES tasks(id) ON UPDATE CASCADE ON DELETE CASCADE +); + +CREATE INDEX IF NOT EXISTS idx_task_events_task_id ON task_events(task_id); +CREATE INDEX IF NOT EXISTS idx_task_events_event_type ON task_events(event_type); + +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conrelid = 'async_jobs'::regclass AND conname = 'fk_async_jobs_task') THEN + ALTER TABLE async_jobs + ADD CONSTRAINT fk_async_jobs_task + FOREIGN KEY (task_id) REFERENCES tasks(id) ON UPDATE CASCADE ON DELETE CASCADE; + END IF; + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conrelid = 'artifacts'::regclass AND conname = 'fk_artifacts_task') THEN + ALTER TABLE artifacts + ADD CONSTRAINT fk_artifacts_task + FOREIGN KEY (task_id) REFERENCES tasks(id) ON UPDATE CASCADE ON DELETE CASCADE; + END IF; + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conrelid = 'task_events'::regclass AND conname = 'fk_task_events_task') THEN + ALTER TABLE task_events + ADD CONSTRAINT fk_task_events_task + FOREIGN KEY (task_id) REFERENCES tasks(id) ON UPDATE CASCADE ON DELETE CASCADE; + END IF; +END +$$; + +COMMIT; diff --git a/docs/sub2api-integration.md b/docs/sub2api-integration.md new file mode 100644 index 000000000000..4d6cb5c074dd --- /dev/null +++ b/docs/sub2api-integration.md @@ -0,0 +1,108 @@ +# Sub2API 账号池集成 + +本集成将 Sub2API 作为独立的内部上游服务，New API 继续负责用户、令牌、计费和公网 API，Sub2API 只负责上游账号、OAuth 凭证、冷却与调度。两者通过兼容 API 连接，不复制或链接 Sub2API 源码。 + +```text +用户 -> New API -> Sub2API -> Claude / Codex / Gemini 账号 +``` + +## 前置条件 + +- Docker Compose v2。 +- New API 已通过 Docker Compose 启动。 +- 已阅读 Sub2API 的 LGPL-3.0-or-later 许可证、合规声明及各上游服务条款。 + +当前部署文件固定使用 Sub2API `0.1.160`，不使用会无提示升级的 `latest` 标签。 + +## 1. 准备配置 + +```bash +cp deploy/sub2api/.env.example deploy/sub2api/.env +chmod 600 deploy/sub2api/.env +``` + +编辑 `deploy/sub2api/.env`，替换所有 `replace_with_...` 占位值。随机密钥可以通过下列命令分别生成： + +```bash +openssl rand -hex 32 +``` + +`NEW_API_DOCKER_NETWORK` 必须是 New API 容器已加入的 Docker 网络。当使用 `deploy/compose.staging.yml` 时，默认值 `new-api-async-staging_backend` 无需修改。其他部署可通过下列命令查询： + +```bash +docker network ls +``` + +## 2. 启动账号池 + +```bash +docker compose \ + --env-file deploy/sub2api/.env \ + -f deploy/sub2api/compose.yml \ + up -d +``` + +检查健康状态： + +```bash +docker compose \ + --env-file deploy/sub2api/.env \ + -f deploy/sub2api/compose.yml \ + ps + +curl --fail http://127.0.0.1:38080/health +``` + +Sub2API 管理端默认只监听 `127.0.0.1:38080`。远程服务器上请使用 SSH 端口转发，不要直接暴露到公网。 + +本地账号测试台默认位于 `http://127.0.0.1:38081`。它通过同源反向代理调用 Sub2API，支持密钥验证、模型列表以及 OpenAI/Gemini 协议的真实生成测试。输入的 API Key 仅保留在当前页面内存中。 + +## 3. 配置 Sub2API + +1. 打开 `http://127.0.0.1:38080`，使用 `.env` 中的管理员账号登录。 +2. 添加你有权使用的上游账号或正式 API Key。 +3. 建立对应平台的分组，并将账号加入分组。 +4. 为 New API 生成一个 Sub2API 下游 API Key。 + +Sub2API 会根据 API Key 绑定分组的平台选择调度逻辑。不要用一个分组混合不相容的平台。建议为 Codex、Claude 和 Gemini 分别生成 Key，并在 New API 中分别建立渠道。 + +## 4. 在 New API 中新建渠道 + +在管理后台的渠道页面新建渠道： + +| 字段 | 值 | +| --- | --- | +| 类型 | `Advanced Custom` | +| API 地址 | `http://sub2api-account-pool:8080` | +| 密钥 | Sub2API 生成的下游 API Key | +| 路由模板 | `Sub2API Gateway` | +| 模型 | 仅填写该 Sub2API Key 所在分组实际可用的模型 | + +打开“高级自定义路由”，选择 `Sub2API Gateway` 后点击“填充模板”。模板会配置以下原生转发路由： + +- OpenAI Chat Completions、Responses、Responses Compact、Embeddings 和 Images。 +- Anthropic Messages。 +- Gemini `generateContent`、`embedContent` 和 `batchEmbedContents`。 + +路由模板统一使用 `Authorization: Bearer `，Sub2API 的 OpenAI、Anthropic 和 Gemini 鉴权中间件均支持此格式。 + +## 5. 验证 + +先直接验证 Sub2API Key： + +```bash +curl --fail \ + -H 'Authorization: Bearer REPLACE_WITH_SUB2API_KEY' \ + http://127.0.0.1:38080/v1/models +``` + +然后在 New API 后台对新渠道执行“测试”。测试成功后，再用 New API 令牌调用相同模型，确认流式输出、用量记录和扣费结果。 + +## 运维与安全 + +- Sub2API 使用 `simple` 模式，避免与 New API 重复执行用户余额计费；账号和分组限制仍由 Sub2API 调度层执行。 +- Sub2API 拥有独立的 PostgreSQL、Redis 和数据卷，不与 New API 共用表或缓存键空间。 +- 禁止在日志、工单或聊天中粘贴 OAuth Refresh Token、Cookie 或账号密码。 +- 使用会员订阅进行 API 中转可能受到上游条款、转售限制和当地监管影响；技术集成不等于获得上游授权。 + +备份以下命名卷：`new-api-sub2api_sub2api_data`、`new-api-sub2api_sub2api_postgres_data` 和 `new-api-sub2api_sub2api_redis_data`。升级时先阅读 Sub2API 发布说明，再修改 `.env` 中的 `SUB2API_IMAGE`。 diff --git a/docs/upstream-cost-cny.md b/docs/upstream-cost-cny.md new file mode 100644 index 000000000000..49eed8fb05ce --- /dev/null +++ b/docs/upstream-cost-cny.md @@ -0,0 +1,62 @@ +# 多上游人民币成本核算 + +本功能把不同渠道的上游原始成本统一换算为人民币，并为每次实际请求保存不可变成本快照。它只影响管理员成本核算，不改变用户钱包、积分扣除或销售价格。 + +## 渠道配置 + +在“渠道 → 高级设置 → 渠道额外设置”中，为每条上游线路配置： + +- 成本来源： + - `auto`：优先采用上游响应中的实际成本；没有时按网关基础计费单位估算。 + - `response_cost`：只接受上游返回的实际成本；缺失时标记为“未定价”。 + - `billing_units`：按用户分组倍率应用前的网关基础价格估算。 +- 上游原始计费单位：例如 `USD`、`CNY`、`CREDIT`。 +- 每单位人民币成本：人民币 ÷ 上游单位。 +- 价格版本：例如 `yunwu-2026-07`，用于历史对账。 + +云雾示例：100 CREDIT 实付 ¥49.5，因此设置： + +```text +成本来源：billing_units +上游原始计费单位：CREDIT +每单位人民币成本：0.495 +价格版本：yunwu-2026-07 +``` + +若一次请求的上游消耗是 0.011682 CREDIT，则人民币成本为： + +```text +0.011682 × 0.495 = ¥0.00578259 +``` + +账本以微元（1 元 = 1,000,000 微元）保存，本例记为 5,783 微元。 + +## 自动覆盖和准确性 + +文本、音频、实时、绘图和异步任务的消费路径都会生成成本快照。未配置成本资料的渠道不会被静默忽略，而会记录为 `unpriced`，原因是 `missing_channel_cost_profile`。 + +`response_cost` 是上游返回的权威值。`billing_units` 是估算值，其准确性依赖网关基础模型价格与该渠道上游价目表一致。如果同一模型在不同上游的输入、输出或按次价格不同，应优先使用能返回实际成本的上游；当前版本不能用一套全局模型价格准确表达多套不同的渠道价目表，因此不能把这种估算值当作上游账单。 + +## 验证和对账 + +管理员使用日志顶部会显示： + +- 上游成本：当前筛选条件下的人民币成本合计。 +- 未定价：`未定价请求数 / 已跟踪请求总数`。 + +单条日志的计费详情会显示原始成本、换算率、成本来源、是否估算及价格版本。 + +也可以调用管理员接口： + +```http +GET /api/log/upstream-cost/stat +``` + +接口支持 `start_timestamp`、`end_timestamp`、`model_name`、`username`、`token_name`、`group`、`channel`、`request_id` 和 `upstream_request_id` 筛选。 + +上线验收标准： + +1. 每条生产渠道均配置原始单位、换算率和价格版本。 +2. 各渠道分别发送一笔小额请求，日志人民币成本与手工计算一致。 +3. 按渠道和模型筛选时，未定价数为 `0`。 +4. 对 `billing_units` 估算渠道抽样与上游账单核对；出现差异时切换为实际成本模式或修正该渠道的价格口径。 diff --git "a/docs/\345\274\202\346\255\245\345\233\276\347\211\207\344\273\273\345\212\241\344\270\255\350\275\254\347\253\231\346\234\254\345\234\260\350\277\220\350\241\214\344\270\216\350\277\201\347\247\273\350\257\264\346\230\216.md" "b/docs/\345\274\202\346\255\245\345\233\276\347\211\207\344\273\273\345\212\241\344\270\255\350\275\254\347\253\231\346\234\254\345\234\260\350\277\220\350\241\214\344\270\216\350\277\201\347\247\273\350\257\264\346\230\216.md" new file mode 100644 index 000000000000..60580024af9d --- /dev/null +++ "b/docs/\345\274\202\346\255\245\345\233\276\347\211\207\344\273\273\345\212\241\344\270\255\350\275\254\347\253\231\346\234\254\345\234\260\350\277\220\350\241\214\344\270\216\350\277\201\347\247\273\350\257\264\346\230\216.md" @@ -0,0 +1,295 @@ +# 异步图片任务中转站本地运行与迁移说明 + +本说明适用于基于官方 `QuantumNous/new-api` `v1.0.0-rc.21`（提交 `bde9b2f44887d34ec54799ae191d50f97914359e`）开发的异步图片任务中转站。当前交付目标是本地 Docker Compose MVP，不包含服务器部署。 + +## 1. 首版边界 + +- 新增云雾同步图片生成接口的持久化异步包装。 +- 保留原有 `POST /v1/images/generations` 同步语义和原有路由。 +- 异步范围仅限图片生成；不扩展到聊天、语音、文件或视频。 +- 仅允许显式开启异步包装、配置模型白名单并启用归档的云雾渠道。 +- 首版不支持对已发送的上游请求做假取消，也不会自动重试 `UNCERTAIN` 任务。 + +## 2. 本地环境 + +需要： + +- Docker Desktop 或兼容的 Docker Engine。 +- Docker Compose v2。 +- 本地端口 `80`、`443`、`9000` 和 `9001` 未被占用。 +- 可访问容器镜像仓库；首次冷构建会下载 Go、Bun、Debian 和前端依赖。 + +Compose 会启动： + +- `new-api-api`：面板和 HTTP API。 +- `new-api-worker`：独立持久化队列 Worker。 +- `postgres`：主数据库和任务队列。 +- `redis`：New API 缓存。 +- `minio` 与 `minio-init`：私有对象存储和幂等建桶。 +- `caddy`：本地 HTTPS 入口。 + +PostgreSQL 和 Redis 只在 Compose 内部网络中暴露，不绑定宿主机端口。MinIO API 和控制台仅绑定 `127.0.0.1:9000-9001`。 + +## 3. 环境变量 + +1. 复制示例： + + ```bash + cp .env.example .env + ``` + +2. 生成独立随机值： + + ```bash + openssl rand -hex 32 + openssl rand -base64 32 + ``` + +3. 替换 `.env` 中的每一个 `REPLACE_*`。不要复用密码，不要提交 `.env`。 + +必填变量： + +| 变量 | 用途 | +| --- | --- | +| `POSTGRES_PASSWORD` | PostgreSQL 密码 | +| `REDIS_PASSWORD` | Redis 密码 | +| `SESSION_SECRET` | New API 会话签名，至少 32 个字符 | +| `CRYPTO_SECRET` | New API 原有敏感数据加密，不得与 `SESSION_SECRET` 相同 | +| `ASYNC_REQUEST_ENCRYPTION_KEY` | 异步请求体 AES-256-GCM 密钥，为 32 字节的 Base64 或 64 位十六进制 | +| `MINIO_ROOT_USER` | 本地 MinIO 访问键 | +| `MINIO_ROOT_PASSWORD` | 本地 MinIO 密钥 | + +常用可调参数： + +| 变量 | 默认值 | 说明 | +| --- | ---: | --- | +| `SESSION_COOKIE_TRUSTED_URL` | `https://localhost` | 本地安全 Cookie 的可信入口 | +| `S3_PUBLIC_ENDPOINT` | `http://localhost:9000` | 返回给客户端的预签名 URL 端点 | +| `ASYNC_YUNWU_ALLOWED_BASE_URLS` | `https://yunwu.ai` | 允许的云雾基地址，逗号分隔；路径仅允许空或 `/v1` | +| `ASYNC_GRSAI_ALLOWED_BASE_URLS` | `https://grsaiapi.com,https://grsai.dakka.com.cn` | 允许的 GRS AI API 基地址，逗号分隔；不允许控制台域名 | +| `ASYNC_WORKER_CONCURRENCY` | `50` | Worker 全局并发上限 | +| `ASYNC_CHANNEL_DEFAULT_CONCURRENCY` | `10` | 渠道未设置时的默认并发上限 | +| `ASYNC_WORKER_LEASE_SECONDS` | `90` | 任务租约时间，不得小于 15 秒 | +| `ASYNC_WORKER_POLL_MILLISECONDS` | `500` | 队列轮询间隔 | +| `ASYNC_JOB_TIMEOUT_SECONDS` | `1800` | 渠道未覆盖时的执行超时 | +| `ASYNC_RESULT_RETENTION_MINUTES` | `60` | 渠道未覆盖时的产物保留分钟数；最终值会限制在 5–1440 分钟 | +| `ASYNC_ARTIFACT_CLEANUP_INTERVAL_SECONDS` | `60` | Worker 扫描并删除到期产物的间隔；不得小于 10 秒 | +| `ASYNC_SIGNED_URL_TTL_SECONDS` | `900` | 下载预签名 URL 有效期 | +| `ASYNC_MAX_REQUEST_BODY_KB` | `256` | 提交请求体上限 | +| `ASYNC_MAX_PROMPT_CHARS` | `8000` | Prompt 字符上限 | +| `ASYNC_MAX_INPUT_URLS` | `8` | 输入 URL 数量上限 | +| `ASYNC_ARTIFACT_MAX_FILES` | `8` | 单任务归档文件数上限 | +| `ASYNC_ARTIFACT_MAX_FILE_MB` | `25` | 单文件大小上限 | +| `ASYNC_ARTIFACT_MAX_TOTAL_MB` | `100` | 单任务总归档上限 | +| `ASYNC_ARTIFACT_DOWNLOAD_TIMEOUT_SECONDS` | `120` | 远程媒体下载超时 | + +`BATCH_UPDATE_ENABLED` 在 Compose 中固定为 `false`。异步任务需要立即预扣、结算或退款，不支持 New API 的批量更新模式。 + +## 4. 启动和停止 + +校验配置并构建： + +```bash +docker compose config -q +docker compose build +``` + +启动： + +```bash +docker compose up -d +docker compose ps +``` + +健康检查： + +```bash +curl -k https://localhost/api/status +``` + +`-k` 只用于本地 Caddy 内部 CA。不得在生产环境以跳过证书校验代替有效 TLS 证书。 + +查看日志： + +```bash +docker compose logs -f new-api-api new-api-worker +``` + +停止但保留数据： + +```bash +docker compose down +``` + +不要在需要保留数据时使用 `docker compose down -v`；`-v` 会删除 PostgreSQL、Redis、MinIO 和 Caddy 的持久卷。 + +## 5. 首次面板配置 + +1. 访问 `https://localhost`，完成 New API 首次管理员初始化。 +2. 创建用于云雾或 GRS AI 图片生成的 OpenAI 兼容渠道。 +3. 云雾 Base URL 只填 `ASYNC_YUNWU_ALLOWED_BASE_URLS` 中的地址；GRS AI Base URL 只填 `ASYNC_GRSAI_ALLOWED_BASE_URLS` 中的 API 地址（默认推荐 `https://grsaiapi.com`），不要填写控制台域名。 +4. 填写真实上游密钥时只在面板或部署 Secret 中操作，不要写入代码、`.env.example` 或日志。 +5. 在渠道模型列表中加入图片模型，并在「渠道额外设置」中： + - 打开「异步图像中转」。 + - 在「允许的图像模型」填入逗号分隔的精确模型名。 + - 设置渠道最大并发数、任务超时和结果保留天数。 + - 保持「自动归档」开启；首版不接受关闭归档的渠道。 + - GRS AI Worker 固定调用 `POST /v1/api/generate` 并发送 `replyType=json`，不会调用 GRS AI 的异步查询接口。 +6. 为用户创建 New API Token，并确保 Token 的分组和模型限制允许该模型。 + +## 6. 异步 API + +提交： + +```bash +curl -k https://localhost/v1/async/images/generations \ + -H 'Authorization: Bearer REPLACE_WITH_NEW_API_TOKEN' \ + -H 'Idempotency-Key: project-request-0001' \ + -H 'Content-Type: application/json' \ + --data '{ + "model": "REPLACE_WITH_ALLOWED_IMAGE_MODEL", + "prompt": "一座未来城市", + "size": "1728x2304", + "response_format": "url" + }' +``` + +成功返回 HTTP `202`： + +```json +{ + "id": "task_xxx", + "status": "queued", + "status_url": "/v1/async/tasks/task_xxx", + "result_url": "/v1/async/tasks/task_xxx/result" +} +``` + +同一 Token 下重复提交相同 `Idempotency-Key` 和相同请求体会返回原任务，不重复扣费。同一键配合不同请求体返回 HTTP `409`。 + +查询： + +```bash +curl -k https://localhost/v1/async/tasks/task_xxx \ + -H 'Authorization: Bearer REPLACE_WITH_NEW_API_TOKEN' +``` + +获取结果： + +```bash +curl -k https://localhost/v1/async/tasks/task_xxx/result \ + -H 'Authorization: Bearer REPLACE_WITH_NEW_API_TOKEN' +``` + +取消未发送任务： + +```bash +curl -k -X POST https://localhost/v1/async/tasks/task_xxx/cancel \ + -H 'Authorization: Bearer REPLACE_WITH_NEW_API_TOKEN' +``` + +状态为 `queued`、`running`、`success`、`failure`、`uncertain` 或 `cancelled`。`running` 任务无上游取消接口时返回 `409`且不中断 Worker。`uncertain` 表示请求可能已在上游执行，禁止客户端自动重试。 + +原同步入口仍为： + +```http +POST /v1/images/generations +``` + +## 7. 面板运维 + +- 在用量日志的「任务」页面按任务类型 `Async image wrapper` 和状态筛选。 +- 列表显示排队/执行耗时、Worker、尝试次数、稳定错误码和计费状态。 +- 详情中可查看上游原始响应、归档产物、预签名下载链接和任务事件。 +- 用户或管理员可取消 `QUEUED` 任务。 +- 仅管理员可手动重试 `FAILURE` 或 `UNCERTAIN`；面板要求确认「可能重复生成」和「可能再次计费」两项风险。 + +## 8. 数据库迁移 + +### 8.1 全新 Compose 环境 + +API 主节点启动时会通过 GORM AutoMigrate 创建 `async_jobs`、`artifacts` 和 `task_events`，并创建三个指向 `tasks(id)` 的 `ON UPDATE CASCADE ON DELETE CASCADE` 外键。 + +### 8.2 现有 PostgreSQL 环境 + +迁移前先备份： + +```bash +docker compose exec -T postgres sh -c \ + 'pg_dump -U "$POSTGRES_USER" -d "$POSTGRES_DB" -Fc' \ + > newapi-before-async-migration.dump +``` + +停止 API 和 Worker 写入后执行： + +```bash +docker compose stop new-api-api new-api-worker +docker compose exec -T postgres sh -c \ + 'psql -U "$POSTGRES_USER" -d "$POSTGRES_DB" -v ON_ERROR_STOP=1' \ + < docs/migrations/001_async_image_jobs_postgresql.sql +docker compose up -d new-api-api new-api-worker caddy +``` + +迁移 SQL 可重复执行。校验： + +```bash +docker compose exec -T postgres sh -c \ + 'psql -U "$POSTGRES_USER" -d "$POSTGRES_DB" -c \ + "SELECT tablename FROM pg_tables WHERE schemaname='"'"'public'"'"' AND tablename IN ('"'"'async_jobs'"'"','"'"'artifacts'"'"','"'"'task_events'"'"') ORDER BY tablename;"' +``` + +### 8.3 回滚 + +优先做无数据丢失回滚：停止 API/Worker，备份 PostgreSQL，回滚应用镜像，保留新增表。旧版 New API 不会读写这三张表，保留表比立即删表安全。 + +只有在已确认备份可恢复、无需保留异步任务和归档元数据时，才可按 `task_events`、`artifacts`、`async_jobs` 的顺序删表。不要自动删除 `tasks` 中的历史记录。 + +## 9. 备份与恢复演练 + +创建备份： + +```bash +docker compose exec -T postgres sh -c \ + 'pg_dump -U "$POSTGRES_USER" -d "$POSTGRES_DB" -Fc' \ + > newapi.dump +``` + +恢复到独立校验库： + +```bash +docker compose exec -T postgres sh -c \ + 'createdb -U "$POSTGRES_USER" newapi_restore_check' +docker compose exec -T postgres sh -c \ + 'pg_restore -U "$POSTGRES_USER" -d newapi_restore_check' \ + < newapi.dump +docker compose exec -T postgres sh -c \ + 'psql -U "$POSTGRES_USER" -d newapi_restore_check -c "SELECT count(*) FROM async_jobs;"' +``` + +验证后再删除独立校验库。恢复演练必须在不覆盖主库的数据库上进行。 + +## 10. 已执行的测试（2026-07-17） + +| 项目 | 结果 | +| --- | --- | +| `go test ./...` | 通过，包含官方原包回归和新增包 | +| 异步核心单元/集成测试 | 通过：状态机、幂等、租约、崩溃恢复、取消、计费、`UNCERTAIN` 恰好一次结算/对账、手动重试、断线、优雅退出、并发限制和多图归档 | +| 云雾适配器故障注入 | 通过：延迟成功、`429`、`4xx`、`5xx`、读取超时进入 `UNCERTAIN`、禁止重定向 | +| URL/DNS/MIME/容量安全 | 通过：私网与特殊地址拦截、MIME 嗅探、请求体和产物数限制 | +| 真实 MinIO 路径 | 通过：上传、预签名下载、内容校验和删除 | +| 前端 | 通过：TypeScript、定向 Oxfmt、Oxlint `0 warnings / 0 errors`、默认面板生产构建 | +| 完整镜像冷构建 | 通过：默认面板、经典面板和 Go 二进制均在 Dockerfile 中成功 | +| Compose 健康和网络 | 通过：API/PostgreSQL/Redis/MinIO 健康，Worker/Caddy 运行，PostgreSQL/Redis 无宿主机端口 | +| HTTPS 与路由 | 通过：HTTP `301`，HTTPS `200`，新旧图片路由未授权均返回 `401` | +| API 重启与持久性 | 通过：API 单独重启时 Worker 容器 ID/启动时间不变，任务探针保留 | +| PostgreSQL 迁移 | 通过：三张表、三个外键，手工 SQL 幂等重复执行 | +| PostgreSQL 备份恢复 | 通过：`pg_dump -Fc` 管道恢复到独立数据库，三张异步表存在 | +| 密钥检查 | 通过：未写入本地测试凭据或常见真实密钥格式 | + +额外执行的全仓 `go vet ./...` 会报告官方 `v1.0.0-rc.21` 基线中的已有问题，主要是 `common/custom-event.go` 按值复制 `sync.Mutex` 和多个旧渠道适配器的不可达代码。本次新增的 `model`、`service`、`storage` 和 `relay/asyncwrap` 核心包定向 `go vet` 通过。这些基线 vet 问题未在首版中跨范围修改。 + +未使用真实云雾密钥做外部付费调用；云雾协议路径由本地模拟上游完成覆盖。这是为了遵守「不提交真实密钥」和「不擅自产生外部付费」的要求。 + +## 11. 上线前仍需要的外部确认 + +本地 MVP 不等于获得对外商用授权。对外提供服务前，必须另行确认 New API AGPLv3/商业许可义务、云雾转售授权、生产 TLS/域名、备份保留策略和实际服务器权限。 diff --git "a/docs/\345\274\202\346\255\245\345\233\276\347\211\207\347\224\237\346\210\220API\350\260\203\347\224\250\346\226\207\346\241\243.md" "b/docs/\345\274\202\346\255\245\345\233\276\347\211\207\347\224\237\346\210\220API\350\260\203\347\224\250\346\226\207\346\241\243.md" new file mode 100644 index 000000000000..a00145259b6b --- /dev/null +++ "b/docs/\345\274\202\346\255\245\345\233\276\347\211\207\347\224\237\346\210\220API\350\260\203\347\224\250\346\226\207\346\241\243.md" @@ -0,0 +1,598 @@ +# 异步图片生成 API 调用文档 + +## 1. 文档信息 + +- 服务名称：New API 异步图片中转站 +- 公网入口：`https://async-api.nexaapp.cn` +- API Base URL：`https://async-api.nexaapp.cn/v1` +- 当前结果保留时间：任务生成成功后约 60 分钟 +- 单次下载链接有效期：约 15 分钟，可在结果保留期内重新请求结果接口获取新链接 +- 更新时间：2026-08-13 + +本文档面向需要把图片生成接入桌面软件、后端服务或自动化工作流的开发人员。 + +> 这里的异步是本站提供的持久化任务能力。客户端提交后可以立即断开，云端 Worker 会继续调用上游同步生图接口、等待结果并归档图片。本站不会调用 GRS AI 的原生异步轮询接口。 + +## 2. 接入信息 + +### 2.1 应该填写的地址 + +如果软件要求填写 `Base URL`： + +```text +https://async-api.nexaapp.cn/v1 +``` + +提交任务的完整地址： + +```text +POST https://async-api.nexaapp.cn/v1/async/images/generations +``` + +### 2.2 应该使用的密钥 + +使用 New API 后台「API 密钥」页面中名为「异步生图调用」的密钥： + +```text +Authorization: Bearer +``` + +不要把云雾或 GRS AI 的上游密钥填写到客户端。上游密钥只保存在服务器渠道配置中。 + +### 2.3 与标准 OpenAI 图片接口的区别 + +本接口不是立即返回图片的标准同步接口。客户端必须实现以下流程： + +```text +提交任务 → 保存 task_id → 查询任务状态 → 获取结果 → 下载图片 +``` + +如果第三方软件只能固定调用 `/v1/images/generations`，并且要求同一个 HTTP 请求立即返回图片，则不能只靠修改 Base URL 接入本异步接口，需要为该软件增加异步任务适配逻辑。 + +## 3. 通用请求规则 + +所有接口都通过 HTTPS 调用，并携带 New API Token： + +```http +Authorization: Bearer +``` + +提交接口还必须携带： + +```http +Content-Type: application/json +Idempotency-Key: <本次业务请求的唯一键> +``` + +### 3.1 Idempotency-Key 幂等规则 + +- 必填，最长 191 字节。 +- 推荐使用 UUID、ULID，或者软件自身的订单号/任务号。 +- 客户端提交超时或网络断开后，必须使用原来的幂等键重试。 +- 同一个 Token、相同幂等键、相同 JSON 请求会返回原任务，不会重复创建或重复预扣费。 +- 同一个 Token、相同幂等键、不同请求内容会返回 HTTP `409`。 +- 一个任务进入 `failure`、`uncertain` 或 `cancelled` 后，如确实需要重新生成，应在确认风险后使用新的幂等键创建新任务。 + +示例： + +```text +flowpic-20260718-550e8400-e29b-41d4-a716-446655440000 +``` + +### 3.2 超时建议 + +- 提交接口只负责持久化入队，客户端 HTTP 超时建议设置为 15～30 秒。 +- 不要让提交请求一直等待图片生成完成。 +- 状态轮询建议每 2～5 秒一次；长任务可逐步增加到 10 秒一次。 +- 获取到终态后立即停止轮询。 + +## 4. 已配置模型 + +| 模型 | 推荐 `size` | 推荐 `quality` | `n` | +| --- | --- | --- | --- | +| `gemini-3.1-flash-image-preview` | `1:1`、`16:9`、`9:16` | `1K`、`2K`、`4K` | 必须为 `1` | +| `gemini-3-pro-image-preview` | `1:1`、`16:9`、`9:16` | `1K`、`2K`、`4K` | 必须为 `1` | +| `gpt-image-2-vip` | 30 个上游尺寸预设：1K（如 `1280x1280`）、2K（如 `2048x2048`）、4K（如 `2880x2880`） | 不要传 | 必须为 `1` | +| `nano-banana-pro` | `1:1`、`16:9`、`9:16`、`4:3`、`3:4` | `1K`、`2K`、`4K` | 必须为 `1` | +| `nano-banana-2` | `1:1`、`16:9`、`9:16`、`4:3`、`3:4` | `1K`、`2K`、`4K` | 必须为 `1` | + +为获得跨模型一致性，建议客户端默认始终发送： + +```json +"n": 1 +``` + +## 5. 提交异步图片任务 + +### 5.1 请求 + +```http +POST /v1/async/images/generations +Authorization: Bearer +Idempotency-Key: +Content-Type: application/json +``` + +请求体字段： + +| 字段 | 类型 | 必填 | 说明 | +| --- | --- | --- | --- | +| `model` | string | 是 | 必须是 Token 和异步渠道都允许的模型 | +| `prompt` | string | 是 | 图片描述；当前部署最大 8000 个 Unicode 字符 | +| `n` | integer | 建议 | 当前推荐固定为 `1` | +| `size` | string | 否 | 图片尺寸或宽高比，取值见模型表 | +| `quality` | string | 否 | 图片质量，取值见模型表 | +| `response_format` | string | 否 | 推荐使用 `url`；最终结果会统一返回本站归档 URL | +| `image` | string/array | 否 | 参考图 URL 或 Base64/Data URL，是否生效取决于模型 | +| `images` | string/array | 否 | 多张参考图，是否生效取决于模型 | +| `stream` | boolean | 否 | 只能省略或传 `false`，不支持流式图片响应 | + +当前部署限制： + +- 请求体最大 256 KiB。 +- 请求内最多包含 8 个 HTTP/HTTPS 输入 URL。 +- 单任务最多归档 8 个文件。 +- 单文件最大 25 MiB，单任务归档总量最大 100 MiB。 +- Base64 参考图会占用请求体大小，较大参考图建议使用公网 HTTPS URL。 + +### 5.2 GRS AI 示例 + +```bash +curl --request POST \ + 'https://async-api.nexaapp.cn/v1/async/images/generations' \ + --header 'Authorization: Bearer ' \ + --header 'Content-Type: application/json' \ + --header 'Idempotency-Key: 550e8400-e29b-41d4-a716-446655440000' \ + --data '{ + "model": "nano-banana-2", + "prompt": "一只戴着蓝色围巾的橘猫坐在窗边，柔和晨光，简洁背景", + "n": 1, + "size": "1:1", + "quality": "2K", + "response_format": "url" + }' +``` + +### 5.3 云雾 Gemini 示例 + +```bash +curl --request POST \ + 'https://async-api.nexaapp.cn/v1/async/images/generations' \ + --header 'Authorization: Bearer ' \ + --header 'Content-Type: application/json' \ + --header 'Idempotency-Key: project-a-gemini-0001' \ + --data '{ + "model": "gemini-3.1-flash-image-preview", + "prompt": "一座被云海环绕的未来城市", + "n": 1, + "size": "16:9", + "quality": "2K", + "response_format": "url" + }' +``` + +### 5.4 成功响应 + +HTTP 状态码：`202 Accepted` + +```json +{ + "id": "task_dhYtnB2XrXyOOCahrXipl5V7h7LULo0X", + "status": "queued", + "status_url": "/v1/async/tasks/task_dhYtnB2XrXyOOCahrXipl5V7h7LULo0X", + "result_url": "/v1/async/tasks/task_dhYtnB2XrXyOOCahrXipl5V7h7LULo0X/result" +} +``` + +`status_url` 和 `result_url` 是相对路径，需要与以下域名拼接： + +```text +https://async-api.nexaapp.cn +``` + +## 6. 查询任务状态 + +### 6.1 请求 + +```http +GET /v1/async/tasks/{task_id} +Authorization: Bearer +``` + +```bash +curl \ + 'https://async-api.nexaapp.cn/v1/async/tasks/task_xxx' \ + --header 'Authorization: Bearer ' +``` + +必须使用创建该任务时的同一个 Token 查询。其他 Token 查询会返回 `task_not_found`。 + +### 6.2 响应 + +```json +{ + "id": "task_xxx", + "status": "running", + "progress": 0, + "created_at": 1784351116, + "started_at": 1784351116, + "finished_at": null, + "error": null +} +``` + +时间字段均为 Unix 时间戳，单位为秒。 + +### 6.3 状态说明 + +| 状态 | 是否终态 | 客户端处理方式 | +| --- | --- | --- | +| `queued` | 否 | 继续轮询；任务已安全持久化 | +| `running` | 否 | 继续轮询；云端正在调用上游或归档产物 | +| `success` | 是 | 调用结果接口 | +| `failure` | 是 | 查看 `error`；不要无限自动重试 | +| `uncertain` | 是 | 上游可能已经执行，禁止自动重新生成 | +| `cancelled` | 是 | 任务已取消；如需重新生成请使用新幂等键 | + +失败状态示例： + +```json +{ + "id": "task_xxx", + "status": "failure", + "progress": 0, + "created_at": 1784351116, + "started_at": 1784351117, + "finished_at": 1784351122, + "error": { + "phase": "upstream_response", + "code": "upstream_rate_limited", + "message": "GRS AI rate limit retry budget was exhausted" + } +} +``` + +## 7. 获取任务结果 + +### 7.1 请求 + +```http +GET /v1/async/tasks/{task_id}/result +Authorization: Bearer +``` + +推荐在生产客户端添加 `include_upstream=false`，避免传输不需要的上游原始响应或内嵌图片数据： + +```bash +curl \ + 'https://async-api.nexaapp.cn/v1/async/tasks/task_xxx/result?include_upstream=false' \ + --header 'Authorization: Bearer ' +``` + +### 7.2 成功响应 + +```json +{ + "id": "task_xxx", + "status": "success", + "response": { + "data": [ + { + "url": "https://async-files.nexaapp.cn/new-api-staging-artifacts/async/task_xxx/00-example.jpg?..." + } + ] + }, + "artifacts": [ + { + "content_type": "image/jpeg", + "size_bytes": 802436, + "sha256": "50b20ce3294d85bb9985ba11493e017eb7aa6a9a803a1da7c7710123ca413fca", + "expires_at": 1784354716, + "url": "https://async-files.nexaapp.cn/new-api-staging-artifacts/async/task_xxx/00-example.jpg?..." + } + ] +} +``` + +如果没有传 `include_upstream=false`，响应还会包含： + +```json +"upstream_response": { + "status": "succeeded" +} +``` + +字段说明： + +| 字段 | 说明 | +| --- | --- | +| `response` | 统一后的 OpenAI 风格结果，优先读取 `response.data[].url` | +| `upstream_response` | 上游原始响应；可通过查询参数关闭 | +| `artifacts` | 本站对象存储中的归档文件元数据 | +| `artifacts[].url` | 短期签名下载地址 | +| `artifacts[].expires_at` | 归档产物永久删除时间，不是签名 URL 的失效时间 | +| `artifacts[].sha256` | 文件 SHA-256，可用于下载完整性校验 | + +### 7.3 下载注意事项 + +- 签名 URL 当前约 15 分钟后失效。 +- 在结果保留期内，可重新调用结果接口获取新的签名 URL。 +- 当前任务结果约保留 60 分钟，过期后返回 HTTP `410`。 +- 客户端应尽快把图片下载到自己的存储，不要长期保存签名 URL。 +- 下载时不要删除或重新编码 URL 中的查询参数。 +- 不需要给 `async-files.nexaapp.cn` 的下载请求添加 Authorization 头，签名已经包含临时授权。 + +## 8. 取消排队任务 + +### 8.1 请求 + +```http +POST /v1/async/tasks/{task_id}/cancel +Authorization: Bearer +``` + +```bash +curl --request POST \ + 'https://async-api.nexaapp.cn/v1/async/tasks/task_xxx/cancel' \ + --header 'Authorization: Bearer ' +``` + +规则： + +- `queued` 任务可以取消，并进入退款流程。 +- `running` 任务不能取消，因为同步上游没有取消接口；返回 HTTP `409`。 +- 已经进入终态的任务不会改变状态，接口返回当前状态。 + +## 9. 错误响应 + +异步接口业务错误通常采用以下格式： + +```json +{ + "error": { + "message": "Idempotency-Key header is required", + "type": "async_task_error", + "code": "idempotency_key_required" + } +} +``` + +Token 鉴权、模型路由或额度检查产生的错误可能使用 New API 通用错误类型 `new_api_error`，客户端应主要依据 HTTP 状态码和 `error.code`/`error.message` 处理。 + +### 9.1 常见 HTTP 状态码 + +| HTTP | 常见错误码 | 说明 | +| --- | --- | --- | +| `202` | — | 任务已入队或命中原幂等任务 | +| `400` | `idempotency_key_required`、`invalid_request`、`invalid_provider_request` | 请求字段、模型参数或幂等头错误 | +| `401` | `invalid_token` 或通用鉴权错误 | Token 缺失、错误或已失效 | +| `403` | 模型权限、用户状态或额度错误 | Token 无模型权限、账户不可用或额度不足 | +| `404` | `task_not_found` | 任务不存在，或不属于当前 Token | +| `409` | `idempotency_key_conflict` | 相同幂等键对应了不同请求 | +| `409` | `task_not_ready` | 任务仍在排队或运行 | +| `409` | `task_uncertain` | 任务执行结果无法确认，禁止自动重试 | +| `409` | `task_cancelled` | 任务已取消 | +| `409` | `upstream_cancel_unsupported` | 运行中的同步上游请求无法取消 | +| `410` | `result_expired` | 归档结果已过期并删除 | +| `422` | 任务的稳定错误码 | 任务已明确失败 | +| `429` | 速率限制错误 | 当前 Token 或模型请求过快 | +| `500` | `create_task_failed`、`query_task_failed` | 服务器内部错误 | +| `503` | `artifact_store_unavailable`、`artifact_sign_failed` | 对象存储或签名服务暂时不可用 | + +### 9.2 常见任务错误码 + +| 错误码 | 含义 | 建议 | +| --- | --- | --- | +| `upstream_rate_limited` | 上游 429 重试预算耗尽 | 延迟后使用新幂等键人工重试 | +| `upstream_http_400` 等 | 上游明确返回 HTTP 错误 | 检查提示词、尺寸、质量和账户状态 | +| `upstream_generation_failed` | 上游明确拒绝或生成失败 | 检查上游返回信息后决定是否重试 | +| `upstream_connect_failed` | 请求体发送前无法连接上游 | 通常可延迟重试 | +| `upstream_result_uncertain` | 请求可能已发送，但结果无法确认 | 禁止自动重试，避免重复生成和计费 | +| `upstream_sync_result_pending` | GRS 同步模式未返回最终结果 | 本站不会转为上游异步轮询，需人工处理 | +| `invalid_upstream_response` | 上游响应结构不符合预期 | 保留任务 ID 并交由管理员排查 | +| `upstream_response_too_large` | 上游响应超过安全上限 | 交由管理员检查模型响应 | +| `artifact_archive_failed` | 生成完成但归档失败 | 不要自动重复生成，交由管理员处理 | + +## 10. 完整 Python 示例 + +依赖： + +```bash +pip install requests +``` + +```python +import time +import uuid +from pathlib import Path + +import requests + +BASE_URL = "https://async-api.nexaapp.cn/v1" +API_TOKEN = "" + +session = requests.Session() +session.headers.update({"Authorization": f"Bearer {API_TOKEN}"}) + +payload = { + "model": "nano-banana-2", + "prompt": "一只戴着蓝色围巾的橘猫坐在窗边，柔和晨光", + "n": 1, + "size": "1:1", + "quality": "2K", + "response_format": "url", +} + +submit_response = session.post( + f"{BASE_URL}/async/images/generations", + headers={ + "Content-Type": "application/json", + "Idempotency-Key": str(uuid.uuid4()), + }, + json=payload, + timeout=30, +) +submit_response.raise_for_status() +task = submit_response.json() +task_id = task["id"] +print("task_id:", task_id) + +while True: + status_response = session.get( + f"{BASE_URL}/async/tasks/{task_id}", + timeout=15, + ) + status_response.raise_for_status() + status_data = status_response.json() + status = status_data["status"] + print("status:", status, "progress:", status_data["progress"]) + + if status == "success": + break + if status in {"failure", "uncertain", "cancelled"}: + raise RuntimeError(status_data) + + time.sleep(2) + +result_response = session.get( + f"{BASE_URL}/async/tasks/{task_id}/result", + params={"include_upstream": "false"}, + timeout=30, +) +result_response.raise_for_status() +result = result_response.json() + +for index, artifact in enumerate(result["artifacts"]): + image_response = requests.get(artifact["url"], timeout=120) + image_response.raise_for_status() + + suffix = ".png" if artifact["content_type"] == "image/png" else ".jpg" + output = Path(f"{task_id}-{index}{suffix}") + output.write_bytes(image_response.content) + print("saved:", output) +``` + +生产软件应把 `Idempotency-Key` 和 `task_id` 持久化到数据库，避免软件重启后丢失任务关联。 + +## 11. 完整 JavaScript/TypeScript 示例 + +适用于 Node.js 18 及以上版本： + +```javascript +import { randomUUID } from 'node:crypto' +import { writeFile } from 'node:fs/promises' + +const baseUrl = 'https://async-api.nexaapp.cn/v1' +const apiToken = '' + +const authHeaders = { + Authorization: `Bearer ${apiToken}`, +} + +const submitResponse = await fetch(`${baseUrl}/async/images/generations`, { + method: 'POST', + headers: { + ...authHeaders, + 'Content-Type': 'application/json', + 'Idempotency-Key': randomUUID(), + }, + body: JSON.stringify({ + model: 'nano-banana-2', + prompt: '一只戴着蓝色围巾的橘猫坐在窗边，柔和晨光', + n: 1, + size: '1:1', + quality: '2K', + response_format: 'url', + }), +}) + +if (!submitResponse.ok) { + throw new Error(await submitResponse.text()) +} + +const task = await submitResponse.json() +console.log('task_id:', task.id) + +let status +while (true) { + const statusResponse = await fetch( + `${baseUrl}/async/tasks/${task.id}`, + { headers: authHeaders }, + ) + + if (!statusResponse.ok) { + throw new Error(await statusResponse.text()) + } + + status = await statusResponse.json() + console.log('status:', status.status, 'progress:', status.progress) + + if (status.status === 'success') break + if (['failure', 'uncertain', 'cancelled'].includes(status.status)) { + throw new Error(JSON.stringify(status)) + } + + await new Promise((resolve) => setTimeout(resolve, 2000)) +} + +const resultResponse = await fetch( + `${baseUrl}/async/tasks/${task.id}/result?include_upstream=false`, + { headers: authHeaders }, +) + +if (!resultResponse.ok) { + throw new Error(await resultResponse.text()) +} + +const result = await resultResponse.json() +const imageResponse = await fetch(result.artifacts[0].url) + +if (!imageResponse.ok) { + throw new Error(`image download failed: ${imageResponse.status}`) +} + +await writeFile( + `${task.id}.jpg`, + Buffer.from(await imageResponse.arrayBuffer()), +) +``` + +## 12. 推荐的客户端状态保存 + +客户端至少保存以下字段： + +| 字段 | 用途 | +| --- | --- | +| `local_request_id` | 软件自己的业务任务 ID | +| `idempotency_key` | 提交超时或断线时安全重试 | +| `remote_task_id` | 本站返回的 `task_id` | +| `status` | 最近一次查询状态 | +| `last_polled_at` | 控制轮询频率 | +| `result_downloaded_at` | 判断是否已保存到自己的存储 | +| `error_code`、`error_message` | 失败诊断和人工处理 | + +建议的软件恢复流程： + +1. 软件启动时加载所有非终态任务。 +2. 已有 `remote_task_id` 的任务只查询状态，不重新提交。 +3. 只有提交请求超时且尚未获得 `remote_task_id` 时，才使用原 `idempotency_key` 重新提交。 +4. `uncertain` 任务进入人工处理队列，禁止自动重新生成。 +5. `success` 后立即下载归档图片并保存到软件自己的长期存储。 + +## 13. 安全要求 + +- 不要把 New API Token 写入前端网页、公开仓库、截图或日志。 +- 桌面软件应使用系统安全凭据存储；服务端应使用环境变量或 Secret Manager。 +- 不要把 Token 放在 URL 查询参数中。 +- 不要向客户端分发云雾或 GRS AI 上游密钥。 +- 日志中可以记录 `task_id` 和幂等键，但应脱敏 Authorization 请求头和签名下载 URL。 +- 如果 Token 泄露，应立即在 New API 后台禁用并重新创建。 + +## 14. 当前不支持的能力 + +- 不支持 Webhook/回调通知，客户端需要轮询。 +- 不支持流式图片响应。 +- 不支持取消已经进入 `running` 的同步上游请求。 +- 不支持使用标准 OpenAI 图片 SDK 方法自动完成异步轮询。 +- 不保证长期保存图片；客户端必须在保留期内下载。 +- `uncertain` 状态不会自动向上游重试，以避免重复生成和重复计费。 diff --git a/dto/async_job.go b/dto/async_job.go new file mode 100644 index 000000000000..f672aa1ca53c --- /dev/null +++ b/dto/async_job.go @@ -0,0 +1,42 @@ +package dto + +import "encoding/json" + +type AsyncSubmitResponse struct { + ID string `json:"id"` + Status string `json:"status"` + StatusURL string `json:"status_url"` + ResultURL string `json:"result_url"` +} + +type AsyncTaskError struct { + Phase string `json:"phase"` + Code string `json:"code"` + Message string `json:"message"` +} + +type AsyncTaskStatusResponse struct { + ID string `json:"id"` + Status string `json:"status"` + Progress int `json:"progress"` + CreatedAt int64 `json:"created_at"` + StartedAt *int64 `json:"started_at"` + FinishedAt *int64 `json:"finished_at"` + Error *AsyncTaskError `json:"error"` +} + +type AsyncArtifactResponse struct { + ContentType string `json:"content_type"` + SizeBytes int64 `json:"size_bytes"` + SHA256 string `json:"sha256"` + ExpiresAt int64 `json:"expires_at"` + URL string `json:"url"` +} + +type AsyncTaskResultResponse struct { + ID string `json:"id"` + Status string `json:"status"` + Response json.RawMessage `json:"response"` + UpstreamResponse json.RawMessage `json:"upstream_response,omitempty"` + Artifacts []AsyncArtifactResponse `json:"artifacts"` +} diff --git a/dto/channel_settings.go b/dto/channel_settings.go index c92a3f988a3a..209c464ede00 100644 --- a/dto/channel_settings.go +++ b/dto/channel_settings.go @@ -11,12 +11,68 @@ import ( ) type ChannelSettings struct { - ForceFormat bool `json:"force_format,omitempty"` - ThinkingToContent bool `json:"thinking_to_content,omitempty"` - Proxy string `json:"proxy"` - PassThroughBodyEnabled bool `json:"pass_through_body_enabled,omitempty"` - SystemPrompt string `json:"system_prompt,omitempty"` - SystemPromptOverride bool `json:"system_prompt_override,omitempty"` + ForceFormat bool `json:"force_format,omitempty"` + ThinkingToContent bool `json:"thinking_to_content,omitempty"` + Proxy string `json:"proxy"` + PassThroughBodyEnabled bool `json:"pass_through_body_enabled,omitempty"` + SystemPrompt string `json:"system_prompt,omitempty"` + SystemPromptOverride bool `json:"system_prompt_override,omitempty"` + AsyncImageEnabled bool `json:"async_image_enabled,omitempty"` + AsyncImageModels []string `json:"async_image_models,omitempty"` + AsyncMaxConcurrency int `json:"async_max_concurrency,omitempty"` + AsyncJobTimeoutSeconds int `json:"async_job_timeout_seconds,omitempty"` + AsyncRetentionMinutes int `json:"async_retention_minutes,omitempty"` + // AsyncRetentionDays is retained only for channels saved before minute-level retention was introduced. + AsyncRetentionDays int `json:"async_retention_days,omitempty"` + AsyncAutoArchive *bool `json:"async_auto_archive,omitempty"` +} + +const ( + AsyncRetentionMinMinutes = 5 + AsyncRetentionMaxMinutes = 24 * 60 + AsyncRetentionDefaultMinutes = 60 +) + +func NormalizeAsyncRetentionMinutes(minutes int) int { + if minutes < AsyncRetentionMinMinutes { + return AsyncRetentionMinMinutes + } + if minutes > AsyncRetentionMaxMinutes { + return AsyncRetentionMaxMinutes + } + return minutes +} + +func (s ChannelSettings) EffectiveAsyncRetentionMinutes(fallback int) int { + minutes := s.AsyncRetentionMinutes + if minutes <= 0 && s.AsyncRetentionDays > 0 { + // The old setting was measured in whole days, so every valid legacy value + // is at least the new 24-hour upper bound. + minutes = AsyncRetentionMaxMinutes + } + if minutes <= 0 { + minutes = fallback + } + if minutes <= 0 { + minutes = AsyncRetentionDefaultMinutes + } + return NormalizeAsyncRetentionMinutes(minutes) +} + +func (s ChannelSettings) AllowsAsyncImageModel(model string) bool { + if !s.AsyncImageEnabled || strings.TrimSpace(model) == "" { + return false + } + for _, allowed := range s.AsyncImageModels { + if strings.TrimSpace(allowed) == model { + return true + } + } + return false +} + +func (s ChannelSettings) AsyncArchiveEnabled() bool { + return s.AsyncAutoArchive == nil || *s.AsyncAutoArchive } type VertexKeyType string @@ -37,14 +93,18 @@ type ChannelOtherSettings struct { AzureResponsesVersion string `json:"azure_responses_version,omitempty"` VertexKeyType VertexKeyType `json:"vertex_key_type,omitempty"` // "json" or "api_key" OpenRouterEnterprise *bool `json:"openrouter_enterprise,omitempty"` - ClaudeBetaQuery bool `json:"claude_beta_query,omitempty"` // Claude 渠道是否强制追加 ?beta=true - AllowServiceTier bool `json:"allow_service_tier,omitempty"` // 是否允许 service_tier 透传（默认过滤以避免额外计费） - AllowInferenceGeo bool `json:"allow_inference_geo,omitempty"` // 是否允许 inference_geo 透传（仅 Claude，默认过滤以满足数据驻留合规 - AllowSpeed bool `json:"allow_speed,omitempty"` // 是否允许 speed 透传（仅 Claude，默认过滤以避免意外切换推理速度模式） - AllowSafetyIdentifier bool `json:"allow_safety_identifier,omitempty"` // 是否允许 safety_identifier 透传（默认过滤以保护用户隐私） - DisableStore bool `json:"disable_store,omitempty"` // 是否禁用 store 透传（默认允许透传，禁用后可能导致 Codex 无法使用） - AllowIncludeObfuscation bool `json:"allow_include_obfuscation,omitempty"` // 是否允许 stream_options.include_obfuscation 透传（默认过滤以避免关闭流混淆保护） - DisableTaskPollingSleep bool `json:"disable_task_polling_sleep,omitempty"` // 是否跳过异步任务轮询间隔 + UpstreamCostMode string `json:"upstream_cost_mode,omitempty"` // auto / response_cost / billing_units + UpstreamCostUnit string `json:"upstream_cost_unit,omitempty"` // 上游原始计费单位，例如 USD、CNY、CREDIT + UpstreamCostRateCNY *float64 `json:"upstream_cost_rate_cny,omitempty"` // 人民币/上游计费单位；仅用于管理员成本核算 + UpstreamCostPriceVersion string `json:"upstream_cost_price_version,omitempty"` // 人工价格版本，写入每次成本快照 + ClaudeBetaQuery bool `json:"claude_beta_query,omitempty"` // Claude 渠道是否强制追加 ?beta=true + AllowServiceTier bool `json:"allow_service_tier,omitempty"` // 是否允许 service_tier 透传（默认过滤以避免额外计费） + AllowInferenceGeo bool `json:"allow_inference_geo,omitempty"` // 是否允许 inference_geo 透传（仅 Claude，默认过滤以满足数据驻留合规 + AllowSpeed bool `json:"allow_speed,omitempty"` // 是否允许 speed 透传（仅 Claude，默认过滤以避免意外切换推理速度模式） + AllowSafetyIdentifier bool `json:"allow_safety_identifier,omitempty"` // 是否允许 safety_identifier 透传（默认过滤以保护用户隐私） + DisableStore bool `json:"disable_store,omitempty"` // 是否禁用 store 透传（默认允许透传，禁用后可能导致 Codex 无法使用） + AllowIncludeObfuscation bool `json:"allow_include_obfuscation,omitempty"` // 是否允许 stream_options.include_obfuscation 透传（默认过滤以避免关闭流混淆保护） + DisableTaskPollingSleep bool `json:"disable_task_polling_sleep,omitempty"` // 是否跳过异步任务轮询间隔 AwsKeyType AwsKeyType `json:"aws_key_type,omitempty"` UpstreamModelUpdateCheckEnabled bool `json:"upstream_model_update_check_enabled,omitempty"` // 是否检测上游模型更新 UpstreamModelUpdateAutoSyncEnabled bool `json:"upstream_model_update_auto_sync_enabled,omitempty"` // 是否自动同步上游模型更新 @@ -55,6 +115,16 @@ type ChannelOtherSettings struct { AdvancedCustom *AdvancedCustomConfig `json:"advanced_custom,omitempty"` } +const ( + UpstreamCostModeAuto = "auto" + UpstreamCostModeResponseCost = "response_cost" + UpstreamCostModeBillingUnits = "billing_units" + + MaxUpstreamCostRateCNY = 1_000_000 + MaxUpstreamCostUnitLength = 32 + MaxUpstreamPriceVersionSize = 64 +) + func (s *ChannelOtherSettings) IsOpenRouterEnterprise() bool { if s == nil || s.OpenRouterEnterprise == nil { return false diff --git a/dto/channel_settings_test.go b/dto/channel_settings_test.go index 080863bedf5e..fe2949a73a07 100644 --- a/dto/channel_settings_test.go +++ b/dto/channel_settings_test.go @@ -477,3 +477,27 @@ func TestAdvancedCustomSupportedEndpointTypesForModel(t *testing.T) { constant.EndpointTypeAnthropic, }, config.SupportedEndpointTypesForModel("other-model")) } + +func TestEffectiveAsyncRetentionMinutes(t *testing.T) { + tests := []struct { + name string + setting ChannelSettings + fallback int + want int + }{ + {name: "configured minimum", setting: ChannelSettings{AsyncRetentionMinutes: 5}, fallback: 60, want: 5}, + {name: "configured maximum", setting: ChannelSettings{AsyncRetentionMinutes: 1440}, fallback: 60, want: 1440}, + {name: "configured value below minimum is bounded", setting: ChannelSettings{AsyncRetentionMinutes: 1}, fallback: 60, want: 5}, + {name: "configured value above maximum is bounded", setting: ChannelSettings{AsyncRetentionMinutes: 2000}, fallback: 60, want: 1440}, + {name: "legacy one day setting", setting: ChannelSettings{AsyncRetentionDays: 1}, fallback: 60, want: 1440}, + {name: "legacy multi-day setting is bounded", setting: ChannelSettings{AsyncRetentionDays: 30}, fallback: 60, want: 1440}, + {name: "environment fallback", setting: ChannelSettings{}, fallback: 90, want: 90}, + {name: "invalid fallback uses default", setting: ChannelSettings{}, fallback: 0, want: AsyncRetentionDefaultMinutes}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + assert.Equal(t, tt.want, tt.setting.EffectiveAsyncRetentionMinutes(tt.fallback)) + }) + } +} diff --git a/dto/openai_image.go b/dto/openai_image.go index 275fd5559080..012a84cc7b2e 100644 --- a/dto/openai_image.go +++ b/dto/openai_image.go @@ -6,6 +6,7 @@ import ( "strings" "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/setting/model_setting" "github.com/QuantumNous/new-api/types" "github.com/gin-gonic/gin" @@ -153,6 +154,8 @@ func (i *ImageRequest) GetTokenCountMeta() *types.TokenCountMeta { qualityRatio = 1.5 } } + } else { + qualityRatio = model_setting.GetImageGenerationPriceMultiplier(i.Model, i.Quality, i.Size) } imageN := uint(1) diff --git a/dto/openai_request.go b/dto/openai_request.go index 3bb2b34c6456..5374aaafc6e2 100644 --- a/dto/openai_request.go +++ b/dto/openai_request.go @@ -32,6 +32,7 @@ type GeneralOpenAIRequest struct { Prompt any `json:"prompt,omitempty"` Prefix any `json:"prefix,omitempty"` Suffix any `json:"suffix,omitempty"` + IncludeBilling *bool `json:"include_billing,omitempty"` Stream *bool `json:"stream,omitempty"` StreamOptions *StreamOptions `json:"stream_options,omitempty"` MaxTokens *uint `json:"max_tokens,omitempty"` @@ -838,9 +839,10 @@ type WebSearchOptions struct { // https://platform.openai.com/docs/api-reference/responses/create type OpenAIResponsesRequest struct { - Model string `json:"model"` - Input json.RawMessage `json:"input,omitempty"` - Include json.RawMessage `json:"include,omitempty"` + Model string `json:"model"` + Input json.RawMessage `json:"input,omitempty"` + Include json.RawMessage `json:"include,omitempty"` + IncludeBilling *bool `json:"include_billing,omitempty"` // 在后台运行推理，暂时还不支持依赖的接口 // Background json.RawMessage `json:"background,omitempty"` Conversation json.RawMessage `json:"conversation,omitempty"` diff --git a/dto/openai_response.go b/dto/openai_response.go index 2de6014f4d05..fecd99dd2859 100644 --- a/dto/openai_response.go +++ b/dto/openai_response.go @@ -221,13 +221,14 @@ type CompletionsStreamResponse struct { } type Usage struct { - PromptTokens int `json:"prompt_tokens"` - CompletionTokens int `json:"completion_tokens"` - TotalTokens int `json:"total_tokens"` - PromptCacheHitTokens int `json:"prompt_cache_hit_tokens,omitempty"` - UsageSemantic string `json:"usage_semantic,omitempty"` - UsageSource string `json:"usage_source,omitempty"` - BillingUsage *BillingUsage `json:"billing_usage,omitempty"` + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + TotalTokens int `json:"total_tokens"` + PromptCacheHitTokens int `json:"prompt_cache_hit_tokens,omitempty"` + UsageSemantic string `json:"usage_semantic,omitempty"` + UsageSource string `json:"usage_source,omitempty"` + BillingUsage *BillingUsage `json:"billing_usage,omitempty"` + Billing *ResponseBilling `json:"billing,omitempty"` PromptTokensDetails InputTokenDetails `json:"prompt_tokens_details"` CompletionTokenDetails OutputTokenDetails `json:"completion_tokens_details"` diff --git a/dto/response_billing.go b/dto/response_billing.go new file mode 100644 index 000000000000..90f685f9bd2d --- /dev/null +++ b/dto/response_billing.go @@ -0,0 +1,16 @@ +package dto + +// ResponseBilling is the gateway-calculated charge for one completed request. +// Unit prices are present for standard token billing. Dynamic expression and +// fixed-per-request billing return the authoritative total cost instead. +type ResponseBilling struct { + Currency string `json:"currency"` + TotalCost float64 `json:"total_cost"` + BillingMode string `json:"billing_mode"` + BillingSource string `json:"billing_source,omitempty"` + GroupRatio float64 `json:"group_ratio"` + InputUnitPricePerMillion *float64 `json:"input_unit_price_per_million,omitempty"` + OutputUnitPricePerMillion *float64 `json:"output_unit_price_per_million,omitempty"` + RequestPrice *float64 `json:"request_price,omitempty"` + MatchedTier string `json:"matched_tier,omitempty"` +} diff --git a/dto/task.go b/dto/task.go index 4a9a8e2e6d18..bb583e388f7d 100644 --- a/dto/task.go +++ b/dto/task.go @@ -50,6 +50,20 @@ type TaskDto struct { Properties any `json:"properties"` Username string `json:"username,omitempty"` Data json.RawMessage `json:"data"` + Async *AsyncTaskMeta `json:"async,omitempty"` +} + +// AsyncTaskMeta contains operational metadata exposed only through the +// authenticated task-management APIs. The encrypted request body and channel +// credentials are deliberately excluded. +type AsyncTaskMeta struct { + ExecutionStatus string `json:"execution_status"` + WorkerID string `json:"worker_id,omitempty"` + Attempt int `json:"attempt"` + RequestSentAt int64 `json:"request_sent_at,omitempty"` + ErrorPhase string `json:"error_phase,omitempty"` + ErrorCode string `json:"error_code,omitempty"` + BillingStatus string `json:"billing_status"` } type FetchReq struct { diff --git a/dto/upstream_cost.go b/dto/upstream_cost.go new file mode 100644 index 000000000000..52529e4d29fe --- /dev/null +++ b/dto/upstream_cost.go @@ -0,0 +1,30 @@ +package dto + +const ( + UpstreamCostStatusSettled = "settled" + UpstreamCostStatusUnpriced = "unpriced" + + UpstreamCostSourceResponseCost = "response_cost" + UpstreamCostSourceBillingUnits = "billing_units" +) + +// UpstreamCostSnapshot is the immutable, admin-only cost result for one +// completed request. AmountCNYMicros is the accounting value; float fields are +// retained for log display and compatibility with existing clients. +type UpstreamCostSnapshot struct { + Status string `json:"status"` + Mode string `json:"mode"` + Source string `json:"source,omitempty"` + Reason string `json:"reason,omitempty"` + NativeUnit string `json:"native_unit"` + NativeAmount float64 `json:"native_amount,omitempty"` + NativeAmountDecimal string `json:"native_amount_decimal,omitempty"` + Units float64 `json:"units,omitempty"` + RateCNYPerUnit float64 `json:"rate_cny_per_unit"` + RateCNYPerUnitDecimal string `json:"rate_cny_per_unit_decimal,omitempty"` + AmountCNY float64 `json:"amount_cny,omitempty"` + AmountCNYMicros int64 `json:"amount_cny_micros,omitempty"` + Estimated bool `json:"estimated"` + PriceVersion string `json:"price_version,omitempty"` + SettlementCurrency string `json:"settlement_currency"` +} diff --git a/go.mod b/go.mod index 98f291f64510..2919e53c38a4 100644 --- a/go.mod +++ b/go.mod @@ -8,10 +8,10 @@ require ( github.com/abema/go-mp4 v1.4.1 github.com/andybalholm/brotli v1.1.1 github.com/anknown/ahocorasick v0.0.0-20190904063843-d75dbd5169c0 - github.com/aws/aws-sdk-go-v2 v1.41.5 - github.com/aws/aws-sdk-go-v2/credentials v1.19.10 + github.com/aws/aws-sdk-go-v2 v1.42.1 + github.com/aws/aws-sdk-go-v2/credentials v1.19.29 github.com/aws/aws-sdk-go-v2/service/bedrockruntime v1.50.4 - github.com/aws/smithy-go v1.24.2 + github.com/aws/smithy-go v1.27.3 github.com/bytedance/gopkg v0.1.3 github.com/casbin/casbin/v2 v2.135.0 github.com/gin-contrib/cors v1.7.2 @@ -68,6 +68,18 @@ require ( require ( github.com/ClickHouse/ch-go v0.65.0 // indirect github.com/ClickHouse/clickhouse-go/v2 v2.32.0 // indirect + github.com/aws/aws-sdk-go-v2/config v1.32.30 // indirect + github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.30 // indirect + github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.31 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.13 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.9.23 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.30 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.31 // indirect + github.com/aws/aws-sdk-go-v2/service/s3 v1.105.2 // indirect + github.com/aws/aws-sdk-go-v2/service/signin v1.4.1 // indirect + github.com/aws/aws-sdk-go-v2/service/sso v1.32.1 // indirect + github.com/aws/aws-sdk-go-v2/service/ssooidc v1.37.1 // indirect + github.com/aws/aws-sdk-go-v2/service/sts v1.44.1 // indirect github.com/bmatcuk/doublestar/v4 v4.6.1 // indirect github.com/casbin/govaluate v1.10.0 // indirect github.com/go-faster/city v1.0.1 // indirect @@ -90,9 +102,9 @@ require ( require ( github.com/DmitriyVTitov/size v1.5.0 // indirect github.com/anknown/darts v0.0.0-20151216065714-83ff685239e6 // indirect - github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.8 // indirect - github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.21 // indirect - github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.21 // indirect + github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.14 // indirect + github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.30 // indirect + github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.30 // indirect github.com/beorn7/perks v1.0.1 // indirect github.com/boombuler/barcode v1.1.0 // indirect github.com/bytedance/sonic v1.14.1 // indirect diff --git a/go.sum b/go.sum index ca6bd091f831..964f9ebc514d 100644 --- a/go.sum +++ b/go.sum @@ -724,18 +724,54 @@ github.com/aws/aws-sdk-go v1.15.11/go.mod h1:mFuSZ37Z9YOHbQEwBWztmVzqXrEkub65tZo github.com/aws/aws-sdk-go v1.43.16/go.mod h1:y4AeaBuwd2Lk+GepC1E9v0qOiTws0MIWAX4oIKwKHZo= github.com/aws/aws-sdk-go-v2 v1.41.5 h1:dj5kopbwUsVUVFgO4Fi5BIT3t4WyqIDjGKCangnV/yY= github.com/aws/aws-sdk-go-v2 v1.41.5/go.mod h1:mwsPRE8ceUUpiTgF7QmQIJ7lgsKUPQOUl3o72QBrE1o= +github.com/aws/aws-sdk-go-v2 v1.42.1 h1:9eOTgu1z/dVtYpNZ3/8/XbbaX0x/BqE3HUzAzs6K0ek= +github.com/aws/aws-sdk-go-v2 v1.42.1/go.mod h1:5pKeft2eJj+gElQ38Jqg4ibCqh+/AK33/0X3hip7IjM= github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.8 h1:eBMB84YGghSocM7PsjmmPffTa+1FBUeNvGvFou6V/4o= github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.8/go.mod h1:lyw7GFp3qENLh7kwzf7iMzAxDn+NzjXEAGjKS2UOKqI= +github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.14 h1:3IZY0XAJquT3aHzbkHfPzy4ACPcEjVG0x87KOwtpqGY= +github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.14/go.mod h1:zwM6veDkhGgQFqkBy+uT28AAYpLu+uFMlPl+rCg/73E= +github.com/aws/aws-sdk-go-v2/config v1.32.30 h1:XwsEzpTJfQYJbFicz/QMLwAZdyeNVVoOEkbF7R3gPJk= +github.com/aws/aws-sdk-go-v2/config v1.32.30/go.mod h1:Ud32SuMc+/9BGxfpSVld7HrE2o05JwKmXY4M3jOQNZU= github.com/aws/aws-sdk-go-v2/credentials v1.19.10 h1:EEhmEUFCE1Yhl7vDhNOI5OCL/iKMdkkYFTRpZXNw7m8= github.com/aws/aws-sdk-go-v2/credentials v1.19.10/go.mod h1:RnnlFCAlxQCkN2Q379B67USkBMu1PipEEiibzYN5UTE= +github.com/aws/aws-sdk-go-v2/credentials v1.19.29 h1:WHZGssHH887cO0ox07SIQZsFx3MKD4ps6w0xUEmnKYQ= +github.com/aws/aws-sdk-go-v2/credentials v1.19.29/go.mod h1:Mhl0xR6zjguiuj00XRx2wMx22sAltk7oya39sT7fdg8= +github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.30 h1:/hi1JADLEW9YYryEz1w4GQu0EtP23pP553Cf9KgsDV4= +github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.30/go.mod h1:/3AOgy4K17Dm4ucMZVC/MJkzy5kmfKUcINRHZyo0koQ= github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.21 h1:Rgg6wvjjtX8bNHcvi9OnXWwcE0a2vGpbwmtICOsvcf4= github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.21/go.mod h1:A/kJFst/nm//cyqonihbdpQZwiUhhzpqTsdbhDdRF9c= +github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.30 h1:xM/Is9cKMHa8Jj8zkvWhvrFkZsXJV9E+BB4g0HW0duQ= +github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.30/go.mod h1:WueJeNDZvK1fMYEWJIkcivBfEzUkTpBhzlrUKKY8EuA= github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.21 h1:PEgGVtPoB6NTpPrBgqSE5hE/o47Ij9qk/SEZFbUOe9A= github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.21/go.mod h1:p+hz+PRAYlY3zcpJhPwXlLC4C+kqn70WIHwnzAfs6ps= +github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.30 h1:jn46zC9LdsVR/ZpMIJqMqb8hHv31BlLx3ulVqNspUOk= +github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.30/go.mod h1:1hTMsAgbdS/AtUi4bw8+gUuh1pceo+eXRLfpSuSQj3M= +github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.31 h1:3GUprIsfmGcC5SACIyB0e7E0BM1O1b3Erl5CePYIAeQ= +github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.31/go.mod h1:7PuV1yl5e2xnUbm+RqvVg5i2iBM8EyijZNoI9wsOoOc= github.com/aws/aws-sdk-go-v2/service/bedrockruntime v1.50.4 h1:W6tKfa/s37faUnwJ71pGqsBO7/wfUX1L7tVprupQGo4= github.com/aws/aws-sdk-go-v2/service/bedrockruntime v1.50.4/go.mod h1:BZ+9thH0QOTDUwE8KAv/ZwUzsNC7CSMJXj/wtnZMs5k= +github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.13 h1:mbRIur/BiHK6SKPjoBIXSE/hJ6g6JGRLuxQy1jGjlN4= +github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.13/go.mod h1:ITg9em2KbJx1s0y4aqRX5OYWG6HBZ5TVR//OdpEZ2CQ= +github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.9.23 h1:9Fjh6fi/U5JEStVZijmaMpUwE/gvBJj7x2B/PjbO9To= +github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.9.23/go.mod h1:iMoT2f1tClxrWAAnKCXjZQ6LOmfLrMG14wmnWpM+F14= +github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.30 h1:/Z5jmNrKsSD7EmDjzAPsm/3L9IuOkzaynklJZ1qX7S4= +github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.30/go.mod h1:lEzEZnOosE7zi8Z6royW1cFJTD9fpab4Ul1SBrllewk= +github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.31 h1:uao4A3QZ5UmB326V6KF+qRpv9Tjz7IlnlnTbbANntlU= +github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.31/go.mod h1:I/1+z0VwL1GhQyLgkoHDlygpUZ+iTAwOQ/NsftiUL2I= +github.com/aws/aws-sdk-go-v2/service/s3 v1.105.2 h1:5C00eQYpTrgQXnp6V3P6P7zPElna3AXvlukbANE6nJI= +github.com/aws/aws-sdk-go-v2/service/s3 v1.105.2/go.mod h1:zdmCoFO/dSI7GlrwsPqFJI+WlFnSU4Tc8TJnlXrM1Do= +github.com/aws/aws-sdk-go-v2/service/signin v1.4.1 h1:V7ZZ300WPXGjvkyore5DGe0ljVPOxCXie/thWdtSBXE= +github.com/aws/aws-sdk-go-v2/service/signin v1.4.1/go.mod h1:mxC0nT/C8wMMS97DemZPzvUZxvIt+2Iq+eS3JdFZGgg= +github.com/aws/aws-sdk-go-v2/service/sso v1.32.1 h1:gYFYh4iLLcAOJRLNPY2aD2g9DIhKn4eof8UkIrr1rTk= +github.com/aws/aws-sdk-go-v2/service/sso v1.32.1/go.mod h1:u8af9Nqkmqnr96f7v9nHqzZT9XBwbXEkTiqT4ROuJSE= +github.com/aws/aws-sdk-go-v2/service/ssooidc v1.37.1 h1:arjT9Cm3/WYbGmD5TUZHk4UQn4Lle1fUNZs5FC6CtF0= +github.com/aws/aws-sdk-go-v2/service/ssooidc v1.37.1/go.mod h1:DMPWJBjYs6+3+f/qhBFEFPPlQ6NlhWjai3dJNvipJ84= +github.com/aws/aws-sdk-go-v2/service/sts v1.44.1 h1:RvfHDg+xvAeZ+5741vUEjpOVtYSIm93W2zhx10Xtydw= +github.com/aws/aws-sdk-go-v2/service/sts v1.44.1/go.mod h1:9gdl4RrflIdpDb2TlXshWgR1F9TeCkvqDx77Vpr4Z/Q= github.com/aws/smithy-go v1.24.2 h1:FzA3bu/nt/vDvmnkg+R8Xl46gmzEDam6mZ1hzmwXFng= github.com/aws/smithy-go v1.24.2/go.mod h1:YE2RhdIuDbA5E5bTdciG9KrW3+TiEONeUWCqxX9i1Fc= +github.com/aws/smithy-go v1.27.3 h1:F3Zb497UhhskkfpJmfkXswyo+t0sh9OTBnIHjogWbVY= +github.com/aws/smithy-go v1.27.3/go.mod h1:YE2RhdIuDbA5E5bTdciG9KrW3+TiEONeUWCqxX9i1Fc= github.com/benbjohnson/clock v1.0.3/go.mod h1:bGMdMPoPVvcYyt1gHDf4J2KE153Yf9BuiUKYMaxlTDM= github.com/benbjohnson/clock v1.1.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA= github.com/benbjohnson/clock v1.3.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA= diff --git a/main.go b/main.go index 548034a19307..5d83e9c8c142 100644 --- a/main.go +++ b/main.go @@ -25,6 +25,7 @@ import ( "github.com/QuantumNous/new-api/oauth" perfmetrics "github.com/QuantumNous/new-api/pkg/perf_metrics" "github.com/QuantumNous/new-api/relay" + "github.com/QuantumNous/new-api/relay/asyncwrap" "github.com/QuantumNous/new-api/router" "github.com/QuantumNous/new-api/service" "github.com/QuantumNous/new-api/service/authz" @@ -68,6 +69,29 @@ func main() { } }() + appRole := strings.ToLower(strings.TrimSpace(os.Getenv("APP_ROLE"))) + if appRole == "" { + appRole = "api" + } + if appRole == "worker" { + runAsyncWorkerRole() + return + } + if appRole != "api" { + common.FatalLog("invalid APP_ROLE: " + appRole + " (expected api or worker)") + return + } + if processed, err := model.ReconcileAsyncBilling(context.Background(), 100); err != nil { + common.SysError("async billing startup reconciliation failed: " + err.Error()) + } else if processed > 0 { + common.SysLog(fmt.Sprintf("reconciled %d async billing records", processed)) + } + if processed, err := service.ReconcileAsyncUpstreamCosts(context.Background(), 100); err != nil { + common.SysError("async upstream cost startup reconciliation failed: " + err.Error()) + } else if processed > 0 { + common.SysLog(fmt.Sprintf("reconciled %d async upstream cost records", processed)) + } + if common.RedisEnabled { // for compatibility with old versions common.MemoryCacheEnabled = true @@ -231,6 +255,36 @@ func main() { common.SysLog("server exited") } +func runAsyncWorkerRole() { + service.NewAsyncImageExecutor = func(channel *model.Channel, apiKey string, timeout time.Duration) (service.AsyncImageExecutor, error) { + provider, ok := common.AsyncImageProviderForBaseURL(channel.GetBaseURL()) + if !ok { + return nil, fmt.Errorf("channel %d does not use an allowed synchronous image provider", channel.Id) + } + switch provider { + case common.AsyncImageProviderYunwu: + return asyncwrap.NewYunwuExecutor(channel.GetBaseURL(), apiKey, timeout) + case common.AsyncImageProviderGRSAI: + return asyncwrap.NewGRSAIExecutor(channel.GetBaseURL(), apiKey, timeout) + default: + return nil, fmt.Errorf("channel %d uses an unsupported synchronous image provider", channel.Id) + } + } + ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM) + defer stop() + worker, err := service.NewAsyncWorkerFromEnv(ctx) + if err != nil { + common.FatalLog("failed to initialize async worker: " + err.Error()) + return + } + common.SysLog(fmt.Sprintf("async worker %s started with concurrency %d", worker.ID, worker.Concurrency)) + if err := worker.Run(ctx); err != nil { + common.FatalLog("async worker stopped with error: " + err.Error()) + return + } + common.SysLog("async worker stopped") +} + func InjectUmamiAnalytics() { analyticsInjectBuilder := &strings.Builder{} if os.Getenv("UMAMI_WEBSITE_ID") != "" { diff --git a/middleware/async_distributor.go b/middleware/async_distributor.go new file mode 100644 index 000000000000..fce2169e9802 --- /dev/null +++ b/middleware/async_distributor.go @@ -0,0 +1,76 @@ +package middleware + +import ( + "fmt" + "net/http" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/service" + "github.com/QuantumNous/new-api/setting/ratio_setting" + "github.com/QuantumNous/new-api/types" + "github.com/gin-gonic/gin" +) + +// AsyncImageDistribute selects only channels that explicitly opt in to the +// synchronous-to-asynchronous image wrapper. Synchronous relay selection is +// deliberately left unchanged. +func AsyncImageDistribute() func(c *gin.Context) { + return func(c *gin.Context) { + request, _, err := getModelRequest(c) + if err != nil { + abortWithOpenAiMessage(c, http.StatusBadRequest, "invalid async image request: "+err.Error(), types.ErrorCodeInvalidRequest) + return + } + if request.Model == "" { + abortWithOpenAiMessage(c, http.StatusBadRequest, "model is required", types.ErrorCodeInvalidRequest) + return + } + + if common.GetContextKeyBool(c, constant.ContextKeyTokenModelLimitEnabled) { + limits, ok := common.GetContextKey(c, constant.ContextKeyTokenModelLimit) + allowed, valid := limits.(map[string]bool) + if !ok || !valid || !allowed[ratio_setting.FormatMatchingModelName(request.Model)] { + abortWithOpenAiMessage(c, http.StatusForbidden, fmt.Sprintf("token cannot access model %s", request.Model), types.ErrorCodeAccessDenied) + return + } + } + + usingGroup := common.GetContextKeyString(c, constant.ContextKeyUsingGroup) + selectedGroup := usingGroup + var channel *model.Channel + if usingGroup == "auto" { + userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup) + for _, group := range service.GetUserAutoGroup(userGroup) { + channel, err = model.GetAsyncImageChannel(group, request.Model) + if err != nil { + break + } + if channel != nil { + selectedGroup = group + common.SetContextKey(c, constant.ContextKeyAutoGroup, group) + break + } + } + } else { + channel, err = model.GetAsyncImageChannel(usingGroup, request.Model) + } + if err != nil { + abortWithOpenAiMessage(c, http.StatusInternalServerError, "failed to select async image channel", types.ErrorCodeGetChannelFailed) + return + } + if channel == nil { + abortWithOpenAiMessage(c, http.StatusServiceUnavailable, fmt.Sprintf("no async image channel is enabled for model %s", request.Model), types.ErrorCodeModelNotFound) + return + } + common.SetContextKey(c, constant.ContextKeyUsingGroup, selectedGroup) + common.SetContextKey(c, constant.ContextKeyRequestStartTime, time.Now()) + if apiErr := SetupContextForSelectedChannel(c, channel, request.Model); apiErr != nil { + abortWithOpenAiMessage(c, apiErr.StatusCode, apiErr.Error(), apiErr.GetErrorCode()) + return + } + c.Next() + } +} diff --git a/middleware/email-verification-rate-limit.go b/middleware/email-verification-rate-limit.go index 49081b55d466..b26238ff729f 100644 --- a/middleware/email-verification-rate-limit.go +++ b/middleware/email-verification-rate-limit.go @@ -1,8 +1,11 @@ package middleware import ( + "crypto/sha256" "fmt" "net/http" + "strconv" + "strings" "github.com/QuantumNous/new-api/common" @@ -10,48 +13,86 @@ import ( ) const ( - EmailVerificationRateLimitMark = "EV" - EmailVerificationMaxRequests = 2 // 30秒内最多2次 - EmailVerificationDuration = 30 // 30秒时间窗口 + EmailVerificationRateLimitMark = "EV" + EmailVerificationIPMaxRequests = 2 + EmailVerificationIPDuration = 30 + EmailVerificationEmailMaxRequests = 1 + EmailVerificationEmailCooldownSecs = 60 ) +func emailVerificationRateLimitKey(email string) string { + normalizedEmail := strings.ToLower(strings.TrimSpace(email)) + digest := sha256.Sum256([]byte(normalizedEmail)) + return fmt.Sprintf("%s:email:%s:%x", redisRateLimitNamespace, EmailVerificationRateLimitMark, digest) +} + +func writeEmailVerificationRateLimited(c *gin.Context, waitSeconds int64) { + if waitSeconds <= 0 { + waitSeconds = EmailVerificationEmailCooldownSecs + } + c.Header("Retry-After", strconv.FormatInt(waitSeconds, 10)) + c.JSON(http.StatusTooManyRequests, gin.H{ + "success": false, + "message": fmt.Sprintf("发送过于频繁，请等待 %d 秒后再试", waitSeconds), + }) + c.Abort() +} + func redisEmailVerificationRateLimiter(c *gin.Context) { allowed, _, ttlSeconds, err := redisFixedWindowTake( c.Request.Context(), redisIPRateLimitKey(EmailVerificationRateLimitMark, c.ClientIP()), - EmailVerificationMaxRequests, - EmailVerificationDuration, + EmailVerificationIPMaxRequests, + EmailVerificationIPDuration, ) if err != nil { memoryEmailVerificationRateLimiter(c) return } - if allowed { - c.Next() + if !allowed { + waitSeconds := int64(EmailVerificationIPDuration) + if ttlSeconds > 0 { + waitSeconds = ttlSeconds + } + writeEmailVerificationRateLimited(c, waitSeconds) return } - waitSeconds := int64(EmailVerificationDuration) - if ttlSeconds > 0 { - waitSeconds = ttlSeconds + email := strings.TrimSpace(c.Query("email")) + if email != "" { + allowed, _, ttlSeconds, err = redisFixedWindowTake( + c.Request.Context(), + emailVerificationRateLimitKey(email), + EmailVerificationEmailMaxRequests, + EmailVerificationEmailCooldownSecs, + ) + if err != nil { + memoryEmailVerificationRateLimiter(c) + return + } + if !allowed { + writeEmailVerificationRateLimited(c, ttlSeconds) + return + } } - c.JSON(http.StatusTooManyRequests, gin.H{ - "success": false, - "message": fmt.Sprintf("发送过于频繁，请等待 %d 秒后再试", waitSeconds), - }) - c.Abort() + c.Next() } func memoryEmailVerificationRateLimiter(c *gin.Context) { key := EmailVerificationRateLimitMark + ":" + c.ClientIP() + if !inMemoryRateLimiter.Request(key, EmailVerificationIPMaxRequests, EmailVerificationIPDuration) { + writeEmailVerificationRateLimited(c, EmailVerificationIPDuration) + return + } - if !inMemoryRateLimiter.Request(key, EmailVerificationMaxRequests, EmailVerificationDuration) { - c.JSON(http.StatusTooManyRequests, gin.H{ - "success": false, - "message": "发送过于频繁，请稍后再试", - }) - c.Abort() + email := strings.TrimSpace(c.Query("email")) + if email != "" && !inMemoryRateLimiter.Request( + emailVerificationRateLimitKey(email), + EmailVerificationEmailMaxRequests, + EmailVerificationEmailCooldownSecs, + ) { + writeEmailVerificationRateLimited(c, EmailVerificationEmailCooldownSecs) return } diff --git a/middleware/rate_limit_test.go b/middleware/rate_limit_test.go index 1ca48372b494..8f446b56aeb8 100644 --- a/middleware/rate_limit_test.go +++ b/middleware/rate_limit_test.go @@ -2,6 +2,7 @@ package middleware import ( "context" + "fmt" "net/http" "net/http/httptest" "sync" @@ -104,15 +105,71 @@ func TestRedisEmailVerificationRateLimiterPreservesResponseAndTTL(t *testing.T) }) remoteAddr := "192.0.2.30:12345" - assert.Equal(t, http.StatusNoContent, performRateLimitRequest(router, "/verify", remoteAddr).Code) - assert.Equal(t, http.StatusNoContent, performRateLimitRequest(router, "/verify", remoteAddr).Code) - response := performRateLimitRequest(router, "/verify", remoteAddr) + path := "/verify?email=Test%40Example.com" + assert.Equal(t, http.StatusNoContent, performRateLimitRequest(router, path, remoteAddr).Code) + response := performRateLimitRequest(router, "/verify?email=%20test%40example.com%20", remoteAddr) assert.Equal(t, http.StatusTooManyRequests, response.Code) - assert.JSONEq(t, `{"success":false,"message":"发送过于频繁，请等待 30 秒后再试"}`, response.Body.String()) + assert.Equal(t, "60", response.Header().Get("Retry-After")) + assert.JSONEq(t, `{"success":false,"message":"发送过于频繁，请等待 60 秒后再试"}`, response.Body.String()) - key := redisIPRateLimitKey(EmailVerificationRateLimitMark, "192.0.2.30") - assert.True(t, redisServer.Exists(key)) - assert.Equal(t, time.Duration(EmailVerificationDuration)*time.Second, redisServer.TTL(key)) + emailKey := emailVerificationRateLimitKey("test@example.com") + assert.True(t, redisServer.Exists(emailKey)) + assert.Equal(t, time.Duration(EmailVerificationEmailCooldownSecs)*time.Second, redisServer.TTL(emailKey)) + + redisServer.FastForward(time.Duration(EmailVerificationEmailCooldownSecs) * time.Second) + assert.Equal(t, http.StatusNoContent, performRateLimitRequest(router, path, remoteAddr).Code) +} + +func TestRedisEmailVerificationRateLimiterAllowsDifferentEmailsWithinIPLimit(t *testing.T) { + gin.SetMode(gin.TestMode) + _, _ = useRateLimitMiniRedis(t) + + router := gin.New() + require.NoError(t, router.SetTrustedProxies(nil)) + router.GET("/verify", EmailVerificationRateLimit(), func(c *gin.Context) { + c.Status(http.StatusNoContent) + }) + + remoteAddr := "192.0.2.31:12345" + assert.Equal(t, http.StatusNoContent, performRateLimitRequest(router, "/verify?email=one%40example.com", remoteAddr).Code) + assert.Equal(t, http.StatusNoContent, performRateLimitRequest(router, "/verify?email=two%40example.com", remoteAddr).Code) + response := performRateLimitRequest(router, "/verify?email=three%40example.com", remoteAddr) + assert.Equal(t, http.StatusTooManyRequests, response.Code) + assert.Equal(t, "30", response.Header().Get("Retry-After")) +} + +func TestRedisEmailVerificationRateLimiterAllowsOnlyOneConcurrentRequestPerEmail(t *testing.T) { + gin.SetMode(gin.TestMode) + _, _ = useRateLimitMiniRedis(t) + + router := gin.New() + require.NoError(t, router.SetTrustedProxies(nil)) + router.GET("/verify", EmailVerificationRateLimit(), func(c *gin.Context) { + c.Status(http.StatusNoContent) + }) + + const requestCount = 12 + var allowedCount atomic.Int64 + var limitedCount atomic.Int64 + var waitGroup sync.WaitGroup + waitGroup.Add(requestCount) + for requestIndex := range requestCount { + go func() { + defer waitGroup.Done() + remoteAddr := fmt.Sprintf("192.0.2.%d:12345", 100+requestIndex) + response := performRateLimitRequest(router, "/verify?email=parallel%40example.com", remoteAddr) + switch response.Code { + case http.StatusNoContent: + allowedCount.Add(1) + case http.StatusTooManyRequests: + limitedCount.Add(1) + } + }() + } + waitGroup.Wait() + + assert.Equal(t, int64(1), allowedCount.Load()) + assert.Equal(t, int64(requestCount-1), limitedCount.Load()) } func TestRedisFixedWindowIsAtomicUnderConcurrency(t *testing.T) { diff --git a/model/async_billing.go b/model/async_billing.go new file mode 100644 index 000000000000..e7778083c369 --- /dev/null +++ b/model/async_billing.go @@ -0,0 +1,335 @@ +package model + +import ( + "context" + "errors" + "fmt" + "time" + + "github.com/QuantumNous/new-api/common" + "gorm.io/gorm" +) + +var ErrAsyncRetryQuotaInsufficient = errors.New("insufficient quota for async task retry") + +// SettleAsyncJobBilling commits the already-reserved quota exactly once and +// updates usage counters in the same transaction as the billing status guard. +// UNCERTAIN is settleable because the request crossed the send boundary and +// the upstream may already have executed and charged it. +func SettleAsyncJobBilling(ctx context.Context, jobID int64) (bool, error) { + changed := false + err := DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + var job AsyncJob + result := lockForUpdate(tx).Where("id = ?", jobID).First(&job) + if result.Error != nil { + return result.Error + } + if job.BillingStatus != AsyncBillingReserved { + return nil + } + var task Task + if err := tx.First(&task, job.TaskID).Error; err != nil { + return err + } + if task.Status != TaskStatusSuccess && task.Status != TaskStatusFailure && task.Status != TaskStatusUncertain { + return fmt.Errorf("cannot settle async billing while task status is %s", task.Status) + } + if err := tx.Model(&User{}).Where("id = ?", task.UserId).Updates(map[string]any{ + "used_quota": gorm.Expr("used_quota + ?", task.Quota), + "request_count": gorm.Expr("request_count + 1"), + }).Error; err != nil { + return err + } + if task.Quota > 0 { + if err := tx.Model(&Channel{}).Where("id = ?", task.ChannelId). + Update("used_quota", gorm.Expr("used_quota + ?", task.Quota)).Error; err != nil { + return err + } + } + if err := tx.Model(&AsyncJob{}).Where("id = ? AND billing_status = ?", job.ID, AsyncBillingReserved). + Update("billing_status", AsyncBillingSettled).Error; err != nil { + return err + } + details := []byte(fmt.Sprintf(`{"quota":%d}`, task.Quota)) + if err := tx.Create(&TaskEvent{TaskID: task.ID, EventType: "billing_settled", ActorType: "system", Details: details}).Error; err != nil { + return err + } + changed = true + return nil + }) + return changed, err +} + +// RefundAsyncJobBilling refunds a queued cancellation or confirmed failure +// exactly once. User, token, subscription and the status guard are committed in +// one database transaction; caches are invalidated after commit. +func RefundAsyncJobBilling(ctx context.Context, jobID int64) (bool, error) { + changed := false + userID := 0 + err := DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + var job AsyncJob + if err := lockForUpdate(tx).Where("id = ?", jobID).First(&job).Error; err != nil { + return err + } + if job.BillingStatus != AsyncBillingReserved { + return nil + } + var task Task + if err := tx.First(&task, job.TaskID).Error; err != nil { + return err + } + if task.Status != TaskStatusFailure && task.Status != TaskStatusCancelled { + return fmt.Errorf("cannot refund async billing while task status is %s", task.Status) + } + quota := task.Quota + userID = task.UserId + if quota > 0 { + if task.PrivateData.BillingSource == "subscription" && task.PrivateData.SubscriptionId > 0 { + var subscription UserSubscription + if err := lockForUpdate(tx).Where("id = ?", task.PrivateData.SubscriptionId).First(&subscription).Error; err != nil { + return err + } + subscription.AmountUsed -= int64(quota) + if subscription.AmountUsed < 0 { + subscription.AmountUsed = 0 + } + if err := tx.Save(&subscription).Error; err != nil { + return err + } + if job.BillingRequestID != "" { + if err := tx.Model(&SubscriptionPreConsumeRecord{}). + Where("request_id = ? AND status != ?", job.BillingRequestID, "refunded"). + Update("status", "refunded").Error; err != nil { + return err + } + } + } else { + if err := tx.Model(&User{}).Where("id = ?", task.UserId). + Update("quota", gorm.Expr("quota + ?", quota)).Error; err != nil { + return err + } + } + if task.PrivateData.TokenId > 0 { + if err := tx.Model(&Token{}).Where("id = ?", task.PrivateData.TokenId).Updates(map[string]any{ + "remain_quota": gorm.Expr("remain_quota + ?", quota), + "used_quota": gorm.Expr("used_quota - ?", quota), + }).Error; err != nil { + return err + } + } + } + result := tx.Model(&AsyncJob{}).Where("id = ? AND billing_status = ?", job.ID, AsyncBillingReserved). + Update("billing_status", AsyncBillingRefunded) + if result.Error != nil { + return result.Error + } + if result.RowsAffected != 1 { + return nil + } + details := []byte(fmt.Sprintf(`{"quota":%d}`, quota)) + if err := tx.Create(&TaskEvent{TaskID: task.ID, EventType: "billing_refunded", ActorType: "system", Details: details}).Error; err != nil { + return err + } + changed = true + return nil + }) + if err != nil { + return false, err + } + if changed && userID > 0 { + _ = InvalidateUserCache(userID) + _ = InvalidateUserTokensCache(userID) + } + return changed, nil +} + +func ReconcileAsyncBilling(ctx context.Context, limit int) (int, error) { + if limit <= 0 { + limit = 100 + } + var jobs []AsyncJob + err := DB.WithContext(ctx). + Joins("JOIN tasks ON tasks.id = async_jobs.task_id"). + Where("async_jobs.billing_status = ? AND tasks.status IN ?", AsyncBillingReserved, []TaskStatus{TaskStatusSuccess, TaskStatusFailure, TaskStatusCancelled, TaskStatusUncertain}). + Order("async_jobs.id ASC").Limit(limit).Find(&jobs).Error + if err != nil { + return 0, err + } + processed := 0 + for _, job := range jobs { + var changed bool + if err := DB.WithContext(ctx).First(&job.Task, job.TaskID).Error; err != nil { + return processed, err + } + switch job.Task.Status { + case TaskStatusSuccess, TaskStatusUncertain: + changed, err = SettleAsyncJobBilling(ctx, job.ID) + case TaskStatusFailure: + if job.RefundEligible { + changed, err = RefundAsyncJobBilling(ctx, job.ID) + } else { + changed, err = SettleAsyncJobBilling(ctx, job.ID) + } + case TaskStatusCancelled: + changed, err = RefundAsyncJobBilling(ctx, job.ID) + } + if err != nil && !errors.Is(err, gorm.ErrRecordNotFound) { + return processed, err + } + if changed { + processed++ + } + } + return processed, nil +} + +// RetryAsyncJob requeues an administrator-approved terminal task and reserves +// quota for the new attempt. For UNCERTAIN (or another non-refundable +// attempt), the previous reservation is first committed because the upstream +// may already have charged it; the new attempt is then reserved separately. +func RetryAsyncJob(ctx context.Context, jobID int64, actorID int) (*AsyncJob, bool, error) { + var retried AsyncJob + changed := false + userID := 0 + err := DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + var job AsyncJob + if err := lockForUpdate(tx).Where("id = ?", jobID).First(&job).Error; err != nil { + return err + } + if job.ExecutionStatus != AsyncStatusFailure && job.ExecutionStatus != AsyncStatusUncertain { + retried = job + return loadAsyncJobTask(tx, &retried) + } + if err := ValidateAsyncTransition(job.ExecutionStatus, AsyncStatusQueued); err != nil { + return err + } + + var task Task + if err := tx.First(&task, job.TaskID).Error; err != nil { + return err + } + userID = task.UserId + priorStatus := job.ExecutionStatus + priorBilling := job.BillingStatus + priorMustSettle := job.BillingStatus == AsyncBillingReserved && + (job.ExecutionStatus == AsyncStatusUncertain || !job.RefundEligible) + if priorMustSettle { + if err := recordAsyncAttemptUsageTx(tx, &task); err != nil { + return err + } + } + + needsNewReservation := priorMustSettle || job.BillingStatus != AsyncBillingReserved + billingRequestID := job.BillingRequestID + if needsNewReservation && task.Quota > 0 { + if task.PrivateData.BillingSource == "subscription" && task.PrivateData.SubscriptionId > 0 { + var subscription UserSubscription + now := time.Now().Unix() + if err := lockForUpdate(tx).Where("id = ? AND user_id = ? AND status = ? AND end_time > ?", task.PrivateData.SubscriptionId, task.UserId, "active", now).First(&subscription).Error; err != nil { + return fmt.Errorf("%w: active subscription unavailable", ErrAsyncRetryQuotaInsufficient) + } + if subscription.AmountTotal > 0 && subscription.AmountTotal-subscription.AmountUsed < int64(task.Quota) { + return fmt.Errorf("%w: subscription quota", ErrAsyncRetryQuotaInsufficient) + } + subscription.AmountUsed += int64(task.Quota) + if err := tx.Save(&subscription).Error; err != nil { + return err + } + billingRequestID = fmt.Sprintf("async-retry-%d-%d", job.ID, time.Now().UnixNano()) + if err := tx.Create(&SubscriptionPreConsumeRecord{ + RequestId: billingRequestID, + UserId: task.UserId, + UserSubscriptionId: subscription.Id, + PreConsumed: int64(task.Quota), + Status: "consumed", + }).Error; err != nil { + return err + } + } else { + result := tx.Model(&User{}). + Where("id = ? AND status = ? AND quota >= ?", task.UserId, common.UserStatusEnabled, task.Quota). + Update("quota", gorm.Expr("quota - ?", task.Quota)) + if result.Error != nil { + return result.Error + } + if result.RowsAffected != 1 { + return fmt.Errorf("%w: wallet quota", ErrAsyncRetryQuotaInsufficient) + } + } + + var token Token + if err := lockForUpdate(tx).Where("id = ? AND user_id = ? AND status = ?", task.PrivateData.TokenId, task.UserId, common.TokenStatusEnabled).First(&token).Error; err != nil { + return fmt.Errorf("%w: token unavailable", ErrAsyncRetryQuotaInsufficient) + } + if !token.UnlimitedQuota && token.RemainQuota < task.Quota { + return fmt.Errorf("%w: token quota", ErrAsyncRetryQuotaInsufficient) + } + if err := tx.Model(&Token{}).Where("id = ?", token.Id).Updates(map[string]any{ + "remain_quota": gorm.Expr("remain_quota - ?", task.Quota), + "used_quota": gorm.Expr("used_quota + ?", task.Quota), + "accessed_time": time.Now().Unix(), + }).Error; err != nil { + return err + } + } + + now := time.Now().Unix() + if err := tx.Model(&AsyncJob{}).Where("id = ? AND execution_status IN ?", job.ID, []AsyncExecutionStatus{AsyncStatusFailure, AsyncStatusUncertain}).Updates(map[string]any{ + "execution_status": AsyncStatusQueued, + "worker_id": "", + "lease_until": 0, + "request_sent_at": 0, + "result_payload": nil, + "error_phase": "", + "error_code": "", + "refund_eligible": false, + "billing_status": AsyncBillingReserved, + "billing_request_id": billingRequestID, + "updated_at": now, + }).Error; err != nil { + return err + } + if err := tx.Model(&Task{}).Where("id = ?", task.ID).Updates(map[string]any{ + "status": TaskStatusQueued, + "progress": "0%", + "start_time": 0, + "finish_time": 0, + "fail_reason": "", + "data": []byte("{}"), + "updated_at": now, + }).Error; err != nil { + return err + } + details := []byte(fmt.Sprintf(`{"previous_status":%q,"previous_billing_status":%q,"duplicate_risk":%t}`, priorStatus, priorBilling, priorStatus == AsyncStatusUncertain)) + if err := tx.Create(&TaskEvent{TaskID: task.ID, EventType: "manual_retry", FromStatus: string(priorStatus), ToStatus: string(AsyncStatusQueued), ActorType: "admin", ActorID: actorID, Details: details}).Error; err != nil { + return err + } + if err := tx.First(&retried, job.ID).Error; err != nil { + return err + } + changed = true + return loadAsyncJobTask(tx, &retried) + }) + if err != nil { + return nil, false, err + } + if changed && userID > 0 { + _ = InvalidateUserCache(userID) + _ = InvalidateUserTokensCache(userID) + } + return &retried, changed, nil +} + +func recordAsyncAttemptUsageTx(tx *gorm.DB, task *Task) error { + if err := tx.Model(&User{}).Where("id = ?", task.UserId).Updates(map[string]any{ + "used_quota": gorm.Expr("used_quota + ?", task.Quota), + "request_count": gorm.Expr("request_count + 1"), + }).Error; err != nil { + return err + } + if task.Quota <= 0 { + return nil + } + return tx.Model(&Channel{}).Where("id = ?", task.ChannelId). + Update("used_quota", gorm.Expr("used_quota + ?", task.Quota)).Error +} diff --git a/model/async_job.go b/model/async_job.go new file mode 100644 index 000000000000..d83272bc0b89 --- /dev/null +++ b/model/async_job.go @@ -0,0 +1,657 @@ +package model + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "hash/fnv" + "sync" + "time" + + "github.com/QuantumNous/new-api/common" + "gorm.io/gorm" +) + +type AsyncExecutionStatus string + +const ( + AsyncStatusQueued AsyncExecutionStatus = "QUEUED" + AsyncStatusRunning AsyncExecutionStatus = "RUNNING" + AsyncStatusSuccess AsyncExecutionStatus = "SUCCESS" + AsyncStatusFailure AsyncExecutionStatus = "FAILURE" + AsyncStatusUncertain AsyncExecutionStatus = "UNCERTAIN" + AsyncStatusCancelled AsyncExecutionStatus = "CANCELLED" +) + +const ( + AsyncEndpointImageGeneration = "image_generation" + AsyncBillingReserved = "RESERVED" + AsyncBillingSettled = "SETTLED" + AsyncBillingRefunded = "REFUNDED" +) + +var ErrInvalidAsyncTransition = errors.New("invalid async job status transition") + +var asyncIdempotencyLocks [256]sync.Mutex + +type AsyncJob struct { + ID int64 `json:"id" gorm:"primaryKey;autoIncrement"` + TaskID int64 `json:"task_id" gorm:"not null;uniqueIndex"` + TokenID int `json:"token_id" gorm:"not null;uniqueIndex:idx_async_token_idempotency,priority:1;index"` + ChannelID int `json:"channel_id" gorm:"not null;index"` + EndpointType string `json:"endpoint_type" gorm:"type:varchar(40);not null;index"` + RequestPayload []byte `json:"-" gorm:"not null"` + RequestHash string `json:"request_hash" gorm:"type:char(64);not null"` + IdempotencyKey string `json:"idempotency_key" gorm:"type:varchar(191);not null;uniqueIndex:idx_async_token_idempotency,priority:2"` + ExecutionStatus AsyncExecutionStatus `json:"execution_status" gorm:"type:varchar(20);not null;index:idx_async_status_lease,priority:1"` + WorkerID string `json:"worker_id,omitempty" gorm:"type:varchar(128);index"` + LeaseUntil int64 `json:"lease_until,omitempty" gorm:"index:idx_async_status_lease,priority:2"` + Attempt int `json:"attempt" gorm:"not null;default:0"` + RequestSentAt int64 `json:"request_sent_at,omitempty" gorm:"index"` + ResultPayload JSONValue `json:"result_payload,omitempty" gorm:"type:text"` + ErrorPhase string `json:"error_phase,omitempty" gorm:"type:varchar(40);index"` + ErrorCode string `json:"error_code,omitempty" gorm:"type:varchar(80);index"` + RefundEligible bool `json:"refund_eligible" gorm:"not null;default:false"` + BillingStatus string `json:"billing_status" gorm:"type:varchar(20);not null;default:'RESERVED';index"` + BillingRequestID string `json:"-" gorm:"type:varchar(64);index"` + CreatedAt int64 `json:"created_at" gorm:"autoCreateTime;index"` + UpdatedAt int64 `json:"updated_at" gorm:"autoUpdateTime"` + + Task Task `json:"task" gorm:"belongsTo:true;foreignKey:TaskID;references:ID;constraint:OnUpdate:CASCADE,OnDelete:CASCADE"` +} + +type Artifact struct { + ID int64 `json:"id" gorm:"primaryKey;autoIncrement"` + TaskID int64 `json:"task_id" gorm:"not null;index:idx_artifact_task_object,priority:1"` + ObjectKey string `json:"object_key" gorm:"type:varchar(512);not null;uniqueIndex;index:idx_artifact_task_object,priority:2"` + ContentType string `json:"content_type" gorm:"type:varchar(128);not null"` + SizeBytes int64 `json:"size_bytes" gorm:"not null"` + SHA256 string `json:"sha256" gorm:"type:char(64);not null;index"` + SourceURLHash string `json:"source_url_hash" gorm:"type:char(64);not null"` + CreatedAt int64 `json:"created_at" gorm:"autoCreateTime;index"` + ExpiresAt int64 `json:"expires_at" gorm:"not null;index"` + + Task Task `json:"-" gorm:"belongsTo:true;foreignKey:TaskID;references:ID;constraint:OnUpdate:CASCADE,OnDelete:CASCADE"` +} + +type TaskEvent struct { + ID int64 `json:"id" gorm:"primaryKey;autoIncrement"` + TaskID int64 `json:"task_id" gorm:"not null;index"` + EventType string `json:"event_type" gorm:"type:varchar(40);not null;index"` + FromStatus string `json:"from_status,omitempty" gorm:"type:varchar(20)"` + ToStatus string `json:"to_status,omitempty" gorm:"type:varchar(20)"` + WorkerID string `json:"worker_id,omitempty" gorm:"type:varchar(128)"` + ErrorPhase string `json:"error_phase,omitempty" gorm:"type:varchar(40)"` + ErrorCode string `json:"error_code,omitempty" gorm:"type:varchar(80)"` + ActorType string `json:"actor_type,omitempty" gorm:"type:varchar(20)"` + ActorID int `json:"actor_id,omitempty"` + Details JSONValue `json:"details,omitempty" gorm:"type:text"` + CreatedAt int64 `json:"created_at" gorm:"autoCreateTime;index"` + + Task Task `json:"-" gorm:"belongsTo:true;foreignKey:TaskID;references:ID;constraint:OnUpdate:CASCADE,OnDelete:CASCADE"` +} + +type AsyncTaskRecord struct { + Job AsyncJob + Task Task +} + +func AsyncStatusIsTerminal(status AsyncExecutionStatus) bool { + switch status { + case AsyncStatusSuccess, AsyncStatusFailure, AsyncStatusUncertain, AsyncStatusCancelled: + return true + default: + return false + } +} + +func ValidateAsyncTransition(from, to AsyncExecutionStatus) error { + allowed := false + switch from { + case AsyncStatusQueued: + allowed = to == AsyncStatusRunning || to == AsyncStatusCancelled + case AsyncStatusRunning: + allowed = to == AsyncStatusSuccess || to == AsyncStatusFailure || to == AsyncStatusUncertain + case AsyncStatusFailure, AsyncStatusUncertain: + allowed = to == AsyncStatusQueued + } + if !allowed { + return fmt.Errorf("%w: %s -> %s", ErrInvalidAsyncTransition, from, to) + } + return nil +} + +func asyncTaskStatus(status AsyncExecutionStatus) TaskStatus { + switch status { + case AsyncStatusQueued: + return TaskStatusQueued + case AsyncStatusRunning: + return TaskStatusInProgress + case AsyncStatusSuccess: + return TaskStatusSuccess + case AsyncStatusFailure: + return TaskStatusFailure + case AsyncStatusUncertain: + return TaskStatusUncertain + case AsyncStatusCancelled: + return TaskStatusCancelled + default: + return TaskStatusUnknown + } +} + +func CreateAsyncTask(task *Task, job *AsyncJob) error { + if task == nil || job == nil { + return errors.New("task and async job are required") + } + return DB.Transaction(func(tx *gorm.DB) error { + if err := tx.Create(task).Error; err != nil { + return err + } + job.TaskID = task.ID + if err := tx.Create(job).Error; err != nil { + return err + } + return tx.Create(&TaskEvent{ + TaskID: task.ID, + EventType: "created", + ToStatus: string(AsyncStatusQueued), + ActorType: "token", + ActorID: job.TokenID, + }).Error + }) +} + +// WithAsyncIdempotencyLock serializes a token/idempotency-key pair in this +// process and, on PostgreSQL, across API replicas using a transaction-scoped +// advisory lock. The database unique index remains the final invariant. +func WithAsyncIdempotencyLock(ctx context.Context, tokenID int, key string, fn func() error) error { + hasher := fnv.New64a() + _, _ = fmt.Fprintf(hasher, "%d:%s", tokenID, key) + lockKey := hasher.Sum64() + local := &asyncIdempotencyLocks[lockKey%uint64(len(asyncIdempotencyLocks))] + local.Lock() + defer local.Unlock() + + if !common.UsingMainDatabase(common.DatabaseTypePostgreSQL) { + return fn() + } + return DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + if err := tx.Exec("SELECT pg_advisory_xact_lock(?)", int64(lockKey)).Error; err != nil { + return err + } + return fn() + }) +} + +func GetAsyncJobByTokenAndKey(ctx context.Context, tokenID int, key string) (*AsyncJob, error) { + var job AsyncJob + err := DB.WithContext(ctx).Where("token_id = ? AND idempotency_key = ?", tokenID, key).First(&job).Error + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil, nil + } + if err != nil { + return nil, err + } + return &job, loadAsyncJobTask(DB.WithContext(ctx), &job) +} + +func GetAsyncJobByPublicTaskID(ctx context.Context, publicTaskID string, tokenID int) (*AsyncJob, error) { + var job AsyncJob + query := DB.WithContext(ctx). + Joins("JOIN tasks ON tasks.id = async_jobs.task_id"). + Where("tasks.task_id = ?", publicTaskID) + if tokenID > 0 { + query = query.Where("async_jobs.token_id = ?", tokenID) + } + err := query.First(&job).Error + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil, nil + } + if err != nil { + return nil, err + } + return &job, loadAsyncJobTask(DB.WithContext(ctx), &job) +} + +func GetAsyncJobForSession(ctx context.Context, publicTaskID string, userID int, administrator bool) (*AsyncJob, error) { + var job AsyncJob + query := DB.WithContext(ctx). + Joins("JOIN tasks ON tasks.id = async_jobs.task_id"). + Where("tasks.task_id = ?", publicTaskID) + if !administrator { + query = query.Where("tasks.user_id = ?", userID) + } + err := query.First(&job).Error + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil, nil + } + if err != nil { + return nil, err + } + return &job, loadAsyncJobTask(DB.WithContext(ctx), &job) +} + +func ListAsyncJobsByTaskIDs(ctx context.Context, taskIDs []int64) (map[int64]AsyncJob, error) { + result := make(map[int64]AsyncJob, len(taskIDs)) + if len(taskIDs) == 0 { + return result, nil + } + var jobs []AsyncJob + if err := DB.WithContext(ctx).Where("task_id IN ?", taskIDs).Find(&jobs).Error; err != nil { + return nil, err + } + for _, job := range jobs { + result[job.TaskID] = job + } + return result, nil +} + +func ListTaskEvents(ctx context.Context, taskID int64) ([]TaskEvent, error) { + var events []TaskEvent + err := DB.WithContext(ctx).Where("task_id = ?", taskID).Order("id ASC").Find(&events).Error + return events, err +} + +func ListQueuedAsyncJobs(ctx context.Context, limit int) ([]AsyncJob, error) { + if limit <= 0 { + limit = 50 + } + var jobs []AsyncJob + err := DB.WithContext(ctx). + Where("execution_status = ?", AsyncStatusQueued). + Order("created_at ASC, id ASC").Limit(limit).Find(&jobs).Error + if err == nil { + for i := range jobs { + if loadErr := loadAsyncJobTask(DB.WithContext(ctx), &jobs[i]); loadErr != nil { + return nil, loadErr + } + } + } + return jobs, err +} + +// ListSettledAsyncJobsMissingUpstreamCost returns settled async attempts whose +// canonical cost ledger entry has not been persisted yet. Joining channels +// excludes historical tasks whose deleted channel can no longer provide an +// auditable conversion profile. +func ListSettledAsyncJobsMissingUpstreamCost(ctx context.Context, limit int) ([]AsyncJob, error) { + if limit <= 0 { + limit = 100 + } + var jobs []AsyncJob + err := DB.WithContext(ctx). + Table("async_jobs"). + Select("async_jobs.*"). + Joins("JOIN tasks ON tasks.id = async_jobs.task_id"). + Joins("JOIN channels ON channels.id = async_jobs.channel_id"). + Joins("LEFT JOIN upstream_cost_records ON upstream_cost_records.request_id = async_jobs.billing_request_id"). + Where("async_jobs.billing_status = ?", AsyncBillingSettled). + Where("async_jobs.billing_request_id <> ?", ""). + Where("upstream_cost_records.id IS NULL"). + Order("async_jobs.id ASC"). + Limit(limit). + Find(&jobs).Error + if err != nil { + return nil, err + } + for i := range jobs { + if err := loadAsyncJobTask(DB.WithContext(ctx), &jobs[i]); err != nil { + return nil, err + } + } + return jobs, nil +} + +func loadAsyncJobTask(db *gorm.DB, job *AsyncJob) error { + if job == nil || job.TaskID == 0 { + return nil + } + return db.First(&job.Task, job.TaskID).Error +} + +func ClaimAsyncJob(ctx context.Context, jobID int64, workerID string, leaseUntil int64) (*AsyncJob, bool, error) { + var claimed AsyncJob + err := DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + var job AsyncJob + result := lockForUpdate(tx).Where("id = ? AND execution_status = ?", jobID, AsyncStatusQueued).First(&job) + if errors.Is(result.Error, gorm.ErrRecordNotFound) { + return nil + } + if result.Error != nil { + return result.Error + } + if err := ValidateAsyncTransition(job.ExecutionStatus, AsyncStatusRunning); err != nil { + return err + } + now := time.Now().Unix() + updates := map[string]any{ + "execution_status": AsyncStatusRunning, + "worker_id": workerID, + "lease_until": leaseUntil, + "attempt": gorm.Expr("attempt + 1"), + "updated_at": now, + } + jobUpdate := tx.Model(&AsyncJob{}).Where("id = ? AND execution_status = ?", job.ID, AsyncStatusQueued).Updates(updates) + if jobUpdate.Error != nil { + return jobUpdate.Error + } + if jobUpdate.RowsAffected != 1 { + return nil + } + taskUpdate := tx.Model(&Task{}).Where("id = ? AND status = ?", job.TaskID, TaskStatusQueued).Updates(map[string]any{ + "status": TaskStatusInProgress, + "progress": "1%", + "start_time": now, + "updated_at": now, + }) + if taskUpdate.Error != nil { + return taskUpdate.Error + } + if taskUpdate.RowsAffected != 1 { + return errors.New("async task state changed while claiming job") + } + if err := tx.Create(&TaskEvent{TaskID: job.TaskID, EventType: "claimed", FromStatus: string(AsyncStatusQueued), ToStatus: string(AsyncStatusRunning), WorkerID: workerID}).Error; err != nil { + return err + } + if err := tx.First(&claimed, job.ID).Error; err != nil { + return err + } + return loadAsyncJobTask(tx, &claimed) + }) + if err != nil { + return nil, false, err + } + if claimed.ID == 0 { + return nil, false, nil + } + return &claimed, true, nil +} + +func RenewAsyncJobLease(ctx context.Context, jobID int64, workerID string, leaseUntil int64) (bool, error) { + result := DB.WithContext(ctx).Model(&AsyncJob{}). + Where("id = ? AND execution_status = ? AND worker_id = ?", jobID, AsyncStatusRunning, workerID). + Updates(map[string]any{"lease_until": leaseUntil, "updated_at": time.Now().Unix()}) + return result.RowsAffected == 1, result.Error +} + +func MarkAsyncRequestSent(ctx context.Context, jobID int64, workerID string, sentAt int64) (bool, error) { + result := DB.WithContext(ctx).Model(&AsyncJob{}). + Where("id = ? AND execution_status = ? AND worker_id = ? AND request_sent_at = 0", jobID, AsyncStatusRunning, workerID). + Updates(map[string]any{"request_sent_at": sentAt, "updated_at": sentAt}) + return result.RowsAffected == 1, result.Error +} + +func CompleteAsyncJob(ctx context.Context, jobID int64, workerID string, status AsyncExecutionStatus, resultPayload json.RawMessage, errorPhase, errorCode, failReason string, refundEligible bool) (bool, error) { + if status != AsyncStatusSuccess && status != AsyncStatusFailure && status != AsyncStatusUncertain { + return false, fmt.Errorf("unsupported completion status %s", status) + } + changed := false + err := DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + var job AsyncJob + result := lockForUpdate(tx).Where("id = ? AND execution_status = ? AND worker_id = ?", jobID, AsyncStatusRunning, workerID).First(&job) + if errors.Is(result.Error, gorm.ErrRecordNotFound) { + return nil + } + if result.Error != nil { + return result.Error + } + if err := ValidateAsyncTransition(job.ExecutionStatus, status); err != nil { + return err + } + now := time.Now().Unix() + jobUpdate := tx.Model(&AsyncJob{}).Where("id = ? AND execution_status = ? AND worker_id = ?", job.ID, AsyncStatusRunning, workerID).Updates(map[string]any{ + "execution_status": status, + "worker_id": "", + "lease_until": 0, + "result_payload": JSONValue(resultPayload), + "error_phase": errorPhase, + "error_code": errorCode, + "refund_eligible": refundEligible, + "updated_at": now, + }) + if jobUpdate.Error != nil { + return jobUpdate.Error + } + if jobUpdate.RowsAffected != 1 { + return nil + } + taskUpdates := map[string]any{ + "status": asyncTaskStatus(status), + "finish_time": now, + "updated_at": now, + "fail_reason": failReason, + } + if status == AsyncStatusSuccess { + taskUpdates["progress"] = "100%" + } else { + taskUpdates["progress"] = "0%" + } + if len(resultPayload) > 0 { + taskUpdates["data"] = resultPayload + } + taskUpdate := tx.Model(&Task{}).Where("id = ? AND status = ?", job.TaskID, TaskStatusInProgress).Updates(taskUpdates) + if taskUpdate.Error != nil { + return taskUpdate.Error + } + if taskUpdate.RowsAffected != 1 { + return errors.New("async task state changed while completing job") + } + if err := tx.Create(&TaskEvent{TaskID: job.TaskID, EventType: "completed", FromStatus: string(AsyncStatusRunning), ToStatus: string(status), WorkerID: workerID, ErrorPhase: errorPhase, ErrorCode: errorCode}).Error; err != nil { + return err + } + changed = true + return nil + }) + return changed, err +} + +func CancelQueuedAsyncJob(ctx context.Context, publicTaskID string, tokenID int) (*AsyncJob, bool, error) { + job, err := GetAsyncJobByPublicTaskID(ctx, publicTaskID, tokenID) + if err != nil || job == nil { + return job, false, err + } + return CancelQueuedAsyncJobByID(ctx, job.ID, "token", tokenID) +} + +func CancelQueuedAsyncJobByID(ctx context.Context, jobID int64, actorType string, actorID int) (*AsyncJob, bool, error) { + var cancelled AsyncJob + changed := false + err := DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + var job AsyncJob + result := lockForUpdate(tx).Where("id = ?", jobID).First(&job) + if errors.Is(result.Error, gorm.ErrRecordNotFound) { + return nil + } + if result.Error != nil { + return result.Error + } + if job.ExecutionStatus != AsyncStatusQueued { + cancelled = job + return loadAsyncJobTask(tx, &cancelled) + } + if err := ValidateAsyncTransition(job.ExecutionStatus, AsyncStatusCancelled); err != nil { + return err + } + now := time.Now().Unix() + jobUpdate := tx.Model(&AsyncJob{}).Where("id = ? AND execution_status = ?", job.ID, AsyncStatusQueued).Updates(map[string]any{ + "execution_status": AsyncStatusCancelled, + "refund_eligible": true, + "updated_at": now, + }) + if jobUpdate.Error != nil { + return jobUpdate.Error + } + if jobUpdate.RowsAffected != 1 { + cancelled = job + return loadAsyncJobTask(tx, &cancelled) + } + taskUpdate := tx.Model(&Task{}).Where("id = ? AND status = ?", job.TaskID, TaskStatusQueued).Updates(map[string]any{ + "status": TaskStatusCancelled, + "finish_time": now, + "updated_at": now, + "fail_reason": "cancelled before upstream request", + }) + if taskUpdate.Error != nil { + return taskUpdate.Error + } + if taskUpdate.RowsAffected != 1 { + return errors.New("async task state changed while cancelling job") + } + if err := tx.Create(&TaskEvent{TaskID: job.TaskID, EventType: "cancelled", FromStatus: string(AsyncStatusQueued), ToStatus: string(AsyncStatusCancelled), ActorType: actorType, ActorID: actorID}).Error; err != nil { + return err + } + changed = true + if err := tx.First(&cancelled, job.ID).Error; err != nil { + return err + } + return loadAsyncJobTask(tx, &cancelled) + }) + if err != nil { + return nil, false, err + } + if cancelled.ID == 0 { + return nil, false, nil + } + return &cancelled, changed, nil +} + +type AsyncRecoverySummary struct { + Requeued int `json:"requeued"` + Uncertain int `json:"uncertain"` +} + +func RecoverExpiredAsyncJobs(ctx context.Context, now int64, limit int) (AsyncRecoverySummary, error) { + if limit <= 0 { + limit = 100 + } + var ids []int64 + if err := DB.WithContext(ctx).Model(&AsyncJob{}). + Where("execution_status = ? AND lease_until > 0 AND lease_until < ?", AsyncStatusRunning, now). + Order("lease_until ASC").Limit(limit).Pluck("id", &ids).Error; err != nil { + return AsyncRecoverySummary{}, err + } + summary := AsyncRecoverySummary{} + for _, id := range ids { + err := DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + var job AsyncJob + result := lockForUpdate(tx).Where("id = ? AND execution_status = ? AND lease_until < ?", id, AsyncStatusRunning, now).First(&job) + if errors.Is(result.Error, gorm.ErrRecordNotFound) { + return nil + } + if result.Error != nil { + return result.Error + } + if job.RequestSentAt == 0 { + if err := tx.Model(&AsyncJob{}).Where("id = ?", job.ID).Updates(map[string]any{ + "execution_status": AsyncStatusQueued, + "worker_id": "", + "lease_until": 0, + "updated_at": now, + }).Error; err != nil { + return err + } + if err := tx.Model(&Task{}).Where("id = ?", job.TaskID).Updates(map[string]any{ + "status": TaskStatusQueued, + "progress": "0%", + "start_time": 0, + "updated_at": now, + }).Error; err != nil { + return err + } + if err := tx.Create(&TaskEvent{TaskID: job.TaskID, EventType: "lease_recovered", FromStatus: string(AsyncStatusRunning), ToStatus: string(AsyncStatusQueued), WorkerID: job.WorkerID}).Error; err != nil { + return err + } + summary.Requeued++ + return nil + } + if err := tx.Model(&AsyncJob{}).Where("id = ?", job.ID).Updates(map[string]any{ + "execution_status": AsyncStatusUncertain, + "worker_id": "", + "lease_until": 0, + "error_phase": "worker_recovery", + "error_code": "lease_expired_after_send", + "updated_at": now, + }).Error; err != nil { + return err + } + if err := tx.Model(&Task{}).Where("id = ?", job.TaskID).Updates(map[string]any{ + "status": TaskStatusUncertain, + "finish_time": now, + "fail_reason": "worker lease expired after the upstream request may have been sent", + "updated_at": now, + }).Error; err != nil { + return err + } + if err := tx.Create(&TaskEvent{TaskID: job.TaskID, EventType: "lease_uncertain", FromStatus: string(AsyncStatusRunning), ToStatus: string(AsyncStatusUncertain), WorkerID: job.WorkerID, ErrorPhase: "worker_recovery", ErrorCode: "lease_expired_after_send"}).Error; err != nil { + return err + } + summary.Uncertain++ + return nil + }) + if err != nil { + return summary, err + } + } + return summary, nil +} + +func CreateArtifacts(ctx context.Context, artifacts []Artifact) error { + if len(artifacts) == 0 { + return nil + } + return DB.WithContext(ctx).Create(&artifacts).Error +} + +func ListArtifactsByTaskID(ctx context.Context, taskID int64) ([]Artifact, error) { + var artifacts []Artifact + err := DB.WithContext(ctx).Where("task_id = ?", taskID).Order("id ASC").Find(&artifacts).Error + return artifacts, err +} + +func ListExpiredArtifacts(ctx context.Context, before int64, limit int) ([]Artifact, error) { + if limit <= 0 { + limit = 100 + } + var artifacts []Artifact + err := DB.WithContext(ctx).Where("expires_at <= ?", before).Order("expires_at ASC").Limit(limit).Find(&artifacts).Error + return artifacts, err +} + +func DeleteArtifactAndClearResultIfLast(ctx context.Context, artifactID, taskID int64) (bool, error) { + resultCleared := false + err := DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + var job AsyncJob + jobQuery := lockForUpdate(tx).Select("id").Where("task_id = ?", taskID).First(&job) + if jobQuery.Error != nil && !errors.Is(jobQuery.Error, gorm.ErrRecordNotFound) { + return jobQuery.Error + } + + deleted := tx.Where("id = ? AND task_id = ?", artifactID, taskID).Delete(&Artifact{}) + if deleted.Error != nil || deleted.RowsAffected == 0 { + return deleted.Error + } + if job.ID == 0 { + return nil + } + + var remaining int64 + if err := tx.Model(&Artifact{}).Where("task_id = ?", taskID).Count(&remaining).Error; err != nil { + return err + } + if remaining > 0 { + return nil + } + update := tx.Model(&AsyncJob{}).Where("id = ?", job.ID).Updates(map[string]any{ + "result_payload": nil, + "updated_at": time.Now().Unix(), + }) + if update.Error != nil { + return update.Error + } + resultCleared = update.RowsAffected == 1 + return nil + }) + return resultCleared, err +} diff --git a/model/async_job_test.go b/model/async_job_test.go new file mode 100644 index 000000000000..4a5b675f0baa --- /dev/null +++ b/model/async_job_test.go @@ -0,0 +1,373 @@ +package model + +import ( + "context" + "encoding/json" + "sync" + "testing" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/dto" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm/schema" +) + +func TestAsyncModelsBelongToTask(t *testing.T) { + models := map[string]any{ + "async job": &AsyncJob{}, + "artifact": &Artifact{}, + "task event": &TaskEvent{}, + } + for name, value := range models { + t.Run(name, func(t *testing.T) { + parsed, err := schema.Parse(value, &sync.Map{}, schema.NamingStrategy{}) + require.NoError(t, err) + relation := parsed.Relationships.Relations["Task"] + require.NotNil(t, relation) + assert.Equal(t, schema.BelongsTo, relation.Type) + require.Len(t, relation.References, 1) + assert.Equal(t, "task_id", relation.References[0].ForeignKey.DBName) + assert.Equal(t, "id", relation.References[0].PrimaryKey.DBName) + }) + } +} + +func createAsyncFixture(t *testing.T, suffix string) (*Task, *AsyncJob) { + t.Helper() + task := &Task{ + TaskID: "task_async_" + suffix, + Platform: constant.TaskPlatformAsyncImage, + UserId: 1, + ChannelId: 2, + Status: TaskStatusQueued, + Progress: "0%", + Data: json.RawMessage(`{"model":"image-model"}`), + } + job := &AsyncJob{ + TokenID: 3, + ChannelID: 2, + EndpointType: AsyncEndpointImageGeneration, + RequestPayload: []byte("encrypted"), + RequestHash: "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", + IdempotencyKey: "idem-" + suffix, + ExecutionStatus: AsyncStatusQueued, + BillingStatus: AsyncBillingReserved, + BillingRequestID: "req-" + suffix, + } + require.NoError(t, CreateAsyncTask(task, job)) + return task, job +} + +func TestValidateAsyncTransition(t *testing.T) { + valid := [][2]AsyncExecutionStatus{ + {AsyncStatusQueued, AsyncStatusRunning}, + {AsyncStatusQueued, AsyncStatusCancelled}, + {AsyncStatusRunning, AsyncStatusSuccess}, + {AsyncStatusRunning, AsyncStatusFailure}, + {AsyncStatusRunning, AsyncStatusUncertain}, + {AsyncStatusFailure, AsyncStatusQueued}, + {AsyncStatusUncertain, AsyncStatusQueued}, + } + for _, transition := range valid { + require.NoError(t, ValidateAsyncTransition(transition[0], transition[1])) + } + require.ErrorIs(t, ValidateAsyncTransition(AsyncStatusSuccess, AsyncStatusQueued), ErrInvalidAsyncTransition) + require.ErrorIs(t, ValidateAsyncTransition(AsyncStatusRunning, AsyncStatusCancelled), ErrInvalidAsyncTransition) +} + +func TestAsyncIdempotencyUniquePerToken(t *testing.T) { + truncateTables(t) + _, first := createAsyncFixture(t, "unique") + duplicateTask := &Task{TaskID: "task_async_duplicate", Platform: constant.TaskPlatformAsyncImage, Status: TaskStatusQueued, Data: json.RawMessage(`{}`)} + duplicate := *first + duplicate.ID = 0 + duplicate.TaskID = 0 + require.Error(t, CreateAsyncTask(duplicateTask, &duplicate)) + + loaded, err := GetAsyncJobByTokenAndKey(context.Background(), first.TokenID, first.IdempotencyKey) + require.NoError(t, err) + require.NotNil(t, loaded) + assert.Equal(t, first.RequestHash, loaded.RequestHash) + assert.Equal(t, "task_async_unique", loaded.Task.TaskID) +} + +func TestAsyncClaimHasSingleWinnerAndRenewsLease(t *testing.T) { + truncateTables(t) + _, job := createAsyncFixture(t, "claim") + + const workers = 5 + wins := make([]bool, workers) + var wg sync.WaitGroup + for i := 0; i < workers; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + _, claimed, err := ClaimAsyncJob(context.Background(), job.ID, "worker", time.Now().Add(time.Minute).Unix()) + if err == nil { + wins[i] = claimed + } + }(i) + } + wg.Wait() + winnerCount := 0 + for _, won := range wins { + if won { + winnerCount++ + } + } + assert.Equal(t, 1, winnerCount) + + renewed, err := RenewAsyncJobLease(context.Background(), job.ID, "worker", time.Now().Add(2*time.Minute).Unix()) + require.NoError(t, err) + assert.True(t, renewed) + renewed, err = RenewAsyncJobLease(context.Background(), job.ID, "other-worker", time.Now().Add(3*time.Minute).Unix()) + require.NoError(t, err) + assert.False(t, renewed) +} + +func TestRecoverExpiredAsyncJobs(t *testing.T) { + truncateTables(t) + _, beforeSend := createAsyncFixture(t, "recover-before") + _, claimed, err := ClaimAsyncJob(context.Background(), beforeSend.ID, "worker-a", time.Now().Add(-time.Minute).Unix()) + require.NoError(t, err) + require.True(t, claimed) + + _, afterSend := createAsyncFixture(t, "recover-after") + _, claimed, err = ClaimAsyncJob(context.Background(), afterSend.ID, "worker-b", time.Now().Add(-time.Minute).Unix()) + require.NoError(t, err) + require.True(t, claimed) + marked, err := MarkAsyncRequestSent(context.Background(), afterSend.ID, "worker-b", time.Now().Add(-2*time.Minute).Unix()) + require.NoError(t, err) + require.True(t, marked) + + summary, err := RecoverExpiredAsyncJobs(context.Background(), time.Now().Unix(), 10) + require.NoError(t, err) + assert.Equal(t, 1, summary.Requeued) + assert.Equal(t, 1, summary.Uncertain) + + var requeued, uncertain AsyncJob + require.NoError(t, DB.First(&requeued, beforeSend.ID).Error) + require.NoError(t, DB.First(&uncertain, afterSend.ID).Error) + assert.Equal(t, AsyncStatusQueued, requeued.ExecutionStatus) + assert.Equal(t, AsyncStatusUncertain, uncertain.ExecutionStatus) +} + +func TestCancelQueuedAsyncJob(t *testing.T) { + truncateTables(t) + task, _ := createAsyncFixture(t, "cancel") + + job, changed, err := CancelQueuedAsyncJob(context.Background(), task.TaskID, 3) + require.NoError(t, err) + require.True(t, changed) + assert.Equal(t, AsyncStatusCancelled, job.ExecutionStatus) + assert.EqualValues(t, TaskStatusCancelled, job.Task.Status) + + _, changed, err = CancelQueuedAsyncJob(context.Background(), task.TaskID, 3) + require.NoError(t, err) + assert.False(t, changed) +} + +func TestAsyncBillingRefundIsIdempotent(t *testing.T) { + truncateTables(t) + require.NoError(t, DB.Create(&User{Id: 11, Username: "async-user", Quota: 900, Status: common.UserStatusEnabled}).Error) + require.NoError(t, DB.Create(&Token{Id: 12, UserId: 11, Key: "async-token", Name: "async", Status: common.TokenStatusEnabled, RemainQuota: 900, UsedQuota: 100}).Error) + require.NoError(t, DB.Create(&Channel{Id: 13, Name: "async-channel", Status: common.ChannelStatusEnabled}).Error) + + task := &Task{ + TaskID: "task_async_refund", + Platform: constant.TaskPlatformAsyncImage, + UserId: 11, + ChannelId: 13, + Quota: 100, + Status: TaskStatusQueued, + Progress: "0%", + PrivateData: TaskPrivateData{ + BillingSource: "wallet", + TokenId: 12, + }, + Data: json.RawMessage(`{}`), + } + job := &AsyncJob{TokenID: 12, ChannelID: 13, EndpointType: AsyncEndpointImageGeneration, RequestPayload: []byte("encrypted"), RequestHash: "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", IdempotencyKey: "refund-once", ExecutionStatus: AsyncStatusQueued, BillingStatus: AsyncBillingReserved} + require.NoError(t, CreateAsyncTask(task, job)) + _, changed, err := CancelQueuedAsyncJob(context.Background(), task.TaskID, 12) + require.NoError(t, err) + require.True(t, changed) + + refunded, err := RefundAsyncJobBilling(context.Background(), job.ID) + require.NoError(t, err) + require.True(t, refunded) + refunded, err = RefundAsyncJobBilling(context.Background(), job.ID) + require.NoError(t, err) + assert.False(t, refunded) + + var user User + var token Token + require.NoError(t, DB.First(&user, 11).Error) + require.NoError(t, DB.First(&token, 12).Error) + assert.Equal(t, 1000, user.Quota) + assert.Equal(t, 1000, token.RemainQuota) + assert.Equal(t, 0, token.UsedQuota) +} + +func TestAsyncBillingSettlementIsIdempotent(t *testing.T) { + truncateTables(t) + require.NoError(t, DB.Create(&User{Id: 21, Username: "settle-user", Quota: 900, Status: common.UserStatusEnabled}).Error) + require.NoError(t, DB.Create(&Token{Id: 22, UserId: 21, Key: "settle-token", Name: "async", Status: common.TokenStatusEnabled, RemainQuota: 900, UsedQuota: 100}).Error) + require.NoError(t, DB.Create(&Channel{Id: 23, Name: "settle-channel", Status: common.ChannelStatusEnabled}).Error) + + task := &Task{TaskID: "task_async_settle", Platform: constant.TaskPlatformAsyncImage, UserId: 21, ChannelId: 23, Quota: 100, Status: TaskStatusQueued, Progress: "0%", PrivateData: TaskPrivateData{BillingSource: "wallet", TokenId: 22}, Data: json.RawMessage(`{}`)} + job := &AsyncJob{TokenID: 22, ChannelID: 23, EndpointType: AsyncEndpointImageGeneration, RequestPayload: []byte("encrypted"), RequestHash: "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb", IdempotencyKey: "settle-once", ExecutionStatus: AsyncStatusQueued, BillingStatus: AsyncBillingReserved} + require.NoError(t, CreateAsyncTask(task, job)) + _, claimed, err := ClaimAsyncJob(context.Background(), job.ID, "worker", time.Now().Add(time.Minute).Unix()) + require.NoError(t, err) + require.True(t, claimed) + completed, err := CompleteAsyncJob(context.Background(), job.ID, "worker", AsyncStatusSuccess, json.RawMessage(`{"data":[]}`), "", "", "", false) + require.NoError(t, err) + require.True(t, completed) + + settled, err := SettleAsyncJobBilling(context.Background(), job.ID) + require.NoError(t, err) + require.True(t, settled) + settled, err = SettleAsyncJobBilling(context.Background(), job.ID) + require.NoError(t, err) + assert.False(t, settled) + + var user User + var channel Channel + require.NoError(t, DB.First(&user, 21).Error) + require.NoError(t, DB.First(&channel, 23).Error) + assert.Equal(t, 100, user.UsedQuota) + assert.Equal(t, 1, user.RequestCount) + assert.Equal(t, int64(100), channel.UsedQuota) +} + +func TestUncertainAsyncBillingReconciliationSettlesExactlyOnce(t *testing.T) { + truncateTables(t) + require.NoError(t, DB.Create(&User{Id: 24, Username: "uncertain-settle-user", Quota: 900, Status: common.UserStatusEnabled}).Error) + require.NoError(t, DB.Create(&Token{Id: 25, UserId: 24, Key: "uncertain-settle-token", Name: "async", Status: common.TokenStatusEnabled, RemainQuota: 900, UsedQuota: 100}).Error) + require.NoError(t, DB.Create(&Channel{Id: 26, Name: "uncertain-settle-channel", Status: common.ChannelStatusEnabled}).Error) + + task := &Task{TaskID: "task_async_uncertain_settle", Platform: constant.TaskPlatformAsyncImage, UserId: 24, ChannelId: 26, Quota: 100, Status: TaskStatusUncertain, Progress: "0%", FinishTime: time.Now().Unix(), PrivateData: TaskPrivateData{BillingSource: "wallet", TokenId: 25}, Data: json.RawMessage(`{}`)} + job := &AsyncJob{TokenID: 25, ChannelID: 26, EndpointType: AsyncEndpointImageGeneration, RequestPayload: []byte("encrypted"), RequestHash: "abababababababababababababababababababababababababababababababab", IdempotencyKey: "uncertain-settle-once", ExecutionStatus: AsyncStatusUncertain, BillingStatus: AsyncBillingReserved, RequestSentAt: time.Now().Unix()} + require.NoError(t, CreateAsyncTask(task, job)) + + processed, err := ReconcileAsyncBilling(context.Background(), 10) + require.NoError(t, err) + assert.Equal(t, 1, processed) + processed, err = ReconcileAsyncBilling(context.Background(), 10) + require.NoError(t, err) + assert.Zero(t, processed) + + var user User + var channel Channel + var loadedJob AsyncJob + require.NoError(t, DB.First(&user, 24).Error) + require.NoError(t, DB.First(&channel, 26).Error) + require.NoError(t, DB.First(&loadedJob, job.ID).Error) + assert.Equal(t, 100, user.UsedQuota) + assert.Equal(t, 1, user.RequestCount) + assert.Equal(t, int64(100), channel.UsedQuota) + assert.Equal(t, AsyncBillingSettled, loadedJob.BillingStatus) +} + +func TestManualRetryReservesRefundedFailureAgain(t *testing.T) { + truncateTables(t) + require.NoError(t, DB.Create(&User{Id: 31, Username: "retry-user", Quota: 1000, Status: common.UserStatusEnabled}).Error) + require.NoError(t, DB.Create(&Token{Id: 32, UserId: 31, Key: "retry-token", Name: "async", Status: common.TokenStatusEnabled, RemainQuota: 1000}).Error) + require.NoError(t, DB.Create(&Channel{Id: 33, Name: "retry-channel", Status: common.ChannelStatusEnabled}).Error) + task := &Task{TaskID: "task_async_retry_failure", Platform: constant.TaskPlatformAsyncImage, UserId: 31, ChannelId: 33, Quota: 100, Status: TaskStatusFailure, Progress: "0%", FinishTime: time.Now().Unix(), PrivateData: TaskPrivateData{BillingSource: "wallet", TokenId: 32}, Data: json.RawMessage(`{}`)} + job := &AsyncJob{TokenID: 32, ChannelID: 33, EndpointType: AsyncEndpointImageGeneration, RequestPayload: []byte("encrypted"), RequestHash: "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc", IdempotencyKey: "retry-refunded", ExecutionStatus: AsyncStatusFailure, BillingStatus: AsyncBillingRefunded, RefundEligible: true} + require.NoError(t, CreateAsyncTask(task, job)) + + retried, changed, err := RetryAsyncJob(context.Background(), job.ID, 99) + require.NoError(t, err) + require.True(t, changed) + assert.Equal(t, AsyncStatusQueued, retried.ExecutionStatus) + assert.Equal(t, AsyncBillingReserved, retried.BillingStatus) + assert.EqualValues(t, TaskStatusQueued, retried.Task.Status) + + var user User + var token Token + require.NoError(t, DB.First(&user, 31).Error) + require.NoError(t, DB.First(&token, 32).Error) + assert.Equal(t, 900, user.Quota) + assert.Equal(t, 900, token.RemainQuota) + assert.Equal(t, 100, token.UsedQuota) +} + +func TestManualRetryUncertainSettlesPriorAndReservesNextAttempt(t *testing.T) { + truncateTables(t) + require.NoError(t, DB.Create(&User{Id: 41, Username: "uncertain-user", Quota: 900, Status: common.UserStatusEnabled}).Error) + require.NoError(t, DB.Create(&Token{Id: 42, UserId: 41, Key: "uncertain-token", Name: "async", Status: common.TokenStatusEnabled, RemainQuota: 900, UsedQuota: 100}).Error) + require.NoError(t, DB.Create(&Channel{Id: 43, Name: "uncertain-channel", Status: common.ChannelStatusEnabled}).Error) + task := &Task{TaskID: "task_async_retry_uncertain", Platform: constant.TaskPlatformAsyncImage, UserId: 41, ChannelId: 43, Quota: 100, Status: TaskStatusUncertain, Progress: "0%", FinishTime: time.Now().Unix(), PrivateData: TaskPrivateData{BillingSource: "wallet", TokenId: 42}, Data: json.RawMessage(`{}`)} + job := &AsyncJob{TokenID: 42, ChannelID: 43, EndpointType: AsyncEndpointImageGeneration, RequestPayload: []byte("encrypted"), RequestHash: "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd", IdempotencyKey: "retry-uncertain", ExecutionStatus: AsyncStatusUncertain, BillingStatus: AsyncBillingReserved, RequestSentAt: time.Now().Unix()} + require.NoError(t, CreateAsyncTask(task, job)) + + retried, changed, err := RetryAsyncJob(context.Background(), job.ID, 99) + require.NoError(t, err) + require.True(t, changed) + assert.Equal(t, AsyncStatusQueued, retried.ExecutionStatus) + assert.Zero(t, retried.RequestSentAt) + + var user User + var token Token + var channel Channel + require.NoError(t, DB.First(&user, 41).Error) + require.NoError(t, DB.First(&token, 42).Error) + require.NoError(t, DB.First(&channel, 43).Error) + assert.Equal(t, 800, user.Quota) + assert.Equal(t, 100, user.UsedQuota) + assert.Equal(t, 1, user.RequestCount) + assert.Equal(t, 800, token.RemainQuota) + assert.Equal(t, 200, token.UsedQuota) + assert.Equal(t, int64(100), channel.UsedQuota) + + _, changed, err = RetryAsyncJob(context.Background(), job.ID, 99) + require.NoError(t, err) + assert.False(t, changed) +} + +func TestAsyncChannelSelectionSkipsNonYunwuOptIn(t *testing.T) { + truncateTables(t) + t.Setenv("ASYNC_YUNWU_ALLOWED_BASE_URLS", "https://yunwu.ai") + archive := true + invalidBase := "https://example.com" + validBase := "https://yunwu.ai" + invalid := &Channel{Id: 51, Name: "invalid-async-origin", BaseURL: &invalidBase, Status: common.ChannelStatusEnabled, Models: "image-model", Group: "default"} + valid := &Channel{Id: 52, Name: "valid-yunwu-origin", BaseURL: &validBase, Status: common.ChannelStatusEnabled, Models: "image-model", Group: "default"} + setting := dto.ChannelSettings{AsyncImageEnabled: true, AsyncImageModels: []string{"image-model"}, AsyncAutoArchive: &archive} + invalid.SetSetting(setting) + valid.SetSetting(setting) + require.NoError(t, DB.Create(invalid).Error) + require.NoError(t, DB.Create(valid).Error) + invalidPriority := int64(100) + validPriority := int64(10) + require.NoError(t, DB.Create(&Ability{Group: "default", Model: "image-model", ChannelId: invalid.Id, Enabled: true, Priority: &invalidPriority, Weight: 1}).Error) + require.NoError(t, DB.Create(&Ability{Group: "default", Model: "image-model", ChannelId: valid.Id, Enabled: true, Priority: &validPriority, Weight: 1}).Error) + + selected, err := GetAsyncImageChannel("default", "image-model") + require.NoError(t, err) + require.NotNil(t, selected) + assert.Equal(t, valid.Id, selected.Id) +} + +func TestAsyncChannelSelectionAcceptsAllowedGRSAIProvider(t *testing.T) { + truncateTables(t) + t.Setenv("ASYNC_GRSAI_ALLOWED_BASE_URLS", "https://grsaiapi.com") + archive := true + baseURL := "https://grsaiapi.com/v1" + channel := &Channel{Id: 53, Name: "valid-grsai-origin", BaseURL: &baseURL, Status: common.ChannelStatusEnabled, Models: "nano-banana-2", Group: "default"} + channel.SetSetting(dto.ChannelSettings{AsyncImageEnabled: true, AsyncImageModels: []string{"nano-banana-2"}, AsyncAutoArchive: &archive}) + require.NoError(t, DB.Create(channel).Error) + priority := int64(20) + require.NoError(t, DB.Create(&Ability{Group: "default", Model: "nano-banana-2", ChannelId: channel.Id, Enabled: true, Priority: &priority, Weight: 1}).Error) + + selected, err := GetAsyncImageChannel("default", "nano-banana-2") + require.NoError(t, err) + require.NotNil(t, selected) + assert.Equal(t, channel.Id, selected.Id) +} diff --git a/model/channel.go b/model/channel.go index 1de876877559..b83e325a2ab1 100644 --- a/model/channel.go +++ b/model/channel.go @@ -5,14 +5,17 @@ import ( "encoding/json" "errors" "fmt" + "math" "math/rand" "strings" "sync" + "unicode" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/dto" "github.com/QuantumNous/new-api/logger" + "github.com/QuantumNous/new-api/setting/ratio_setting" "github.com/QuantumNous/new-api/types" "github.com/samber/lo" @@ -293,6 +296,82 @@ func (channel *Channel) GetModels() []string { return strings.Split(strings.Trim(channel.Models, ","), ",") } +// GetAsyncImageChannel selects an enabled channel at the highest configured +// priority that explicitly enables the asynchronous image wrapper for model. +// It intentionally queries PostgreSQL rather than the general relay cache so a +// synchronous-only channel can never be selected for an asynchronous job. +func GetAsyncImageChannel(group string, modelName string) (*Channel, error) { + type candidate struct { + Channel + AbilityPriority int64 `gorm:"column:ability_priority"` + AbilityWeight uint `gorm:"column:ability_weight"` + } + + load := func(candidateModel string) ([]candidate, error) { + var candidates []candidate + err := DB.Table("channels"). + Select("channels.*, abilities.priority AS ability_priority, abilities.weight AS ability_weight"). + Joins("JOIN abilities ON abilities.channel_id = channels.id"). + Where("abilities."+commonGroupCol+" = ? AND abilities.model = ? AND abilities.enabled = ?", group, candidateModel, true). + Where("channels.status = ?", common.ChannelStatusEnabled). + Order("abilities.priority DESC, channels.id ASC"). + Scan(&candidates).Error + return candidates, err + } + + candidates, err := load(modelName) + if err != nil { + return nil, err + } + if len(candidates) == 0 { + normalized := ratio_setting.FormatMatchingModelName(modelName) + if normalized != modelName { + candidates, err = load(normalized) + if err != nil { + return nil, err + } + } + } + + filtered := make([]candidate, 0, len(candidates)) + var highestPriority int64 + for _, item := range candidates { + setting := item.GetSetting() + if !setting.AllowsAsyncImageModel(modelName) || !setting.AsyncArchiveEnabled() || !common.IsAllowedAsyncImageBaseURL(item.GetBaseURL()) { + continue + } + if len(filtered) == 0 { + highestPriority = item.AbilityPriority + } + if item.AbilityPriority != highestPriority { + break + } + filtered = append(filtered, item) + } + if len(filtered) == 0 { + return nil, nil + } + if len(filtered) == 1 { + channel := filtered[0].Channel + return &channel, nil + } + + weightSum := 0 + for _, item := range filtered { + weightSum += int(item.AbilityWeight) + 10 + } + target := common.GetRandomInt(weightSum) + for _, item := range filtered { + target -= int(item.AbilityWeight) + 10 + if target <= 0 { + channel := item.Channel + return &channel, nil + } + } + channel := filtered[len(filtered)-1].Channel + return &channel, nil +} + func (channel *Channel) GetGroups() []string { if channel.Group == "" { return []string{} @@ -971,6 +1050,30 @@ func (channel *Channel) ValidateSettings() error { return err } } + if rate := channelOtherSettings.UpstreamCostRateCNY; rate != nil { + if *rate <= 0 || + math.IsNaN(*rate) || + math.IsInf(*rate, 0) || + *rate > dto.MaxUpstreamCostRateCNY { + return fmt.Errorf("upstream_cost_rate_cny must be greater than 0 and no more than %d", dto.MaxUpstreamCostRateCNY) + } + } + switch channelOtherSettings.UpstreamCostMode { + case "", dto.UpstreamCostModeAuto, dto.UpstreamCostModeResponseCost, dto.UpstreamCostModeBillingUnits: + default: + return fmt.Errorf("unsupported upstream_cost_mode: %s", channelOtherSettings.UpstreamCostMode) + } + if channelOtherSettings.UpstreamCostMode != "" && channelOtherSettings.UpstreamCostRateCNY == nil { + return fmt.Errorf("upstream_cost_rate_cny is required when upstream cost tracking is enabled") + } + unit := strings.TrimSpace(channelOtherSettings.UpstreamCostUnit) + if len(unit) > dto.MaxUpstreamCostUnitLength || strings.IndexFunc(unit, unicode.IsControl) >= 0 { + return fmt.Errorf("upstream_cost_unit must not exceed %d characters or contain control characters", dto.MaxUpstreamCostUnitLength) + } + priceVersion := strings.TrimSpace(channelOtherSettings.UpstreamCostPriceVersion) + if len(priceVersion) > dto.MaxUpstreamPriceVersionSize || strings.IndexFunc(priceVersion, unicode.IsControl) >= 0 { + return fmt.Errorf("upstream_cost_price_version must not exceed %d characters or contain control characters", dto.MaxUpstreamPriceVersionSize) + } if channel.Type == constant.ChannelTypeAdvancedCustom && channelOtherSettings.UpstreamModelUpdateCheckEnabled { if _, ok := channelOtherSettings.AdvancedCustom.ModelListRoute(); !ok { return fmt.Errorf("advanced custom channels require a %s route when upstream model update checks are enabled", dto.AdvancedCustomModelListPath) diff --git a/model/channel_settings_test.go b/model/channel_settings_test.go index c4974faf4f18..04ebae5107b7 100644 --- a/model/channel_settings_test.go +++ b/model/channel_settings_test.go @@ -66,3 +66,93 @@ func TestAdvancedCustomChannelRequiresModelListRouteOnlyWhenUpdateChecksEnabled( }) } } + +func TestChannelUpstreamCostRateValidation(t *testing.T) { + validRate := 0.495 + zeroRate := 0.0 + tooLargeRate := float64(dto.MaxUpstreamCostRateCNY + 1) + + tests := []struct { + name string + rate *float64 + wantErr bool + }{ + {name: "unset disables cost tracking"}, + {name: "positive fractional rate", rate: &validRate}, + {name: "zero is rejected", rate: &zeroRate, wantErr: true}, + {name: "excessive rate is rejected", rate: &tooLargeRate, wantErr: true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + channel := &Channel{Type: constant.ChannelTypeOpenAI} + channel.SetOtherSettings(dto.ChannelOtherSettings{UpstreamCostRateCNY: tt.rate}) + + err := channel.ValidateSettings() + if tt.wantErr { + require.Error(t, err) + assert.Contains(t, err.Error(), "upstream_cost_rate_cny") + return + } + require.NoError(t, err) + }) + } +} + +func TestChannelUpstreamCostProfileValidation(t *testing.T) { + validRate := 0.495 + tests := []struct { + name string + settings dto.ChannelOtherSettings + wantErr string + }{ + { + name: "automatic profile is accepted", + settings: dto.ChannelOtherSettings{ + UpstreamCostMode: dto.UpstreamCostModeAuto, + UpstreamCostUnit: "CREDIT", + UpstreamCostRateCNY: &validRate, + UpstreamCostPriceVersion: "yunwu-2026-07", + }, + }, + { + name: "mode requires a rate", + settings: dto.ChannelOtherSettings{ + UpstreamCostMode: dto.UpstreamCostModeBillingUnits, + }, + wantErr: "upstream_cost_rate_cny is required", + }, + { + name: "unknown mode is rejected", + settings: dto.ChannelOtherSettings{ + UpstreamCostMode: "guess", + UpstreamCostRateCNY: &validRate, + }, + wantErr: "unsupported upstream_cost_mode", + }, + { + name: "control characters in unit are rejected", + settings: dto.ChannelOtherSettings{ + UpstreamCostMode: dto.UpstreamCostModeAuto, + UpstreamCostUnit: "USD\nCNY", + UpstreamCostRateCNY: &validRate, + }, + wantErr: "upstream_cost_unit", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + channel := &Channel{Type: constant.ChannelTypeOpenAI} + channel.SetOtherSettings(tt.settings) + + err := channel.ValidateSettings() + if tt.wantErr == "" { + require.NoError(t, err) + return + } + require.Error(t, err) + assert.Contains(t, err.Error(), tt.wantErr) + }) + } +} diff --git a/model/log.go b/model/log.go index 401d53c435a5..7be174370976 100644 --- a/model/log.go +++ b/model/log.go @@ -341,10 +341,10 @@ type RecordConsumeLogParams struct { } func RecordConsumeLog(c *gin.Context, userId int, params RecordConsumeLogParams) { - if !common.LogConsumeEnabled { - return + recordUsageLog := common.LogConsumeEnabled + if recordUsageLog { + logger.LogInfo(c, fmt.Sprintf("record consume log: userId=%d, params=%s", userId, common.GetJsonString(params))) } - logger.LogInfo(c, fmt.Sprintf("record consume log: userId=%d, params=%s", userId, common.GetJsonString(params))) username := c.GetString("username") requestId := c.GetString(common.RequestIdKey) upstreamRequestId := c.GetString(common.UpstreamRequestIdKey) @@ -352,9 +352,11 @@ func RecordConsumeLog(c *gin.Context, userId int, params RecordConsumeLogParams) otherStr := common.MapToJsonStr(params.Other) // 判断是否需要记录 IP needRecordIp := false - if settingMap, err := GetUserSetting(userId, false); err == nil { - if settingMap.RecordIpLog { - needRecordIp = true + if recordUsageLog { + if settingMap, err := GetUserSetting(userId, false); err == nil { + if settingMap.RecordIpLog { + needRecordIp = true + } } } log := &Log{ @@ -383,9 +385,18 @@ func RecordConsumeLog(c *gin.Context, userId int, params RecordConsumeLogParams) UpstreamRequestId: upstreamRequestId, Other: otherStr, } - err := createLog(log) - if err != nil { - logger.LogError(c, "failed to record log: "+err.Error()) + if recordUsageLog { + if err := createLog(log); err != nil { + logger.LogError(c, "failed to record log: "+err.Error()) + } + } else { + ensureLogRequestId(log) + } + if costErr := RecordUpstreamCostFromLog(log, params.Other); costErr != nil { + logger.LogError(c, "failed to record upstream cost: "+costErr.Error()) + } + if !recordUsageLog { + return } if common.DataExportEnabled { LogQuotaData(QuotaDataLogParams{ @@ -412,14 +423,13 @@ type RecordTaskBillingLogParams struct { Quota int TokenId int Group string + RequestId string Other map[string]interface{} NodeName string // 任务发起节点；为空时回退当前节点 } func RecordTaskBillingLog(params RecordTaskBillingLogParams) { - if params.LogType == LogTypeConsume && !common.LogConsumeEnabled { - return - } + recordUsageLog := params.LogType != LogTypeConsume || common.LogConsumeEnabled username, _ := GetUsernameById(params.UserId, false) tokenName := "" if params.TokenId > 0 { @@ -440,11 +450,21 @@ func RecordTaskBillingLog(params RecordTaskBillingLogParams) { ChannelId: params.ChannelId, TokenId: params.TokenId, Group: params.Group, + RequestId: params.RequestId, Other: common.MapToJsonStr(params.Other), } - err := createLog(log) - if err != nil { - common.SysLog("failed to record task billing log: " + err.Error()) + if recordUsageLog { + if err := createLog(log); err != nil { + common.SysLog("failed to record task billing log: " + err.Error()) + } + } else { + ensureLogRequestId(log) + } + if costErr := RecordUpstreamCostFromLog(log, params.Other); costErr != nil { + common.SysLog("failed to record task upstream cost: " + costErr.Error()) + } + if !recordUsageLog { + return } if params.LogType == LogTypeConsume && common.DataExportEnabled { nodeName := params.NodeName diff --git a/model/main.go b/model/main.go index ac63d1f59c12..e69e28095993 100644 --- a/model/main.go +++ b/model/main.go @@ -283,7 +283,11 @@ func migrateDB() error { &Midjourney{}, &TopUp{}, &QuotaData{}, + &UpstreamCostRecord{}, &Task{}, + &AsyncJob{}, + &Artifact{}, + &TaskEvent{}, &Model{}, &Vendor{}, &PrefillGroup{}, @@ -346,7 +350,11 @@ func migrateDBFast() error { {&Midjourney{}, "Midjourney"}, {&TopUp{}, "TopUp"}, {&QuotaData{}, "QuotaData"}, + {&UpstreamCostRecord{}, "UpstreamCostRecord"}, {&Task{}, "Task"}, + {&AsyncJob{}, "AsyncJob"}, + {&Artifact{}, "Artifact"}, + {&TaskEvent{}, "TaskEvent"}, {&Model{}, "Model"}, {&Vendor{}, "Vendor"}, {&PrefillGroup{}, "PrefillGroup"}, diff --git a/model/pricing.go b/model/pricing.go index 440e1e0999b9..5e9e08a3c2af 100644 --- a/model/pricing.go +++ b/model/pricing.go @@ -11,31 +11,33 @@ import ( "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/dto" "github.com/QuantumNous/new-api/setting/billing_setting" + "github.com/QuantumNous/new-api/setting/model_setting" "github.com/QuantumNous/new-api/setting/ratio_setting" "github.com/QuantumNous/new-api/types" ) type Pricing struct { - ModelName string `json:"model_name"` - Description string `json:"description,omitempty"` - Icon string `json:"icon,omitempty"` - Tags string `json:"tags,omitempty"` - VendorID int `json:"vendor_id,omitempty"` - QuotaType int `json:"quota_type"` - ModelRatio float64 `json:"model_ratio"` - ModelPrice float64 `json:"model_price"` - OwnerBy string `json:"owner_by"` - CompletionRatio float64 `json:"completion_ratio"` - CacheRatio *float64 `json:"cache_ratio,omitempty"` - CreateCacheRatio *float64 `json:"create_cache_ratio,omitempty"` - ImageRatio *float64 `json:"image_ratio,omitempty"` - AudioRatio *float64 `json:"audio_ratio,omitempty"` - AudioCompletionRatio *float64 `json:"audio_completion_ratio,omitempty"` - EnableGroup []string `json:"enable_groups"` - SupportedEndpointTypes []constant.EndpointType `json:"supported_endpoint_types"` - BillingMode string `json:"billing_mode,omitempty"` - BillingExpr string `json:"billing_expr,omitempty"` - PricingVersion string `json:"pricing_version,omitempty"` + ModelName string `json:"model_name"` + Description string `json:"description,omitempty"` + Icon string `json:"icon,omitempty"` + Tags string `json:"tags,omitempty"` + VendorID int `json:"vendor_id,omitempty"` + QuotaType int `json:"quota_type"` + ModelRatio float64 `json:"model_ratio"` + ModelPrice float64 `json:"model_price"` + OwnerBy string `json:"owner_by"` + CompletionRatio float64 `json:"completion_ratio"` + CacheRatio *float64 `json:"cache_ratio,omitempty"` + CreateCacheRatio *float64 `json:"create_cache_ratio,omitempty"` + ImageRatio *float64 `json:"image_ratio,omitempty"` + AudioRatio *float64 `json:"audio_ratio,omitempty"` + AudioCompletionRatio *float64 `json:"audio_completion_ratio,omitempty"` + EnableGroup []string `json:"enable_groups"` + SupportedEndpointTypes []constant.EndpointType `json:"supported_endpoint_types"` + BillingMode string `json:"billing_mode,omitempty"` + BillingExpr string `json:"billing_expr,omitempty"` + PricingVersion string `json:"pricing_version,omitempty"` + ImageGeneration *model_setting.ImageGenerationCapabilities `json:"image_generation,omitempty"` } type PricingVendor struct { @@ -360,6 +362,7 @@ func updatePricing() { ModelName: model, EnableGroup: groups.Items(), SupportedEndpointTypes: modelSupportEndpointTypes[model], + ImageGeneration: model_setting.GetImageGenerationCapabilities(model), } // 补充模型元数据（描述、标签、供应商、状态） @@ -411,7 +414,7 @@ func updatePricing() { // 防止大更新后数据不通用 if len(pricingMap) > 0 { - pricingMap[0].PricingVersion = "5a90f2b86c08bd983a9a2e6d66c255f4eaef9c4bc934386d2b6ae84ef0ff1f1f" + pricingMap[0].PricingVersion = "c93f4990684023eef6ec35670795311442418348207af6cd284468468f812f69" } // 刷新缓存映射，供高并发快速查询 diff --git a/model/task.go b/model/task.go index ecaf70f3eff3..945a3fc2a00b 100644 --- a/model/task.go +++ b/model/task.go @@ -38,6 +38,8 @@ const ( TaskStatusInProgress = "IN_PROGRESS" TaskStatusFailure = "FAILURE" TaskStatusSuccess = "SUCCESS" + TaskStatusUncertain = "UNCERTAIN" + TaskStatusCancelled = "CANCELLED" TaskStatusUnknown = "UNKNOWN" ) @@ -298,6 +300,7 @@ func GetTimedOutUnfinishedTasks(cutoffUnix int64, limit int) []*Task { var tasks []*Task err := DB.Where("progress != ?", "100%"). Where("status NOT IN ?", []string{TaskStatusFailure, TaskStatusSuccess}). + Where("platform != ?", constant.TaskPlatformAsyncImage). Where("submit_time < ?", cutoffUnix). Order("submit_time"). Limit(limit). @@ -334,7 +337,8 @@ func GetAllUnFinishSyncTasks(limit int) []*Task { var tasks []*Task var err error // get all tasks progress is not 100% - err = DB.Where("progress != ?", "100%").Where("status != ?", TaskStatusFailure).Where("status != ?", TaskStatusSuccess).Limit(limit).Order("id").Find(&tasks).Error + err = DB.Where("progress != ?", "100%").Where("status != ?", TaskStatusFailure).Where("status != ?", TaskStatusSuccess). + Where("platform != ?", constant.TaskPlatformAsyncImage).Limit(limit).Order("id").Find(&tasks).Error if err != nil { return nil } @@ -351,6 +355,7 @@ func HasUnfinishedSyncTasks() bool { Where("progress != ?", "100%"). Where("status != ?", TaskStatusFailure). Where("status != ?", TaskStatusSuccess). + Where("platform != ?", constant.TaskPlatformAsyncImage). Limit(1). Pluck("id", &id).Error return err == nil && id != 0 diff --git a/model/task_cas_test.go b/model/task_cas_test.go index d99a8d89eb51..9076f77e822e 100644 --- a/model/task_cas_test.go +++ b/model/task_cas_test.go @@ -36,6 +36,9 @@ func TestMain(m *testing.M) { if err := db.AutoMigrate( &Task{}, + &AsyncJob{}, + &Artifact{}, + &TaskEvent{}, &User{}, &UserSession{}, &AuthFlow{}, @@ -52,6 +55,7 @@ func TestMain(m *testing.M) { &SubscriptionPlan{}, &SubscriptionOrder{}, &UserSubscription{}, + &SubscriptionPreConsumeRecord{}, &UserOAuthBinding{}, &PerfMetric{}, &SystemInstance{}, @@ -68,6 +72,9 @@ func truncateTables(t *testing.T) { t.Helper() t.Cleanup(func() { DB.Exec("DELETE FROM tasks") + DB.Exec("DELETE FROM task_events") + DB.Exec("DELETE FROM artifacts") + DB.Exec("DELETE FROM async_jobs") DB.Exec("DELETE FROM auth_flows") DB.Exec("DELETE FROM external_identity_claims") DB.Exec("DELETE FROM user_sessions") @@ -83,6 +90,7 @@ func truncateTables(t *testing.T) { DB.Exec("DELETE FROM abilities") DB.Exec("DELETE FROM top_ups") DB.Exec("DELETE FROM subscription_orders") + DB.Exec("DELETE FROM subscription_pre_consume_records") DB.Exec("DELETE FROM subscription_plans") DB.Exec("DELETE FROM user_subscriptions") DB.Exec("DELETE FROM perf_metrics") diff --git a/model/upstream_cost.go b/model/upstream_cost.go new file mode 100644 index 000000000000..a801c84e0686 --- /dev/null +++ b/model/upstream_cost.go @@ -0,0 +1,168 @@ +package model + +import ( + "errors" + "fmt" + "strconv" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/dto" + + "gorm.io/gorm/clause" +) + +// UpstreamCostRecord is the canonical CNY cost ledger. Monetary accounting is +// stored as integer micro-yuan; native values remain decimal strings so the +// original upstream unit can be audited without floating-point round trips. +type UpstreamCostRecord struct { + Id int64 `json:"id"` + RequestId string `json:"request_id" gorm:"type:varchar(64);uniqueIndex"` + CreatedAt int64 `json:"created_at" gorm:"bigint;index"` + UserId int `json:"user_id" gorm:"index"` + Username string `json:"username" gorm:"type:varchar(191);index"` + TokenName string `json:"token_name" gorm:"type:varchar(191);index"` + UseGroup string `json:"group" gorm:"column:use_group;type:varchar(64);index"` + ChannelId int `json:"channel_id" gorm:"index"` + ModelName string `json:"model_name" gorm:"type:varchar(191);index"` + UpstreamRequestId string `json:"upstream_request_id" gorm:"type:varchar(128);index"` + Status string `json:"status" gorm:"type:varchar(16);index"` + Mode string `json:"mode" gorm:"type:varchar(32)"` + Source string `json:"source" gorm:"type:varchar(32)"` + Reason string `json:"reason" gorm:"type:varchar(64)"` + NativeUnit string `json:"native_unit" gorm:"type:varchar(32)"` + NativeAmount string `json:"native_amount" gorm:"type:varchar(64)"` + RateCNYPerUnit string `json:"rate_cny_per_unit" gorm:"type:varchar(64)"` + AmountCNYMicros int64 `json:"amount_cny_micros" gorm:"bigint"` + Estimated bool `json:"estimated"` + PriceVersion string `json:"price_version" gorm:"type:varchar(64)"` + SettlementCurrency string `json:"settlement_currency" gorm:"type:varchar(8)"` +} + +func upstreamCostSnapshotFromOther(other map[string]interface{}) (*dto.UpstreamCostSnapshot, bool) { + if other == nil { + return nil, false + } + adminInfo, ok := other["admin_info"].(map[string]interface{}) + if !ok || adminInfo == nil { + return nil, false + } + value, ok := adminInfo["upstream_cost"] + if !ok || value == nil { + return nil, false + } + if snapshot, ok := value.(*dto.UpstreamCostSnapshot); ok && snapshot != nil { + return snapshot, true + } + raw, err := common.Marshal(value) + if err != nil { + return nil, false + } + var snapshot dto.UpstreamCostSnapshot + if err := common.Unmarshal(raw, &snapshot); err != nil { + return nil, false + } + return &snapshot, true +} + +// RecordUpstreamCostFromLog persists the admin-only snapshot attached to a +// consume log. The request ID is unique, making retries idempotent. +func RecordUpstreamCostFromLog(log *Log, other map[string]interface{}) error { + if log == nil { + return errors.New("log is required") + } + snapshot, ok := upstreamCostSnapshotFromOther(other) + if !ok { + return nil + } + if snapshot.Status != dto.UpstreamCostStatusSettled && snapshot.Status != dto.UpstreamCostStatusUnpriced { + return fmt.Errorf("unsupported upstream cost status: %s", snapshot.Status) + } + ensureLogRequestId(log) + record := &UpstreamCostRecord{ + RequestId: log.RequestId, + CreatedAt: log.CreatedAt, + UserId: log.UserId, + Username: log.Username, + TokenName: log.TokenName, + UseGroup: log.Group, + ChannelId: log.ChannelId, + ModelName: log.ModelName, + UpstreamRequestId: log.UpstreamRequestId, + Status: snapshot.Status, + Mode: snapshot.Mode, + Source: snapshot.Source, + Reason: snapshot.Reason, + NativeUnit: snapshot.NativeUnit, + NativeAmount: snapshot.NativeAmountDecimal, + RateCNYPerUnit: snapshot.RateCNYPerUnitDecimal, + AmountCNYMicros: snapshot.AmountCNYMicros, + Estimated: snapshot.Estimated, + PriceVersion: snapshot.PriceVersion, + SettlementCurrency: snapshot.SettlementCurrency, + } + if record.NativeAmount == "" { + record.NativeAmount = strconv.FormatFloat(snapshot.NativeAmount, 'f', -1, 64) + } + if record.RateCNYPerUnit == "" { + record.RateCNYPerUnit = strconv.FormatFloat(snapshot.RateCNYPerUnit, 'f', -1, 64) + } + return DB.Clauses(clause.OnConflict{ + Columns: []clause.Column{{Name: "request_id"}}, + DoNothing: true, + }).Create(record).Error +} + +type UpstreamCostStat struct { + SettledRequests int64 `json:"settled_requests"` + UnpricedRequests int64 `json:"unpriced_requests"` + EstimatedRequests int64 `json:"estimated_requests"` + AmountCNYMicros int64 `json:"amount_cny_micros"` +} + +func GetUpstreamCostStat( + startTimestamp, endTimestamp int64, + modelName, username, tokenName, useGroup, requestId, upstreamRequestId string, + channelId int, +) (UpstreamCostStat, error) { + query := DB.Model(&UpstreamCostRecord{}) + if startTimestamp > 0 { + query = query.Where("created_at >= ?", startTimestamp) + } + if endTimestamp > 0 { + query = query.Where("created_at <= ?", endTimestamp) + } + if modelName != "" { + query = query.Where("model_name = ?", modelName) + } + if username != "" { + query = query.Where("username = ?", username) + } + if tokenName != "" { + query = query.Where("token_name = ?", tokenName) + } + if useGroup != "" { + query = query.Where("use_group = ?", useGroup) + } + if requestId != "" { + query = query.Where("request_id = ?", requestId) + } + if upstreamRequestId != "" { + query = query.Where("upstream_request_id = ?", upstreamRequestId) + } + if channelId > 0 { + query = query.Where("channel_id = ?", channelId) + } + + var stat UpstreamCostStat + err := query.Select( + "COALESCE(SUM(CASE WHEN status = ? THEN 1 ELSE 0 END), 0) AS settled_requests, "+ + "COALESCE(SUM(CASE WHEN status = ? THEN 1 ELSE 0 END), 0) AS unpriced_requests, "+ + "COALESCE(SUM(CASE WHEN status = ? AND source = ? THEN 1 ELSE 0 END), 0) AS estimated_requests, "+ + "COALESCE(SUM(amount_cny_micros), 0) AS amount_cny_micros", + dto.UpstreamCostStatusSettled, + dto.UpstreamCostStatusUnpriced, + dto.UpstreamCostStatusSettled, + dto.UpstreamCostSourceBillingUnits, + ).Scan(&stat).Error + return stat, err +} diff --git a/model/upstream_cost_test.go b/model/upstream_cost_test.go new file mode 100644 index 000000000000..3f46528598eb --- /dev/null +++ b/model/upstream_cost_test.go @@ -0,0 +1,149 @@ +package model + +import ( + "net/http/httptest" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/dto" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" + "gorm.io/gorm" +) + +func TestRecordUpstreamCostFromLogIsIdempotentAndAggregated(t *testing.T) { + previousDB := DB + previousMainDatabaseType := common.MainDatabaseType() + previousLogDatabaseType := common.LogDatabaseType() + db, err := gorm.Open(sqlite.Open("file:upstream_cost_ledger?mode=memory&cache=shared"), &gorm.Config{}) + require.NoError(t, err) + DB = db + common.SetDatabaseTypes(common.DatabaseTypeSQLite, common.DatabaseTypeSQLite) + initCol() + t.Cleanup(func() { + DB = previousDB + common.SetDatabaseTypes(previousMainDatabaseType, previousLogDatabaseType) + initCol() + }) + require.NoError(t, db.AutoMigrate(&UpstreamCostRecord{})) + + log := &Log{ + RequestId: "req-cost-1", + UpstreamRequestId: "upstream-1", + CreatedAt: 100, + UserId: 9, + Username: "admin", + TokenName: "Nexa", + Group: "codex", + ChannelId: 36, + ModelName: "gpt-5.6-sol", + } + snapshot := &dto.UpstreamCostSnapshot{ + Status: dto.UpstreamCostStatusSettled, + Mode: dto.UpstreamCostModeBillingUnits, + Source: dto.UpstreamCostSourceBillingUnits, + NativeUnit: "CREDIT", + NativeAmount: 0.011682, + NativeAmountDecimal: "0.011682", + RateCNYPerUnit: 0.495, + RateCNYPerUnitDecimal: "0.495", + AmountCNYMicros: 5783, + Estimated: true, + PriceVersion: "yunwu-2026-07", + SettlementCurrency: "CNY", + } + other := map[string]interface{}{ + "admin_info": map[string]interface{}{"upstream_cost": snapshot}, + } + + require.NoError(t, RecordUpstreamCostFromLog(log, other)) + require.NoError(t, RecordUpstreamCostFromLog(log, other)) + + unpricedLog := *log + unpricedLog.RequestId = "req-cost-2" + unpricedSnapshot := &dto.UpstreamCostSnapshot{ + Status: dto.UpstreamCostStatusUnpriced, + Mode: dto.UpstreamCostModeBillingUnits, + Reason: "missing_channel_cost_profile", + NativeUnit: "UNIT", + PriceVersion: "manual", + SettlementCurrency: "CNY", + } + require.NoError(t, RecordUpstreamCostFromLog(&unpricedLog, map[string]interface{}{ + "admin_info": map[string]interface{}{"upstream_cost": unpricedSnapshot}, + })) + + var recordCount int64 + require.NoError(t, db.Model(&UpstreamCostRecord{}).Count(&recordCount).Error) + assert.Equal(t, int64(2), recordCount) + var settledRecord UpstreamCostRecord + require.NoError(t, db.Where("request_id = ?", "req-cost-1").First(&settledRecord).Error) + assert.Equal(t, int64(5783), settledRecord.AmountCNYMicros) + assert.Equal(t, "0.011682", settledRecord.NativeAmount) + assert.Equal(t, "0.495", settledRecord.RateCNYPerUnit) + + stat, err := GetUpstreamCostStat( + 0, + 0, + "gpt-5.6-sol", + "admin", + "Nexa", + "codex", + "", + "upstream-1", + 36, + ) + require.NoError(t, err) + assert.Equal(t, int64(1), stat.SettledRequests) + assert.Equal(t, int64(1), stat.UnpricedRequests) + assert.Equal(t, int64(1), stat.EstimatedRequests) + assert.Equal(t, int64(5783), stat.AmountCNYMicros) +} + +func TestRecordConsumeLogKeepsCostLedgerWhenUsageLogsAreDisabled(t *testing.T) { + previousDB := DB + previousMainDatabaseType := common.MainDatabaseType() + previousLogDatabaseType := common.LogDatabaseType() + previousLogConsumeEnabled := common.LogConsumeEnabled + db, err := gorm.Open(sqlite.Open("file:upstream_cost_without_usage_log?mode=memory&cache=shared"), &gorm.Config{}) + require.NoError(t, err) + DB = db + common.SetDatabaseTypes(common.DatabaseTypeSQLite, common.DatabaseTypeSQLite) + common.LogConsumeEnabled = false + initCol() + t.Cleanup(func() { + DB = previousDB + common.SetDatabaseTypes(previousMainDatabaseType, previousLogDatabaseType) + common.LogConsumeEnabled = previousLogConsumeEnabled + initCol() + }) + require.NoError(t, db.AutoMigrate(&UpstreamCostRecord{})) + + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set("username", "admin") + snapshot := &dto.UpstreamCostSnapshot{ + Status: dto.UpstreamCostStatusUnpriced, + Mode: dto.UpstreamCostModeBillingUnits, + Reason: "missing_channel_cost_profile", + NativeUnit: "UNIT", + PriceVersion: "manual", + SettlementCurrency: "CNY", + } + RecordConsumeLog(ctx, 9, RecordConsumeLogParams{ + ChannelId: 36, + ModelName: "gpt-5.6-sol", + TokenName: "Nexa", + Group: "codex", + Other: map[string]interface{}{ + "admin_info": map[string]interface{}{"upstream_cost": snapshot}, + }, + }) + + var record UpstreamCostRecord + require.NoError(t, db.First(&record).Error) + assert.Equal(t, dto.UpstreamCostStatusUnpriced, record.Status) + assert.Equal(t, "admin", record.Username) +} diff --git a/relay/asyncwrap/executor.go b/relay/asyncwrap/executor.go new file mode 100644 index 000000000000..ce086ee96fca --- /dev/null +++ b/relay/asyncwrap/executor.go @@ -0,0 +1,228 @@ +package asyncwrap + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net" + "net/http" + "strconv" + "sync" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/dto" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/service" +) + +type synchronousImageRequest struct { + endpoint string + payload []byte + parse func([]byte) service.AsyncExecutionOutcome +} + +type synchronousImageRequestPreparer func(dto.ImageRequest, []byte) (synchronousImageRequest, error) + +func newSynchronousImageHTTPClient(timeout time.Duration) *http.Client { + if timeout <= 0 { + timeout = 30 * time.Minute + } + transport := &http.Transport{ + Proxy: http.ProxyFromEnvironment, + ForceAttemptHTTP2: true, + DialContext: (&net.Dialer{Timeout: 10 * time.Second, KeepAlive: 30 * time.Second}).DialContext, + TLSHandshakeTimeout: 10 * time.Second, + ResponseHeaderTimeout: minDuration(60*time.Second, timeout), + ExpectContinueTimeout: 2 * time.Second, + IdleConnTimeout: 90 * time.Second, + MaxIdleConns: 100, + MaxIdleConnsPerHost: 50, + } + return &http.Client{ + Transport: transport, + Timeout: timeout, + CheckRedirect: func(*http.Request, []*http.Request) error { + return http.ErrUseLastResponse + }, + } +} + +func executeSynchronousImage( + ctx context.Context, + provider string, + client *http.Client, + apiKey string, + payload []byte, + markRequestSent func() error, + prepare synchronousImageRequestPreparer, +) service.AsyncExecutionOutcome { + if client == nil || prepare == nil { + return executorFailure("executor", "executor_unavailable", provider+" executor is unavailable", true) + } + var imageRequest dto.ImageRequest + if err := common.Unmarshal(payload, &imageRequest); err != nil || imageRequest.Model == "" { + return executorFailure("request_validate", "invalid_stored_request", "stored async image request is invalid", true) + } + prepared, err := prepare(imageRequest, payload) + if err != nil { + return executorFailure("request_validate", "invalid_stored_request", err.Error(), true) + } + if prepared.endpoint == "" || prepared.parse == nil { + return executorFailure("executor", "executor_unavailable", provider+" executor did not prepare a valid synchronous request", true) + } + + var markOnce sync.Once + var markErr error + mark := func() error { + markOnce.Do(func() { + if markRequestSent != nil { + markErr = markRequestSent() + } + }) + return markErr + } + + const maxAttempts = 3 + for attempt := 1; attempt <= maxAttempts; attempt++ { + tracker := &sentTrackingReader{reader: bytes.NewReader(prepared.payload), mark: mark} + request, err := http.NewRequestWithContext(ctx, http.MethodPost, prepared.endpoint, tracker) + if err != nil { + return executorFailure("request_build", "request_build_failed", "failed to construct "+provider+" request", true) + } + request.ContentLength = int64(len(prepared.payload)) + request.Header.Set("Authorization", "Bearer "+apiKey) + request.Header.Set("Content-Type", "application/json") + request.Header.Set("Accept", "application/json") + request.Header.Set("User-Agent", "new-api-async-image-worker/1") + + response, requestErr := client.Do(request) + if requestErr != nil { + if tracker.sent { + return service.AsyncExecutionOutcome{ + Status: model.AsyncStatusUncertain, + ErrorPhase: "upstream_read", + ErrorCode: "upstream_result_uncertain", + ErrorMessage: "the " + provider + " request may have executed but its result could not be confirmed", + } + } + if attempt < maxAttempts && ctx.Err() == nil { + if !waitRetry(ctx, time.Duration(attempt)*200*time.Millisecond) { + return executorFailure("upstream_connect", "upstream_connect_cancelled", provider+" connection attempt was cancelled before sending", true) + } + continue + } + return executorFailure("upstream_connect", "upstream_connect_failed", "failed to connect to "+provider+" before sending the request body", true) + } + + if response.StatusCode == http.StatusTooManyRequests { + _, _ = io.Copy(io.Discard, io.LimitReader(response.Body, 64*1024)) + _ = response.Body.Close() + if attempt < maxAttempts { + delay := service.ParseRetryAfter(response.Header.Get("Retry-After"), 30*time.Second) + if delay == 0 { + delay = time.Duration(attempt) * time.Second + } + if waitRetry(ctx, delay) { + continue + } + } + return executorFailure("upstream_response", "upstream_rate_limited", provider+" rate limit retry budget was exhausted", true) + } + if response.StatusCode < 200 || response.StatusCode >= 300 { + _, _ = io.Copy(io.Discard, io.LimitReader(response.Body, 64*1024)) + _ = response.Body.Close() + return executorFailure( + "upstream_response", + "upstream_http_"+strconv.Itoa(response.StatusCode), + fmt.Sprintf("%s returned HTTP %d", provider, response.StatusCode), + synchronousImageRefundEligibleStatus(response.StatusCode), + ) + } + + maxResponseBytes := int64(common.GetEnvOrDefault("ASYNC_UPSTREAM_MAX_RESPONSE_MB", 64)) * 1024 * 1024 + if maxResponseBytes <= 0 { + maxResponseBytes = 64 * 1024 * 1024 + } + body, readErr := io.ReadAll(io.LimitReader(response.Body, maxResponseBytes+1)) + _ = response.Body.Close() + if readErr != nil { + return service.AsyncExecutionOutcome{Status: model.AsyncStatusUncertain, ErrorPhase: "upstream_read", ErrorCode: "upstream_result_uncertain", ErrorMessage: provider + " returned an incomplete response after accepting the request"} + } + if int64(len(body)) > maxResponseBytes { + return executorFailure("upstream_parse", "upstream_response_too_large", provider+" response exceeded the configured safety limit", false) + } + outcome := prepared.parse(body) + outcome.Payload = json.RawMessage(body) + if outcome.Status == "" { + return executorFailure("upstream_parse", "invalid_upstream_response", provider+" returned an invalid response state", false) + } + return outcome + } + return executorFailure("upstream_connect", "upstream_connect_failed", "failed to connect to "+provider, true) +} + +func synchronousImageRefundEligibleStatus(status int) bool { + switch status { + case http.StatusBadRequest, + http.StatusUnauthorized, + http.StatusPaymentRequired, + http.StatusForbidden, + http.StatusNotFound, + http.StatusMethodNotAllowed, + http.StatusRequestEntityTooLarge, + http.StatusUnsupportedMediaType, + http.StatusUnprocessableEntity: + return true + default: + return false + } +} + +type sentTrackingReader struct { + reader io.Reader + mark func() error + sent bool +} + +func (r *sentTrackingReader) Read(buffer []byte) (int, error) { + if !r.sent { + if r.mark != nil { + if err := r.mark(); err != nil { + return 0, err + } + } + r.sent = true + } + return r.reader.Read(buffer) +} + +func executorFailure(phase, code, message string, refundable bool) service.AsyncExecutionOutcome { + return service.AsyncExecutionOutcome{ + Status: model.AsyncStatusFailure, + ErrorPhase: phase, + ErrorCode: code, + ErrorMessage: message, + RefundEligible: refundable, + } +} + +func waitRetry(ctx context.Context, delay time.Duration) bool { + timer := time.NewTimer(delay) + defer timer.Stop() + select { + case <-ctx.Done(): + return false + case <-timer.C: + return true + } +} + +func minDuration(first, second time.Duration) time.Duration { + if first < second { + return first + } + return second +} diff --git a/relay/asyncwrap/grsai.go b/relay/asyncwrap/grsai.go new file mode 100644 index 000000000000..0d3c760a8d9d --- /dev/null +++ b/relay/asyncwrap/grsai.go @@ -0,0 +1,198 @@ +package asyncwrap + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "net/url" + "strings" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/dto" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/service" + "github.com/QuantumNous/new-api/setting/model_setting" +) + +const grsaiSynchronousImagePath = "/v1/api/generate" + +type GRSAIExecutor struct { + baseURL string + apiKey string + client *http.Client +} + +func NewGRSAIExecutor(baseURL, apiKey string, timeout time.Duration) (*GRSAIExecutor, error) { + if _, err := grsaiSynchronousImageEndpoint(baseURL); err != nil { + return nil, err + } + if strings.TrimSpace(apiKey) == "" { + return nil, errors.New("GRS AI API key is empty") + } + return &GRSAIExecutor{ + baseURL: strings.TrimSpace(baseURL), + apiKey: apiKey, + client: newSynchronousImageHTTPClient(timeout), + }, nil +} + +func (e *GRSAIExecutor) Execute(ctx context.Context, payload []byte, markRequestSent func() error) service.AsyncExecutionOutcome { + if e == nil || e.client == nil { + return executorFailure("executor", "executor_unavailable", "GRS AI executor is unavailable", true) + } + return executeSynchronousImage(ctx, "GRS AI", e.client, e.apiKey, payload, markRequestSent, e.prepareRequest) +} + +func (e *GRSAIExecutor) prepareRequest(request dto.ImageRequest, _ []byte) (synchronousImageRequest, error) { + if request.N != nil && *request.N != 1 { + return synchronousImageRequest{}, errors.New("GRS AI synchronous image generation supports exactly one image per task") + } + endpoint, err := grsaiSynchronousImageEndpoint(e.baseURL) + if err != nil { + return synchronousImageRequest{}, err + } + images, err := grsaiReferenceImages(request.Images, request.Image) + if err != nil { + return synchronousImageRequest{}, err + } + payload := struct { + Model string `json:"model"` + Prompt string `json:"prompt"` + Images []string `json:"images,omitempty"` + AspectRatio string `json:"aspectRatio,omitempty"` + ImageSize string `json:"imageSize,omitempty"` + ReplyType string `json:"replyType"` + }{ + Model: request.Model, + Prompt: request.Prompt, + Images: images, + AspectRatio: strings.TrimSpace(request.Size), + ReplyType: "json", + } + normalizedModel := strings.ToLower(strings.TrimSpace(request.Model)) + capabilities := model_setting.GetImageGenerationCapabilities(normalizedModel) + if capabilities != nil && capabilities.ResolutionParameter == model_setting.ImageResolutionParameterQuality { + payload.ImageSize, err = grsaiImageSize(request.Quality) + if err != nil { + return synchronousImageRequest{}, err + } + } + requestPayload, err := common.Marshal(payload) + if err != nil { + return synchronousImageRequest{}, err + } + return synchronousImageRequest{ + endpoint: endpoint, + payload: requestPayload, + parse: parseGRSAISynchronousImageResponse, + }, nil +} + +func grsaiSynchronousImageEndpoint(baseURL string) (string, error) { + parsed, err := parseGRSAIBaseURL(baseURL) + if err != nil { + return "", err + } + parsed.Path = grsaiSynchronousImagePath + parsed.RawPath = "" + return parsed.String(), nil +} + +func parseGRSAIBaseURL(baseURL string) (*url.URL, error) { + parsed, err := url.Parse(strings.TrimSpace(baseURL)) + if err != nil || parsed.Hostname() == "" || (parsed.Scheme != "http" && parsed.Scheme != "https") { + return nil, errors.New("GRS AI base URL must be an absolute http or https URL") + } + if parsed.User != nil || parsed.RawQuery != "" || parsed.Fragment != "" { + return nil, errors.New("GRS AI base URL must not contain credentials, query parameters or fragments") + } + basePath := strings.TrimRight(parsed.Path, "/") + if basePath != "" && basePath != "/v1" { + return nil, errors.New("GRS AI base URL path must be empty or /v1") + } + parsed.Path = "" + parsed.RawPath = "" + return parsed, nil +} + +func grsaiReferenceImages(values ...json.RawMessage) ([]string, error) { + images := make([]string, 0) + for _, value := range values { + if len(value) == 0 || string(value) == "null" { + continue + } + var list []string + if err := common.Unmarshal(value, &list); err == nil { + images = append(images, list...) + continue + } + var single string + if err := common.Unmarshal(value, &single); err == nil && strings.TrimSpace(single) != "" { + images = append(images, single) + continue + } + return nil, errors.New("GRS AI reference images must be a URL/base64 string or an array of strings") + } + return images, nil +} + +func grsaiImageSize(quality string) (string, error) { + switch strings.ToLower(strings.TrimSpace(quality)) { + case "", "auto", "standard", "medium", "1k": + return "1K", nil + case "hd", "high", "2k": + return "2K", nil + case "4k": + return "4K", nil + default: + return "", errors.New("unsupported GRS AI image quality; use 1K, 2K or 4K") + } +} + +func parseGRSAISynchronousImageResponse(body []byte) service.AsyncExecutionOutcome { + var response struct { + Status string `json:"status"` + Results []struct { + URL string `json:"url"` + } `json:"results"` + Error string `json:"error"` + } + if err := common.Unmarshal(body, &response); err != nil { + return executorFailure("upstream_parse", "invalid_upstream_response", "GRS AI returned invalid JSON", false) + } + switch strings.ToLower(strings.TrimSpace(response.Status)) { + case "succeeded": + media := make([]service.AsyncMediaSource, 0, len(response.Results)) + for _, result := range response.Results { + if strings.TrimSpace(result.URL) == "" { + continue + } + if source, ok := service.ParseDataURLSource(result.URL); ok { + media = append(media, source) + } else { + media = append(media, service.AsyncMediaSource{URL: result.URL}) + } + } + if len(media) == 0 { + return executorFailure("upstream_parse", "invalid_upstream_response", "GRS AI synchronous response did not contain an image", false) + } + return service.AsyncExecutionOutcome{Status: model.AsyncStatusSuccess, Media: media} + case "failed", "violation": + message := strings.TrimSpace(response.Error) + if message == "" { + message = "GRS AI rejected the synchronous image request" + } + return executorFailure("upstream_response", "upstream_generation_failed", message, true) + case "running": + return service.AsyncExecutionOutcome{ + Status: model.AsyncStatusUncertain, + ErrorPhase: "upstream_response", + ErrorCode: "upstream_sync_result_pending", + ErrorMessage: "GRS AI did not return a final result in synchronous mode; no upstream async polling was performed", + } + default: + return executorFailure("upstream_parse", "invalid_upstream_response", "GRS AI returned an unknown synchronous response status", false) + } +} diff --git a/relay/asyncwrap/grsai_test.go b/relay/asyncwrap/grsai_test.go new file mode 100644 index 000000000000..5cc67caeb8bb --- /dev/null +++ b/relay/asyncwrap/grsai_test.go @@ -0,0 +1,125 @@ +package asyncwrap + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + "time" + + "github.com/QuantumNous/new-api/model" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestGRSAIExecutorUsesSynchronousJSONMode(t *testing.T) { + var requests atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + requests.Add(1) + assert.Equal(t, "/v1/api/generate", request.URL.Path) + assert.Equal(t, "Bearer test-placeholder-key", request.Header.Get("Authorization")) + body, err := io.ReadAll(request.Body) + require.NoError(t, err) + assert.JSONEq(t, `{ + "model":"nano-banana-2", + "prompt":"draw a lighthouse", + "images":["https://example.com/reference.png"], + "aspectRatio":"16:9", + "imageSize":"2K", + "replyType":"json" + }`, string(body)) + _, _ = writer.Write([]byte(`{"id":"sync-1","status":"succeeded","results":[{"url":"https://example.com/result.png"}]}`)) + })) + defer server.Close() + + executor, err := NewGRSAIExecutor(server.URL+"/v1", "test-placeholder-key", time.Second) + require.NoError(t, err) + var marked atomic.Int32 + outcome := executor.Execute(context.Background(), []byte(`{ + "model":"nano-banana-2", + "prompt":"draw a lighthouse", + "image":["https://example.com/reference.png"], + "size":"16:9", + "quality":"2K", + "n":1, + "replyType":"async" + }`), func() error { + marked.Add(1) + return nil + }) + + assert.Equal(t, model.AsyncStatusSuccess, outcome.Status) + require.Len(t, outcome.Media, 1) + assert.Equal(t, "https://example.com/result.png", outcome.Media[0].URL) + assert.Equal(t, int32(1), requests.Load()) + assert.Equal(t, int32(1), marked.Load()) +} + +func TestGRSAIExecutorNeverPollsPendingSynchronousResponse(t *testing.T) { + var requestedPaths []string + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + requestedPaths = append(requestedPaths, request.URL.Path) + _, _ = writer.Write([]byte(`{"id":"unexpected-pending","status":"running","progress":10}`)) + })) + defer server.Close() + + executor, err := NewGRSAIExecutor(server.URL, "test-placeholder-key", time.Second) + require.NoError(t, err) + outcome := executor.Execute(context.Background(), []byte(`{"model":"gpt-image-2","prompt":"draw","n":1}`), func() error { return nil }) + + assert.Equal(t, model.AsyncStatusUncertain, outcome.Status) + assert.Equal(t, "upstream_sync_result_pending", outcome.ErrorCode) + assert.Equal(t, []string{"/v1/api/generate"}, requestedPaths) +} + +func TestGRSAILiteModelUsesProviderDefaultResolution(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + body, err := io.ReadAll(request.Body) + require.NoError(t, err) + assert.JSONEq(t, `{ + "model":"nano-banana-2-lite", + "prompt":"draw", + "aspectRatio":"1:1", + "replyType":"json" + }`, string(body)) + _, _ = writer.Write([]byte(`{"status":"succeeded","results":[{"url":"https://example.com/result.png"}]}`)) + })) + defer server.Close() + + executor, err := NewGRSAIExecutor(server.URL, "test-placeholder-key", time.Second) + require.NoError(t, err) + outcome := executor.Execute(context.Background(), []byte(`{"model":"nano-banana-2-lite","prompt":"draw","size":"1:1","quality":"auto"}`), func() error { return nil }) + assert.Equal(t, model.AsyncStatusSuccess, outcome.Status) +} + +func TestGRSAIGPTImageVIPForwardsExplicitSizeWithoutQuality(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + body, err := io.ReadAll(request.Body) + require.NoError(t, err) + assert.JSONEq(t, `{ + "model":"gpt-image-2-vip", + "prompt":"draw", + "aspectRatio":"3840x2160", + "replyType":"json" + }`, string(body)) + _, _ = writer.Write([]byte(`{"status":"succeeded","results":[{"url":"https://example.com/result.png"}]}`)) + })) + defer server.Close() + + executor, err := NewGRSAIExecutor(server.URL, "test-placeholder-key", time.Second) + require.NoError(t, err) + outcome := executor.Execute(context.Background(), []byte(`{"model":"gpt-image-2-vip","prompt":"draw","size":"3840x2160","n":1}`), func() error { return nil }) + assert.Equal(t, model.AsyncStatusSuccess, outcome.Status) +} + +func TestGRSAISynchronousImageEndpointRejectsUnapprovedPaths(t *testing.T) { + endpoint, err := grsaiSynchronousImageEndpoint("https://grsaiapi.com/v1") + require.NoError(t, err) + assert.Equal(t, "https://grsaiapi.com/v1/api/generate", endpoint) + _, err = grsaiSynchronousImageEndpoint("https://grsaiapi.com/dashboard") + require.Error(t, err) + _, err = grsaiSynchronousImageEndpoint("file:///tmp/socket") + require.Error(t, err) +} diff --git a/relay/asyncwrap/yunwu.go b/relay/asyncwrap/yunwu.go new file mode 100644 index 000000000000..cd582661224b --- /dev/null +++ b/relay/asyncwrap/yunwu.go @@ -0,0 +1,266 @@ +package asyncwrap + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "net/url" + "os" + "strings" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/dto" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/service" + "github.com/QuantumNous/new-api/setting/model_setting" +) + +const yunwuImagePath = "/v1/images/generations" + +type YunwuExecutor struct { + baseURL string + apiKey string + client *http.Client +} + +func NewYunwuExecutor(baseURL, apiKey string, timeout time.Duration) (*YunwuExecutor, error) { + if _, err := yunwuImageEndpoint(baseURL); err != nil { + return nil, err + } + if strings.TrimSpace(apiKey) == "" { + return nil, errors.New("yunwu API key is empty") + } + return &YunwuExecutor{ + baseURL: strings.TrimSpace(baseURL), + apiKey: apiKey, + client: newSynchronousImageHTTPClient(timeout), + }, nil +} + +func yunwuGeminiImageEndpoint(baseURL, model string) (string, error) { + parsed, err := parseYunwuBaseURL(baseURL) + if err != nil { + return "", err + } + if strings.TrimSpace(model) == "" || strings.ContainsAny(model, "/?#") { + return "", errors.New("yunwu gemini model name is invalid") + } + parsed.Path = "/v1beta/models/" + model + ":generateContent" + parsed.RawPath = "" + return parsed.String(), nil +} + +func yunwuImageEndpoint(baseURL string) (string, error) { + parsed, err := parseYunwuBaseURL(baseURL) + if err != nil { + return "", err + } + parsed.Path = yunwuImagePath + parsed.RawPath = "" + return parsed.String(), nil +} + +func parseYunwuBaseURL(baseURL string) (*url.URL, error) { + parsed, err := url.Parse(strings.TrimSpace(baseURL)) + if err != nil || parsed.Hostname() == "" || (parsed.Scheme != "http" && parsed.Scheme != "https") { + return nil, errors.New("yunwu base URL must be an absolute http or https URL") + } + if parsed.User != nil || parsed.RawQuery != "" || parsed.Fragment != "" { + return nil, errors.New("yunwu base URL must not contain credentials, query parameters or fragments") + } + basePath := strings.TrimRight(parsed.Path, "/") + if basePath != "" && basePath != "/v1" { + return nil, errors.New("yunwu base URL path must be empty or /v1") + } + parsed.Path = "" + parsed.RawPath = "" + return parsed, nil +} + +func (e *YunwuExecutor) Execute(ctx context.Context, payload []byte, markRequestSent func() error) service.AsyncExecutionOutcome { + if e == nil || e.client == nil { + return executorFailure("executor", "executor_unavailable", "yunwu executor is unavailable", true) + } + return executeSynchronousImage(ctx, "yunwu", e.client, e.apiKey, payload, markRequestSent, func(imageRequest dto.ImageRequest, original []byte) (synchronousImageRequest, error) { + endpoint, requestPayload, geminiResponse, err := e.prepareRequest(imageRequest, original) + if err != nil { + return synchronousImageRequest{}, err + } + return synchronousImageRequest{ + endpoint: endpoint, + payload: requestPayload, + parse: func(body []byte) service.AsyncExecutionOutcome { + var media []service.AsyncMediaSource + var parseErr error + if geminiResponse { + media, parseErr = parseYunwuGeminiImageMedia(body) + } else { + media, parseErr = parseYunwuImageMedia(body) + } + if parseErr != nil { + return executorFailure("upstream_parse", "invalid_upstream_response", "yunwu returned a response without valid image artifacts", false) + } + return service.AsyncExecutionOutcome{Status: model.AsyncStatusSuccess, Media: media} + }, + }, nil + }) +} + +func (e *YunwuExecutor) prepareRequest(imageRequest dto.ImageRequest, original []byte) (string, []byte, bool, error) { + model := imageRequest.Model + if suffix := strings.TrimSpace(os.Getenv("ASYNC_YUNWU_ROUTE_SUFFIX")); suffix != "" { + switch suffix { + case "floor", "nitro", "stable": + model += ":" + suffix + default: + return "", nil, false, errors.New("ASYNC_YUNWU_ROUTE_SUFFIX must be floor, nitro or stable") + } + } + if model_setting.IsGeminiModelSupportImagine(imageRequest.Model) { + endpoint, err := yunwuGeminiImageEndpoint(e.baseURL, model) + if err != nil { + return "", nil, false, err + } + body, err := yunwuGeminiImagePayload(imageRequest) + return endpoint, body, true, err + } + endpoint, err := yunwuImageEndpoint(e.baseURL) + if err != nil { + return "", nil, false, err + } + if model == imageRequest.Model { + return endpoint, original, false, nil + } + var requestMap map[string]json.RawMessage + if err := common.Unmarshal(original, &requestMap); err != nil { + return "", nil, false, err + } + routedModel, err := common.Marshal(model) + if err != nil { + return "", nil, false, err + } + requestMap["model"] = routedModel + body, err := common.Marshal(requestMap) + return endpoint, body, false, err +} + +func yunwuGeminiImagePayload(request dto.ImageRequest) ([]byte, error) { + if request.N != nil && *request.N != 1 { + return nil, errors.New("yunwu gemini image models support exactly one image per async task") + } + aspectRatio, err := yunwuGeminiAspectRatio(request.Size) + if err != nil { + return nil, err + } + imageSize, err := yunwuGeminiImageSize(request.Quality) + if err != nil { + return nil, err + } + imageConfig, err := common.Marshal(map[string]string{ + "aspectRatio": aspectRatio, + "imageSize": imageSize, + }) + if err != nil { + return nil, err + } + payload := dto.GeminiChatRequest{ + Contents: []dto.GeminiChatContent{{ + Role: "user", + Parts: []dto.GeminiPart{{Text: request.Prompt}}, + }}, + GenerationConfig: dto.GeminiChatGenerationConfig{ + ResponseModalities: []string{"IMAGE", "TEXT"}, + ImageConfig: imageConfig, + }, + } + return common.Marshal(payload) +} + +func yunwuGeminiAspectRatio(size string) (string, error) { + switch strings.TrimSpace(size) { + case "", "1:1", "256x256", "512x512", "1024x1024": + return "1:1", nil + case "3:2", "1536x1024": + return "3:2", nil + case "2:3", "1024x1536": + return "2:3", nil + case "9:16", "1024x1792": + return "9:16", nil + case "16:9", "1792x1024": + return "16:9", nil + case "4:3", "3:4", "4:5", "5:4", "21:9": + return strings.TrimSpace(size), nil + default: + return "", errors.New("unsupported yunwu gemini image aspect ratio") + } +} + +func yunwuGeminiImageSize(quality string) (string, error) { + switch strings.ToLower(strings.TrimSpace(quality)) { + case "", "auto", "standard", "medium", "1k": + return "1K", nil + case "hd", "high", "2k": + return "2K", nil + case "4k": + return "4K", nil + default: + return "", errors.New("unsupported yunwu gemini image quality; use 1K, 2K or 4K") + } +} + +func parseYunwuImageMedia(body []byte) ([]service.AsyncMediaSource, error) { + var response struct { + Data []struct { + URL string `json:"url"` + B64JSON string `json:"b64_json"` + } `json:"data"` + } + if err := common.Unmarshal(body, &response); err != nil { + return nil, err + } + if len(response.Data) == 0 { + return nil, errors.New("image response data is empty") + } + media := make([]service.AsyncMediaSource, 0, len(response.Data)) + for _, item := range response.Data { + if source, ok := service.ParseDataURLSource(item.URL); ok { + media = append(media, source) + continue + } + if strings.TrimSpace(item.URL) != "" { + media = append(media, service.AsyncMediaSource{URL: item.URL}) + continue + } + if strings.TrimSpace(item.B64JSON) != "" { + media = append(media, service.AsyncMediaSource{Base64: item.B64JSON}) + continue + } + return nil, errors.New("image response item has no media") + } + return media, nil +} + +func parseYunwuGeminiImageMedia(body []byte) ([]service.AsyncMediaSource, error) { + var response dto.GeminiChatResponse + if err := common.Unmarshal(body, &response); err != nil { + return nil, err + } + media := make([]service.AsyncMediaSource, 0, len(response.Candidates)) + for _, candidate := range response.Candidates { + for _, part := range candidate.Content.Parts { + if part.InlineData == nil || strings.TrimSpace(part.InlineData.Data) == "" { + continue + } + media = append(media, service.AsyncMediaSource{ + Base64: part.InlineData.Data, + ContentType: part.InlineData.MimeType, + }) + } + } + if len(media) == 0 { + return nil, errors.New("gemini response has no image data") + } + return media, nil +} diff --git a/relay/asyncwrap/yunwu_test.go b/relay/asyncwrap/yunwu_test.go new file mode 100644 index 000000000000..4588341d4f13 --- /dev/null +++ b/relay/asyncwrap/yunwu_test.go @@ -0,0 +1,182 @@ +package asyncwrap + +import ( + "context" + "encoding/base64" + "fmt" + "io" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + "time" + + "github.com/QuantumNous/new-api/model" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestYunwuExecutorDelayedSuccess(t *testing.T) { + png := base64.StdEncoding.EncodeToString([]byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n', 0, 0, 0, 0}) + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + time.Sleep(20 * time.Millisecond) + assert.Equal(t, "/v1/images/generations", request.URL.Path) + assert.Equal(t, "Bearer test-placeholder-key", request.Header.Get("Authorization")) + writer.Header().Set("Content-Type", "application/json") + _, _ = fmt.Fprintf(writer, `{"created":1,"data":[{"b64_json":%q}]}`, png) + })) + defer server.Close() + + executor, err := NewYunwuExecutor(server.URL, "test-placeholder-key", time.Second) + require.NoError(t, err) + var marked atomic.Int32 + outcome := executor.Execute(context.Background(), []byte(`{"model":"m","prompt":"p"}`), func() error { + marked.Add(1) + return nil + }) + assert.Equal(t, model.AsyncStatusSuccess, outcome.Status) + assert.Len(t, outcome.Media, 1) + assert.Equal(t, int32(1), marked.Load()) +} + +func TestYunwuExecutorConvertsGeminiImageRequest(t *testing.T) { + png := base64.StdEncoding.EncodeToString([]byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n', 0, 0, 0, 0}) + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + assert.Equal(t, "/v1beta/models/gemini-3.1-flash-image-preview:generateContent", request.URL.Path) + body, err := io.ReadAll(request.Body) + require.NoError(t, err) + assert.JSONEq(t, `{ + "contents":[{"role":"user","parts":[{"text":"draw a blue square"}]}], + "generationConfig":{"responseModalities":["IMAGE","TEXT"],"imageConfig":{"aspectRatio":"16:9","imageSize":"2K"}} + }`, string(body)) + writer.Header().Set("Content-Type", "application/json") + _, _ = fmt.Fprintf(writer, `{"candidates":[{"content":{"parts":[{"inlineData":{"mimeType":"image/png","data":%q}}]}}]}`, png) + })) + defer server.Close() + + executor, err := NewYunwuExecutor(server.URL+"/v1", "test-placeholder-key", time.Second) + require.NoError(t, err) + outcome := executor.Execute(context.Background(), []byte(`{ + "model":"gemini-3.1-flash-image-preview", + "prompt":"draw a blue square", + "size":"1792x1024", + "quality":"2K", + "n":1 + }`), func() error { return nil }) + require.Equal(t, model.AsyncStatusSuccess, outcome.Status) + require.Len(t, outcome.Media, 1) + assert.Equal(t, "image/png", outcome.Media[0].ContentType) + assert.Equal(t, png, outcome.Media[0].Base64) +} + +func TestYunwuExecutorAppliesStableRouteWithoutChangingClientModel(t *testing.T) { + t.Setenv("ASYNC_YUNWU_ROUTE_SUFFIX", "stable") + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + body, err := io.ReadAll(request.Body) + require.NoError(t, err) + assert.JSONEq(t, `{"model":"gpt-image-2:stable","prompt":"p"}`, string(body)) + _, _ = writer.Write([]byte(`{"data":[{"b64_json":"iVBORw0KGgo="}]}`)) + })) + defer server.Close() + + executor, err := NewYunwuExecutor(server.URL, "test-placeholder-key", time.Second) + require.NoError(t, err) + outcome := executor.Execute(context.Background(), []byte(`{"model":"gpt-image-2","prompt":"p"}`), func() error { return nil }) + assert.Equal(t, model.AsyncStatusSuccess, outcome.Status) +} + +func TestYunwuExecutorRetries429(t *testing.T) { + var requests atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + attempt := requests.Add(1) + if attempt < 3 { + writer.Header().Set("Retry-After", "0") + writer.WriteHeader(http.StatusTooManyRequests) + return + } + _, _ = writer.Write([]byte(`{"data":[{"b64_json":"iVBORw0KGgo="}]}`)) + })) + defer server.Close() + + executor, err := NewYunwuExecutor(server.URL, "test-placeholder-key", 5*time.Second) + require.NoError(t, err) + var marked atomic.Int32 + outcome := executor.Execute(context.Background(), []byte(`{"model":"m","prompt":"p"}`), func() error { + marked.Add(1) + return nil + }) + assert.Equal(t, model.AsyncStatusSuccess, outcome.Status) + assert.Equal(t, int32(3), requests.Load()) + assert.Equal(t, int32(1), marked.Load()) +} + +func TestYunwuExecutorReadTimeoutIsUncertain(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + time.Sleep(200 * time.Millisecond) + _, _ = writer.Write([]byte(`{"data":[{"b64_json":"iVBORw0KGgo="}]}`)) + })) + defer server.Close() + + executor, err := NewYunwuExecutor(server.URL, "test-placeholder-key", 50*time.Millisecond) + require.NoError(t, err) + outcome := executor.Execute(context.Background(), []byte(`{"model":"m","prompt":"p"}`), func() error { return nil }) + assert.Equal(t, model.AsyncStatusUncertain, outcome.Status) + assert.Equal(t, "upstream_result_uncertain", outcome.ErrorCode) +} + +func TestYunwuImageEndpointWhitelistPath(t *testing.T) { + endpoint, err := yunwuImageEndpoint("https://yunwu.ai/v1") + require.NoError(t, err) + assert.Equal(t, "https://yunwu.ai/v1/images/generations", endpoint) + _, err = yunwuImageEndpoint("file:///tmp/socket") + require.Error(t, err) + _, err = yunwuImageEndpoint("https://yunwu.ai/unapproved-prefix") + require.Error(t, err) + endpoint, err = yunwuGeminiImageEndpoint("https://yunwu.ai/v1", "gemini-3-pro-image-preview:stable") + require.NoError(t, err) + assert.Equal(t, "https://yunwu.ai/v1beta/models/gemini-3-pro-image-preview:stable:generateContent", endpoint) +} + +func TestYunwuExecutorClassifiesExplicitFailureAndServerError(t *testing.T) { + tests := []struct { + name string + status int + refundable bool + }{ + {name: "rejected request", status: http.StatusBadRequest, refundable: true}, + {name: "server failure", status: http.StatusInternalServerError, refundable: false}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + writer.WriteHeader(test.status) + })) + defer server.Close() + executor, err := NewYunwuExecutor(server.URL, "test-placeholder-key", time.Second) + require.NoError(t, err) + outcome := executor.Execute(context.Background(), []byte(`{"model":"m","prompt":"p"}`), func() error { return nil }) + assert.Equal(t, model.AsyncStatusFailure, outcome.Status) + assert.Equal(t, test.refundable, outcome.RefundEligible) + assert.Equal(t, fmt.Sprintf("upstream_http_%d", test.status), outcome.ErrorCode) + }) + } +} + +func TestYunwuExecutorDoesNotFollowRedirects(t *testing.T) { + var targetRequests atomic.Int32 + target := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + targetRequests.Add(1) + writer.WriteHeader(http.StatusOK) + })) + defer target.Close() + source := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + http.Redirect(writer, &http.Request{}, target.URL, http.StatusTemporaryRedirect) + })) + defer source.Close() + + executor, err := NewYunwuExecutor(source.URL, "test-placeholder-key", time.Second) + require.NoError(t, err) + outcome := executor.Execute(context.Background(), []byte(`{"model":"m","prompt":"p"}`), func() error { return nil }) + assert.Equal(t, model.AsyncStatusFailure, outcome.Status) + assert.Zero(t, targetRequests.Load()) +} diff --git a/relay/channel/openai/chat_via_responses.go b/relay/channel/openai/chat_via_responses.go index 18758e728d41..2ce734d24f2c 100644 --- a/relay/channel/openai/chat_via_responses.go +++ b/relay/channel/openai/chat_via_responses.go @@ -59,6 +59,8 @@ func OaiResponsesToChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp usage = service.ResponseText2Usage(c, text, info.UpstreamModelName, info.GetEstimatePromptTokens()) chatResp.Usage = *usage } + service.AttachResponseBilling(c, info, usage) + chatResp.Usage.Billing = usage.Billing responseValue := any(chatResp) if info.RelayFormat != types.RelayFormatOpenAI { @@ -167,6 +169,8 @@ func OaiResponsesToChatBufferedStreamHandler(c *gin.Context, info *relaycommon.R usage = service.ResponseText2Usage(c, text, info.UpstreamModelName, info.GetEstimatePromptTokens()) chatResp.Usage = *usage } + service.AttachResponseBilling(c, info, usage) + chatResp.Usage.Billing = usage.Billing responseValue := any(chatResp) if info.RelayFormat != types.RelayFormatOpenAI { @@ -316,6 +320,7 @@ func OaiResponsesToChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo usage = service.ResponseText2Usage(c, state.UsageText(), info.UpstreamModelName, info.GetEstimatePromptTokens()) state.SetUsage(usage) } + service.AttachResponseBilling(c, info, usage) if info.RelayFormat == types.RelayFormatClaude && info.ClaudeConvertInfo != nil { info.ClaudeConvertInfo.Usage = usage diff --git a/relay/channel/openai/relay-openai.go b/relay/channel/openai/relay-openai.go index 50415c8b3533..e29437ed6eec 100644 --- a/relay/channel/openai/relay-openai.go +++ b/relay/channel/openai/relay-openai.go @@ -169,18 +169,26 @@ func OaiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Re logger.LogError(c, fmt.Sprintf("error handling last response: %s, lastStreamData: [%s]", err.Error(), lastStreamData)) } - if info.RelayFormat == types.RelayFormatOpenAI { - if shouldSendLastResp { - _ = sendStreamData(c, info, lastStreamData, info.ChannelSetting.ForceFormat, info.ChannelSetting.ThinkingToContent) - } - } - if !containStreamUsage { usage = service.ResponseText2Usage(c, responseTextBuilder.String(), info.UpstreamModelName, info.GetEstimatePromptTokens()) usage.CompletionTokens += toolCount * 7 } applyUsagePostProcessing(info, usage, common.StringToByteSlice(lastStreamData)) + service.AttachResponseBilling(c, info, usage) + + if info.RelayFormat == types.RelayFormatOpenAI && shouldSendLastResp { + if containStreamUsage && usage.Billing != nil { + enrichedData, err := addBillingToUsageJSON(common.StringToByteSlice(lastStreamData), usage.Billing) + if err != nil { + logger.LogError(c, "failed to add billing to stream usage: "+err.Error()) + containStreamUsage = false + } else { + lastStreamData = string(enrichedData) + } + } + _ = sendStreamData(c, info, lastStreamData, info.ChannelSetting.ForceFormat, info.ChannelSetting.ThinkingToContent) + } HandleFinalResponse(c, info, lastStreamData, responseId, createAt, model, systemFingerprint, usage, containStreamUsage) @@ -251,6 +259,7 @@ func OpenaiHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Respo } applyUsagePostProcessing(info, &simpleResponse.Usage, responseBody) + service.AttachResponseBilling(c, info, &simpleResponse.Usage) switch info.RelayFormat { case types.RelayFormatOpenAI: @@ -262,6 +271,11 @@ func OpenaiHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Respo } bodyMap["usage"] = simpleResponse.Usage responseBody, _ = common.Marshal(bodyMap) + } else if simpleResponse.Usage.Billing != nil { + responseBody, err = addBillingToUsageJSON(responseBody, simpleResponse.Usage.Billing) + if err != nil { + return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError) + } } if forceFormat { responseBody, err = common.Marshal(simpleResponse) diff --git a/relay/channel/openai/relay_responses.go b/relay/channel/openai/relay_responses.go index 9293183168ee..adac7f782e15 100644 --- a/relay/channel/openai/relay_responses.go +++ b/relay/channel/openai/relay_responses.go @@ -40,9 +40,6 @@ func OaiResponsesHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http c.Set("image_generation_call_size", responsesResponse.GetSize()) } - // 写入新的 response body - service.IOCopyBytesGracefully(c, resp, responseBody) - // compute usage usage := dto.Usage{} if responsesResponse.Usage != nil { @@ -54,18 +51,27 @@ func OaiResponsesHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http usage.PromptTokensDetails.CacheWriteTokens = responsesResponse.Usage.InputTokensDetails.CacheWriteTokens } } - if info == nil || info.ResponsesUsageInfo == nil || info.ResponsesUsageInfo.BuiltInTools == nil { - return &usage, nil + if info != nil && info.ResponsesUsageInfo != nil && info.ResponsesUsageInfo.BuiltInTools != nil { + // 解析 Tools 用量 + for _, tool := range responsesResponse.Tools { + buildToolinfo, ok := info.ResponsesUsageInfo.BuiltInTools[common.Interface2String(tool["type"])] + if !ok || buildToolinfo == nil { + logger.LogError(c, fmt.Sprintf("BuiltInTools not found for tool type: %v", tool["type"])) + continue + } + buildToolinfo.CallCount++ + } } - // 解析 Tools 用量 - for _, tool := range responsesResponse.Tools { - buildToolinfo, ok := info.ResponsesUsageInfo.BuiltInTools[common.Interface2String(tool["type"])] - if !ok || buildToolinfo == nil { - logger.LogError(c, fmt.Sprintf("BuiltInTools not found for tool type: %v", tool["type"])) - continue + + service.AttachResponseBilling(c, info, &usage) + if usage.Billing != nil { + responseBody, err = addBillingToUsageJSON(responseBody, usage.Billing) + if err != nil { + return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError) } - buildToolinfo.CallCount++ } + + service.IOCopyBytesGracefully(c, resp, responseBody) return &usage, nil } @@ -89,7 +95,6 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp sr.Error(err) return } - sendResponsesStreamData(c, streamResponse, data) switch streamResponse.Type { case "response.completed": if streamResponse.Response != nil { @@ -114,6 +119,26 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp c.Set("image_generation_call_size", streamResponse.Response.GetSize()) } } + if usage.CompletionTokens == 0 { + text := responseTextBuilder.String() + if text != "" { + usage.CompletionTokens = service.CountTextToken(text, info.UpstreamModelName) + } + } + if usage.PromptTokens == 0 && usage.CompletionTokens != 0 { + usage.PromptTokens = info.GetEstimatePromptTokens() + } + usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens + service.AttachResponseBilling(c, info, usage) + if usage.Billing != nil { + enrichedData, err := addBillingToResponsesEventJSON(common.StringToByteSlice(data), usage.Billing) + if err != nil { + logger.LogError(c, "failed to add billing to responses stream: "+err.Error()) + sr.Error(err) + return + } + data = string(enrichedData) + } case "response.output_text.delta": // 处理输出文本 responseTextBuilder.WriteString(streamResponse.Delta) @@ -130,6 +155,7 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp } } } + sendResponsesStreamData(c, streamResponse, data) }) if usage.CompletionTokens == 0 { @@ -147,6 +173,7 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp } usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens + service.AttachResponseBilling(c, info, usage) return usage, nil } diff --git a/relay/channel/openai/response_billing.go b/relay/channel/openai/response_billing.go new file mode 100644 index 000000000000..c432fc7c2613 --- /dev/null +++ b/relay/channel/openai/response_billing.go @@ -0,0 +1,64 @@ +package openai + +import ( + "encoding/json" + "fmt" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/dto" +) + +func addBillingToUsageJSON(data []byte, billing *dto.ResponseBilling) ([]byte, error) { + if billing == nil { + return data, nil + } + + var response map[string]json.RawMessage + if err := common.Unmarshal(data, &response); err != nil { + return nil, err + } + + rawUsage, ok := response["usage"] + if !ok || string(rawUsage) == "null" { + rawUsage = []byte(`{}`) + } + var usage map[string]json.RawMessage + if err := common.Unmarshal(rawUsage, &usage); err != nil { + return nil, fmt.Errorf("decode response usage: %w", err) + } + + rawBilling, err := common.Marshal(billing) + if err != nil { + return nil, err + } + usage["billing"] = rawBilling + + rawUsage, err = common.Marshal(usage) + if err != nil { + return nil, err + } + response["usage"] = rawUsage + return common.Marshal(response) +} + +func addBillingToResponsesEventJSON(data []byte, billing *dto.ResponseBilling) ([]byte, error) { + if billing == nil { + return data, nil + } + + var event map[string]json.RawMessage + if err := common.Unmarshal(data, &event); err != nil { + return nil, err + } + rawResponse, ok := event["response"] + if !ok || string(rawResponse) == "null" { + return data, nil + } + + enrichedResponse, err := addBillingToUsageJSON(rawResponse, billing) + if err != nil { + return nil, err + } + event["response"] = enrichedResponse + return common.Marshal(event) +} diff --git a/relay/channel/openai/response_billing_integration_test.go b/relay/channel/openai/response_billing_integration_test.go new file mode 100644 index 000000000000..b82f3ffcab6e --- /dev/null +++ b/relay/channel/openai/response_billing_integration_test.go @@ -0,0 +1,101 @@ +package openai + +import ( + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/dto" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/types" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func responseBillingRelayInfo(stream bool) *relaycommon.RelayInfo { + return &relaycommon.RelayInfo{ + ChannelMeta: &relaycommon.ChannelMeta{ + ChannelType: constant.ChannelTypeOpenAI, + UpstreamModelName: "test-model", + }, + IsStream: stream, + RelayFormat: types.RelayFormatOpenAI, + OriginModelName: "test-model", + ShouldIncludeUsage: true, + ShouldIncludeBilling: true, + DisablePing: true, + StartTime: time.Now(), + PriceData: types.PriceData{ + ModelRatio: 1, + CompletionRatio: 2, + GroupRatioInfo: types.GroupRatioInfo{GroupRatio: 0.5}, + }, + } +} + +func TestOpenaiHandlerIncludesRequestedBilling(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + + body := `{"id":"chatcmpl_1","object":"chat.completion","created":1710000000,"model":"test-model","choices":[{"index":0,"message":{"role":"assistant","content":"hello"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1000,"completion_tokens":500,"total_tokens":1500}}` + resp := &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader(body)), + Header: http.Header{"Content-Type": []string{"application/json"}}, + } + + usage, apiErr := OpenaiHandler(ctx, responseBillingRelayInfo(false), resp) + + require.Nil(t, apiErr) + require.NotNil(t, usage) + require.NotNil(t, usage.Billing) + assert.InDelta(t, 0.002, usage.Billing.TotalCost, 1e-12) + + var result dto.OpenAITextResponse + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &result)) + require.NotNil(t, result.Usage.Billing) + assert.Equal(t, "USD", result.Usage.Billing.Currency) + assert.InDelta(t, 0.002, result.Usage.Billing.TotalCost, 1e-12) +} + +func TestOaiStreamHandlerIncludesBillingInFinalUsageChunk(t *testing.T) { + gin.SetMode(gin.TestMode) + oldTimeout := constant.StreamingTimeout + constant.StreamingTimeout = 30 + t.Cleanup(func() { constant.StreamingTimeout = oldTimeout }) + + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + ctx.Set(common.RequestIdKey, "billing-stream-test") + + body := strings.Join([]string{ + `data: {"id":"chatcmpl_1","object":"chat.completion.chunk","created":1710000000,"model":"test-model","choices":[{"index":0,"delta":{"content":"hello"},"finish_reason":null}]}`, + `data: {"id":"chatcmpl_1","object":"chat.completion.chunk","created":1710000000,"model":"test-model","choices":[],"usage":{"prompt_tokens":1000,"completion_tokens":500,"total_tokens":1500}}`, + `data: [DONE]`, + ``, + }, "\n") + resp := &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader(body)), + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + } + + usage, apiErr := OaiStreamHandler(ctx, responseBillingRelayInfo(true), resp) + + require.Nil(t, apiErr) + require.NotNil(t, usage) + require.NotNil(t, usage.Billing) + got := recorder.Body.String() + assert.Contains(t, got, `"billing":{"currency":"USD","total_cost":0.002`) + assert.Contains(t, got, `data: [DONE]`) +} diff --git a/relay/common/local_request_fields.go b/relay/common/local_request_fields.go new file mode 100644 index 000000000000..7b9ec546ff21 --- /dev/null +++ b/relay/common/local_request_fields.go @@ -0,0 +1,28 @@ +package common + +import ( + "encoding/json" + + appcommon "github.com/QuantumNous/new-api/common" +) + +// RemoveLocalRequestFields removes gateway-only top-level fields while +// preserving all other passthrough JSON values verbatim. +func RemoveLocalRequestFields(data []byte, fields ...string) ([]byte, error) { + var body map[string]json.RawMessage + if err := appcommon.Unmarshal(data, &body); err != nil { + return nil, err + } + + changed := false + for _, field := range fields { + if _, ok := body[field]; ok { + delete(body, field) + changed = true + } + } + if !changed { + return data, nil + } + return appcommon.Marshal(body) +} diff --git a/relay/common/local_request_fields_test.go b/relay/common/local_request_fields_test.go new file mode 100644 index 000000000000..6fea8c159284 --- /dev/null +++ b/relay/common/local_request_fields_test.go @@ -0,0 +1,23 @@ +package common + +import ( + "testing" + + appcommon "github.com/QuantumNous/new-api/common" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestRemoveLocalRequestFieldsPreservesPassthroughPayload(t *testing.T) { + input := []byte(`{"model":"gpt-test","include_billing":true,"large_integer":18446744073686646784,"vendor":{"flag":true}}`) + + got, err := RemoveLocalRequestFields(input, "include_billing") + + require.NoError(t, err) + var decoded map[string]any + require.NoError(t, appcommon.Unmarshal(got, &decoded)) + assert.NotContains(t, decoded, "include_billing") + assert.Contains(t, string(got), "18446744073686646784") + assert.Contains(t, string(got), `"vendor":{"flag":true}`) +} diff --git a/relay/common/relay_info.go b/relay/common/relay_info.go index 9f460ce5c6a7..ab059bdadafc 100644 --- a/relay/common/relay_info.go +++ b/relay/common/relay_info.go @@ -106,6 +106,7 @@ type RelayInfo struct { RequestURLPath string RequestHeaders map[string]string ShouldIncludeUsage bool + ShouldIncludeBilling bool DisablePing bool // 是否禁止向下游发送自定义 Ping ClientWs *websocket.Conn TargetWs *websocket.Conn @@ -262,6 +263,7 @@ func (info *RelayInfo) ToString() string { fmt.Fprintf(b, "OriginModelName: %q, ", info.OriginModelName) fmt.Fprintf(b, "EstimatePromptTokens: %d, ", info.estimatePromptTokens) fmt.Fprintf(b, "ShouldIncludeUsage: %t, ", info.ShouldIncludeUsage) + fmt.Fprintf(b, "ShouldIncludeBilling: %t, ", info.ShouldIncludeBilling) fmt.Fprintf(b, "DisablePing: %t, ", info.DisablePing) fmt.Fprintf(b, "SendResponseCount: %d, ", info.SendResponseCount) fmt.Fprintf(b, "FinalPreConsumedQuota: %d, ", info.FinalPreConsumedQuota) diff --git a/relay/compatible_handler.go b/relay/compatible_handler.go index a68cfe730f60..329de927d22a 100644 --- a/relay/compatible_handler.go +++ b/relay/compatible_handler.go @@ -35,6 +35,10 @@ func TextHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types return types.NewError(fmt.Errorf("failed to copy request to GeneralOpenAIRequest: %w", err), types.ErrorCodeInvalidRequest, types.ErrOptionWithSkipRetry()) } + hasIncludeBilling := request.IncludeBilling != nil + info.ShouldIncludeBilling = lo.FromPtrOr(request.IncludeBilling, false) + request.IncludeBilling = nil + if request.WebSearchOptions != nil { c.Set("chat_completion_web_search_context_size", request.WebSearchOptions.SearchContextSize) } @@ -49,6 +53,9 @@ func TextHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types if request.StreamOptions != nil { includeUsage = request.StreamOptions.IncludeUsage } + if info.ShouldIncludeBilling { + includeUsage = true + } // 如果不支持StreamOptions，将StreamOptions设置为nil if !info.SupportStreamOptions || !lo.FromPtrOr(request.Stream, false) { @@ -104,7 +111,25 @@ func TextHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types logger.LogDebug(c, "requestBody: %s", debugBytes) } } - requestBody = common.ReaderOnly(storage) + if hasIncludeBilling { + jsonData, err := storage.Bytes() + if err != nil { + return types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry()) + } + jsonData, err = relaycommon.RemoveLocalRequestFields(jsonData, "include_billing") + if err != nil { + return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) + } + body, size, closer, err := relaycommon.NewOutboundJSONBody(jsonData) + if err != nil { + return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) + } + defer closer.Close() + info.UpstreamRequestBodySize = size + requestBody = body + } else { + requestBody = common.ReaderOnly(storage) + } } else { convertedRequest, err := adaptor.ConvertOpenAIRequest(c, info, request) if err != nil { diff --git a/relay/helper/openai_image_request_test.go b/relay/helper/openai_image_request_test.go index e9fb1b9992d9..c468496fd88e 100644 --- a/relay/helper/openai_image_request_test.go +++ b/relay/helper/openai_image_request_test.go @@ -158,3 +158,26 @@ func TestGetAndValidOpenAIImageRequestNBounds(t *testing.T) { require.Contains(t, err.Error(), boundErr) }) } + +func TestAsyncImageResolutionPriceRatio(t *testing.T) { + tests := []struct { + name string + model string + quality string + size string + want float64 + }{ + {name: "flash 1K", model: "gemini-3.1-flash-image-preview", quality: "1K", want: 1}, + {name: "flash 4K", model: "gemini-3.1-flash-image-preview", quality: "4K", want: 2}, + {name: "pro 4K", model: "gemini-3-pro-image-preview", quality: "4K", want: 1}, + {name: "gpt 2K", model: "gpt-image-2-vip", size: "2048x2048", want: 1}, + {name: "gpt 4K", model: "gpt-image-2-vip", size: "3840x2160", want: 2}, + {name: "nano 4K", model: "nano-banana-pro", quality: "4K", want: 1}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + request := dto.ImageRequest{Model: test.model, Prompt: "a cat", Quality: test.quality, Size: test.size} + require.InDelta(t, test.want, request.GetTokenCountMeta().ImagePriceRatio, 0.000001) + }) + } +} diff --git a/relay/responses_handler.go b/relay/responses_handler.go index 5fa23d099623..86b406bc9b61 100644 --- a/relay/responses_handler.go +++ b/relay/responses_handler.go @@ -18,6 +18,7 @@ import ( "github.com/QuantumNous/new-api/types" "github.com/gin-gonic/gin" + "github.com/samber/lo" ) func ResponsesHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types.NewAPIError) { @@ -70,6 +71,10 @@ func ResponsesHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError * return types.NewError(fmt.Errorf("failed to copy request to GeneralOpenAIRequest: %w", err), types.ErrorCodeInvalidRequest, types.ErrOptionWithSkipRetry()) } + hasIncludeBilling := request.IncludeBilling != nil + info.ShouldIncludeBilling = lo.FromPtrOr(request.IncludeBilling, false) + request.IncludeBilling = nil + err = helper.ModelMappedHelper(c, info, request) if err != nil { return types.NewError(err, types.ErrorCodeChannelModelMappedError, types.ErrOptionWithSkipRetry()) @@ -86,7 +91,25 @@ func ResponsesHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError * if err != nil { return types.NewError(err, types.ErrorCodeReadRequestBodyFailed, types.ErrOptionWithSkipRetry()) } - requestBody = common.ReaderOnly(storage) + if hasIncludeBilling { + jsonData, err := storage.Bytes() + if err != nil { + return types.NewError(err, types.ErrorCodeReadRequestBodyFailed, types.ErrOptionWithSkipRetry()) + } + jsonData, err = relaycommon.RemoveLocalRequestFields(jsonData, "include_billing") + if err != nil { + return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) + } + body, size, closer, err := relaycommon.NewOutboundJSONBody(jsonData) + if err != nil { + return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) + } + defer closer.Close() + info.UpstreamRequestBodySize = size + requestBody = body + } else { + requestBody = common.ReaderOnly(storage) + } } else { convertedRequest, err := adaptor.ConvertOpenAIResponsesRequest(c, info, *request) if err != nil { diff --git a/router/api-router.go b/router/api-router.go index 80fd65178c44..8bf79dd2b305 100644 --- a/router/api-router.go +++ b/router/api-router.go @@ -271,6 +271,7 @@ func SetApiRouter(router *gin.Engine) { logRoute := apiRouter.Group("/log") logRoute.GET("/", middleware.AdminAuth(), controller.GetAllLogs) logRoute.GET("/stat", middleware.AdminAuth(), controller.GetLogsStat) + logRoute.GET("/upstream-cost/stat", middleware.AdminAuth(), controller.GetUpstreamCostStat) logRoute.GET("/self/stat", middleware.UserAuth(), controller.GetLogsSelfStat) logRoute.GET("/channel_affinity_usage_cache", middleware.AdminAuth(), controller.GetChannelAffinityUsageCacheStats) logRoute.GET("/search", middleware.AdminAuth(), controller.SearchAllLogs) @@ -326,7 +327,12 @@ func SetApiRouter(router *gin.Engine) { taskRoute := apiRouter.Group("/task") { taskRoute.GET("/self", middleware.UserAuth(), controller.GetUserTask) + taskRoute.GET("/self/async/:task_id", middleware.UserAuth(), controller.GetUserAsyncTaskDetail) + taskRoute.POST("/self/async/:task_id/cancel", middleware.UserAuth(), controller.CancelUserAsyncTask) taskRoute.GET("/", middleware.AdminAuth(), controller.GetAllTask) + taskRoute.GET("/async/:task_id", middleware.AdminAuth(), controller.GetAdminAsyncTaskDetail) + taskRoute.POST("/async/:task_id/cancel", middleware.AdminAuth(), controller.CancelAdminAsyncTask) + taskRoute.POST("/async/:task_id/retry", middleware.AdminAuth(), controller.RetryAdminAsyncTask) } vendorRoute := apiRouter.Group("/vendors") diff --git a/router/relay-router.go b/router/relay-router.go index 17a13cad7fd6..5bba740a3d27 100644 --- a/router/relay-router.go +++ b/router/relay-router.go @@ -59,6 +59,18 @@ func SetRelayRouter(router *gin.Engine) { }) } + asyncRouter := router.Group("/v1/async") + asyncRouter.Use(middleware.RouteTag("relay")) + asyncRouter.Use(middleware.SystemPerformanceCheck()) + asyncRouter.Use(middleware.TokenAuth()) + asyncRouter.Use(middleware.ModelRequestRateLimit()) + { + asyncRouter.POST("/images/generations", middleware.AsyncImageDistribute(), controller.SubmitAsyncImageTask) + asyncRouter.GET("/tasks/:task_id", controller.GetAsyncTask) + asyncRouter.GET("/tasks/:task_id/result", controller.GetAsyncTaskResult) + asyncRouter.POST("/tasks/:task_id/cancel", controller.CancelAsyncTask) + } + playgroundRouter := router.Group("/pg") playgroundRouter.Use(middleware.RouteTag("relay")) playgroundRouter.Use(middleware.SystemPerformanceCheck()) diff --git a/service/async_artifact.go b/service/async_artifact.go new file mode 100644 index 000000000000..d0bac44277ab --- /dev/null +++ b/service/async_artifact.go @@ -0,0 +1,385 @@ +package service + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/base64" + "encoding/hex" + "errors" + "fmt" + "io" + "mime" + "net" + "net/http" + "net/url" + "os" + "path/filepath" + "strconv" + "strings" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/dto" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/storage" +) + +type AsyncMediaSource struct { + URL string + Base64 string + ContentType string +} + +var blockedArtifactNetworks = mustParseArtifactNetworks( + "0.0.0.0/8", + "100.64.0.0/10", + "192.0.0.0/24", + "198.18.0.0/15", + "240.0.0.0/4", + "::/128", +) + +func mustParseArtifactNetworks(values ...string) []*net.IPNet { + networks := make([]*net.IPNet, 0, len(values)) + for _, value := range values { + _, network, err := net.ParseCIDR(value) + if err != nil { + panic(err) + } + networks = append(networks, network) + } + return networks +} + +func IsPublicArtifactIP(ip net.IP) bool { + if ip == nil || ip.IsUnspecified() || ip.IsLoopback() || ip.IsPrivate() || ip.IsLinkLocalUnicast() || ip.IsLinkLocalMulticast() || ip.IsMulticast() { + return false + } + for _, network := range blockedArtifactNetworks { + if network.Contains(ip) { + return false + } + } + return true +} + +func validateArtifactURL(raw string) (*url.URL, error) { + parsed, err := url.Parse(strings.TrimSpace(raw)) + if err != nil || parsed.Hostname() == "" || (parsed.Scheme != "http" && parsed.Scheme != "https") { + return nil, errors.New("artifact URL must be an absolute http or https URL") + } + if parsed.User != nil { + return nil, errors.New("artifact URL credentials are forbidden") + } + return parsed, nil +} + +func safeArtifactHTTPClient() *http.Client { + dialer := &net.Dialer{Timeout: 10 * time.Second, KeepAlive: 30 * time.Second} + transport := &http.Transport{ + Proxy: nil, + ForceAttemptHTTP2: true, + TLSHandshakeTimeout: 10 * time.Second, + ResponseHeaderTimeout: 30 * time.Second, + IdleConnTimeout: 30 * time.Second, + MaxIdleConns: 20, + DialContext: func(ctx context.Context, network, address string) (net.Conn, error) { + host, port, err := net.SplitHostPort(address) + if err != nil { + return nil, err + } + addresses, err := net.DefaultResolver.LookupIPAddr(ctx, host) + if err != nil { + return nil, err + } + for _, address := range addresses { + if !IsPublicArtifactIP(address.IP) { + continue + } + return dialer.DialContext(ctx, network, net.JoinHostPort(address.IP.String(), port)) + } + return nil, errors.New("artifact host resolved only to blocked addresses") + }, + } + return &http.Client{ + Transport: transport, + Timeout: time.Duration(common.GetEnvOrDefault("ASYNC_ARTIFACT_DOWNLOAD_TIMEOUT_SECONDS", 120)) * time.Second, + CheckRedirect: func(request *http.Request, via []*http.Request) error { + if len(via) >= 5 { + return errors.New("too many artifact redirects") + } + _, err := validateArtifactURL(request.URL.String()) + return err + }, + } +} + +type materializedArtifact struct { + file *os.File + path string + contentType string + size int64 + sha256 string + sourceURLHash string +} + +func (m *materializedArtifact) Close() { + if m == nil { + return + } + if m.file != nil { + _ = m.file.Close() + } + if m.path != "" { + _ = os.Remove(m.path) + } +} + +func ArchiveAsyncMedia(ctx context.Context, task *model.Task, sources []AsyncMediaSource, store storage.ArtifactStore, retentionMinutes int) ([]model.Artifact, error) { + if task == nil || task.ID == 0 || task.TaskID == "" { + return nil, errors.New("persisted task is required for artifact archiving") + } + if store == nil { + return nil, errors.New("artifact store is required") + } + existing, err := model.ListArtifactsByTaskID(ctx, task.ID) + if err != nil { + return nil, err + } + if len(existing) > 0 { + return existing, nil + } + maxFiles := common.GetEnvOrDefault("ASYNC_ARTIFACT_MAX_FILES", 8) + if len(sources) == 0 || len(sources) > maxFiles { + return nil, fmt.Errorf("artifact count must be between 1 and %d", maxFiles) + } + if retentionMinutes <= 0 { + retentionMinutes = dto.AsyncRetentionDefaultMinutes + } + retentionMinutes = dto.NormalizeAsyncRetentionMinutes(retentionMinutes) + maxSingle := int64(common.GetEnvOrDefault("ASYNC_ARTIFACT_MAX_FILE_MB", 25)) * 1024 * 1024 + maxTotal := int64(common.GetEnvOrDefault("ASYNC_ARTIFACT_MAX_TOTAL_MB", 100)) * 1024 * 1024 + if maxSingle <= 0 || maxTotal <= 0 { + return nil, errors.New("artifact size limits must be positive") + } + + uploaded := make([]string, 0, len(sources)) + artifacts := make([]model.Artifact, 0, len(sources)) + cleanupUploads := func() { + for _, key := range uploaded { + _ = store.Delete(context.Background(), key) + } + } + total := int64(0) + for index, source := range sources { + materialized, err := materializeAsyncArtifact(ctx, source, maxSingle, index) + if err != nil { + cleanupUploads() + return nil, err + } + total += materialized.size + if total > maxTotal { + materialized.Close() + cleanupUploads() + return nil, fmt.Errorf("artifact total exceeds %d bytes", maxTotal) + } + extension := extensionForArtifactMIME(materialized.contentType) + objectKey := fmt.Sprintf("async/%s/%02d-%s%s", task.TaskID, index, materialized.sha256[:16], extension) + if _, err := materialized.file.Seek(0, io.SeekStart); err != nil { + materialized.Close() + cleanupUploads() + return nil, err + } + if err := store.Put(ctx, objectKey, materialized.file, materialized.contentType); err != nil { + materialized.Close() + cleanupUploads() + return nil, err + } + uploaded = append(uploaded, objectKey) + artifacts = append(artifacts, model.Artifact{ + TaskID: task.ID, + ObjectKey: objectKey, + ContentType: materialized.contentType, + SizeBytes: materialized.size, + SHA256: materialized.sha256, + SourceURLHash: materialized.sourceURLHash, + ExpiresAt: time.Now().Add(time.Duration(retentionMinutes) * time.Minute).Unix(), + }) + materialized.Close() + } + if err := model.CreateArtifacts(ctx, artifacts); err != nil { + cleanupUploads() + return nil, err + } + return artifacts, nil +} + +func materializeAsyncArtifact(ctx context.Context, source AsyncMediaSource, maxBytes int64, index int) (*materializedArtifact, error) { + file, err := os.CreateTemp("", "new-api-async-artifact-*") + if err != nil { + return nil, err + } + result := &materializedArtifact{file: file, path: file.Name()} + failed := true + defer func() { + if failed { + result.Close() + } + }() + + var reader io.ReadCloser + declaredType := strings.TrimSpace(source.ContentType) + sourceIdentity := fmt.Sprintf("inline:%d", index) + if source.URL != "" { + parsed, err := validateArtifactURL(source.URL) + if err != nil { + return nil, err + } + request, err := http.NewRequestWithContext(ctx, http.MethodGet, parsed.String(), nil) + if err != nil { + return nil, err + } + request.Header.Set("Accept", "image/avif,image/webp,image/png,image/jpeg,image/gif") + response, err := safeArtifactHTTPClient().Do(request) + if err != nil { + return nil, err + } + if response.StatusCode != http.StatusOK { + _ = response.Body.Close() + return nil, fmt.Errorf("artifact download returned HTTP %d", response.StatusCode) + } + if response.ContentLength > maxBytes { + _ = response.Body.Close() + return nil, fmt.Errorf("artifact exceeds %d bytes", maxBytes) + } + reader = response.Body + declaredType = response.Header.Get("Content-Type") + sourceIdentity = parsed.String() + } else if source.Base64 != "" { + encoded := strings.TrimSpace(source.Base64) + decoder := base64.NewDecoder(base64.StdEncoding, strings.NewReader(encoded)) + reader = io.NopCloser(decoder) + } else { + return nil, errors.New("artifact source has neither URL nor base64 data") + } + defer reader.Close() + + hasher := sha256.New() + written, err := io.Copy(io.MultiWriter(file, hasher), io.LimitReader(reader, maxBytes+1)) + if err != nil { + return nil, err + } + if written > maxBytes { + return nil, fmt.Errorf("artifact exceeds %d bytes", maxBytes) + } + if written == 0 { + return nil, errors.New("artifact is empty") + } + if _, err := file.Seek(0, io.SeekStart); err != nil { + return nil, err + } + probe := make([]byte, 512) + n, readErr := file.Read(probe) + if readErr != nil && !errors.Is(readErr, io.EOF) { + return nil, readErr + } + detected := detectArtifactMIME(probe[:n]) + contentType, _, _ := mime.ParseMediaType(declaredType) + if !allowedArtifactMIME(detected) { + return nil, fmt.Errorf("artifact content is not an allowed raster image (%s)", detected) + } + if contentType != "" && contentType != "application/octet-stream" && contentType != detected { + return nil, fmt.Errorf("artifact MIME type %s does not match detected content %s", contentType, detected) + } + contentType = detected + + sourceHash := sha256.Sum256([]byte(sourceIdentity)) + result.contentType = contentType + result.size = written + result.sha256 = hex.EncodeToString(hasher.Sum(nil)) + result.sourceURLHash = hex.EncodeToString(sourceHash[:]) + if _, err := file.Seek(0, io.SeekStart); err != nil { + return nil, err + } + failed = false + return result, nil +} + +func detectArtifactMIME(probe []byte) string { + detected := http.DetectContentType(probe) + if detected == "application/octet-stream" && len(probe) >= 12 && bytes.Equal(probe[4:8], []byte("ftyp")) { + brand := string(probe[8:12]) + if brand == "avif" || brand == "avis" { + return "image/avif" + } + } + return detected +} + +func allowedArtifactMIME(value string) bool { + switch strings.ToLower(strings.TrimSpace(value)) { + case "image/jpeg", "image/png", "image/webp", "image/gif", "image/avif": + return true + default: + return false + } +} + +func extensionForArtifactMIME(value string) string { + switch value { + case "image/jpeg": + return ".jpg" + case "image/png": + return ".png" + case "image/webp": + return ".webp" + case "image/gif": + return ".gif" + case "image/avif": + return ".avif" + default: + return filepath.Ext("artifact") + } +} + +func ParseDataURLSource(raw string) (AsyncMediaSource, bool) { + if !strings.HasPrefix(raw, "data:") { + return AsyncMediaSource{}, false + } + comma := strings.IndexByte(raw, ',') + if comma < 0 { + return AsyncMediaSource{}, false + } + metadata := strings.TrimPrefix(raw[:comma], "data:") + parts := strings.Split(metadata, ";") + if len(parts) < 2 || parts[len(parts)-1] != "base64" || !allowedArtifactMIME(parts[0]) { + return AsyncMediaSource{}, false + } + return AsyncMediaSource{Base64: raw[comma+1:], ContentType: parts[0]}, true +} + +func ParseRetryAfter(value string, maximum time.Duration) time.Duration { + value = strings.TrimSpace(value) + if value == "" { + return 0 + } + if seconds, err := strconv.Atoi(value); err == nil && seconds >= 0 { + delay := time.Duration(seconds) * time.Second + if delay > maximum { + return maximum + } + return delay + } + if parsed, err := http.ParseTime(value); err == nil { + delay := time.Until(parsed) + if delay < 0 { + return 0 + } + if delay > maximum { + return maximum + } + return delay + } + return 0 +} diff --git a/service/async_artifact_test.go b/service/async_artifact_test.go new file mode 100644 index 000000000000..20c49fa84a4f --- /dev/null +++ b/service/async_artifact_test.go @@ -0,0 +1,198 @@ +package service + +import ( + "bytes" + "context" + "encoding/base64" + "encoding/json" + "io" + "net" + "strings" + "sync" + "testing" + "time" + + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/model" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +type memoryArtifactStore struct { + mu sync.Mutex + objects map[string][]byte +} + +func (s *memoryArtifactStore) Put(_ context.Context, key string, body io.Reader, _ string) error { + s.mu.Lock() + defer s.mu.Unlock() + data, err := io.ReadAll(body) + if err != nil { + return err + } + s.objects[key] = data + return nil +} + +func (s *memoryArtifactStore) SignedURL(_ context.Context, key string, _ time.Duration) (string, error) { + return "https://objects.example/" + key, nil +} + +func (s *memoryArtifactStore) Delete(_ context.Context, key string) error { + s.mu.Lock() + defer s.mu.Unlock() + delete(s.objects, key) + return nil +} + +func TestIsPublicArtifactIP(t *testing.T) { + assert.False(t, IsPublicArtifactIP(net.ParseIP("127.0.0.1"))) + assert.False(t, IsPublicArtifactIP(net.ParseIP("169.254.169.254"))) + assert.False(t, IsPublicArtifactIP(net.ParseIP("10.0.0.1"))) + assert.False(t, IsPublicArtifactIP(net.ParseIP("100.64.0.1"))) + assert.True(t, IsPublicArtifactIP(net.ParseIP("1.1.1.1"))) +} + +func TestArchiveInlineMediaStreamsAndPersists(t *testing.T) { + task := &model.Task{TaskID: "task_inline_artifacts", Status: model.TaskStatusInProgress} + require.NoError(t, model.DB.Create(task).Error) + t.Cleanup(func() { + model.DB.Where("task_id = ?", task.ID).Delete(&model.Artifact{}) + model.DB.Delete(task) + }) + + png := append([]byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n'}, bytes.Repeat([]byte{0}, 64)...) + store := &memoryArtifactStore{objects: map[string][]byte{}} + artifacts, err := ArchiveAsyncMedia(context.Background(), task, []AsyncMediaSource{{ + Base64: base64.StdEncoding.EncodeToString(png), + ContentType: "image/png", + }}, store, 30) + require.NoError(t, err) + require.Len(t, artifacts, 1) + assert.Equal(t, "image/png", artifacts[0].ContentType) + assert.Equal(t, int64(len(png)), artifacts[0].SizeBytes) + assert.WithinDuration(t, time.Now().Add(30*time.Minute), time.Unix(artifacts[0].ExpiresAt, 0), 2*time.Second) + assert.Len(t, store.objects, 1) +} + +func TestCleanupExpiredArtifactAlsoRemovesEmbeddedImagePayload(t *testing.T) { + task := &model.Task{ + TaskID: "task_expired_artifact_cleanup", + Platform: constant.TaskPlatformAsyncImage, + UserId: 1, + ChannelId: 2, + Status: model.TaskStatusSuccess, + Progress: "100%", + SubmitTime: time.Now().Add(-time.Hour).Unix(), + FinishTime: time.Now().Add(-time.Minute).Unix(), + Data: json.RawMessage(`{}`), + } + job := &model.AsyncJob{ + TokenID: 17, + ChannelID: 2, + EndpointType: model.AsyncEndpointImageGeneration, + RequestPayload: []byte("encrypted"), + RequestHash: strings.Repeat("a", 64), + IdempotencyKey: "expired-artifact-cleanup", + ExecutionStatus: model.AsyncStatusSuccess, + BillingStatus: model.AsyncBillingSettled, + ResultPayload: model.JSONValue(`{"data":[{"b64_json":"embedded-image-data"}]}`), + } + require.NoError(t, model.CreateAsyncTask(task, job)) + artifact := &model.Artifact{ + TaskID: task.ID, + ObjectKey: "async/task_expired_artifact_cleanup/image.png", + ContentType: "image/png", + SizeBytes: 12, + SHA256: strings.Repeat("b", 64), + SourceURLHash: strings.Repeat("c", 64), + ExpiresAt: time.Now().Add(-time.Minute).Unix(), + } + require.NoError(t, model.DB.Create(artifact).Error) + t.Cleanup(func() { + model.DB.Where("task_id = ?", task.ID).Delete(&model.Artifact{}) + model.DB.Where("task_id = ?", task.ID).Delete(&model.TaskEvent{}) + model.DB.Where("task_id = ?", task.ID).Delete(&model.AsyncJob{}) + model.DB.Delete(task) + }) + + store := &memoryArtifactStore{objects: map[string][]byte{artifact.ObjectKey: []byte("image")}} + deleted, err := CleanupExpiredAsyncArtifacts(context.Background(), store, 100) + require.NoError(t, err) + assert.Equal(t, 1, deleted) + assert.NotContains(t, store.objects, artifact.ObjectKey) + + var artifactCount int64 + require.NoError(t, model.DB.Model(&model.Artifact{}).Where("id = ?", artifact.ID).Count(&artifactCount).Error) + assert.Zero(t, artifactCount) + var storedJob model.AsyncJob + require.NoError(t, model.DB.First(&storedJob, job.ID).Error) + assert.Empty(t, storedJob.ResultPayload) + var taskCount int64 + require.NoError(t, model.DB.Model(&model.Task{}).Where("id = ?", task.ID).Count(&taskCount).Error) + assert.Equal(t, int64(1), taskCount) +} + +func TestArchivePersistsEveryImageInMultiImageResponse(t *testing.T) { + task := &model.Task{TaskID: "task_multiple_artifacts", Status: model.TaskStatusInProgress} + require.NoError(t, model.DB.Create(task).Error) + t.Cleanup(func() { + model.DB.Where("task_id = ?", task.ID).Delete(&model.Artifact{}) + model.DB.Delete(task) + }) + + png := append([]byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n'}, bytes.Repeat([]byte{0}, 64)...) + encoded := base64.StdEncoding.EncodeToString(png) + store := &memoryArtifactStore{objects: map[string][]byte{}} + artifacts, err := ArchiveAsyncMedia(context.Background(), task, []AsyncMediaSource{ + {Base64: encoded, ContentType: "image/png"}, + {Base64: encoded, ContentType: "image/png"}, + }, store, 30) + require.NoError(t, err) + assert.Len(t, artifacts, 2) + assert.Len(t, store.objects, 2) +} + +func TestArchiveReusesExistingArtifactsAfterWorkerRecovery(t *testing.T) { + task := &model.Task{TaskID: "task_recovered_artifacts", Status: model.TaskStatusInProgress} + require.NoError(t, model.DB.Create(task).Error) + t.Cleanup(func() { + model.DB.Where("task_id = ?", task.ID).Delete(&model.Artifact{}) + model.DB.Delete(task) + }) + + png := append([]byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n'}, bytes.Repeat([]byte{0}, 64)...) + store := &memoryArtifactStore{objects: map[string][]byte{}} + first, err := ArchiveAsyncMedia(context.Background(), task, []AsyncMediaSource{{ + Base64: base64.StdEncoding.EncodeToString(png), + ContentType: "image/png", + }}, store, 30) + require.NoError(t, err) + require.Len(t, first, 1) + + second, err := ArchiveAsyncMedia(context.Background(), task, []AsyncMediaSource{{ + Base64: base64.StdEncoding.EncodeToString(append(png, 1)), + ContentType: "image/png", + }}, store, 30) + require.NoError(t, err) + require.Len(t, second, 1) + assert.Equal(t, first[0].ObjectKey, second[0].ObjectKey) + assert.Len(t, store.objects, 1) +} + +func TestMaterializeArtifactRejectsSpoofedImageMIME(t *testing.T) { + _, err := materializeAsyncArtifact(context.Background(), AsyncMediaSource{ + Base64: base64.StdEncoding.EncodeToString([]byte("not an image")), + ContentType: "image/png", + }, 1024, 0) + require.ErrorContains(t, err, "allowed raster image") +} + +func TestParseDataURLSource(t *testing.T) { + source, ok := ParseDataURLSource("data:image/png;base64,aGVsbG8=") + assert.True(t, ok) + assert.Equal(t, "image/png", source.ContentType) + assert.Equal(t, "aGVsbG8=", source.Base64) + _, ok = ParseDataURLSource("data:image/svg+xml;base64,PHN2Zz4=") + assert.False(t, ok) +} diff --git a/service/async_security.go b/service/async_security.go new file mode 100644 index 000000000000..166316a39714 --- /dev/null +++ b/service/async_security.go @@ -0,0 +1,248 @@ +package service + +import ( + "crypto/aes" + "crypto/cipher" + "crypto/rand" + "crypto/sha256" + "encoding/base64" + "encoding/hex" + "errors" + "fmt" + "net/url" + "os" + "strconv" + "strings" + "unicode/utf8" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/dto" + "github.com/QuantumNous/new-api/setting/model_setting" +) + +const asyncPayloadVersion byte = 1 + +func CanonicalAsyncJSON(body []byte) ([]byte, error) { + var value any + if err := common.Unmarshal(body, &value); err != nil { + return nil, err + } + return common.Marshal(value) +} + +func HashAsyncRequest(body []byte) (string, error) { + canonical, err := CanonicalAsyncJSON(body) + if err != nil { + return "", err + } + sum := sha256.Sum256(canonical) + return hex.EncodeToString(sum[:]), nil +} + +func asyncEncryptionKey() ([]byte, error) { + if encoded := strings.TrimSpace(os.Getenv("ASYNC_REQUEST_ENCRYPTION_KEY")); encoded != "" { + if decoded, err := base64.StdEncoding.DecodeString(encoded); err == nil && len(decoded) == 32 { + return decoded, nil + } + if decoded, err := base64.RawStdEncoding.DecodeString(encoded); err == nil && len(decoded) == 32 { + return decoded, nil + } + if decoded, err := hex.DecodeString(encoded); err == nil && len(decoded) == 32 { + return decoded, nil + } + return nil, errors.New("ASYNC_REQUEST_ENCRYPTION_KEY must encode exactly 32 bytes") + } + if secret := os.Getenv("CRYPTO_SECRET"); len(secret) >= 32 { + sum := sha256.Sum256([]byte(secret)) + return sum[:], nil + } + return nil, errors.New("ASYNC_REQUEST_ENCRYPTION_KEY or a CRYPTO_SECRET of at least 32 characters is required") +} + +func EncryptAsyncPayload(plaintext []byte) ([]byte, error) { + key, err := asyncEncryptionKey() + if err != nil { + return nil, err + } + block, err := aes.NewCipher(key) + if err != nil { + return nil, err + } + gcm, err := cipher.NewGCM(block) + if err != nil { + return nil, err + } + nonce := make([]byte, gcm.NonceSize()) + if _, err := rand.Read(nonce); err != nil { + return nil, err + } + sealed := gcm.Seal(nil, nonce, plaintext, nil) + out := make([]byte, 1, 1+len(nonce)+len(sealed)) + out[0] = asyncPayloadVersion + out = append(out, nonce...) + out = append(out, sealed...) + return out, nil +} + +func DecryptAsyncPayload(payload []byte) ([]byte, error) { + if len(payload) < 2 || payload[0] != asyncPayloadVersion { + return nil, errors.New("unsupported encrypted async payload") + } + key, err := asyncEncryptionKey() + if err != nil { + return nil, err + } + block, err := aes.NewCipher(key) + if err != nil { + return nil, err + } + gcm, err := cipher.NewGCM(block) + if err != nil { + return nil, err + } + if len(payload) < 1+gcm.NonceSize()+gcm.Overhead() { + return nil, errors.New("encrypted async payload is truncated") + } + nonce := payload[1 : 1+gcm.NonceSize()] + return gcm.Open(nil, nonce, payload[1+gcm.NonceSize():], nil) +} + +func ValidateAsyncImageRequest(request *dto.ImageRequest, raw []byte) error { + if request == nil { + return errors.New("image request is required") + } + if strings.TrimSpace(request.Model) == "" { + return errors.New("model is required") + } + if strings.TrimSpace(request.Prompt) == "" { + return errors.New("prompt is required") + } + maxPromptChars := 20000 + if configured, err := strconv.Atoi(os.Getenv("ASYNC_MAX_PROMPT_CHARS")); err == nil && configured > 0 { + maxPromptChars = configured + } + if utf8.RuneCountInString(request.Prompt) > maxPromptChars { + return fmt.Errorf("prompt exceeds %d characters", maxPromptChars) + } + if request.Stream != nil && *request.Stream { + return errors.New("streaming image responses are not supported by the async wrapper") + } + if len(raw) == 0 { + return errors.New("request body is empty") + } + maxRequestKB := 256 + if configured, err := strconv.Atoi(os.Getenv("ASYNC_MAX_REQUEST_BODY_KB")); err == nil && configured > 0 { + maxRequestKB = configured + } + if len(raw) > maxRequestKB*1024 { + return fmt.Errorf("request body exceeds %d KiB", maxRequestKB) + } + maxImages := 8 + if configured, err := strconv.Atoi(os.Getenv("ASYNC_ARTIFACT_MAX_FILES")); err == nil && configured > 0 { + maxImages = configured + } + if request.N != nil && int(*request.N) > maxImages { + return fmt.Errorf("image count exceeds %d", maxImages) + } + if model_setting.IsGeminiModelSupportImagine(request.Model) { + if request.N != nil && *request.N != 1 { + return errors.New("yunwu gemini image models support exactly one image per async task") + } + if !validAsyncGeminiAspectRatio(request.Size) { + return errors.New("unsupported yunwu gemini image aspect ratio") + } + if !validAsyncGeminiImageSize(request.Quality) { + return errors.New("unsupported yunwu gemini image quality; use 1K, 2K or 4K") + } + } + maxInputURLs := 8 + if configured, err := strconv.Atoi(os.Getenv("ASYNC_MAX_INPUT_URLS")); err == nil && configured > 0 { + maxInputURLs = configured + } + canonical, err := CanonicalAsyncJSON(raw) + if err != nil { + return err + } + var value any + if err := common.Unmarshal(canonical, &value); err != nil { + return err + } + if countHTTPURLs(value) > maxInputURLs { + return fmt.Errorf("request contains more than %d input URLs", maxInputURLs) + } + return nil +} + +func ValidateAsyncImageProviderRequest(request *dto.ImageRequest, provider common.AsyncImageProvider) error { + if request == nil { + return errors.New("image request is required") + } + if provider == common.AsyncImageProviderGRSAI && request.N != nil && *request.N != 1 { + return errors.New("GRS AI synchronous image generation supports exactly one image per task") + } + if provider == common.AsyncImageProviderGRSAI && strings.EqualFold(strings.TrimSpace(request.Model), "gpt-image-2-vip") { + capabilities := model_setting.GetImageGenerationCapabilities(request.Model) + size := strings.TrimSpace(request.Size) + if size != "" && (capabilities == nil || !common.StringsContains(capabilities.Sizes, size)) { + return errors.New("unsupported gpt-image-2-vip size; use one of the documented 1K, 2K or 4K presets") + } + if strings.TrimSpace(request.Quality) != "" { + return errors.New("gpt-image-2-vip does not support the quality parameter; select resolution with size") + } + } + return nil +} + +func validAsyncGeminiAspectRatio(size string) bool { + switch strings.TrimSpace(size) { + case "", "1:1", "3:2", "2:3", "9:16", "16:9", "4:3", "3:4", "4:5", "5:4", "21:9", + "256x256", "512x512", "1024x1024", "1536x1024", "1024x1536", "1024x1792", "1792x1024": + return true + default: + return false + } +} + +func validAsyncGeminiImageSize(quality string) bool { + switch strings.ToLower(strings.TrimSpace(quality)) { + case "", "auto", "standard", "medium", "hd", "high", "1k", "2k", "4k": + return true + default: + return false + } +} + +func countHTTPURLs(value any) int { + switch typed := value.(type) { + case string: + parsed, err := url.Parse(typed) + if err == nil && (parsed.Scheme == "http" || parsed.Scheme == "https") && parsed.Host != "" { + return 1 + } + case []any: + count := 0 + for _, item := range typed { + count += countHTTPURLs(item) + } + return count + case map[string]any: + count := 0 + for _, item := range typed { + count += countHTTPURLs(item) + } + return count + } + return 0 +} + +func IsAllowedYunwuBaseURL(raw string) bool { + return common.IsAllowedYunwuBaseURL(raw) +} + +func AsyncImageProviderForBaseURL(raw string) (common.AsyncImageProvider, bool) { + return common.AsyncImageProviderForBaseURL(raw) +} + +func IsAllowedAsyncImageBaseURL(raw string) bool { + return common.IsAllowedAsyncImageBaseURL(raw) +} diff --git a/service/async_security_test.go b/service/async_security_test.go new file mode 100644 index 000000000000..e95f13177067 --- /dev/null +++ b/service/async_security_test.go @@ -0,0 +1,121 @@ +package service + +import ( + "bytes" + "encoding/base64" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/dto" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestAsyncPayloadEncryptionRoundTrip(t *testing.T) { + key := bytes.Repeat([]byte{0x42}, 32) + t.Setenv("ASYNC_REQUEST_ENCRYPTION_KEY", base64.StdEncoding.EncodeToString(key)) + plaintext := []byte(`{"model":"test","prompt":"secret prompt"}`) + encrypted, err := EncryptAsyncPayload(plaintext) + require.NoError(t, err) + assert.NotContains(t, string(encrypted), "secret prompt") + decrypted, err := DecryptAsyncPayload(encrypted) + require.NoError(t, err) + assert.Equal(t, plaintext, decrypted) +} + +func TestHashAsyncRequestCanonicalizesObjectKeys(t *testing.T) { + first, err := HashAsyncRequest([]byte(`{"model":"m","prompt":"p"}`)) + require.NoError(t, err) + second, err := HashAsyncRequest([]byte(" { \"prompt\" : \"p\", \"model\": \"m\" } ")) + require.NoError(t, err) + assert.Equal(t, first, second) +} + +func TestValidateAsyncImageRequest(t *testing.T) { + request := &dto.ImageRequest{Model: "m", Prompt: "hello"} + require.NoError(t, ValidateAsyncImageRequest(request, []byte(`{"model":"m","prompt":"hello"}`))) + + stream := true + request.Stream = &stream + require.ErrorContains(t, ValidateAsyncImageRequest(request, []byte(`{"model":"m","prompt":"hello","stream":true}`)), "streaming") + + stream = false + count := uint(9) + request.N = &count + require.ErrorContains(t, ValidateAsyncImageRequest(request, []byte(`{"model":"m","prompt":"hello","n":9}`)), "image count") + + t.Setenv("ASYNC_MAX_REQUEST_BODY_KB", "1") + request.N = nil + require.ErrorContains(t, ValidateAsyncImageRequest(request, bytes.Repeat([]byte{'x'}, 1025)), "request body") +} + +func TestValidateAsyncGeminiImageRequest(t *testing.T) { + count := uint(1) + request := &dto.ImageRequest{ + Model: "gemini-3.1-flash-image-preview", + Prompt: "draw a square", + N: &count, + Size: "16:9", + Quality: "2K", + } + require.NoError(t, ValidateAsyncImageRequest(request, []byte(`{"model":"gemini-3.1-flash-image-preview","prompt":"draw a square","n":1,"size":"16:9","quality":"2K"}`))) + + count = 2 + require.ErrorContains(t, ValidateAsyncImageRequest(request, []byte(`{"model":"gemini-3.1-flash-image-preview","prompt":"draw a square","n":2}`)), "exactly one") + count = 1 + request.Quality = "ultra" + require.ErrorContains(t, ValidateAsyncImageRequest(request, []byte(`{"model":"gemini-3.1-flash-image-preview","prompt":"draw a square","quality":"ultra"}`)), "use 1K, 2K or 4K") +} + +func TestAllowedYunwuBaseURL(t *testing.T) { + t.Setenv("ASYNC_YUNWU_ALLOWED_BASE_URLS", "") + assert.True(t, IsAllowedYunwuBaseURL("https://yunwu.ai")) + assert.False(t, IsAllowedYunwuBaseURL("https://yunwu.ai.attacker.example")) + assert.False(t, IsAllowedYunwuBaseURL("http://yunwu.ai")) + assert.False(t, IsAllowedYunwuBaseURL("https://yunwu.ai/unapproved-prefix")) + + t.Setenv("ASYNC_YUNWU_ALLOWED_BASE_URLS", "http://mock-upstream:8080") + assert.True(t, IsAllowedYunwuBaseURL("http://mock-upstream:8080/v1")) +} + +func TestAllowedGRSAIBaseURL(t *testing.T) { + t.Setenv("ASYNC_GRSAI_ALLOWED_BASE_URLS", "") + provider, allowed := AsyncImageProviderForBaseURL("https://grsaiapi.com/v1") + assert.True(t, allowed) + assert.Equal(t, common.AsyncImageProviderGRSAI, provider) + assert.True(t, IsAllowedAsyncImageBaseURL("https://grsai.dakka.com.cn")) + assert.False(t, IsAllowedAsyncImageBaseURL("https://grsai.com")) + assert.False(t, IsAllowedAsyncImageBaseURL("https://grsaiapi.com.attacker.example")) + + t.Setenv("ASYNC_GRSAI_ALLOWED_BASE_URLS", "http://mock-grsai:8080") + assert.True(t, IsAllowedAsyncImageBaseURL("http://mock-grsai:8080/v1")) +} + +func TestValidateGRSAISynchronousImageRequest(t *testing.T) { + count := uint(1) + request := &dto.ImageRequest{Model: "nano-banana-2", Prompt: "draw", N: &count} + require.NoError(t, ValidateAsyncImageProviderRequest(request, common.AsyncImageProviderGRSAI)) + + count = 2 + require.ErrorContains(t, ValidateAsyncImageProviderRequest(request, common.AsyncImageProviderGRSAI), "exactly one") +} + +func TestValidateGRSAIGPTImageVIPResolution(t *testing.T) { + count := uint(1) + request := &dto.ImageRequest{ + Model: "gpt-image-2-vip", + Prompt: "draw", + N: &count, + Size: "3840x2160", + } + require.NoError(t, ValidateAsyncImageProviderRequest(request, common.AsyncImageProviderGRSAI)) + request.Size = "" + require.NoError(t, ValidateAsyncImageProviderRequest(request, common.AsyncImageProviderGRSAI)) + + request.Size = "1024x1024" + require.ErrorContains(t, ValidateAsyncImageProviderRequest(request, common.AsyncImageProviderGRSAI), "1K, 2K or 4K presets") + + request.Size = "2048x2048" + request.Quality = "high" + require.ErrorContains(t, ValidateAsyncImageProviderRequest(request, common.AsyncImageProviderGRSAI), "does not support the quality parameter") +} diff --git a/service/async_worker.go b/service/async_worker.go new file mode 100644 index 000000000000..d40d313e7ae7 --- /dev/null +++ b/service/async_worker.go @@ -0,0 +1,433 @@ +package service + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "os" + "strconv" + "sync" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/dto" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/storage" +) + +type AsyncExecutionOutcome struct { + Status model.AsyncExecutionStatus + Payload json.RawMessage + Media []AsyncMediaSource + ErrorPhase string + ErrorCode string + ErrorMessage string + RefundEligible bool +} + +type AsyncImageExecutor interface { + Execute(ctx context.Context, payload []byte, markRequestSent func() error) AsyncExecutionOutcome +} + +type AsyncImageExecutorFactory func(channel *model.Channel, apiKey string, timeout time.Duration) (AsyncImageExecutor, error) + +var NewAsyncImageExecutor AsyncImageExecutorFactory + +type AsyncWorker struct { + ID string + Concurrency int + LeaseDuration time.Duration + PollInterval time.Duration + JobTimeout time.Duration + Store storage.ArtifactStore + Factory AsyncImageExecutorFactory + + semaphores *asyncSemaphoreRegistry + wg sync.WaitGroup +} + +func NewAsyncWorkerFromEnv(ctx context.Context) (*AsyncWorker, error) { + store, err := storage.NewS3ArtifactStore(ctx) + if err != nil { + return nil, err + } + workerID := os.Getenv("ASYNC_WORKER_ID") + if workerID == "" { + workerID = fmt.Sprintf("%s-%d", common.NodeName, os.Getpid()) + } + concurrency := common.GetEnvOrDefault("ASYNC_WORKER_CONCURRENCY", 50) + if concurrency <= 0 { + return nil, errors.New("ASYNC_WORKER_CONCURRENCY must be positive") + } + leaseSeconds := common.GetEnvOrDefault("ASYNC_WORKER_LEASE_SECONDS", 90) + if leaseSeconds < 15 { + return nil, errors.New("ASYNC_WORKER_LEASE_SECONDS must be at least 15") + } + jobTimeout := common.GetEnvOrDefault("ASYNC_JOB_TIMEOUT_SECONDS", 1800) + if jobTimeout <= 0 { + return nil, errors.New("ASYNC_JOB_TIMEOUT_SECONDS must be positive") + } + if NewAsyncImageExecutor == nil { + return nil, errors.New("async image executor factory is not configured") + } + return &AsyncWorker{ + ID: workerID, + Concurrency: concurrency, + LeaseDuration: time.Duration(leaseSeconds) * time.Second, + PollInterval: time.Duration(common.GetEnvOrDefault("ASYNC_WORKER_POLL_MILLISECONDS", 500)) * time.Millisecond, + JobTimeout: time.Duration(jobTimeout) * time.Second, + Store: store, + Factory: NewAsyncImageExecutor, + semaphores: newAsyncSemaphoreRegistry(), + }, nil +} + +func (w *AsyncWorker) Run(ctx context.Context) error { + if w == nil || w.Store == nil || w.Factory == nil { + return errors.New("async worker is not fully configured") + } + if w.PollInterval <= 0 { + w.PollInterval = 500 * time.Millisecond + } + global := make(chan struct{}, w.Concurrency) + recoveryTicker := time.NewTicker(15 * time.Second) + cleanupIntervalSeconds := common.GetEnvOrDefault("ASYNC_ARTIFACT_CLEANUP_INTERVAL_SECONDS", 60) + if cleanupIntervalSeconds < 10 { + cleanupIntervalSeconds = 10 + } + cleanupTicker := time.NewTicker(time.Duration(cleanupIntervalSeconds) * time.Second) + defer recoveryTicker.Stop() + defer cleanupTicker.Stop() + + if _, err := model.RecoverExpiredAsyncJobs(ctx, time.Now().Unix(), w.Concurrency*4); err != nil { + common.SysError("initial async lease recovery failed: " + err.Error()) + } + if _, err := model.ReconcileAsyncBilling(ctx, w.Concurrency*4); err != nil { + common.SysError("initial async billing reconciliation failed: " + err.Error()) + } + if _, err := ReconcileAsyncUpstreamCosts(ctx, w.Concurrency*4); err != nil { + common.SysError("initial async upstream cost reconciliation failed: " + err.Error()) + } + if _, err := CleanupExpiredAsyncArtifacts(ctx, w.Store, 100); err != nil { + common.SysError("initial async artifact cleanup failed: " + err.Error()) + } + + for { + select { + case <-ctx.Done(): + w.wg.Wait() + return nil + case <-recoveryTicker.C: + if _, err := model.RecoverExpiredAsyncJobs(ctx, time.Now().Unix(), w.Concurrency*4); err != nil && !errors.Is(err, context.Canceled) { + common.SysError("async lease recovery failed: " + err.Error()) + } + if _, err := model.ReconcileAsyncBilling(ctx, w.Concurrency*4); err != nil && !errors.Is(err, context.Canceled) { + common.SysError("async billing reconciliation failed: " + err.Error()) + } + if _, err := ReconcileAsyncUpstreamCosts(ctx, w.Concurrency*4); err != nil && !errors.Is(err, context.Canceled) { + common.SysError("async upstream cost reconciliation failed: " + err.Error()) + } + case <-cleanupTicker.C: + if _, err := CleanupExpiredAsyncArtifacts(ctx, w.Store, 100); err != nil && !errors.Is(err, context.Canceled) { + common.SysError("async artifact cleanup failed: " + err.Error()) + } + default: + } + + available := w.Concurrency - len(global) + if available <= 0 { + if !waitAsyncPoll(ctx, w.PollInterval) { + w.wg.Wait() + return nil + } + continue + } + candidates, err := model.ListQueuedAsyncJobs(ctx, available*4) + if err != nil { + if errors.Is(err, context.Canceled) { + w.wg.Wait() + return nil + } + common.SysError("list async jobs failed: " + err.Error()) + waitAsyncPoll(ctx, w.PollInterval) + continue + } + claimedCount := 0 + for _, candidate := range candidates { + if len(global) >= cap(global) { + break + } + channel, channelErr := model.GetChannelById(candidate.ChannelID, true) + limit := 1 + if channelErr == nil && channel != nil { + limit = channel.GetSetting().AsyncMaxConcurrency + if limit <= 0 { + limit = common.GetEnvOrDefault("ASYNC_CHANNEL_DEFAULT_CONCURRENCY", 10) + } + } + modelName := candidate.Task.Properties.OriginModelName + release, acquired := w.semaphores.TryAcquire(candidate.ChannelID, modelName, limit) + if !acquired { + continue + } + global <- struct{}{} + claimed, won, claimErr := model.ClaimAsyncJob(ctx, candidate.ID, w.ID, time.Now().Add(w.LeaseDuration).Unix()) + if claimErr != nil || !won { + <-global + release() + if claimErr != nil { + common.SysError("claim async job failed: " + claimErr.Error()) + } + continue + } + claimedCount++ + w.wg.Add(1) + go func(job *model.AsyncJob, releaseLimits func()) { + defer w.wg.Done() + defer func() { <-global }() + defer releaseLimits() + // Shutdown stops new claims but deliberately does not cancel an + // already-sent upstream request. Each job drains under its own + // configured timeout before Run returns. + w.process(context.Background(), job) + }(claimed, release) + } + if claimedCount == 0 && !waitAsyncPoll(ctx, w.PollInterval) { + w.wg.Wait() + return nil + } + } +} + +func waitAsyncPoll(ctx context.Context, duration time.Duration) bool { + timer := time.NewTimer(duration) + defer timer.Stop() + select { + case <-ctx.Done(): + return false + case <-timer.C: + return true + } +} + +func (w *AsyncWorker) process(parent context.Context, job *model.AsyncJob) { + jobCtx, cancel := context.WithCancel(parent) + defer cancel() + leaseDone := make(chan struct{}) + go w.renewLease(jobCtx, cancel, job.ID, leaseDone) + defer close(leaseDone) + + fail := func(phase, code, message string, refundable bool) { + changed, err := model.CompleteAsyncJob(context.Background(), job.ID, w.ID, model.AsyncStatusFailure, nil, phase, code, message, refundable) + if err != nil { + common.SysError(fmt.Sprintf("complete async job %d as failure: %v", job.ID, err)) + return + } + if !changed { + return + } + if refundable { + _, err = model.RefundAsyncJobBilling(context.Background(), job.ID) + } else { + err = settleAsyncJobBilling(context.Background(), job) + } + if err != nil { + common.SysError(fmt.Sprintf("finalize async billing for job %d: %v", job.ID, err)) + } + } + + payload, err := DecryptAsyncPayload(job.RequestPayload) + if err != nil { + fail("request_decrypt", "request_decrypt_failed", "encrypted task payload could not be decrypted", true) + return + } + channel, err := model.GetChannelById(job.ChannelID, true) + if err != nil || channel == nil || channel.Status != common.ChannelStatusEnabled { + fail("channel_load", "channel_unavailable", "the selected async channel is unavailable", true) + return + } + setting := channel.GetSetting() + modelName := job.Task.Properties.OriginModelName + if !setting.AllowsAsyncImageModel(modelName) || !setting.AsyncArchiveEnabled() || !IsAllowedAsyncImageBaseURL(channel.GetBaseURL()) { + fail("channel_validate", "channel_async_disabled", "the selected channel no longer permits this async image model", true) + return + } + apiKey, _, apiErr := channel.GetNextEnabledKey() + if apiErr != nil { + fail("channel_key", "channel_key_unavailable", "the selected channel has no available credential", true) + return + } + timeout := w.JobTimeout + if setting.AsyncJobTimeoutSeconds > 0 { + timeout = time.Duration(setting.AsyncJobTimeoutSeconds) * time.Second + } + if timeout <= 0 { + timeout = 30 * time.Minute + } + executor, err := w.Factory(channel, apiKey, timeout) + if err != nil { + fail("executor_init", "executor_init_failed", "failed to initialize the synchronous image executor", true) + return + } + executionCtx, timeoutCancel := context.WithTimeout(jobCtx, timeout) + defer timeoutCancel() + outcome := executor.Execute(executionCtx, payload, func() error { + marked, markErr := model.MarkAsyncRequestSent(context.Background(), job.ID, w.ID, time.Now().Unix()) + if markErr != nil { + return markErr + } + if !marked { + return errors.New("async job lease was lost before request send") + } + return nil + }) + + switch outcome.Status { + case model.AsyncStatusSuccess: + retention := setting.EffectiveAsyncRetentionMinutes( + common.GetEnvOrDefault("ASYNC_RESULT_RETENTION_MINUTES", dto.AsyncRetentionDefaultMinutes), + ) + archiveTimeout := time.Duration(common.GetEnvOrDefault("ASYNC_ARTIFACT_ARCHIVE_TIMEOUT_SECONDS", 300)) * time.Second + if archiveTimeout <= 0 { + archiveTimeout = 5 * time.Minute + } + archiveCtx, archiveCancel := context.WithTimeout(jobCtx, archiveTimeout) + _, archiveErr := ArchiveAsyncMedia(archiveCtx, &job.Task, outcome.Media, w.Store, retention) + archiveCancel() + if archiveErr != nil { + fail("artifact_archive", "artifact_archive_failed", "upstream succeeded but one or more artifacts could not be archived", false) + return + } + changed, err := model.CompleteAsyncJob(context.Background(), job.ID, w.ID, model.AsyncStatusSuccess, outcome.Payload, "", "", "", false) + if err != nil { + common.SysError(fmt.Sprintf("complete async job %d as success: %v", job.ID, err)) + return + } + if changed { + if err := settleAsyncJobBilling(context.Background(), job); err != nil { + common.SysError(fmt.Sprintf("settle async job %d: %v", job.ID, err)) + } + } + case model.AsyncStatusUncertain: + changed, err := model.CompleteAsyncJob(context.Background(), job.ID, w.ID, model.AsyncStatusUncertain, outcome.Payload, outcome.ErrorPhase, outcome.ErrorCode, outcome.ErrorMessage, false) + if err != nil { + common.SysError(fmt.Sprintf("complete async job %d as uncertain: %v", job.ID, err)) + return + } + if changed { + if err := settleAsyncJobBilling(context.Background(), job); err != nil { + common.SysError(fmt.Sprintf("settle uncertain async job %d: %v", job.ID, err)) + } + } + case model.AsyncStatusFailure: + fail(outcome.ErrorPhase, outcome.ErrorCode, outcome.ErrorMessage, outcome.RefundEligible) + default: + fail("executor", "invalid_executor_outcome", "synchronous image executor returned an invalid state", false) + } +} + +func settleAsyncJobBilling(ctx context.Context, job *model.AsyncJob) error { + if job == nil { + return errors.New("async job is required") + } + settled, err := model.SettleAsyncJobBilling(ctx, job.ID) + if err != nil || !settled { + return err + } + channel, err := model.GetChannelById(job.ChannelID, false) + if err != nil { + return err + } + if channel == nil { + return fmt.Errorf("async channel %d not found after settlement", job.ChannelID) + } + LogAsyncTaskSettlement(job, channel) + return nil +} + +func (w *AsyncWorker) renewLease(ctx context.Context, cancel context.CancelFunc, jobID int64, done <-chan struct{}) { + interval := w.LeaseDuration / 3 + if interval < 5*time.Second { + interval = 5 * time.Second + } + ticker := time.NewTicker(interval) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-done: + return + case <-ticker.C: + renewed, err := model.RenewAsyncJobLease(context.Background(), jobID, w.ID, time.Now().Add(w.LeaseDuration).Unix()) + if err != nil || !renewed { + cancel() + return + } + } + } +} + +func CleanupExpiredAsyncArtifacts(ctx context.Context, store storage.ArtifactStore, limit int) (int, error) { + artifacts, err := model.ListExpiredArtifacts(ctx, time.Now().Unix(), limit) + if err != nil { + return 0, err + } + deleted := 0 + for _, artifact := range artifacts { + if err := store.Delete(ctx, artifact.ObjectKey); err != nil { + return deleted, err + } + if _, err := model.DeleteArtifactAndClearResultIfLast(ctx, artifact.ID, artifact.TaskID); err != nil { + return deleted, err + } + deleted++ + } + return deleted, nil +} + +type asyncSemaphoreRegistry struct { + mu sync.Mutex + items map[string]chan struct{} +} + +func newAsyncSemaphoreRegistry() *asyncSemaphoreRegistry { + return &asyncSemaphoreRegistry{items: make(map[string]chan struct{})} +} + +func (r *asyncSemaphoreRegistry) TryAcquire(channelID int, modelName string, limit int) (func(), bool) { + if limit <= 0 { + limit = 1 + } + channelKey := "channel:" + strconv.Itoa(channelID) + modelKey := channelKey + ":model:" + modelName + channelSemaphore := r.get(channelKey, limit) + modelSemaphore := r.get(modelKey, limit) + select { + case channelSemaphore <- struct{}{}: + default: + return func() {}, false + } + select { + case modelSemaphore <- struct{}{}: + return func() { + <-modelSemaphore + <-channelSemaphore + }, true + default: + <-channelSemaphore + return func() {}, false + } +} + +func (r *asyncSemaphoreRegistry) get(key string, limit int) chan struct{} { + r.mu.Lock() + defer r.mu.Unlock() + if existing, ok := r.items[key]; ok { + return existing + } + created := make(chan struct{}, limit) + r.items[key] = created + return created +} diff --git a/service/async_worker_test.go b/service/async_worker_test.go new file mode 100644 index 000000000000..9f318b7ec811 --- /dev/null +++ b/service/async_worker_test.go @@ -0,0 +1,201 @@ +package service + +import ( + "context" + "encoding/base64" + "encoding/json" + "sync" + "testing" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/dto" + "github.com/QuantumNous/new-api/model" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +type successfulAsyncExecutor struct{} + +func (successfulAsyncExecutor) Execute(_ context.Context, _ []byte, markRequestSent func() error) AsyncExecutionOutcome { + if err := markRequestSent(); err != nil { + return AsyncExecutionOutcome{Status: model.AsyncStatusFailure, ErrorPhase: "mark", ErrorCode: "mark_failed", ErrorMessage: err.Error(), RefundEligible: true} + } + png := []byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n', 0, 0, 0, 0} + return AsyncExecutionOutcome{ + Status: model.AsyncStatusSuccess, + Payload: json.RawMessage(`{"created":1,"data":[{"b64_json":"iVBORw0KGgoAAAAA"}]}`), + Media: []AsyncMediaSource{{Base64: base64.StdEncoding.EncodeToString(png), ContentType: "image/png"}}, + } +} + +type drainingAsyncExecutor struct { + started chan struct{} + release chan struct{} + once sync.Once +} + +func (e *drainingAsyncExecutor) Execute(ctx context.Context, payload []byte, markRequestSent func() error) AsyncExecutionOutcome { + if err := markRequestSent(); err != nil { + return AsyncExecutionOutcome{Status: model.AsyncStatusFailure, ErrorPhase: "mark", ErrorCode: "mark_failed", ErrorMessage: err.Error(), RefundEligible: true} + } + e.once.Do(func() { close(e.started) }) + select { + case <-e.release: + return successfulAsyncExecutor{}.Execute(context.Background(), payload, func() error { return nil }) + case <-ctx.Done(): + return AsyncExecutionOutcome{Status: model.AsyncStatusUncertain, ErrorPhase: "shutdown", ErrorCode: "cancelled_after_send", ErrorMessage: "request context was cancelled"} + } +} + +func TestAsyncWorkerCompletesDisconnectedClientTask(t *testing.T) { + t.Setenv("ASYNC_REQUEST_ENCRYPTION_KEY", base64.StdEncoding.EncodeToString(make([]byte, 32))) + t.Setenv("ASYNC_YUNWU_ALLOWED_BASE_URLS", "https://yunwu.ai") + model.DB.Exec("DELETE FROM task_events") + model.DB.Exec("DELETE FROM artifacts") + model.DB.Exec("DELETE FROM async_jobs") + model.DB.Exec("DELETE FROM tasks") + model.DB.Exec("DELETE FROM upstream_cost_records") + model.DB.Exec("DELETE FROM logs") + model.DB.Exec("DELETE FROM tokens") + model.DB.Exec("DELETE FROM users") + model.DB.Exec("DELETE FROM channels") + + user := &model.User{Id: 301, Username: "worker-user", Quota: 900, Status: common.UserStatusEnabled} + token := &model.Token{Id: 302, UserId: user.Id, Key: "worker-token-placeholder", Name: "worker", Status: common.TokenStatusEnabled, RemainQuota: 900, UsedQuota: 100} + baseURL := "https://yunwu.ai" + archive := true + channel := &model.Channel{Id: 303, Name: "yunwu-worker", Key: "upstream-placeholder", BaseURL: &baseURL, Status: common.ChannelStatusEnabled, Models: "image-model", Group: "default"} + channel.SetSetting(dto.ChannelSettings{AsyncImageEnabled: true, AsyncImageModels: []string{"image-model"}, AsyncMaxConcurrency: 1, AsyncAutoArchive: &archive}) + rate := 0.495 + channel.SetOtherSettings(dto.ChannelOtherSettings{ + UpstreamCostMode: dto.UpstreamCostModeBillingUnits, + UpstreamCostUnit: "CREDIT", + UpstreamCostRateCNY: &rate, + UpstreamCostPriceVersion: "yunwu-test", + }) + require.NoError(t, model.DB.Create(user).Error) + require.NoError(t, model.DB.Create(token).Error) + require.NoError(t, model.DB.Create(channel).Error) + + payload, err := EncryptAsyncPayload([]byte(`{"model":"image-model","prompt":"kept after disconnect"}`)) + require.NoError(t, err) + task := &model.Task{TaskID: "task_worker_disconnected", Platform: constant.TaskPlatformAsyncImage, UserId: user.Id, ChannelId: channel.Id, Quota: 100, Status: model.TaskStatusQueued, Progress: "0%", Properties: model.Properties{OriginModelName: "image-model"}, PrivateData: model.TaskPrivateData{BillingSource: BillingSourceWallet, TokenId: token.Id}, Data: json.RawMessage(`{}`)} + job := &model.AsyncJob{TokenID: token.Id, ChannelID: channel.Id, EndpointType: model.AsyncEndpointImageGeneration, RequestPayload: payload, RequestHash: "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd", IdempotencyKey: "worker-disconnected", ExecutionStatus: model.AsyncStatusQueued, BillingStatus: model.AsyncBillingReserved, BillingRequestID: "async-cost-request"} + require.NoError(t, model.CreateAsyncTask(task, job)) + claimed, won, err := model.ClaimAsyncJob(context.Background(), job.ID, "test-worker", time.Now().Add(time.Minute).Unix()) + require.NoError(t, err) + require.True(t, won) + + store := &memoryArtifactStore{objects: map[string][]byte{}} + worker := &AsyncWorker{ID: "test-worker", LeaseDuration: time.Minute, JobTimeout: time.Minute, Store: store, Factory: func(*model.Channel, string, time.Duration) (AsyncImageExecutor, error) { + return successfulAsyncExecutor{}, nil + }} + worker.process(context.Background(), claimed) + + loaded, err := model.GetAsyncJobByPublicTaskID(context.Background(), task.TaskID, token.Id) + require.NoError(t, err) + require.NotNil(t, loaded) + assert.Equal(t, model.AsyncStatusSuccess, loaded.ExecutionStatus) + assert.Equal(t, model.AsyncBillingSettled, loaded.BillingStatus) + artifacts, err := model.ListArtifactsByTaskID(context.Background(), task.ID) + require.NoError(t, err) + assert.Len(t, artifacts, 1) + assert.Len(t, store.objects, 1) + var cost model.UpstreamCostRecord + require.NoError(t, model.DB.Where("request_id = ?", job.BillingRequestID).First(&cost).Error) + assert.Equal(t, channel.Id, cost.ChannelId) + assert.Equal(t, "image-model", cost.ModelName) + assert.Equal(t, "CREDIT", cost.NativeUnit) + assert.Equal(t, "0.0002", cost.NativeAmount) + assert.Equal(t, "0.495", cost.RateCNYPerUnit) + assert.EqualValues(t, 99, cost.AmountCNYMicros) + assert.True(t, cost.Estimated) + + require.NoError(t, model.DB.Where("request_id = ?", job.BillingRequestID).Delete(&model.UpstreamCostRecord{}).Error) + require.NoError(t, model.LOG_DB.Where("request_id = ?", job.BillingRequestID).Delete(&model.Log{}).Error) + processed, err := ReconcileAsyncUpstreamCosts(context.Background(), 10) + require.NoError(t, err) + assert.Equal(t, 1, processed) + require.NoError(t, model.DB.Where("request_id = ?", job.BillingRequestID).First(&cost).Error) + processed, err = ReconcileAsyncUpstreamCosts(context.Background(), 10) + require.NoError(t, err) + assert.Equal(t, 0, processed) +} + +func TestAsyncChannelAndModelSemaphoreKeepsExcessQueued(t *testing.T) { + registry := newAsyncSemaphoreRegistry() + release, acquired := registry.TryAcquire(1, "image-model", 1) + require.True(t, acquired) + _, acquired = registry.TryAcquire(1, "image-model", 1) + assert.False(t, acquired) + release() + _, acquired = registry.TryAcquire(1, "image-model", 1) + assert.True(t, acquired) +} + +func TestAsyncWorkerShutdownDrainsRunningRequest(t *testing.T) { + t.Setenv("ASYNC_REQUEST_ENCRYPTION_KEY", base64.StdEncoding.EncodeToString(make([]byte, 32))) + t.Setenv("ASYNC_YUNWU_ALLOWED_BASE_URLS", "https://yunwu.ai") + model.DB.Exec("DELETE FROM task_events") + model.DB.Exec("DELETE FROM artifacts") + model.DB.Exec("DELETE FROM async_jobs") + model.DB.Exec("DELETE FROM tasks") + model.DB.Exec("DELETE FROM tokens") + model.DB.Exec("DELETE FROM users") + model.DB.Exec("DELETE FROM channels") + + user := &model.User{Id: 401, Username: "drain-user", Quota: 900, Status: common.UserStatusEnabled} + token := &model.Token{Id: 402, UserId: user.Id, Key: "drain-token-placeholder", Name: "worker", Status: common.TokenStatusEnabled, RemainQuota: 900, UsedQuota: 100} + baseURL := "https://yunwu.ai" + archive := true + channel := &model.Channel{Id: 403, Name: "yunwu-drain", Key: "upstream-placeholder", BaseURL: &baseURL, Status: common.ChannelStatusEnabled, Models: "image-model", Group: "default"} + channel.SetSetting(dto.ChannelSettings{AsyncImageEnabled: true, AsyncImageModels: []string{"image-model"}, AsyncMaxConcurrency: 1, AsyncAutoArchive: &archive}) + require.NoError(t, model.DB.Create(user).Error) + require.NoError(t, model.DB.Create(token).Error) + require.NoError(t, model.DB.Create(channel).Error) + payload, err := EncryptAsyncPayload([]byte(`{"model":"image-model","prompt":"drain on shutdown"}`)) + require.NoError(t, err) + task := &model.Task{TaskID: "task_worker_drain", Platform: constant.TaskPlatformAsyncImage, UserId: user.Id, ChannelId: channel.Id, Quota: 100, Status: model.TaskStatusQueued, Progress: "0%", Properties: model.Properties{OriginModelName: "image-model"}, PrivateData: model.TaskPrivateData{BillingSource: BillingSourceWallet, TokenId: token.Id}, Data: json.RawMessage(`{}`)} + job := &model.AsyncJob{TokenID: token.Id, ChannelID: channel.Id, EndpointType: model.AsyncEndpointImageGeneration, RequestPayload: payload, RequestHash: "eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee", IdempotencyKey: "worker-drain", ExecutionStatus: model.AsyncStatusQueued, BillingStatus: model.AsyncBillingReserved} + require.NoError(t, model.CreateAsyncTask(task, job)) + + executor := &drainingAsyncExecutor{started: make(chan struct{}), release: make(chan struct{})} + worker := &AsyncWorker{ + ID: "draining-worker", + Concurrency: 1, + LeaseDuration: 30 * time.Second, + PollInterval: 5 * time.Millisecond, + JobTimeout: 2 * time.Second, + Store: &memoryArtifactStore{objects: map[string][]byte{}}, + Factory: func(*model.Channel, string, time.Duration) (AsyncImageExecutor, error) { return executor, nil }, + semaphores: newAsyncSemaphoreRegistry(), + } + runCtx, cancelRun := context.WithCancel(context.Background()) + runDone := make(chan error, 1) + go func() { runDone <- worker.Run(runCtx) }() + select { + case <-executor.started: + case <-time.After(2 * time.Second): + t.Fatal("worker did not start the queued task") + } + cancelRun() + select { + case err := <-runDone: + t.Fatalf("worker returned before the running request drained: %v", err) + case <-time.After(75 * time.Millisecond): + } + close(executor.release) + select { + case err := <-runDone: + require.NoError(t, err) + case <-time.After(2 * time.Second): + t.Fatal("worker did not finish after the running request drained") + } + + loaded, err := model.GetAsyncJobByPublicTaskID(context.Background(), task.TaskID, token.Id) + require.NoError(t, err) + require.NotNil(t, loaded) + assert.Equal(t, model.AsyncStatusSuccess, loaded.ExecutionStatus) +} diff --git a/service/log_info_generate.go b/service/log_info_generate.go index 207b0af5bfef..a0b1d166973c 100644 --- a/service/log_info_generate.go +++ b/service/log_info_generate.go @@ -297,6 +297,7 @@ func GenerateMjOtherInfo(relayInfo *relaycommon.RelayInfo, priceData types.Price other["user_group_ratio"] = priceData.GroupRatioInfo.GroupSpecialRatio } appendRequestPath(nil, relayInfo, other) + AttachUpstreamCost(relayInfo, priceData.Quota, other) return other } diff --git a/service/quota.go b/service/quota.go index 84ef22b94dff..ff5daf603cff 100644 --- a/service/quota.go +++ b/service/quota.go @@ -241,6 +241,7 @@ func PostWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, mod if tieredResult != nil { InjectTieredBillingInfo(other, relayInfo, tieredResult) } + AttachUpstreamCost(relayInfo, quota, other) attachQuotaSaturation(ctx, relayInfo, other) model.RecordConsumeLog(ctx, relayInfo.UserId, model.RecordConsumeLogParams{ ChannelId: relayInfo.ChannelId, @@ -364,6 +365,7 @@ func PostAudioConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, u if tieredResult != nil { InjectTieredBillingInfo(other, relayInfo, tieredResult) } + AttachUpstreamCostWithUsage(relayInfo, usage, quota, other) attachQuotaSaturation(ctx, relayInfo, other) model.RecordConsumeLog(ctx, relayInfo.UserId, model.RecordConsumeLogParams{ ChannelId: relayInfo.ChannelId, diff --git a/service/response_billing.go b/service/response_billing.go new file mode 100644 index 000000000000..406f0661ea7d --- /dev/null +++ b/service/response_billing.go @@ -0,0 +1,83 @@ +package service + +import ( + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/dto" + "github.com/QuantumNous/new-api/pkg/billingexpr" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/setting/billing_setting" + + "github.com/gin-gonic/gin" + "github.com/shopspring/decimal" +) + +// AttachResponseBilling adds the gateway-calculated charge to an API usage +// object when the client explicitly requested billing details. +func AttachResponseBilling(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage *dto.Usage) { + if relayInfo == nil || !relayInfo.ShouldIncludeBilling || usage == nil { + return + } + usage.Billing = buildResponseBilling(ctx, relayInfo, usage) +} + +func buildResponseBilling(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage *dto.Usage) *dto.ResponseBilling { + billingUsage := effectiveBillingUsage(usage) + summary := calculateTextQuotaSummary(ctx, relayInfo, billingUsage) + mode := billing_setting.BillingModeRatio + matchedTier := "" + quotaPerUnit := common.QuotaPerUnit + + if snap := relayInfo.TieredBillingSnapshot; snap != nil && snap.BillingMode == billing_setting.BillingModeTieredExpr { + mode = billing_setting.BillingModeTieredExpr + if snap.QuotaPerUnit > 0 { + quotaPerUnit = snap.QuotaPerUnit + } + usedVars := billingexpr.UsedVars(snap.ExprString) + ok, tieredQuota, result := TryTieredSettle( + relayInfo, + BuildTieredTokenParams(billingUsage, summary.IsClaudeUsageSemantic, usedVars), + ) + if ok { + summary.Quota = composeTieredTextQuota(relayInfo, summary, tieredQuota, result) + if result != nil { + matchedTier = result.MatchedTier + } + } + } else if relayInfo.PriceData.UsePrice { + mode = "fixed" + } + + totalCost := 0.0 + if quotaPerUnit > 0 { + totalCost, _ = decimal.NewFromInt(int64(summary.Quota)). + Div(decimal.NewFromFloat(quotaPerUnit)). + Float64() + } + + billing := &dto.ResponseBilling{ + Currency: "USD", + TotalCost: totalCost, + BillingMode: mode, + BillingSource: relayInfo.BillingSource, + GroupRatio: summary.GroupRatio, + MatchedTier: matchedTier, + } + + otherRatio := relayInfo.PriceData.OtherRatioMultiplier() + switch mode { + case billing_setting.BillingModeRatio: + unitScale := 0.0 + if common.QuotaPerUnit > 0 { + unitScale = 1_000_000 / common.QuotaPerUnit + } + inputPrice := summary.ModelRatio * summary.GroupRatio * otherRatio * unitScale + outputPrice := inputPrice * summary.CompletionRatio + billing.InputUnitPricePerMillion = &inputPrice + billing.OutputUnitPricePerMillion = &outputPrice + case "fixed": + requestPrice := summary.ModelPrice * summary.GroupRatio * otherRatio + billing.RequestPrice = &requestPrice + } + + return billing +} diff --git a/service/response_billing_test.go b/service/response_billing_test.go new file mode 100644 index 000000000000..a581c295f7d8 --- /dev/null +++ b/service/response_billing_test.go @@ -0,0 +1,96 @@ +package service + +import ( + "net/http/httptest" + "testing" + "time" + + "github.com/QuantumNous/new-api/dto" + "github.com/QuantumNous/new-api/pkg/billingexpr" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/types" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func newResponseBillingTestContext() *gin.Context { + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + return ctx +} + +func TestAttachResponseBillingUsesSettledRatioCharge(t *testing.T) { + gin.SetMode(gin.TestMode) + ctx := newResponseBillingTestContext() + info := &relaycommon.RelayInfo{ + ShouldIncludeBilling: true, + OriginModelName: "test-model", + StartTime: time.Now(), + BillingSource: BillingSourceWallet, + PriceData: types.PriceData{ + ModelRatio: 1, + CompletionRatio: 2, + GroupRatioInfo: types.GroupRatioInfo{ + GroupRatio: 0.5, + }, + }, + } + usage := &dto.Usage{PromptTokens: 1000, CompletionTokens: 500, TotalTokens: 1500} + + AttachResponseBilling(ctx, info, usage) + + require.NotNil(t, usage.Billing) + assert.Equal(t, "USD", usage.Billing.Currency) + assert.Equal(t, "ratio", usage.Billing.BillingMode) + assert.Equal(t, BillingSourceWallet, usage.Billing.BillingSource) + assert.InDelta(t, 0.002, usage.Billing.TotalCost, 1e-12) + require.NotNil(t, usage.Billing.InputUnitPricePerMillion) + require.NotNil(t, usage.Billing.OutputUnitPricePerMillion) + assert.InDelta(t, 1, *usage.Billing.InputUnitPricePerMillion, 1e-12) + assert.InDelta(t, 2, *usage.Billing.OutputUnitPricePerMillion, 1e-12) +} + +func TestAttachResponseBillingUsesTieredExpressionCharge(t *testing.T) { + gin.SetMode(gin.TestMode) + ctx := newResponseBillingTestContext() + expr := `tier("base", p * 2 + c * 4)` + info := &relaycommon.RelayInfo{ + ShouldIncludeBilling: true, + OriginModelName: "tiered-model", + StartTime: time.Now(), + PriceData: types.PriceData{ + GroupRatioInfo: types.GroupRatioInfo{GroupRatio: 0.5}, + }, + TieredBillingSnapshot: &billingexpr.BillingSnapshot{ + BillingMode: "tiered_expr", + ModelName: "tiered-model", + ExprString: expr, + ExprHash: billingexpr.ExprHashString(expr), + GroupRatio: 0.5, + QuotaPerUnit: 500_000, + }, + } + usage := &dto.Usage{PromptTokens: 1000, CompletionTokens: 500, TotalTokens: 1500} + + AttachResponseBilling(ctx, info, usage) + + require.NotNil(t, usage.Billing) + assert.Equal(t, "tiered_expr", usage.Billing.BillingMode) + assert.Equal(t, "base", usage.Billing.MatchedTier) + assert.InDelta(t, 0.002, usage.Billing.TotalCost, 1e-12) + assert.Nil(t, usage.Billing.InputUnitPricePerMillion) + assert.Nil(t, usage.Billing.OutputUnitPricePerMillion) +} + +func TestAttachResponseBillingIsOptIn(t *testing.T) { + gin.SetMode(gin.TestMode) + ctx := newResponseBillingTestContext() + info := &relaycommon.RelayInfo{} + usage := &dto.Usage{PromptTokens: 1, TotalTokens: 1} + + AttachResponseBilling(ctx, info, usage) + + assert.Nil(t, usage.Billing) +} diff --git a/service/task_billing.go b/service/task_billing.go index 51fb77c31543..06ab01b8c8f1 100644 --- a/service/task_billing.go +++ b/service/task_billing.go @@ -51,6 +51,7 @@ func LogTaskConsumption(c *gin.Context, info *relaycommon.RelayInfo) { other["is_model_mapped"] = true other["upstream_model_name"] = info.UpstreamModelName } + AttachUpstreamCost(info, info.PriceData.Quota, other) attachQuotaSaturation(c, info, other) model.RecordConsumeLog(c, info.UserId, model.RecordConsumeLogParams{ ChannelId: info.ChannelId, @@ -152,6 +153,62 @@ func taskBillingContextPriceData(bc *model.TaskBillingContext) *types.PriceData return priceData } +// LogAsyncTaskSettlement writes the consume log and immutable upstream-cost +// snapshot for an async attempt whose reserved quota was committed. +func LogAsyncTaskSettlement(job *model.AsyncJob, channel *model.Channel) { + if job == nil || channel == nil || job.Task.ID == 0 { + return + } + task := &job.Task + other := taskBillingOther(task) + other["task_id"] = task.TaskID + other["is_async"] = true + other["async_execution_status"] = job.ExecutionStatus + relayInfo := &relaycommon.RelayInfo{ + ChannelMeta: &relaycommon.ChannelMeta{ + ChannelId: channel.Id, + ChannelType: channel.Type, + ChannelOtherSettings: channel.GetOtherSettings(), + }, + } + AttachUpstreamCost(relayInfo, task.Quota, other) + model.RecordTaskBillingLog(model.RecordTaskBillingLogParams{ + UserId: task.UserId, + LogType: model.LogTypeConsume, + Content: "异步任务结算", + ChannelId: task.ChannelId, + ModelName: taskModelName(task), + Quota: task.Quota, + TokenId: task.PrivateData.TokenId, + Group: task.Group, + RequestId: job.BillingRequestID, + Other: other, + NodeName: task.PrivateData.NodeName, + }) +} + +// ReconcileAsyncUpstreamCosts backfills the cost ledger if a process stopped +// after committing async billing but before recording its consume log. +func ReconcileAsyncUpstreamCosts(ctx context.Context, limit int) (int, error) { + jobs, err := model.ListSettledAsyncJobsMissingUpstreamCost(ctx, limit) + if err != nil { + return 0, err + } + processed := 0 + for i := range jobs { + channel, err := model.GetChannelById(jobs[i].ChannelID, false) + if err != nil { + return processed, err + } + if channel == nil { + continue + } + LogAsyncTaskSettlement(&jobs[i], channel) + processed++ + } + return processed, nil +} + // taskModelName 从 BillingContext 或 Properties 中获取模型名称。 func taskModelName(task *model.Task) string { if bc := task.PrivateData.BillingContext; bc != nil && bc.OriginModelName != "" { diff --git a/service/task_billing_test.go b/service/task_billing_test.go index 53e3f680d01c..9167c4177681 100644 --- a/service/task_billing_test.go +++ b/service/task_billing_test.go @@ -41,6 +41,9 @@ func TestMain(m *testing.M) { if err := db.AutoMigrate( &model.Task{}, + &model.AsyncJob{}, + &model.Artifact{}, + &model.TaskEvent{}, &model.User{}, &model.Token{}, &model.Log{}, @@ -49,6 +52,7 @@ func TestMain(m *testing.M) { &model.UserSubscription{}, &model.SystemTask{}, &model.SystemTaskLock{}, + &model.UpstreamCostRecord{}, ); err != nil { panic("failed to migrate: " + err.Error()) } diff --git a/service/text_quota.go b/service/text_quota.go index 7da3391206d6..dbb002d47fbb 100644 --- a/service/text_quota.go +++ b/service/text_quota.go @@ -486,6 +486,7 @@ func PostTextConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, us InjectTieredBillingInfo(other, relayInfo, tieredResult) } + AttachUpstreamCostWithUsage(relayInfo, billingUsage, summary.Quota, other) attachQuotaSaturation(ctx, relayInfo, other) model.RecordConsumeLog(ctx, relayInfo.UserId, model.RecordConsumeLogParams{ diff --git a/service/upstream_cost.go b/service/upstream_cost.go new file mode 100644 index 000000000000..349efcc288ca --- /dev/null +++ b/service/upstream_cost.go @@ -0,0 +1,245 @@ +package service + +import ( + "encoding/json" + "math" + "strconv" + "strings" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/dto" + relaycommon "github.com/QuantumNous/new-api/relay/common" + + "github.com/shopspring/decimal" +) + +const upstreamCostMicrosPerCNY = int64(1_000_000) + +type upstreamCostProfile struct { + mode string + nativeUnit string + rate decimal.Decimal + rateFloat float64 + priceVersion string +} + +// AttachUpstreamCost records an estimated CNY acquisition cost from the +// gateway's base billing units. It is retained for billing paths that do not +// expose a normalized upstream usage object. +func AttachUpstreamCost(relayInfo *relaycommon.RelayInfo, quota int, other map[string]interface{}) { + attachUpstreamCost(relayInfo, nil, quota, other) +} + +// AttachUpstreamCostWithUsage prefers an authoritative cost returned by the +// upstream when the channel profile allows it, and otherwise falls back to the +// gateway's base billing units. The snapshot is admin-only and never changes +// the user's quota charge. +func AttachUpstreamCostWithUsage(relayInfo *relaycommon.RelayInfo, usage *dto.Usage, quota int, other map[string]interface{}) { + attachUpstreamCost(relayInfo, usage, quota, other) +} + +func attachUpstreamCost(relayInfo *relaycommon.RelayInfo, usage *dto.Usage, quota int, other map[string]interface{}) { + if other == nil || relayInfo == nil || relayInfo.ChannelMeta == nil { + return + } + profile, ok := resolveUpstreamCostProfile(relayInfo) + if !ok { + settings := relayInfo.ChannelOtherSettings + mode := strings.TrimSpace(settings.UpstreamCostMode) + if mode == "" { + mode = dto.UpstreamCostModeBillingUnits + } + nativeUnit := strings.ToUpper(strings.TrimSpace(settings.UpstreamCostUnit)) + if nativeUnit == "" { + nativeUnit = "UNIT" + } + priceVersion := strings.TrimSpace(settings.UpstreamCostPriceVersion) + if priceVersion == "" { + priceVersion = "manual" + } + attachUpstreamCostToAdminInfo(other, &dto.UpstreamCostSnapshot{ + Status: dto.UpstreamCostStatusUnpriced, + Mode: mode, + Reason: "missing_channel_cost_profile", + NativeUnit: nativeUnit, + PriceVersion: priceVersion, + SettlementCurrency: "CNY", + }) + return + } + + snapshot := &dto.UpstreamCostSnapshot{ + Status: dto.UpstreamCostStatusUnpriced, + Mode: profile.mode, + NativeUnit: profile.nativeUnit, + RateCNYPerUnit: profile.rateFloat, + RateCNYPerUnitDecimal: profile.rate.String(), + PriceVersion: profile.priceVersion, + SettlementCurrency: "CNY", + } + + nativeAmount, source, estimated, reason := resolveUpstreamNativeCost(profile.mode, usage, quota, other) + if reason != "" { + snapshot.Reason = reason + attachUpstreamCostToAdminInfo(other, snapshot) + return + } + + amountCNY := nativeAmount.Mul(profile.rate) + amountMicros := amountCNY.Mul(decimal.NewFromInt(upstreamCostMicrosPerCNY)).Round(0) + maxInt64 := decimal.NewFromInt(math.MaxInt64) + if amountMicros.IsNegative() || amountMicros.GreaterThan(maxInt64) { + snapshot.Reason = "amount_out_of_range" + attachUpstreamCostToAdminInfo(other, snapshot) + return + } + + nativeAmountFloat, _ := nativeAmount.Float64() + amountCNYFloat, _ := amountCNY.Float64() + if math.IsNaN(nativeAmountFloat) || + math.IsInf(nativeAmountFloat, 0) || + math.IsNaN(amountCNYFloat) || + math.IsInf(amountCNYFloat, 0) { + snapshot.Reason = "amount_out_of_range" + attachUpstreamCostToAdminInfo(other, snapshot) + return + } + + snapshot.Status = dto.UpstreamCostStatusSettled + snapshot.Source = source + snapshot.NativeAmount = nativeAmountFloat + snapshot.NativeAmountDecimal = nativeAmount.String() + snapshot.Units = nativeAmountFloat + snapshot.AmountCNY = amountCNYFloat + snapshot.AmountCNYMicros = amountMicros.IntPart() + snapshot.Estimated = estimated + attachUpstreamCostToAdminInfo(other, snapshot) +} + +func resolveUpstreamCostProfile(relayInfo *relaycommon.RelayInfo) (upstreamCostProfile, bool) { + if relayInfo == nil || relayInfo.ChannelMeta == nil { + return upstreamCostProfile{}, false + } + settings := relayInfo.ChannelOtherSettings + rate := settings.UpstreamCostRateCNY + if rate == nil || + *rate <= 0 || + math.IsNaN(*rate) || + math.IsInf(*rate, 0) || + *rate > dto.MaxUpstreamCostRateCNY { + return upstreamCostProfile{}, false + } + + mode := strings.TrimSpace(settings.UpstreamCostMode) + if mode == "" { + // Channels saved before multi-source cost profiles used the local billing + // unit conversion. Preserve that behavior on upgrade. + mode = dto.UpstreamCostModeBillingUnits + } + switch mode { + case dto.UpstreamCostModeAuto, dto.UpstreamCostModeResponseCost, dto.UpstreamCostModeBillingUnits: + default: + return upstreamCostProfile{}, false + } + + nativeUnit := strings.ToUpper(strings.TrimSpace(settings.UpstreamCostUnit)) + if nativeUnit == "" { + nativeUnit = "UNIT" + } + priceVersion := strings.TrimSpace(settings.UpstreamCostPriceVersion) + if priceVersion == "" { + priceVersion = "manual" + } + return upstreamCostProfile{ + mode: mode, + nativeUnit: nativeUnit, + rate: decimal.NewFromFloat(*rate), + rateFloat: *rate, + priceVersion: priceVersion, + }, true +} + +func resolveUpstreamNativeCost(mode string, usage *dto.Usage, quota int, other map[string]interface{}) (decimal.Decimal, string, bool, string) { + if mode == dto.UpstreamCostModeAuto || mode == dto.UpstreamCostModeResponseCost { + if amount, ok := upstreamResponseCost(usage); ok { + return amount, dto.UpstreamCostSourceResponseCost, false, "" + } + if mode == dto.UpstreamCostModeResponseCost { + return decimal.Zero, "", false, "missing_response_cost" + } + } + + if quota < 0 || common.QuotaPerUnit <= 0 { + return decimal.Zero, "", true, "missing_billing_units" + } + units := decimal.NewFromInt(int64(quota)).Div(decimal.NewFromFloat(common.QuotaPerUnit)) + if groupRatio, ok := decimalFromAny(other["group_ratio"]); ok && groupRatio.IsPositive() { + units = units.Div(groupRatio) + } + if units.IsNegative() { + return decimal.Zero, "", true, "invalid_billing_units" + } + return units, dto.UpstreamCostSourceBillingUnits, true, "" +} + +func upstreamResponseCost(usage *dto.Usage) (decimal.Decimal, bool) { + if usage == nil { + return decimal.Zero, false + } + amount, ok := decimalFromAny(usage.Cost) + if !ok || amount.IsNegative() { + return decimal.Zero, false + } + return amount, true +} + +func decimalFromAny(value any) (decimal.Decimal, bool) { + switch number := value.(type) { + case nil: + return decimal.Zero, false + case decimal.Decimal: + return number, true + case json.Number: + result, err := decimal.NewFromString(number.String()) + return result, err == nil + case string: + result, err := decimal.NewFromString(strings.TrimSpace(number)) + return result, err == nil + case float64: + if math.IsNaN(number) || math.IsInf(number, 0) { + return decimal.Zero, false + } + return decimal.NewFromFloat(number), true + case float32: + value64 := float64(number) + if math.IsNaN(value64) || math.IsInf(value64, 0) { + return decimal.Zero, false + } + return decimal.NewFromFloat32(number), true + case int: + return decimal.NewFromInt(int64(number)), true + case int64: + return decimal.NewFromInt(number), true + case int32: + return decimal.NewFromInt(int64(number)), true + case uint: + result, err := decimal.NewFromString(strconv.FormatUint(uint64(number), 10)) + return result, err == nil + case uint64: + result, err := decimal.NewFromString(strconv.FormatUint(number, 10)) + return result, err == nil + case uint32: + return decimal.NewFromInt(int64(number)), true + default: + return decimal.Zero, false + } +} + +func attachUpstreamCostToAdminInfo(other map[string]interface{}, snapshot *dto.UpstreamCostSnapshot) { + adminInfo, ok := other["admin_info"].(map[string]interface{}) + if !ok || adminInfo == nil { + adminInfo = make(map[string]interface{}) + other["admin_info"] = adminInfo + } + adminInfo["upstream_cost"] = snapshot +} diff --git a/service/upstream_cost_test.go b/service/upstream_cost_test.go new file mode 100644 index 000000000000..3dcea21efba3 --- /dev/null +++ b/service/upstream_cost_test.go @@ -0,0 +1,133 @@ +package service + +import ( + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/dto" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestAttachUpstreamCostConvertsBilledUnitsToCNY(t *testing.T) { + rate := 0.495 + info := &relaycommon.RelayInfo{ + ChannelMeta: &relaycommon.ChannelMeta{ + ChannelOtherSettings: dto.ChannelOtherSettings{ + UpstreamCostMode: dto.UpstreamCostModeBillingUnits, + UpstreamCostUnit: "CREDIT", + UpstreamCostRateCNY: &rate, + UpstreamCostPriceVersion: "yunwu-2026-07", + }, + }, + } + quota := int(0.011682 * common.QuotaPerUnit) + other := map[string]interface{}{ + "admin_info": map[string]interface{}{"use_channel": []int{36}}, + } + + AttachUpstreamCost(info, quota, other) + + adminInfo, ok := other["admin_info"].(map[string]interface{}) + require.True(t, ok) + assert.Equal(t, []int{36}, adminInfo["use_channel"]) + cost, ok := adminInfo["upstream_cost"].(*dto.UpstreamCostSnapshot) + require.True(t, ok) + assert.Equal(t, dto.UpstreamCostStatusSettled, cost.Status) + assert.Equal(t, dto.UpstreamCostSourceBillingUnits, cost.Source) + assert.Equal(t, "CREDIT", cost.NativeUnit) + assert.Equal(t, "yunwu-2026-07", cost.PriceVersion) + assert.True(t, cost.Estimated) + assert.InDelta(t, 0.011682, cost.Units, 1e-12) + assert.InDelta(t, 0.495, cost.RateCNYPerUnit, 1e-12) + assert.InDelta(t, 0.00578259, cost.AmountCNY, 1e-12) + assert.Equal(t, int64(5783), cost.AmountCNYMicros) +} + +func TestAttachUpstreamCostMarksMissingChannelProfileUnpriced(t *testing.T) { + other := map[string]interface{}{} + + AttachUpstreamCost(&relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{}}, 100, other) + + adminInfo, ok := other["admin_info"].(map[string]interface{}) + require.True(t, ok) + cost, ok := adminInfo["upstream_cost"].(*dto.UpstreamCostSnapshot) + require.True(t, ok) + assert.Equal(t, dto.UpstreamCostStatusUnpriced, cost.Status) + assert.Equal(t, "missing_channel_cost_profile", cost.Reason) +} + +func TestAttachUpstreamCostExcludesCustomerGroupRatio(t *testing.T) { + rate := 0.495 + info := &relaycommon.RelayInfo{ + ChannelMeta: &relaycommon.ChannelMeta{ + ChannelOtherSettings: dto.ChannelOtherSettings{ + UpstreamCostRateCNY: &rate, + }, + }, + } + other := map[string]interface{}{"group_ratio": 2.0} + quota := int(0.011682 * 2 * common.QuotaPerUnit) + + AttachUpstreamCost(info, quota, other) + + adminInfo, ok := other["admin_info"].(map[string]interface{}) + require.True(t, ok) + cost, ok := adminInfo["upstream_cost"].(*dto.UpstreamCostSnapshot) + require.True(t, ok) + assert.InDelta(t, 0.011682, cost.Units, 1e-12) + assert.InDelta(t, 0.00578259, cost.AmountCNY, 1e-12) +} + +func TestAttachUpstreamCostAutoPrefersAuthoritativeResponseCost(t *testing.T) { + rate := 7.2 + info := &relaycommon.RelayInfo{ + ChannelMeta: &relaycommon.ChannelMeta{ + ChannelOtherSettings: dto.ChannelOtherSettings{ + UpstreamCostMode: dto.UpstreamCostModeAuto, + UpstreamCostUnit: "USD", + UpstreamCostRateCNY: &rate, + }, + }, + } + usage := &dto.Usage{Cost: "0.0125"} + other := map[string]interface{}{"group_ratio": 9.0} + + AttachUpstreamCostWithUsage(info, usage, int(common.QuotaPerUnit), other) + + adminInfo, ok := other["admin_info"].(map[string]interface{}) + require.True(t, ok) + cost, ok := adminInfo["upstream_cost"].(*dto.UpstreamCostSnapshot) + require.True(t, ok) + assert.Equal(t, dto.UpstreamCostStatusSettled, cost.Status) + assert.Equal(t, dto.UpstreamCostSourceResponseCost, cost.Source) + assert.False(t, cost.Estimated) + assert.InDelta(t, 0.0125, cost.NativeAmount, 1e-12) + assert.InDelta(t, 0.09, cost.AmountCNY, 1e-12) + assert.Equal(t, int64(90_000), cost.AmountCNYMicros) +} + +func TestAttachUpstreamCostResponseModeMarksMissingCostUnpriced(t *testing.T) { + rate := 7.2 + info := &relaycommon.RelayInfo{ + ChannelMeta: &relaycommon.ChannelMeta{ + ChannelOtherSettings: dto.ChannelOtherSettings{ + UpstreamCostMode: dto.UpstreamCostModeResponseCost, + UpstreamCostUnit: "USD", + UpstreamCostRateCNY: &rate, + }, + }, + } + other := map[string]interface{}{} + + AttachUpstreamCostWithUsage(info, &dto.Usage{}, 100, other) + + adminInfo, ok := other["admin_info"].(map[string]interface{}) + require.True(t, ok) + cost, ok := adminInfo["upstream_cost"].(*dto.UpstreamCostSnapshot) + require.True(t, ok) + assert.Equal(t, dto.UpstreamCostStatusUnpriced, cost.Status) + assert.Equal(t, "missing_response_cost", cost.Reason) + assert.Zero(t, cost.AmountCNYMicros) +} diff --git a/setting/console_setting/config.go b/setting/console_setting/config.go index 144e95c497be..8dfc043de0fe 100644 --- a/setting/console_setting/config.go +++ b/setting/console_setting/config.go @@ -2,6 +2,8 @@ package console_setting import "github.com/QuantumNous/new-api/setting/config" +const defaultSafetyAnnouncement = `[{"id":1,"content":"## 内容安全与合规提示\n\n请依法合规使用本站 AI 服务。“宽审核”“低审核”“Global”等仅说明上游审核特征，不代表允许违法违规内容，也不保证内容合法、安全或可商用。禁止生成侵权、诈骗、暴力色情、未成年人伤害或侵犯隐私的内容。发布或商用前请人工复核并确认权利；依法需标识 AI 生成合成内容时，不得删除或篡改标识。违规可能导致密钥或账户受限。","publishDate":"2026-08-12T20:28:00+08:00","type":"warning","extra":"本提示仅用于风险告知，不构成法律意见。具体要求以用户协议、隐私政策及适用法律为准。"}]` + type ConsoleSetting struct { ApiInfo string `json:"api_info"` // 控制台 API 信息 (JSON 数组字符串) UptimeKumaGroups string `json:"uptime_kuma_groups"` // Uptime Kuma 分组配置 (JSON 数组字符串) @@ -17,7 +19,7 @@ type ConsoleSetting struct { var defaultConsoleSetting = ConsoleSetting{ ApiInfo: "", UptimeKumaGroups: "", - Announcements: "", + Announcements: defaultSafetyAnnouncement, FAQ: "", ApiInfoEnabled: true, UptimeKumaEnabled: true, diff --git a/setting/console_setting/config_test.go b/setting/console_setting/config_test.go new file mode 100644 index 000000000000..5ffb297c7e9b --- /dev/null +++ b/setting/console_setting/config_test.go @@ -0,0 +1,19 @@ +package console_setting + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestDefaultSafetyAnnouncementIsValid(t *testing.T) { + require.NoError(t, ValidateConsoleSettings(defaultSafetyAnnouncement, "Announcements")) + + announcements := GetAnnouncements() + require.Len(t, announcements, 1) + content, ok := announcements[0]["content"].(string) + require.True(t, ok) + assert.Contains(t, content, "“宽审核”“低审核”“Global”") + assert.Contains(t, content, "AI 生成合成内容") +} diff --git a/setting/model_setting/image_generation.go b/setting/model_setting/image_generation.go new file mode 100644 index 000000000000..03dbfdc2180c --- /dev/null +++ b/setting/model_setting/image_generation.go @@ -0,0 +1,125 @@ +package model_setting + +import "strings" + +const ( + ImageResolutionParameterQuality = "quality" + ImageResolutionParameterSize = "size" +) + +type ImageGenerationCapabilities struct { + Resolutions []string `json:"resolutions"` + ResolutionParameter string `json:"resolution_parameter"` + Sizes []string `json:"sizes"` + DefaultResolution string `json:"default_resolution"` + DefaultSize string `json:"default_size"` + ResolutionPriceMultiplier map[string]float64 `json:"resolution_price_multipliers"` +} + +var imageGenerationCapabilities = map[string]ImageGenerationCapabilities{ + "gemini-3.1-flash-image-preview": { + Resolutions: []string{"1K", "2K", "4K"}, + ResolutionParameter: ImageResolutionParameterQuality, + Sizes: []string{"1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9", "21:9"}, + DefaultResolution: "1K", + DefaultSize: "1:1", + ResolutionPriceMultiplier: map[string]float64{ + "1K": 1, + "2K": 1, + "4K": 2, + }, + }, + "gemini-3-pro-image-preview": { + Resolutions: []string{"1K", "2K", "4K"}, + ResolutionParameter: ImageResolutionParameterQuality, + Sizes: []string{"1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9", "21:9"}, + DefaultResolution: "1K", + DefaultSize: "1:1", + ResolutionPriceMultiplier: map[string]float64{ + "1K": 1, + "2K": 1, + "4K": 1, + }, + }, + "nano-banana-2": { + Resolutions: []string{"1K", "2K", "4K"}, + ResolutionParameter: ImageResolutionParameterQuality, + Sizes: []string{"1:1", "16:9", "9:16", "4:3", "3:4"}, + DefaultResolution: "1K", + DefaultSize: "1:1", + ResolutionPriceMultiplier: map[string]float64{ + "1K": 1, + "2K": 1, + "4K": 1, + }, + }, + "nano-banana-pro": { + Resolutions: []string{"1K", "2K", "4K"}, + ResolutionParameter: ImageResolutionParameterQuality, + Sizes: []string{"1:1", "16:9", "9:16", "4:3", "3:4"}, + DefaultResolution: "1K", + DefaultSize: "1:1", + ResolutionPriceMultiplier: map[string]float64{ + "1K": 1, + "2K": 1, + "4K": 1, + }, + }, + "gpt-image-2-vip": { + Resolutions: []string{"1K", "2K", "4K"}, + ResolutionParameter: ImageResolutionParameterSize, + Sizes: []string{ + "1280x1280", "848x1280", "1280x848", "960x1280", "1280x960", "1024x1280", "1280x1024", "720x1280", "1280x720", "1280x544", + "2048x2048", "1360x2048", "2048x1360", "1536x2048", "2048x1536", "1632x2048", "2048x1632", "1152x2048", "2048x1152", "2048x864", + "2880x2880", "2336x3520", "3520x2336", "2480x3312", "3312x2480", "2560x3216", "3216x2560", "2160x3840", "3840x2160", "3840x1632", + }, + DefaultResolution: "2K", + DefaultSize: "2048x2048", + ResolutionPriceMultiplier: map[string]float64{ + "1K": 1, + "2K": 1, + "4K": 2, + }, + }, +} + +func GetImageGenerationCapabilities(model string) *ImageGenerationCapabilities { + capabilities, ok := imageGenerationCapabilities[strings.ToLower(strings.TrimSpace(model))] + if !ok { + return nil + } + capabilities.Resolutions = append([]string(nil), capabilities.Resolutions...) + capabilities.Sizes = append([]string(nil), capabilities.Sizes...) + capabilities.ResolutionPriceMultiplier = make(map[string]float64, len(capabilities.ResolutionPriceMultiplier)) + for resolution, multiplier := range imageGenerationCapabilities[strings.ToLower(strings.TrimSpace(model))].ResolutionPriceMultiplier { + capabilities.ResolutionPriceMultiplier[resolution] = multiplier + } + return &capabilities +} + +func GetImageGenerationPriceMultiplier(model, quality, size string) float64 { + capabilities := GetImageGenerationCapabilities(model) + if capabilities == nil { + return 1 + } + + resolution := strings.ToUpper(strings.TrimSpace(quality)) + if capabilities.ResolutionParameter == ImageResolutionParameterSize { + resolution = capabilities.DefaultResolution + for index, configuredSize := range capabilities.Sizes { + if strings.EqualFold(configuredSize, strings.TrimSpace(size)) { + resolutionIndex := index * len(capabilities.Resolutions) / len(capabilities.Sizes) + resolution = capabilities.Resolutions[resolutionIndex] + break + } + } + } else if resolution == "" { + resolution = capabilities.DefaultResolution + } + + multiplier, ok := capabilities.ResolutionPriceMultiplier[resolution] + if !ok || multiplier <= 0 { + return 1 + } + return multiplier +} diff --git a/setting/model_setting/image_generation_test.go b/setting/model_setting/image_generation_test.go new file mode 100644 index 000000000000..69a99fe3a918 --- /dev/null +++ b/setting/model_setting/image_generation_test.go @@ -0,0 +1,75 @@ +package model_setting + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestImageGenerationCapabilitiesMatchConfiguredUpstreams(t *testing.T) { + tests := []struct { + model string + resolutionParameter string + defaultResolution string + defaultSize string + sizeCount int + }{ + {model: "gemini-3.1-flash-image-preview", resolutionParameter: ImageResolutionParameterQuality, defaultResolution: "1K", defaultSize: "1:1", sizeCount: 10}, + {model: "gemini-3-pro-image-preview", resolutionParameter: ImageResolutionParameterQuality, defaultResolution: "1K", defaultSize: "1:1", sizeCount: 10}, + {model: "nano-banana-2", resolutionParameter: ImageResolutionParameterQuality, defaultResolution: "1K", defaultSize: "1:1", sizeCount: 5}, + {model: "nano-banana-pro", resolutionParameter: ImageResolutionParameterQuality, defaultResolution: "1K", defaultSize: "1:1", sizeCount: 5}, + {model: "gpt-image-2-vip", resolutionParameter: ImageResolutionParameterSize, defaultResolution: "2K", defaultSize: "2048x2048", sizeCount: 30}, + } + + for _, test := range tests { + t.Run(test.model, func(t *testing.T) { + capabilities := GetImageGenerationCapabilities(test.model) + require.NotNil(t, capabilities) + assert.Equal(t, []string{"1K", "2K", "4K"}, capabilities.Resolutions) + assert.Equal(t, test.resolutionParameter, capabilities.ResolutionParameter) + assert.Equal(t, test.defaultResolution, capabilities.DefaultResolution) + assert.Equal(t, test.defaultSize, capabilities.DefaultSize) + assert.Len(t, capabilities.Sizes, test.sizeCount) + }) + } +} + +func TestImageGenerationCapabilitiesReturnsIndependentSlices(t *testing.T) { + first := GetImageGenerationCapabilities("nano-banana-2") + require.NotNil(t, first) + first.Resolutions[0] = "changed" + first.Sizes[0] = "changed" + first.ResolutionPriceMultiplier["4K"] = 99 + + second := GetImageGenerationCapabilities("NANO-BANANA-2") + require.NotNil(t, second) + assert.Equal(t, "1K", second.Resolutions[0]) + assert.Equal(t, "1:1", second.Sizes[0]) + assert.Equal(t, float64(1), second.ResolutionPriceMultiplier["4K"]) + assert.Nil(t, GetImageGenerationCapabilities("chat-model")) +} + +func TestImageGenerationPriceMultiplierMatchesConfiguredResolution(t *testing.T) { + tests := []struct { + name string + model string + quality string + size string + want float64 + }{ + {name: "flash 1K", model: "gemini-3.1-flash-image-preview", quality: "1K", want: 1}, + {name: "flash 4K", model: "gemini-3.1-flash-image-preview", quality: "4k", want: 2}, + {name: "pro 4K", model: "gemini-3-pro-image-preview", quality: "4K", want: 1}, + {name: "gpt 2K", model: "gpt-image-2-vip", size: "2048x2048", want: 1}, + {name: "gpt 4K", model: "gpt-image-2-vip", size: "3840x2160", want: 2}, + {name: "nano 4K", model: "nano-banana-pro", quality: "4K", want: 1}, + {name: "unknown model", model: "unknown", quality: "4K", want: 1}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + assert.Equal(t, test.want, GetImageGenerationPriceMultiplier(test.model, test.quality, test.size)) + }) + } +} diff --git a/setting/operation_setting/general_setting.go b/setting/operation_setting/general_setting.go index b4a3ccccdaf3..715bdfdc359b 100644 --- a/setting/operation_setting/general_setting.go +++ b/setting/operation_setting/general_setting.go @@ -24,7 +24,7 @@ type GeneralSetting struct { // 默认配置 var generalSetting = GeneralSetting{ - DocsLink: "https://docs.newapi.pro", + DocsLink: "/docs", PingIntervalEnabled: false, PingIntervalSeconds: 60, QuotaDisplayType: QuotaDisplayTypeUSD, diff --git a/setting/ratio_setting/model_ratio.go b/setting/ratio_setting/model_ratio.go index 829e0794a157..d150e90550e2 100644 --- a/setting/ratio_setting/model_ratio.go +++ b/setting/ratio_setting/model_ratio.go @@ -515,7 +515,7 @@ func getHardcodedCompletionModelRatio(name string) (float64, bool) { if strings.HasPrefix(name, "gpt-5.4-nano") { return 6.25, true } - return 6, true + return 6, false } // gpt-5.5 and later models are unlocked return 6, false diff --git a/setting/ratio_setting/model_ratio_test.go b/setting/ratio_setting/model_ratio_test.go new file mode 100644 index 000000000000..3f039bde005f --- /dev/null +++ b/setting/ratio_setting/model_ratio_test.go @@ -0,0 +1,12 @@ +package ratio_setting + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestGPT54CompletionRatioCanBeConfigured(t *testing.T) { + info := GetCompletionRatioInfo("gpt-5.4") + assert.False(t, info.Locked) +} diff --git a/setting/system_setting/legal.go b/setting/system_setting/legal.go index cc84d4085cc2..333d5414e52d 100644 --- a/setting/system_setting/legal.go +++ b/setting/system_setting/legal.go @@ -2,14 +2,197 @@ package system_setting import "github.com/QuantumNous/new-api/setting/config" +// These defaults keep the public legal pages usable on fresh deployments. +// Administrators can replace or clear them through the system settings. +const defaultUserAgreement = `> 更新日期：2026年8月12日 +> 生效日期：2026年8月12日 + +## 特别提示 + +本协议适用于您对本平台网站、控制台、API、异步任务及相关技术服务的访问和使用。平台运营主体、联系方式、实际服务范围、价格与专项规则，以网站届时公示的信息为准。 + +请您在注册、充值或调用 API 前完整阅读本协议，尤其是内容安全、计费退款、第三方服务、责任限制及争议解决条款。您完成注册、勾选同意、充值、调用 API 或继续使用服务，即表示您已理解并同意本协议。若您代表单位使用服务，您确认已获得该单位的有效授权。 + +## 1. 用户资格与账户 + +1. 您应具有与使用行为相适应的民事行为能力。未成年人应在监护人阅读并同意本协议后使用服务；依法需要取得监护人单独同意的，应事先取得该同意。 +2. 您应提供真实、准确、完整且有效的注册资料，并及时更新。平台可基于账户安全、支付风控、监管要求或特定服务需要，要求进行必要的身份或主体核验。 +3. 账户、密码、API 密钥、验证码和访问令牌仅限您本人或经您授权的人员使用。请勿将密钥写入公开代码、前端页面或不受信任的环境。通过您账户或密钥发起的操作，原则上视为您的操作。 +4. 如发现账户或密钥被盗用、泄露或存在异常调用，请立即撤销相关密钥、修改凭据并通过工单联系平台。因您未妥善保管凭据造成的损失，由责任方依法承担。 + +## 2. 服务性质与第三方上游 + +1. 本平台主要提供 AI 接口聚合、请求路由、任务管理、额度计费和相关技术工具。具体模型能力由相应上游服务商提供，平台并非所有模型或生成能力的开发者。 +2. 模型名称、可用区域、并发、速度、审核规则、输出格式及价格可能因上游策略、汇率、网络或维护而变化。平台会尽合理努力通过页面或公告提示重要变化，但无法保证所有服务永久可用或始终保持相同性能。 +3. 您还应遵守所使用上游服务商的适用条款、使用政策和权利限制。上游服务商可能独立处理请求数据，并可能对请求进行审核、拒绝、限流或终止。 +4. 页面、模型名称或公告中出现的“宽审核”“低审核”“Global”等描述，仅用于说明特定上游或路由的审核特征，不代表允许生成违法违规内容，也不构成对内容合法、安全、准确、可商用或必然通过审核的承诺。 + +## 3. 费用、额度与退款 + +1. 服务价格、计费单位、倍率、币种、有效期及优惠以购买页面、模型价格页和实际消费记录为准。您应在调用前核对模型、参数和预计费用。 +2. 部分请求可能先预扣额度，再依据实际用量或上游结算结果进行最终扣费、补扣或退回。异步任务的最终状态、计费结果及失败处理以任务记录和平台规则为准。 +3. 充值金额或额度属于预付服务资源。退款条件、方式、到账时间和手续费以付款页面、专项规则、支付渠道规则及适用法律为准。法律规定应退款，或平台未按约提供服务且依法应退款的，平台将按适用规则处理；本协议不排除用户依法享有的强制性权利。 +4. 因用户主动取消、参数错误、内容违规、上游拒绝或已实际产生上游成本而发生的费用，按页面公示的结算规则处理。您对账单有异议的，应及时通过工单提供请求编号、时间和相关凭证。 +5. 禁止使用来源不明或违法所得资金充值，禁止盗刷、洗钱、恶意拒付、套现或利用计费漏洞牟利。平台可对异常交易采取延迟入账、限制使用或配合调查等必要措施。 + +## 4. 合法合规与可接受使用 + +您不得利用本服务实施、协助或促成下列行为： + +- 违反适用法律法规、监管要求、公共秩序或善良风俗； +- 制作、上传或传播危害国家安全、暴力恐怖、违法色情、涉及未成年人性剥削、仇恨歧视、自残诱导、诈骗、赌博、洗钱或其他有害内容； +- 侵犯他人的著作权、商标权、肖像权、名誉权、隐私权、个人信息权益、商业秘密或其他合法权益； +- 未经合法授权收集、识别、推断、处理或披露个人信息、生物识别信息、医疗信息、金融信息或其他敏感数据； +- 生成或传播虚假信息、冒充他人、伪造身份或来源，或以足以误导公众的方式使用合成内容； +- 绕过安全审核、速率限制、身份验证、计费机制或其他技术保护措施，攻击、探测、干扰平台或上游系统； +- 未经允许转售账户、密钥、额度、接口能力，或为违法业务提供代理、分发和批量调用服务。 + +自动化审核可能出现误判或漏判。无论平台或上游是否审核、是否放行，均不减轻您对输入、输出及实际使用行为承担的责任。依法需要对 AI 生成合成内容添加显式或隐式标识、进行用户提示或主动声明的，您应正确标识并保留来源信息，不得删除、篡改、伪造或隐匿依法要求的标识。 + +平台发现或合理怀疑违规、滥用、安全风险或异常消费时，可视情况拒绝请求、限制模型、暂停密钥、冻结账户、保全必要记录、要求说明或终止服务；依法需要时，平台将配合有权机关处理。您可通过工单提交申诉和证明材料。 + +## 5. 输入、输出与知识产权 + +1. 您应确保对提交的文本、图片、音频、视频、代码、文件及其他内容拥有必要权利或合法授权。为完成请求，您授权平台在必要范围内接收、缓存、转换、传输和向所选上游提供这些内容。 +2. 在法律允许的范围内，您保留对自有输入内容的权利。平台不会因提供接口服务当然取得您内容的所有权；但第三方上游对输入和输出的权利安排可能另有规定，您应自行核对。 +3. AI 输出可能不准确、不完整、具有偏见、与他人内容相似，或包含事实、版权、商标、肖像、隐私及其他风险。平台不保证输出具有独创性、可登记性或可商用性。您在发布、商用或用于重要决策前，应进行人工复核、事实核验和权利审查。 +4. 输出不构成医疗、法律、金融、投资或其他专业意见。高风险场景应由具备相应资质的专业人员作出最终判断。 + +## 6. 隐私与数据安全 + +1. 平台对个人信息的处理遵循另行公示的《隐私政策》，该政策构成本协议的组成部分。 +2. 为提供路由、计费、风控、故障排查和安全审核，平台可能处理账户信息、请求元数据、用量记录、设备与网络信息以及必要的输入输出内容。请求内容还可能被传输至您选择或平台配置的上游服务商。 +3. 请勿提交与请求目的无关的个人信息、国家秘密、商业秘密、未公开源代码、访问凭据或其他高敏感内容。您代表第三方提交个人信息时，应确保具备合法处理依据并履行必要的告知、同意和安全义务。 + +## 7. 服务变更、中断与终止 + +平台可因维护升级、上游故障、网络攻击、不可抗力、法律政策变化或商业调整变更或中断部分服务，并会在合理可行范围内进行通知。平台不保证服务绝对连续、无错误或满足所有特定需求。 + +您可依平台提供的功能申请删除账户。账户删除后将无法继续使用服务；依法应保留的交易、安全、审计或争议处理记录，可能在法定期限内继续保存并限制处理。未使用额度、未完成任务和退款事项按适用规则及法律处理。 + +## 8. 责任限制 + +在法律允许的范围内，平台不对因上游服务变化、网络波动、用户配置错误、第三方侵权、AI 输出错误或用户违法使用造成的间接损失、预期利益损失承担责任。任何责任限制均不适用于依法不得限制或免除的责任，包括因故意或重大过失造成的人身损害等法定情形。 + +## 9. 协议更新与通知 + +平台可根据服务、风险或法律变化更新本协议。重大变更将通过站内公告、页面提示或其他合理方式通知。更新后的协议自公示日期或另行注明的日期生效；若您不同意，应停止使用服务并按规则处理账户。继续使用服务视为接受更新后的协议，但法律要求另行取得同意的除外。 + +## 10. 法律适用、争议与联系 + +本协议的法律适用和争议管辖，以平台运营主体所在地及适用的强制性法律规则为准。发生争议时，双方应先通过工单或平台公示的联系方式友好协商；协商不成的，可向依法具有管辖权的法院或争议解决机构寻求救济。 + +如您对本协议、计费、内容安全或账户处置有疑问，请通过平台工单系统联系我们。` + +const defaultPrivacyPolicy = `> 更新日期：2026年8月12日 +> 生效日期：2026年8月12日 + +## 重要说明 + +本政策说明平台在提供网站、控制台、API、异步任务、计费和技术支持服务时，如何处理与已识别或可识别自然人有关的信息。平台运营主体和联系方式以网站公示信息为准。 + +请您在使用服务前阅读本政策。若您代表单位或通过自己的产品向最终用户提供服务，您还应依法制定面向最终用户的隐私规则，并确保向平台或上游传输个人信息具有合法、正当、必要的处理依据。 + +## 1. 我们处理的信息 + +根据您使用的功能，我们可能处理以下信息： + +1. **账户与身份信息**：用户名、邮箱、联系方式、头像、第三方登录标识、邀请关系，以及您主动提交的身份或单位核验资料。 +2. **认证与安全信息**：密码的安全处理结果、API 密钥及其权限配置、双重验证或通行密钥信息、登录会话、IP 地址、浏览器和设备信息、登录及安全事件记录。 +3. **API 与使用记录**：所用模型、请求时间、令牌或任务编号、用量、延迟、状态码、错误信息、渠道与路由元数据、额度变化和账单明细。 +4. **请求与生成内容**：您提交的提示词、消息、图片、音频、视频、文件、任务参数及生成结果。为完成请求，这些内容可能被临时缓存、格式转换并传输至相应上游；在开启调试、内容安全、错误诊断或任务记录功能时，部分内容可能进入必要日志。 +5. **交易信息**：充值金额、订单号、支付渠道、支付状态、退款或拒付信息。完整银行卡号等支付凭据通常由支付机构处理，平台是否接触相关信息取决于实际接入方式。 +6. **支持与沟通信息**：工单、申诉、投诉、问卷、邮件以及您向客服提供的附件和沟通记录。 +7. **本地与交互数据**：Cookie、本地存储、语言、主题、界面偏好和必要的访问统计信息。 + +请不要在请求中提交与处理目的无关的敏感个人信息、密码、私钥、验证码、国家秘密或其他高敏感内容。 + +## 2. 处理目的 + +我们可能为下列目的处理信息： + +- 创建和管理账户、验证身份、保持登录状态及提供客户支持； +- 验证 API 请求、选择路由、调用上游、返回结果及管理异步任务； +- 计算用量和费用、处理充值退款、生成账单并防止欺诈； +- 监测可用性、排查故障、改进性能和保障账户、平台及上游安全； +- 识别和处置违法违规内容、滥用、攻击、恶意自动化或违反服务协议的行为； +- 履行适用法律规定的备案、审计、内容标识、日志留存、投诉处理和协助监管等义务； +- 在取得有效同意或具备其他合法依据时，提供您选择的其他功能。 + +我们将尽量按照目的明确、直接相关和最小必要的原则处理个人信息，不会仅因技术上可以获取而无限制收集。 + +## 3. 委托处理、共享与上游传输 + +为提供服务，信息可能在必要范围内提供给： + +1. **AI 上游服务商**：接收请求内容、参数和必要标识，用于执行推理、生成、审核和返回结果。不同上游的数据处理地点、保存规则和训练政策可能不同，请同时查阅相应上游政策。 +2. **基础设施服务商**：提供云计算、网络、存储、日志、邮件、短信、验证码、监控和安全防护。 +3. **支付与交易服务商**：处理收款、退款、对账、风控和发票事项。 +4. **登录和集成服务商**：在您选择 OAuth、社交登录、通知或其他第三方集成功能时处理必要信息。 +5. **有权机关及争议相关方**：在适用法律要求、保护用户或公众安全、调查欺诈侵权，或建立、行使、抗辩法律权利所必需时提供必要信息。 + +我们不会以出售个人信息为目的对外提供您的个人信息。若未来处理目的、方式或接收方发生实质变化，我们将依法另行告知，并在需要时取得您的同意。 + +## 4. 跨境处理 + +本平台可接入位于不同国家或地区的上游和基础设施服务。您的请求内容及相关元数据可能因所选模型、渠道或部署配置被传输至境外。平台运营者和使用本服务的业务方应根据实际数据流向履行必要的告知、单独同意、影响评估、合同、认证或安全评估义务。 + +如您不能接受特定上游的数据处理地点，请不要使用对应模型或渠道，并联系平台确认是否存在可选路由。本政策不对任何数据的固定存储地域作出未明确公示的保证。 + +## 5. 保存期限 + +我们仅在实现本政策所述目的所需期限内保存信息，但法律法规、财税审计、安全风控、争议解决或上游规则要求更长时间的除外。具体期限会因信息类型和部署配置不同而变化： + +- 账户资料通常保存至账户删除或服务关系终止后完成必要清理； +- 用量、订单、支付、额度和发票记录按财税、交易及争议处理需要保存； +- 登录、安全、操作和错误日志按安全保护、审计及法定要求保存； +- 请求内容、上传文件和生成结果的保存时间取决于任务类型、日志配置、缓存策略和上游规则，结果链接还可能具有较短有效期。 + +保存期限届满后，我们将根据适用要求删除或匿名化处理相关信息。受备份、灾难恢复或法律保全影响，删除可能存在合理延迟；延迟期间将限制不必要的处理。 + +## 6. 您的权利 + +在适用法律规定的范围内，您可以： + +- 查阅、复制、更正或补充您的个人信息； +- 撤回基于同意的处理，或关闭相应可选功能； +- 请求删除个人信息或注销账户； +- 请求解释本政策和个人信息处理规则； +- 对自动化决策、账户限制或投诉处理结果提出异议； +- 向有权监管机构投诉或依法寻求其他救济。 + +您可优先通过个人资料、账户安全和账户删除功能行使相关权利，也可通过工单提交请求。为保护账户和他人权益，我们可能先核验您的身份。法律要求继续保存的信息，在保存期内可能无法立即删除，但我们会限制与保存目的无关的处理。 + +## 7. Cookie 与本地存储 + +平台使用 Cookie 或浏览器本地存储维持会话、保存语言和主题偏好、记录必要安全状态并改善使用体验。禁用必要 Cookie 或清除本地数据可能导致登录、身份验证或部分功能无法正常使用。第三方登录、支付或嵌入内容可能使用其自身的 Cookie，受相应第三方政策约束。 + +## 8. 信息安全 + +我们将根据实际风险采取访问控制、权限隔离、加密传输、日志审计、备份和安全监测等合理措施。但互联网传输和第三方系统不存在绝对安全，您也应妥善保护账户、API 密钥和终端环境。 + +发生可能影响您权益的个人信息安全事件时，我们将依法采取补救措施，并在适用法律要求时向有关部门报告或向受影响用户告知事件情况、可能影响及建议措施。 + +## 9. 未成年人 + +本服务主要面向具备相应民事行为能力的用户。未成年人应在监护人指导下使用。涉及依法需要监护人同意的未成年人个人信息时，应事先取得监护人同意并采取专门保护措施。监护人可通过工单联系我们，申请查阅、更正或删除相关信息。 + +## 10. 政策更新 + +我们可能根据服务功能、数据流向或法律要求更新本政策。重大变更将通过站内公告、页面提示或其他合理方式告知；法律要求重新取得同意的，我们将依法处理。 + +## 11. 联系我们 + +如您对本政策、个人信息处理或权利请求有疑问，请通过平台工单系统联系我们。我们会在核验请求后，依据适用法律和平台流程进行处理。` + type LegalSettings struct { UserAgreement string `json:"user_agreement"` PrivacyPolicy string `json:"privacy_policy"` } var defaultLegalSettings = LegalSettings{ - UserAgreement: "", - PrivacyPolicy: "", + UserAgreement: defaultUserAgreement, + PrivacyPolicy: defaultPrivacyPolicy, } func init() { diff --git a/setting/system_setting/legal_test.go b/setting/system_setting/legal_test.go new file mode 100644 index 000000000000..36521f182e0b --- /dev/null +++ b/setting/system_setting/legal_test.go @@ -0,0 +1,19 @@ +package system_setting + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestDefaultLegalSettingsProvideUserFacingDocuments(t *testing.T) { + settings := GetLegalSettings() + + require.NotEmpty(t, settings.UserAgreement) + require.NotEmpty(t, settings.PrivacyPolicy) + assert.Contains(t, settings.UserAgreement, "“宽审核”“低审核”“Global”") + assert.Contains(t, settings.UserAgreement, "AI 生成合成内容") + assert.Contains(t, settings.PrivacyPolicy, "请求与生成内容") + assert.Contains(t, settings.PrivacyPolicy, "AI 上游服务商") +} diff --git a/storage/artifact_store.go b/storage/artifact_store.go new file mode 100644 index 000000000000..96f3cee66860 --- /dev/null +++ b/storage/artifact_store.go @@ -0,0 +1,137 @@ +package storage + +import ( + "context" + "errors" + "io" + "net/url" + "os" + "strconv" + "strings" + "time" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/config" + "github.com/aws/aws-sdk-go-v2/credentials" + "github.com/aws/aws-sdk-go-v2/service/s3" +) + +type ArtifactStore interface { + Put(ctx context.Context, key string, body io.Reader, contentType string) error + SignedURL(ctx context.Context, key string, ttl time.Duration) (string, error) + Delete(ctx context.Context, key string) error +} + +type S3ArtifactStore struct { + bucket string + client *s3.Client + presigner *s3.PresignClient +} + +func NewS3ArtifactStore(ctx context.Context) (*S3ArtifactStore, error) { + endpoint := strings.TrimSpace(os.Getenv("S3_ENDPOINT")) + publicEndpoint := strings.TrimSpace(os.Getenv("S3_PUBLIC_ENDPOINT")) + region := strings.TrimSpace(os.Getenv("S3_REGION")) + bucket := strings.TrimSpace(os.Getenv("S3_BUCKET")) + accessKey := os.Getenv("S3_ACCESS_KEY_ID") + secretKey := os.Getenv("S3_SECRET_ACCESS_KEY") + if region == "" { + region = "us-east-1" + } + if bucket == "" || accessKey == "" || secretKey == "" { + return nil, errors.New("S3_BUCKET, S3_ACCESS_KEY_ID and S3_SECRET_ACCESS_KEY are required") + } + if endpoint != "" { + if err := validateEndpoint(endpoint); err != nil { + return nil, err + } + } + if publicEndpoint != "" { + if err := validateEndpoint(publicEndpoint); err != nil { + return nil, err + } + } else { + publicEndpoint = endpoint + } + + awsConfig, err := config.LoadDefaultConfig(ctx, + config.WithRegion(region), + config.WithCredentialsProvider(credentials.NewStaticCredentialsProvider(accessKey, secretKey, "")), + ) + if err != nil { + return nil, err + } + pathStyle, _ := strconv.ParseBool(os.Getenv("S3_USE_PATH_STYLE")) + client := s3.NewFromConfig(awsConfig, func(options *s3.Options) { + options.UsePathStyle = pathStyle + if endpoint != "" { + options.BaseEndpoint = aws.String(strings.TrimRight(endpoint, "/")) + } + }) + presignClient := client + if publicEndpoint != "" && publicEndpoint != endpoint { + presignClient = s3.NewFromConfig(awsConfig, func(options *s3.Options) { + options.UsePathStyle = pathStyle + options.BaseEndpoint = aws.String(strings.TrimRight(publicEndpoint, "/")) + }) + } + return &S3ArtifactStore{ + bucket: bucket, + client: client, + presigner: s3.NewPresignClient(presignClient), + }, nil +} + +func validateEndpoint(raw string) error { + parsed, err := url.Parse(raw) + if err != nil || parsed.Hostname() == "" || (parsed.Scheme != "http" && parsed.Scheme != "https") { + return errors.New("S3 endpoint must be an absolute http or https URL") + } + if parsed.User != nil || parsed.RawQuery != "" || parsed.Fragment != "" { + return errors.New("S3 endpoint must not contain credentials, query parameters or fragments") + } + return nil +} + +func (s *S3ArtifactStore) Put(ctx context.Context, key string, body io.Reader, contentType string) error { + if s == nil || s.client == nil { + return errors.New("artifact store is not initialized") + } + _, err := s.client.PutObject(ctx, &s3.PutObjectInput{ + Bucket: aws.String(s.bucket), + Key: aws.String(key), + Body: body, + ContentType: aws.String(contentType), + }) + return err +} + +func (s *S3ArtifactStore) SignedURL(ctx context.Context, key string, ttl time.Duration) (string, error) { + if s == nil || s.presigner == nil { + return "", errors.New("artifact store is not initialized") + } + if ttl <= 0 { + ttl = 15 * time.Minute + } + result, err := s.presigner.PresignGetObject(ctx, &s3.GetObjectInput{ + Bucket: aws.String(s.bucket), + Key: aws.String(key), + }, func(options *s3.PresignOptions) { + options.Expires = ttl + }) + if err != nil { + return "", err + } + return result.URL, nil +} + +func (s *S3ArtifactStore) Delete(ctx context.Context, key string) error { + if s == nil || s.client == nil { + return errors.New("artifact store is not initialized") + } + _, err := s.client.DeleteObject(ctx, &s3.DeleteObjectInput{ + Bucket: aws.String(s.bucket), + Key: aws.String(key), + }) + return err +} diff --git a/storage/artifact_store_test.go b/storage/artifact_store_test.go new file mode 100644 index 000000000000..0daf002ca5f7 --- /dev/null +++ b/storage/artifact_store_test.go @@ -0,0 +1,52 @@ +package storage + +import ( + "bytes" + "context" + "io" + "net/http" + "os" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestNewS3ArtifactStoreRequiresCredentials(t *testing.T) { + t.Setenv("S3_BUCKET", "") + t.Setenv("S3_ACCESS_KEY_ID", "") + t.Setenv("S3_SECRET_ACCESS_KEY", "") + _, err := NewS3ArtifactStore(context.Background()) + require.Error(t, err) + assert.NotContains(t, err.Error(), "secret-value") +} + +func TestValidateS3Endpoint(t *testing.T) { + require.NoError(t, validateEndpoint("http://minio:9000")) + require.Error(t, validateEndpoint("file:///tmp/bucket")) + require.Error(t, validateEndpoint("https://user:pass@example.com")) +} + +func TestS3ArtifactStoreIntegration(t *testing.T) { + if os.Getenv("S3_INTEGRATION_TEST") != "1" { + t.Skip("set S3_INTEGRATION_TEST=1 to run against an S3-compatible service") + } + store, err := NewS3ArtifactStore(context.Background()) + require.NoError(t, err) + key := "integration/artifact-store-" + time.Now().UTC().Format("20060102T150405.000000000") + ".png" + payload := []byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n', 0, 0, 0, 0} + t.Cleanup(func() { + require.NoError(t, store.Delete(context.Background(), key)) + }) + require.NoError(t, store.Put(context.Background(), key, bytes.NewReader(payload), "image/png")) + signedURL, err := store.SignedURL(context.Background(), key, time.Minute) + require.NoError(t, err) + response, err := http.Get(signedURL) // #nosec G107 -- URL is generated by the configured test object store. + require.NoError(t, err) + defer response.Body.Close() + assert.Equal(t, http.StatusOK, response.StatusCode) + downloaded, err := io.ReadAll(response.Body) + require.NoError(t, err) + assert.Equal(t, payload, downloaded) +} diff --git a/web/index.html b/web/index.html index 7f280b401afd..4acf08fd9d51 100644 --- a/web/index.html +++ b/web/index.html @@ -14,7 +14,7 @@ content="Unified AI API gateway and admin dashboard." /> - + diff --git a/web/src/components/theme-switch.tsx b/web/src/components/theme-switch.tsx index 7a3a69ea1801..77dedc55496c 100644 --- a/web/src/components/theme-switch.tsx +++ b/web/src/components/theme-switch.tsx @@ -32,15 +32,15 @@ import { cn } from '@/lib/utils' export function ThemeSwitch() { const { t } = useTranslation() - const { theme, setTheme } = useTheme() + const { resolvedTheme, theme, setTheme } = useTheme() /* Update theme-color meta tag * when theme is updated */ useEffect(() => { - const themeColor = theme === 'dark' ? '#020817' : '#fff' + const themeColor = resolvedTheme === 'dark' ? '#111817' : '#fff' const metaThemeColor = document.querySelector("meta[name='theme-color']") if (metaThemeColor) metaThemeColor.setAttribute('content', themeColor) - }, [theme]) + }, [resolvedTheme]) return ( diff --git a/web/src/components/ui/slider.tsx b/web/src/components/ui/slider.tsx index 1de6dcc8e8ec..96e30b6790f2 100644 --- a/web/src/components/ui/slider.tsx +++ b/web/src/components/ui/slider.tsx @@ -32,7 +32,11 @@ function Slider({ ? value : Array.isArray(defaultValue) ? defaultValue - : [min, max] + : value !== undefined + ? [value] + : defaultValue !== undefined + ? [defaultValue] + : [min] return ( via data-* attributes so theme-presets.css can // override CSS variables at the right cascade layer. useEffect(() => { - applyAttribute( - 'data-theme-preset', - preset === DEFAULT_THEME_CUSTOMIZATION.preset ? null : preset - ) + applyAttribute('data-theme-preset', preset) }, [preset]) // Font is the one axis where we resolve before writing the attribute: @@ -153,17 +150,11 @@ export function ThemeCustomizationProvider(props: { }, [font, preset]) useEffect(() => { - applyAttribute( - 'data-theme-radius', - radius === DEFAULT_THEME_CUSTOMIZATION.radius ? null : radius - ) + applyAttribute('data-theme-radius', radius === 'default' ? null : radius) }, [radius]) useEffect(() => { - applyAttribute( - 'data-theme-scale', - scale === DEFAULT_THEME_CUSTOMIZATION.scale ? null : scale - ) + applyAttribute('data-theme-scale', scale === 'default' ? null : scale) }, [scale]) useEffect(() => { diff --git a/web/src/context/theme-provider.tsx b/web/src/context/theme-provider.tsx index c93bee199833..ee310eaa0cae 100644 --- a/web/src/context/theme-provider.tsx +++ b/web/src/context/theme-provider.tsx @@ -26,14 +26,18 @@ import { } from 'react' import { getCookie, setCookie, removeCookie } from '@/lib/cookies' +import { + DEFAULT_THEME_MODE, + THEME_MODES, + type ThemeMode, +} from '@/lib/theme-customization' -type Theme = 'dark' | 'light' | 'system' +type Theme = ThemeMode type ResolvedTheme = Exclude -const DEFAULT_THEME = 'system' const THEME_COOKIE_NAME = 'vite-ui-theme' const THEME_COOKIE_MAX_AGE = 60 * 60 * 24 * 365 // 1 year -const THEMES = new Set(['dark', 'light', 'system']) +const THEMES = new Set(THEME_MODES) type ThemeProviderProps = { children: React.ReactNode @@ -50,9 +54,9 @@ type ThemeProviderState = { } const initialState: ThemeProviderState = { - defaultTheme: DEFAULT_THEME, - resolvedTheme: 'light', - theme: DEFAULT_THEME, + defaultTheme: DEFAULT_THEME_MODE, + resolvedTheme: DEFAULT_THEME_MODE, + theme: DEFAULT_THEME_MODE, setTheme: () => null, resetTheme: () => null, } @@ -77,7 +81,7 @@ function getStoredTheme(storageKey: string, fallback: Theme): Theme { export function ThemeProvider({ children, - defaultTheme = DEFAULT_THEME, + defaultTheme = DEFAULT_THEME_MODE, storageKey = THEME_COOKIE_NAME, ...props }: ThemeProviderProps) { diff --git a/web/src/features/api-docs/index.tsx b/web/src/features/api-docs/index.tsx new file mode 100644 index 000000000000..7ee3db8f63f1 --- /dev/null +++ b/web/src/features/api-docs/index.tsx @@ -0,0 +1,111 @@ +/* +Copyright (C) 2023-2026 QuantumNous + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact support@quantumnous.com +*/ +import { + ArrowRight01Icon, + BookOpen01Icon, + Key01Icon, + Store01Icon, +} from '@hugeicons/core-free-icons' +import { HugeiconsIcon } from '@hugeicons/react' +import { Link } from '@tanstack/react-router' +import { useTranslation } from 'react-i18next' + +import { PublicLayout } from '@/components/layout' +import { RichContent } from '@/components/rich-content' +import { Button } from '@/components/ui/button' + +const API_BASE_URL = 'https://async-api.nexaapp.cn/v1' + +export function ApiDocs() { + const { t } = useTranslation() + + return ( + +
+
+
+
+
+ + {t('OpenAI-compatible public API')} +
+

+ {t('API Documentation')} +

+

+ {t( + 'Connect chat and image models through one API. The examples below use the interfaces currently available in production.' + )} +

+ +
+ + + +
+
+
+ +
+
+

+ {t('Base URL')} +

+ + {API_BASE_URL} + +
+
+

+ {t('Authentication')} +

+

Bearer API Key

+
+
+

+ {t('Protocol')} +

+

OpenAI Compatible

+
+
+ +
+ +
+
+ + ) +} diff --git a/web/src/features/async-image-lab/api.ts b/web/src/features/async-image-lab/api.ts new file mode 100644 index 000000000000..e708477687fe --- /dev/null +++ b/web/src/features/async-image-lab/api.ts @@ -0,0 +1,97 @@ +/* +Copyright (C) 2023-2026 QuantumNous + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact support@quantumnous.com +*/ +import axios from 'axios' + +import { api } from '@/lib/api' + +import type { + AsyncApiErrorResponse, + AsyncImageFormValues, + AsyncSubmitResponse, + AsyncTaskResultResponse, + AsyncTaskStatusResponse, +} from './types' + +function tokenRequestConfig(apiKey: string) { + return { + headers: { Authorization: `Bearer ${apiKey}` }, + skipBusinessError: true, + skipErrorHandler: true, + skipAuthRefresh: true, + } +} + +export async function submitAsyncImageTask( + values: AsyncImageFormValues, + apiKey: string +): Promise { + const response = await api.post( + '/v1/async/images/generations', + { + model: values.model, + prompt: values.prompt, + n: 1, + size: values.size, + quality: values.quality, + }, + { + ...tokenRequestConfig(apiKey), + headers: { + ...tokenRequestConfig(apiKey).headers, + 'Idempotency-Key': crypto.randomUUID(), + }, + } + ) + return response.data +} + +export async function getAsyncImageTask( + taskId: string, + apiKey: string +): Promise { + const response = await api.get( + `/v1/async/tasks/${encodeURIComponent(taskId)}`, + { + ...tokenRequestConfig(apiKey), + disableDuplicate: true, + } + ) + return response.data +} + +export async function getAsyncImageResult( + taskId: string, + apiKey: string +): Promise { + const response = await api.get( + `/v1/async/tasks/${encodeURIComponent(taskId)}/result?include_upstream=false`, + { + ...tokenRequestConfig(apiKey), + disableDuplicate: true, + } + ) + return response.data +} + +export function getAsyncApiErrorMessage(error: unknown): string | undefined { + if (!axios.isAxiosError(error)) { + return error instanceof Error ? error.message : undefined + } + return error.response?.data.error?.message || error.message +} diff --git a/web/src/features/async-image-lab/constants.ts b/web/src/features/async-image-lab/constants.ts new file mode 100644 index 000000000000..f0ee09d13ac4 --- /dev/null +++ b/web/src/features/async-image-lab/constants.ts @@ -0,0 +1,137 @@ +/* +Copyright (C) 2023-2026 QuantumNous + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact support@quantumnous.com +*/ +import type { StatusBadgeProps } from '@/components/status-badge' + +import type { + AsyncExecutionStatus, + AsyncImageFormValues, + AsyncImageModel, +} from './types' + +interface AsyncModelConfig { + label: string + description: string + sizeLabel: string + sizes: readonly string[] + qualities: readonly string[] + defaultSize: string + defaultQuality: string +} + +export const ASYNC_MODEL_CONFIGS: Record = { + 'gemini-3.1-flash-image-preview': { + label: 'Gemini 3.1 Flash', + description: 'Fast preview generation', + sizeLabel: 'Aspect ratio', + sizes: ['1:1', '16:9', '9:16'], + qualities: ['1K', '2K', '4K'], + defaultSize: '1:1', + defaultQuality: '1K', + }, + 'gemini-3-pro-image-preview': { + label: 'Gemini 3 Pro', + description: 'Higher-detail preview generation', + sizeLabel: 'Aspect ratio', + sizes: ['1:1', '16:9', '9:16'], + qualities: ['1K', '2K', '4K'], + defaultSize: '1:1', + defaultQuality: '1K', + }, + 'gpt-image-2': { + label: 'GPT Image 2', + description: 'OpenAI-compatible image generation', + sizeLabel: 'Image size', + sizes: ['1024x1024', '1536x1024', '1024x1536'], + qualities: ['low', 'medium', 'high'], + defaultSize: '1024x1024', + defaultQuality: 'low', + }, + 'gpt-image-2-vip': { + label: 'GPT Image 2 VIP', + description: 'GRS AI high-resolution generation', + sizeLabel: 'Image size', + sizes: ['1024x1024'], + qualities: ['standard'], + defaultSize: '1024x1024', + defaultQuality: 'standard', + }, + 'nano-banana-pro': { + label: 'Nano Banana Pro', + description: 'High-detail GRS AI generation', + sizeLabel: 'Aspect ratio', + sizes: ['1:1', '16:9', '9:16', '4:3', '3:4'], + qualities: ['1K', '2K', '4K'], + defaultSize: '1:1', + defaultQuality: '1K', + }, + 'nano-banana-2-lite': { + label: 'Nano Banana 2 Lite', + description: 'Cost-efficient GRS AI generation', + sizeLabel: 'Aspect ratio', + sizes: ['1:1', '16:9', '9:16', '4:3', '3:4'], + qualities: ['auto'], + defaultSize: '1:1', + defaultQuality: 'auto', + }, + 'nano-banana-2': { + label: 'Nano Banana 2', + description: 'Balanced GRS AI generation', + sizeLabel: 'Aspect ratio', + sizes: ['1:1', '16:9', '9:16', '4:3', '3:4'], + qualities: ['1K', '2K', '4K'], + defaultSize: '1:1', + defaultQuality: '1K', + }, + 'nano-banana-fast': { + label: 'Nano Banana Fast', + description: 'Fast GRS AI generation', + sizeLabel: 'Aspect ratio', + sizes: ['1:1', '16:9', '9:16', '4:3', '3:4'], + qualities: ['auto'], + defaultSize: '1:1', + defaultQuality: 'auto', + }, +} + +export const ASYNC_IMAGE_DEFAULT_VALUES: AsyncImageFormValues = { + model: 'gpt-image-2', + tokenId: '', + prompt: '', + size: ASYNC_MODEL_CONFIGS['gpt-image-2'].defaultSize, + quality: ASYNC_MODEL_CONFIGS['gpt-image-2'].defaultQuality, +} + +export const TERMINAL_ASYNC_STATUSES = new Set([ + 'success', + 'failure', + 'uncertain', + 'cancelled', +]) + +export const ASYNC_STATUS_CONFIG: Record< + AsyncExecutionStatus, + { label: string; variant: StatusBadgeProps['variant'] } +> = { + queued: { label: 'Queued', variant: 'neutral' }, + running: { label: 'Running', variant: 'info' }, + success: { label: 'Success', variant: 'success' }, + failure: { label: 'Failed', variant: 'danger' }, + uncertain: { label: 'Uncertain', variant: 'warning' }, + cancelled: { label: 'Cancelled', variant: 'neutral' }, +} diff --git a/web/src/features/async-image-lab/index.tsx b/web/src/features/async-image-lab/index.tsx new file mode 100644 index 000000000000..218f862a8035 --- /dev/null +++ b/web/src/features/async-image-lab/index.tsx @@ -0,0 +1,733 @@ +/* +Copyright (C) 2023-2026 QuantumNous + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact support@quantumnous.com +*/ +import { zodResolver } from '@hookform/resolvers/zod' +import { + AiImageIcon, + AlertCircleIcon, + ArrowRight01Icon, + Clock01Icon, + Download01Icon, + Image01Icon, + PlayCircle02Icon, + Refresh01Icon, +} from '@hugeicons/core-free-icons' +import { HugeiconsIcon } from '@hugeicons/react' +import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query' +import { Link } from '@tanstack/react-router' +import { useEffect, useMemo, useRef, useState } from 'react' +import { Controller, useForm, useWatch } from 'react-hook-form' +import { useTranslation } from 'react-i18next' +import { toast } from 'sonner' + +import { SectionPageLayout } from '@/components/layout' +import { StatusBadge } from '@/components/status-badge' +import { Alert, AlertDescription, AlertTitle } from '@/components/ui/alert' +import { Button } from '@/components/ui/button' +import { + Card, + CardContent, + CardDescription, + CardFooter, + CardHeader, + CardTitle, +} from '@/components/ui/card' +import { + Empty, + EmptyDescription, + EmptyHeader, + EmptyMedia, + EmptyTitle, +} from '@/components/ui/empty' +import { + Field, + FieldDescription, + FieldError, + FieldGroup, + FieldTitle, +} from '@/components/ui/field' +import { + Progress, + ProgressLabel, + ProgressValue, +} from '@/components/ui/progress' +import { + Select, + SelectContent, + SelectGroup, + SelectItem, + SelectTrigger, + SelectValue, +} from '@/components/ui/select' +import { Separator } from '@/components/ui/separator' +import { Spinner } from '@/components/ui/spinner' +import { Textarea } from '@/components/ui/textarea' +import { ToggleGroup, ToggleGroupItem } from '@/components/ui/toggle-group' +import { fetchTokenKey, getApiKeys } from '@/features/keys/api' +import type { ApiKey } from '@/features/keys/types' +import { formatTimestampToDate } from '@/lib/format' + +import { + getAsyncApiErrorMessage, + getAsyncImageResult, + getAsyncImageTask, + submitAsyncImageTask, +} from './api' +import { + ASYNC_IMAGE_DEFAULT_VALUES, + ASYNC_MODEL_CONFIGS, + ASYNC_STATUS_CONFIG, + TERMINAL_ASYNC_STATUSES, +} from './constants' +import { + ASYNC_IMAGE_MODELS, + asyncImageFormSchema, + type ActiveAsyncImageTask, + type AsyncImageFormValues, + type AsyncImageModel, + type AsyncTaskResultResponse, + type AsyncTaskStatusResponse, +} from './types' + +function apiKeyAllowsModel(apiKey: ApiKey, model: AsyncImageModel) { + if (!apiKey.model_limits_enabled || !apiKey.model_limits) return true + return apiKey.model_limits.split(',').includes(model) +} + +interface AsyncImageFormProps { + apiKeys: ApiKey[] + isLoadingKeys: boolean + isSubmitting: boolean + onSubmit: (values: AsyncImageFormValues) => void +} + +function AsyncImageForm(props: AsyncImageFormProps) { + const { t } = useTranslation() + const form = useForm({ + resolver: zodResolver(asyncImageFormSchema), + defaultValues: ASYNC_IMAGE_DEFAULT_VALUES, + }) + const selectedModel = useWatch({ control: form.control, name: 'model' }) + const modelConfig = ASYNC_MODEL_CONFIGS[selectedModel] + const compatibleKeys = useMemo( + () => + props.apiKeys.filter( + (apiKey) => + apiKey.status === 1 && apiKeyAllowsModel(apiKey, selectedModel) + ), + [props.apiKeys, selectedModel] + ) + const apiKeyItems = useMemo( + () => + compatibleKeys.map((apiKey) => ({ + label: apiKey.name, + value: String(apiKey.id), + })), + [compatibleKeys] + ) + + useEffect(() => { + const currentTokenId = form.getValues('tokenId') + if (compatibleKeys.some((apiKey) => String(apiKey.id) === currentTokenId)) { + return + } + form.setValue('tokenId', String(compatibleKeys[0]?.id ?? '')) + }, [compatibleKeys, form]) + + const handleModelChange = (values: string[]) => { + const model = values[0] as AsyncImageModel | undefined + if (!model || model === selectedModel) return + const nextConfig = ASYNC_MODEL_CONFIGS[model] + form.setValue('model', model, { shouldValidate: true }) + form.setValue('size', nextConfig.defaultSize, { shouldValidate: true }) + form.setValue('quality', nextConfig.defaultQuality, { + shouldValidate: true, + }) + const currentTokenId = form.getValues('tokenId') + const currentKey = props.apiKeys.find( + (apiKey) => String(apiKey.id) === currentTokenId + ) + if (!currentKey || !apiKeyAllowsModel(currentKey, model)) { + const nextKey = props.apiKeys.find( + (apiKey) => apiKey.status === 1 && apiKeyAllowsModel(apiKey, model) + ) + form.setValue('tokenId', String(nextKey?.id ?? '')) + } + } + + return ( + + + {t('Generation settings')} + + {t( + 'Choose a model and submit a real task to the isolated staging server.' + )} + + +
+ + + ( + + {t('Image model')} + + {ASYNC_IMAGE_MODELS.map((model) => ( + + + {t(ASYNC_MODEL_CONFIGS[model].label)} + + + {t(ASYNC_MODEL_CONFIGS[model].description)} + + + ))} + + + )} + /> + + ( + + {t('API key')} + + + {compatibleKeys.length > 0 + ? t( + 'Only enabled keys that can access the selected model are shown.' + ) + : t('No enabled API key can access this model.')} + + + {fieldState.error?.message + ? t(fieldState.error.message) + : null} + + + )} + /> + + ( + + {t('Prompt')} +