118 lines
4.3 KiB
Go
118 lines
4.3 KiB
Go
package flow
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"errors"
|
||
"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"
|
||
)
|
||
|
||
// 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 统一异步模型调用入口（spec §5.4）：
|
||
// 提交时把 task_id/msg_topic 落库 flow_async_task，崩溃后重订阅 msg_topic 拿回已完成结果复用，不重复调用。
|
||
// 同步模型直接走 gateway.ModelCallResult，不落库（无恢复语义）。
|
||
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
|
||
}
|