Skip to content

refactor: async task #2985

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Calcium-Ion merged 10 commits into from
Feb 22, 2026
Merged

refactor: async task #2985

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
9e39544
refactor(task): extract billing and polling logic from controller to …
Calcium-Ion Feb 10, 2026
d6e11fd
feat(task): add adaptor billing interface and async settlement framework
Calcium-Ion Feb 10, 2026
809ba92
refactor(logs): add refund logging for asynchronous tasks and update …
Calcium-Ion Feb 21, 2026
a920d1f
refactor(relay): rename RelayTask to RelayTaskFetch and update routing
Calcium-Ion Feb 21, 2026
76892e8
refactor(relay): enhance remix logic for billing context extraction
Calcium-Ion Feb 21, 2026
cda5401
refactor(relay): improve channel locking and retry logic in RelayTask
Calcium-Ion Feb 21, 2026
5ec4633
refactor(task): add CAS-guarded updates to prevent concurrent billing…
Calcium-Ion Feb 21, 2026
9976b31
refactor(task): enhance UpdateWithStatus for CAS updates and add inte…
Calcium-Ion Feb 21, 2026
ec5c6b2
feat(task): add model redirection, per-call billing, and multipart re…
Calcium-Ion Feb 22, 2026
48c9b17
fix(i18n): remove duplicate task ID translations and clean up unused …
Calcium-Ion Feb 22, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 9 additions & 1 deletion common/gin.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -243,7 +243,15 @@ func ParseMultipartFormReusable(c *gin.Context) (*multipart.Form, error) {
return nil, err
}

contentType := c.Request.Header.Get("Content-Type")
// Use the original Content-Type saved on first call to avoid boundary
// mismatch when callers overwrite c.Request.Header after multipart rebuild.
var contentType string
if saved, ok := c.Get("_original_multipart_ct"); ok {
contentType = saved.(string)
} else {
contentType = c.Request.Header.Get("Content-Type")
c.Set("_original_multipart_ct", contentType)
}
Comment on lines +249 to +254

@coderabbitai coderabbitai Bot Feb 22, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Unsafe bare type assertion will panic on unexpected stored type.

saved.(string) at Line 250 panics if anything other than a string is stored under "_original_multipart_ct". Use the comma-ok form and fall back to the header read path on failure.

🛡️ Proposed fix 
-	if saved, ok := c.Get("_original_multipart_ct"); ok {
-		contentType = saved.(string)
-	} else {
+	if saved, ok := c.Get("_original_multipart_ct"); ok {
+		if s, ok := saved.(string); ok {
+			contentType = s
+		} else {
+			contentType = c.Request.Header.Get("Content-Type")
+		}
+	} else {
 		contentType = c.Request.Header.Get("Content-Type")
 		c.Set("_original_multipart_ct", contentType)
 	}
🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@common/gin.go` around lines 249 - 254, The code in the handler using
gin.Context (variable c) unsafely asserts saved.(string) for the key
"_original_multipart_ct" which can panic if the stored value is not a string;
change this to use the comma-ok form (v, ok := saved.(string)) and if ok is
false fall back to reading c.Request.Header.Get("Content-Type") and store that
string into c.Set("_original_multipart_ct", contentType) so the function uses
the header path on failure instead of panicking.

boundary, err := parseBoundary(contentType)
if err != nil {
return nil, err
Expand Down
17 changes: 8 additions & 9 deletions controller/midjourney.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -130,6 +130,7 @@ func UpdateMidjourneyTaskBulk() {
if !checkMjTaskNeedUpdate(task, responseItem) {
continue
}
preStatus := task.Status
task.Code = 1
task.Progress = responseItem.Progress
task.PromptEn = responseItem.PromptEn
Expand Down Expand Up @@ -172,18 +173,16 @@ func UpdateMidjourneyTaskBulk() {
shouldReturnQuota = true
}
}
err = task.Update()
won, err := task.UpdateWithStatus(preStatus)
if err != nil {
logger.LogError(ctx, "UpdateMidjourneyTask task error: "+err.Error())
} else {
if shouldReturnQuota {
err = model.IncreaseUserQuota(task.UserId, task.Quota, false)
if err != nil {
logger.LogError(ctx, "fail to increase user quota: "+err.Error())
}
logContent := fmt.Sprintf("构图失败 %s，补偿 %s", task.MjId, logger.LogQuota(task.Quota))
model.RecordLog(task.UserId, model.LogTypeSystem, logContent)
} else if won && shouldReturnQuota {
err = model.IncreaseUserQuota(task.UserId, task.Quota, false)
if err != nil {
logger.LogError(ctx, "fail to increase user quota: "+err.Error())
}
logContent := fmt.Sprintf("构图失败 %s，补偿 %s", task.MjId, logger.LogQuota(task.Quota))
model.RecordLog(task.UserId, model.LogTypeSystem, logContent)
}
}
}
Expand Down
149 changes: 112 additions & 37 deletions controller/relay.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -450,72 +450,147 @@ func RelayNotFound(c *gin.Context) {
})
}

func RelayTaskFetch(c *gin.Context) {
relayInfo, err := relaycommon.GenRelayInfo(c, types.RelayFormatTask, nil, nil)
if err != nil {
c.JSON(http.StatusInternalServerError, &dto.TaskError{
Code: "gen_relay_info_failed",
Message: err.Error(),
StatusCode: http.StatusInternalServerError,
})
return
}
if taskErr := relay.RelayTaskFetch(c, relayInfo.RelayMode); taskErr != nil {
respondTaskError(c, taskErr)
}
}

func RelayTask(c *gin.Context) {
retryTimes := common.RetryTimes
channelId := c.GetInt("channel_id")
c.Set("use_channel", []string{fmt.Sprintf("%d", channelId)})
relayInfo, err := relaycommon.GenRelayInfo(c, types.RelayFormatTask, nil, nil)
if err != nil {
c.JSON(http.StatusInternalServerError, &dto.TaskError{
Code: "gen_relay_info_failed",
Message: err.Error(),
StatusCode: http.StatusInternalServerError,
})
return
}
taskErr := taskRelayHandler(c, relayInfo)
if taskErr == nil {
retryTimes = 0

if taskErr := relay.ResolveOriginTask(c, relayInfo); taskErr != nil {
respondTaskError(c, taskErr)
return
}

var result *relay.TaskSubmitResult
var taskErr *dto.TaskError
defer func() {
if taskErr != nil && relayInfo.Billing != nil {
relayInfo.Billing.Refund(c)
}
}()

retryParam := &service.RetryParam{
Ctx: c,
TokenGroup: relayInfo.TokenGroup,
ModelName: relayInfo.OriginModelName,
Retry: common.GetPointer(0),
}
for ; shouldRetryTaskRelay(c, channelId, taskErr, retryTimes) && retryParam.GetRetry() < retryTimes; retryParam.IncreaseRetry() {
channel, newAPIError := getChannel(c, relayInfo, retryParam)
if newAPIError != nil {
logger.LogError(c, fmt.Sprintf("CacheGetRandomSatisfiedChannel failed: %s", newAPIError.Error()))
taskErr = service.TaskErrorWrapperLocal(newAPIError.Err, "get_channel_failed", http.StatusInternalServerError)
break

for ; retryParam.GetRetry() <= common.RetryTimes; retryParam.IncreaseRetry() {
var channel *model.Channel

if lockedCh, ok := relayInfo.LockedChannel.(*model.Channel); ok && lockedCh != nil {
channel = lockedCh
if retryParam.GetRetry() > 0 {
if setupErr := middleware.SetupContextForSelectedChannel(c, channel, relayInfo.OriginModelName); setupErr != nil {
taskErr = service.TaskErrorWrapperLocal(setupErr.Err, "setup_locked_channel_failed", http.StatusInternalServerError)
break
}
}
} else {
var channelErr *types.NewAPIError
channel, channelErr = getChannel(c, relayInfo, retryParam)
if channelErr != nil {
logger.LogError(c, channelErr.Error())
taskErr = service.TaskErrorWrapperLocal(channelErr.Err, "get_channel_failed", http.StatusInternalServerError)
break
}
}
channelId = channel.Id
useChannel := c.GetStringSlice("use_channel")
useChannel = append(useChannel, fmt.Sprintf("%d", channelId))
c.Set("use_channel", useChannel)
logger.LogInfo(c, fmt.Sprintf("using channel #%d to retry (remain times %d)", channel.Id, retryParam.GetRetry()))
//middleware.SetupContextForSelectedChannel(c, channel, originalModel)

bodyStorage, err := common.GetBodyStorage(c)
if err != nil {
if common.IsRequestBodyTooLargeError(err) || errors.Is(err, common.ErrRequestBodyTooLarge) {
taskErr = service.TaskErrorWrapperLocal(err, "read_request_body_failed", http.StatusRequestEntityTooLarge)

addUsedChannel(c, channel.Id)
bodyStorage, bodyErr := common.GetBodyStorage(c)
if bodyErr != nil {
if common.IsRequestBodyTooLargeError(bodyErr) || errors.Is(bodyErr, common.ErrRequestBodyTooLarge) {
taskErr = service.TaskErrorWrapperLocal(bodyErr, "read_request_body_failed", http.StatusRequestEntityTooLarge)
} else {
taskErr = service.TaskErrorWrapperLocal(err, "read_request_body_failed", http.StatusBadRequest)
taskErr = service.TaskErrorWrapperLocal(bodyErr, "read_request_body_failed", http.StatusBadRequest)
}
break
}
c.Request.Body = io.NopCloser(bodyStorage)
taskErr = taskRelayHandler(c, relayInfo)

result, taskErr = relay.RelayTaskSubmit(c, relayInfo)
if taskErr == nil {
break
}

if !taskErr.LocalError {
processChannelError(c,
*types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey,
common.GetContextKeyString(c, constant.ContextKeyChannelKey), channel.GetAutoBan()),
types.NewOpenAIError(taskErr.Error, types.ErrorCodeBadResponseStatusCode, taskErr.StatusCode))
}

if !shouldRetryTaskRelay(c, channel.Id, taskErr, common.RetryTimes-retryParam.GetRetry()) {
break
}
}

useChannel := c.GetStringSlice("use_channel")
if len(useChannel) > 1 {
retryLogStr := fmt.Sprintf("重试：%s", strings.Trim(strings.Join(strings.Fields(fmt.Sprint(useChannel)), "->"), "[]"))
logger.LogInfo(c, retryLogStr)
}
if taskErr != nil {
if taskErr.StatusCode == http.StatusTooManyRequests {
taskErr.Message = "当前分组上游负载已饱和，请稍后再试"

// ── 成功：结算 + 日志 + 插入任务 ──
if taskErr == nil {
if settleErr := service.SettleBilling(c, relayInfo, result.Quota); settleErr != nil {
common.SysError("settle task billing error: " + settleErr.Error())
}
c.JSON(taskErr.StatusCode, taskErr)
service.LogTaskConsumption(c, relayInfo)

task := model.InitTask(result.Platform, relayInfo)
task.PrivateData.UpstreamTaskID = result.UpstreamTaskID
task.PrivateData.BillingSource = relayInfo.BillingSource
task.PrivateData.SubscriptionId = relayInfo.SubscriptionId
task.PrivateData.TokenId = relayInfo.TokenId
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: common.StringsContains(constant.TaskPricePatches, relayInfo.OriginModelName),
}
task.Quota = result.Quota
task.Data = result.TaskData
task.Action = relayInfo.Action
if insertErr := task.Insert(); insertErr != nil {
common.SysError("insert task error: " + insertErr.Error())
}
Comment on lines +576 to +580

@coderabbitai coderabbitai Bot Feb 22, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Task record insertion failure is silently swallowed — the user is billed but the task is untracked.

service.SettleBilling (line 557) and the adaptor's DoResponse (which already wrote 200 OK to the client) both succeed before task.Insert() is called. If task.Insert() fails, the quota has been consumed, the upstream task is running, but the task record does not exist — the user can never query or cancel it. The error is only logged to sys-error with no compensating action.

Consider persisting the task record before settling billing (or at least retrying the insert asynchronously), so a billing commit only happens when the task is durably stored.

🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@controller/relay.go` around lines 576 - 580, The task record insertion
(task.Insert()) is performed after service.SettleBilling and adaptor.DoResponse,
so failures can consume quota while leaving the task untracked; move durable
persistence earlier or add a compensating flow: persist the task (call
task.Insert()) before calling service.SettleBilling and before writing the final
response in adaptor.DoResponse, and if immediate pre-billing insert fails
implement a retry/enqueue mechanism (or abort/rollback the billing action) and
surface a server error instead of silently logging via common.SysError; update
the code paths around service.SettleBilling, adaptor.DoResponse, and
task.Insert() to ensure billing only commits after a successful durable insert
(or that failed inserts trigger async retry and billing rollback).

}

if taskErr != nil {
respondTaskError(c, taskErr)
}
}

func taskRelayHandler(c *gin.Context, relayInfo *relaycommon.RelayInfo) *dto.TaskError {
var err *dto.TaskError
switch relayInfo.RelayMode {
case relayconstant.RelayModeSunoFetch, relayconstant.RelayModeSunoFetchByID, relayconstant.RelayModeVideoFetchByID:
err = relay.RelayTaskFetch(c, relayInfo.RelayMode)
default:
err = relay.RelayTaskSubmit(c, relayInfo)
// respondTaskError 统一输出 Task 错误响应（含 429 限流提示改写）
func respondTaskError(c *gin.Context, taskErr *dto.TaskError) {
if taskErr.StatusCode == http.StatusTooManyRequests {
taskErr.Message = "当前分组上游负载已饱和，请稍后再试"
}
return err
c.JSON(taskErr.StatusCode, taskErr)
}

func shouldRetryTaskRelay(c *gin.Context, channelId int, taskErr *dto.TaskError, retryTimes int) bool {
Expand Down
Loading