ai-agent/workflow/service/flow/async.go

package flow

import (
	"context"
	"encoding/json"
	"errors"
	"sync"
	"time"

	"github.com/gogf/gf/v2/frame/g"

	"ai-agent/gateway"
	"ai-agent/workflow/consts/model"
	flowDao "ai-agent/workflow/dao/flow"
	"ai-agent/workflow/model/entity"
)

// 全局等待任务回调的工具
var (
	asyncMu    sync.Mutex
	asyncTasks = make(map[string]chan any)
)

// Wait 阻塞等待回调结果
// 调用后会一直卡住，直到 Notify 唤醒 或 超时/取消
func Wait(ctx context.Context, taskId string) (any, error) {
	asyncMu.Lock()
	ch := make(chan any, 1)
	asyncTasks[taskId] = ch
	asyncMu.Unlock()

	defer close(ch)
	for {
		select {
		case result := <-ch:
			return result, nil
		case <-ctx.Done():
			asyncMu.Lock()
			delete(asyncTasks, taskId)
			asyncMu.Unlock()
			return nil, ctx.Err()
		}
	}
}

// Notify 回调时调用，唤醒等待的任务
func Notify(taskId string, result any) {
	asyncMu.Lock()
	defer asyncMu.Unlock()

	ch, exist := asyncTasks[taskId]
	if !exist {
		return
	}
	ch <- result
	delete(asyncTasks, taskId)
}

// asyncRecoverWaitTimeout in-flight 行重订阅等结果的超时上限。
// 任务可能仍执行中（消息未发布）或提交即失败（不会发布消息）；超时视为结果未知，清记录重提。
const asyncRecoverWaitTimeout = 10 * time.Minute

type asyncAction int

const (
	asyncActionResubmit asyncAction = iota // 无记录 / failed / done 空结果 → 重新提交
	asyncActionReuse                       // done 有结果 → 复用，不重调
	asyncActionFinalize                    // in-flight → 重订阅 msgTopic 收尾
)

// asyncCallAction 异步任务缓存行动决策（纯函数，可单测）
func asyncCallAction(rec *entity.FlowAsyncTask) asyncAction {
	if rec == nil {
		return asyncActionResubmit
	}
	switch rec.State {
	case flowDao.FlowAsyncStateDone:
		if rec.Result != "" && rec.Result != "{}" {
			return asyncActionReuse
		}
		return asyncActionResubmit
	case flowDao.FlowAsyncStateFailed:
		return asyncActionResubmit
	default: // in-flight
		return asyncActionFinalize
	}
}

// AsyncModelCallWithRecovery 统一异步模型调用入口：
// 提交时把 task_id/msg_topic 落库 flow_async_task，崩溃后重订阅 msg_topic 拿回已完成结果复用，不重复调用。
// 同步模型直接走 gateway.ModelCallResult，不落库（无恢复语义）。
// 注意：本函数是节点内阻塞调用（WaitModelCallResult 等回调），不是独立并发触发方，
// 不参与 exec 并发仲裁（谁抢到执行权谁跑）——仲裁语义见《工作流执行并发仲裁设计.md》。
func AsyncModelCallWithRecovery(ctx context.Context, execId int64, nodeId string, segIdx int, modelId int64, responseType model.ResponseType, sessionId string, requestParams map[string]any, businessParams map[string]any) (*gateway.ModelCallRes, error) {
	if responseType == nil || *responseType != *model.ResponseTypeAsync.Code() {
		return gateway.ModelCallResult(ctx, modelId, responseType, sessionId, requestParams, businessParams)
	}

	rec, err := flowDao.FlowAsyncTaskDao.Get(ctx, execId, nodeId, segIdx)
	if err != nil {
		return nil, err
	}
	switch asyncCallAction(rec) {
	case asyncActionReuse:
		return unmarshalModelCallRes(rec.Result)
	case asyncActionFinalize:
		// 重订阅 msgTopic 收尾：拿回已发布结果（消息在 JetStream 保留 7 天）；
		// 成功 → 落 done 复用；超时/订阅失败 → 清记录重提
		waitCtx, cancel := context.WithTimeout(ctx, asyncRecoverWaitTimeout)
		res, waitErr := gateway.WaitModelCallResult(waitCtx, rec.MsgTopic)
		cancel()
		if waitErr != nil {
			if errors.Is(waitErr, context.DeadlineExceeded) || errors.Is(waitErr, context.Canceled) {
				_ = flowDao.FlowAsyncTaskDao.DeleteByKey(ctx, execId, nodeId, segIdx)
				break // 落入下方重新提交
			}
			return nil, waitErr
		}
		_ = flowDao.FlowAsyncTaskDao.UpdateByKey(ctx, execId, nodeId, segIdx, flowDao.FlowAsyncStateDone, marshalModelCallRes(res))
		return res, nil
	case asyncActionResubmit:
		// 清残留（failed 或 done 空结果）
		if rec != nil {
			_ = flowDao.FlowAsyncTaskDao.DeleteByKey(ctx, execId, nodeId, segIdx)
		}
	}

	// 重新提交：先落库 in-flight（task_id/msg_topic），等待结果
	res, msgTopic, err := gateway.SubmitModelCall(ctx, modelId, responseType, sessionId, requestParams, businessParams)
	if err != nil {
		return nil, err
	}
	if err := flowDao.FlowAsyncTaskDao.Upsert(ctx, execId, nodeId, segIdx, modelId, res.TaskId, msgTopic); err != nil {
		return nil, err
	}
	waitRes, waitErr := gateway.WaitModelCallResult(ctx, msgTopic)
	if waitErr != nil {
		// 结果失败（model error/取消）：落 failed，调用方（段重试/恢复）决定后续
		_ = flowDao.FlowAsyncTaskDao.UpdateByKey(ctx, execId, nodeId, segIdx, flowDao.FlowAsyncStateFailed, "")
		return nil, waitErr
	}
	_ = flowDao.FlowAsyncTaskDao.UpdateByKey(ctx, execId, nodeId, segIdx, flowDao.FlowAsyncStateDone, marshalModelCallRes(waitRes))
	return waitRes, nil
}

func marshalModelCallRes(res *gateway.ModelCallRes) string {
	b, err := json.Marshal(res)
	if err != nil {
		g.Log().Warningf(context.Background(), "序列化模型调用结果失败: %v", err)
		return "{}"
	}
	return string(b)
}

func unmarshalModelCallRes(s string) (*gateway.ModelCallRes, error) {
	res := new(gateway.ModelCallRes)
	if err := json.Unmarshal([]byte(s), res); err != nil {
		return nil, err
	}
	return res, nil
}