red-future/model-gateway
4
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: 修复请求头转发与任务状态流转问题

Browse Source 
移除 util.ForwardHeaders，改为从原始请求精确提取 Authorization 或全部请求头；
任务创建时直接设为 Running 状态，避免二次更新与查询；
模型调用使用独立超时上下文，防止外层取消影响回调；
增加 OSS 上传耗时日志，调整数据库连接池参数。
This commit is contained in:
曲皓东
2026-06-18 10:08:36 +08:00
parent fddaf36f48
commit b21d7a8dbf
4 changed files with 90 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
service/task/worker.go
View File

@@ -129,10 +129,14 @@ func (w *asyncWorker) handleOne(ctx context.Context, task *entity.ModelGatewayTa
if attempt > 0 {
g.Log().Infof(ctx, "[执行任务][重试] OSS上传 第%d/%d次 taskId=%s", attempt, maxRetry, task.TaskID)
}
startUpload := time.Now()
oss, err = gateway.UploadByTask(ctx, gjson.New(result).MustToJson(), "json")
if err == nil {
break
}
cost := time.Since(startUpload)
g.Log().Infof(ctx, "本次上传耗时:%s", cost)
g.Log().Errorf(ctx, "[执行任务][失败] OSS上传失败 taskId=%s attempt=%d/%d err=%v", task.TaskID, attempt, maxRetry, err)
if attempt == maxRetry {
task.State = public.TaskStatusFailed
@@ -161,9 +165,9 @@ func (w *asyncWorker) handleOne(ctx context.Context, task *entity.ModelGatewayTa
}
queue.ReleaseQueueSlot(ctx, task.ModelName, task.TaskID)
go gateway.TriggerCallback(context.WithoutCancel(ctx), task)
go gateway.TriggerCallback(util.AsyncCtx(ctx), task)
if req.EpicycleId != 0 {
go gateway.TriggerPromptsCallback(context.WithoutCancel(ctx), task, req.EpicycleId)
go gateway.TriggerPromptsCallback(util.AsyncCtx(ctx), task, req.EpicycleId)
}
g.Log().Infof(ctx, "[执行任务][成功] taskId=%s duration=%ds fileType=%s",
@@ -420,7 +424,7 @@ func injectErrorMessage(payload map[string]any, err error) map[string]any {
// modelKey 用于覆盖/补充模型配置 head_msg（例如每次请求携带不同的 X-API-Key）
func InvokeModel(ctx context.Context, model *entity.ModelGatewayModel, body map[string]any) ([]byte, error) {
// 1) 记录模型调用次数
_ = dao.ModelGatewayLogsStat.IncRequestCount(ctx, time.Now(), model.TenantId, model.Creator, model.ModelName)
//_ = dao.ModelGatewayLogsStat.IncRequestCount(ctx, time.Now(), model.TenantId, model.Creator, model.ModelName)
// 2）请求参数映射：将标准 payload 按模型配置的 requestMapping 转为模型需要的格式
//—— 请求映射实际处理为提示词构建请求，因为有附加字段及其他字段的拼接。这里不方便做请求映射
@@ -447,13 +451,20 @@ func InvokeModel(ctx context.Context, model *entity.ModelGatewayModel, body map[
baseURL = baseURL + "?" + q.Encode()
}
}
req, err = http.NewRequestWithContext(ctx, http.MethodGet, baseURL, nil)
// 改用独立超时ctx，隔绝外层截止
reqCtx, reqCancel := context.WithTimeout(context.Background(), timeout)
defer reqCancel()
req, err = http.NewRequestWithContext(reqCtx, http.MethodGet, baseURL, nil)
//req, err = http.NewRequestWithContext(ctx, http.MethodGet, baseURL, nil)
default:
bodyBytes, err := json.Marshal(body)
if err != nil {
return nil, err
}
req, err = http.NewRequestWithContext(ctx, http.MethodPost, baseURL, bytes.NewReader(bodyBytes))
reqCtx, reqCancel := context.WithTimeout(context.Background(), timeout)
defer reqCancel()
req, err = http.NewRequestWithContext(reqCtx, http.MethodPost, baseURL, bytes.NewReader(bodyBytes))
//req, err = http.NewRequestWithContext(ctx, http.MethodPost, baseURL, bytes.NewReader(bodyBytes))
}
// 5）注入请求头：先模型静态配置，再动态 modelKey（后者可覆盖前者）
@@ -552,5 +563,5 @@ func (w *asyncWorker) failTask(ctx context.Context, t *entity.ModelGatewayTask,
g.Log().Warningf(ctx, "[执行任务][更新数据库失败] taskId=%s err=%v", t.TaskID, err)
}
queue.ReleaseQueueSlot(ctx, t.ModelName, t.TaskID)
go gateway.TriggerCallback(context.WithoutCancel(ctx), t)
go gateway.TriggerCallback(util.AsyncCtx(ctx), t)
}