model-gateway/common/util/pull_task.go

package util

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"regexp"
	"strings"
	"time"

	"github.com/gogf/gf/v2/encoding/gjson"
	"github.com/gogf/gf/v2/frame/g"
	"github.com/gogf/gf/v2/util/gconv"
)

// PullTaskResult 拉取任务结果
func PullTaskResult(ctx context.Context, body map[string]any, queryConfig map[string]any, headMsg map[string]any) (map[string]any, error) {
	taskID, err := extractTaskID(body, queryConfig)
	if err != nil {
		return nil, err
	}
	g.Log().Infof(ctx, "[PullTaskResult] taskID=%s", taskID)

	queryUrl := buildQueryURL(queryConfig, taskID)
	method := gconv.String(queryConfig["method"])
	if method == "" {
		method = "GET"
	}
	interval := gconv.Int(queryConfig["interval_seconds"])
	if interval <= 0 {
		interval = 2
	}

	statusPath := gconv.String(queryConfig["status_path"])
	if statusPath == "" {
		statusPath = "status"
	}
	statusValues, _ := queryConfig["status_values"].(map[string]any)

	// 失败信息路径
	errorPath := gconv.String(queryConfig["error_path"])
	if errorPath == "" {
		errorPath = "error.message"
	}

	reqBodyMap := map[string]any{"task_id": taskID}

	for {
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		default:
		}

		result, err := doQueryRequest(ctx, method, queryUrl, reqBodyMap, headMsg)
		if err != nil {
			g.Log().Warningf(ctx, "[PullTaskResult] 请求失败 taskID=%s err=%v", taskID, err)
			time.Sleep(time.Duration(interval) * time.Second)
			continue
		}

		if result == nil {
			time.Sleep(time.Duration(interval) * time.Second)
			continue
		}

		statusStr := gconv.String(gjson.New(result).Get(statusPath).Val())
		g.Log().Infof(ctx, "[PullTaskResult] 状态 taskID=%s status=%s", taskID, statusStr)

		if matchStatus(statusStr, statusValues["succeeded"]) {
			g.Log().Infof(ctx, "[PullTaskResult] 任务成功 taskID=%s", taskID)
			return result, nil
		}

		if matchStatus(statusStr, statusValues["failed"]) {
			errMsg := gconv.String(gjson.New(result).Get(errorPath).Val())
			if errMsg == "" {
				errMsg = "任务失败"
			}
			g.Log().Errorf(ctx, "[PullTaskResult] 任务失败 taskID=%s err=%s", taskID, errMsg)
			return result, fmt.Errorf("任务失败: %s", errMsg)
		}

		time.Sleep(time.Duration(interval) * time.Second)
	}
}

func extractTaskID(body, queryConfig map[string]any) (string, error) {
	taskIDPath := gconv.String(queryConfig["task_id"])
	taskID := gconv.String(gjson.New(body).Get(taskIDPath).Val())
	if taskID == "" {
		return "", fmt.Errorf("无法从路径 %s 提取 taskID", taskIDPath)
	}
	return taskID, nil
}

func buildQueryURL(queryConfig map[string]any, taskID string) string {
	queryUrl := gconv.String(queryConfig["url"])
	return replaceURLParams(queryUrl, map[string]any{"id": taskID})
}

func doQueryRequest(ctx context.Context, method, queryUrl string, reqBodyMap map[string]any, headMsg map[string]any) (map[string]any, error) {
	var reqBody io.Reader
	if method == "POST" {
		bs, _ := json.Marshal(reqBodyMap)
		reqBody = bytes.NewReader(bs)
	}

	req, err := http.NewRequestWithContext(ctx, method, queryUrl, reqBody)
	if err != nil {
		return nil, fmt.Errorf("创建请求失败: %w", err)
	}

	for hk, hv := range ParseHeadMsgHeaders(headMsg) {
		req.Header.Set(hk, hv)
	}

	client := &http.Client{Timeout: 30 * time.Second}
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	raw, _ := io.ReadAll(resp.Body)

	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return nil, nil
	}

	var result map[string]any
	_ = json.Unmarshal(raw, &result)
	return result, nil
}

func matchStatus(actual string, expected any) bool {
	expectedStr := gconv.String(expected)
	if actual == expectedStr {
		return true
	}
	if arr, ok := expected.([]any); ok {
		for _, item := range arr {
			if actual == gconv.String(item) {
				return true
			}
		}
	}
	return false
}

func replaceURLParams(rawURL string, params map[string]any) string {
	re := regexp.MustCompile(`\{([^}]+)}`)
	return re.ReplaceAllStringFunc(rawURL, func(s string) string {
		key := strings.Trim(s, "{}")
		if val, ok := params[key]; ok {
			return gconv.String(val)
		}
		return s
	})
}