gatway

dao/task_dao_bg.go

@@ -20,7 +20,7 @@ func (d *taskDao) ClaimPendingGlobal(ctx context.Context, batchSize int) (tasks
 	}
 	err = gfdb.DB(ctx).Transaction(ctx, func(ctx context.Context, tx gdb.TX) error {
 		sql := fmt.Sprintf(
-			`SELECT id, tenant_id, creator, model_name, task_id, model_key, input_ref, request_payload, phase, tmp_file
+			`SELECT id, tenant_id, creator, model_name, task_id, biz_name, callback_url, model_key, input_ref, request_payload, phase, tmp_file
 			   FROM %s
 			  WHERE deleted_at IS NULL AND state = 0
 			  ORDER BY enqueue_at ASC
@@ -55,13 +55,51 @@ func (d *taskDao) ClaimPendingGlobal(ctx context.Context, batchSize int) (tasks
 	return
 }
 
-func (d *taskDao) UpdateSuccessGlobal(ctx context.Context, id int64, ossFile, fileType string, fileSize int64, expireAt *gtime.Time) error {
+// ClaimPendingByTaskIDGlobal 按 task_id 定向抢占单个 pending 任务（不加 tenant 过滤）
+// 用于 createTask 创建成功后立即异步尝试执行当前任务，避免只依赖后续 runWork 扫描队列。
+func (d *taskDao) ClaimPendingByTaskIDGlobal(ctx context.Context, taskID string) (task *entity.AsynchTask, err error) {
+	if taskID == "" {
+		return nil, nil
+	}
+	err = gfdb.DB(ctx).Transaction(ctx, func(ctx context.Context, tx gdb.TX) error {
+		sql := fmt.Sprintf(
+			`SELECT id, tenant_id, creator, model_name, task_id, biz_name, callback_url, model_key, input_ref, request_payload, phase, tmp_file
+			   FROM %s
+			  WHERE deleted_at IS NULL AND state = 0 AND task_id = ?
+			  LIMIT 1
+			  FOR UPDATE SKIP LOCKED`,
+			public.TableNameTask,
+		)
+		r, err := tx.GetOne(sql, taskID)
+		if err != nil {
+			return err
+		}
+		if r.IsEmpty() {
+			task = nil
+			return nil
+		}
+		if err := r.Struct(&task); err != nil {
+			return err
+		}
+		now := time.Now()
+		_, err = tx.Exec(
+			fmt.Sprintf(`UPDATE %s SET state=1, started_at=?, updated_at=? WHERE id=?`, public.TableNameTask),
+			now, now, task.Id,
+		)
+		return err
+	})
+	return
+}
+
+func (d *taskDao) UpdateSuccessGlobal(ctx context.Context, id int64, ossFile, fileType, textResult string, fileSize int64, expireAt *gtime.Time, expendTokens int) error {
 	now := gtime.Now()
 	_, err := gfdb.DB(ctx).Exec(ctx,
 		fmt.Sprintf(`UPDATE %s
 SET state=2,
     oss_file=?,
     file_type=?,
+    text_result=?,
+	expend_tokens=?,
     file_size=?,
     error_msg='',
     finished_at=?,
@@ -71,7 +109,7 @@ SET state=2,
     tmp_file='',
     updated_at=?
 WHERE id=?`, public.TableNameTask),
-		ossFile, fileType, fileSize, now, now, now, id,
+		ossFile, fileType, textResult, expendTokens, fileSize, now, now, now, id,
 	)
 	return err
 }