feat(workflow): exec_workflow 增加重试/心跳三列与可恢复查询 DAO

2026-08-25 15:41:50 +08:00
parent b3f4b94b21
commit d5b2d90a27
5 changed files with 203 additions and 0 deletions
@@ -1,10 +1,13 @@
 package session

 import (
+	flow "ai-agent/workflow/consts/flow"
 	"ai-agent/workflow/consts/public"
 	sessionDto "ai-agent/workflow/model/dto/session"
 	"ai-agent/workflow/model/entity"
 	"context"
+	"fmt"
+	"time"

 	"gitea.redpowerfuture.com/red-future/common/beans"
 	"gitea.redpowerfuture.com/red-future/common/db/gfdb"
@@ -117,3 +120,55 @@ func (d *execWorkflowDao) ListBySession(ctx context.Context, sessionId string) (
 	err = r.Structs(&res)
 	return
 }
+
+// ResetRunning 置为运行中并刷新心跳（execute 复用失败记录 / reExecute / 恢复例程共用；
+// 用 map 更新避免 OmitEmpty 省略 0 值，同时写 status、last_heartbeat、node_group_id）
+func (d *execWorkflowDao) ResetRunning(ctx context.Context, id int64, nodeGroupId string) error {
+	_, err := gfdb.DB(ctx, public.DbNameBlackDeacon).Model(ctx, public.TableNameExecWorkflow).
+		Where(entity.ExecWorkflowCol.Id, id).
+		Data(map[string]any{
+			entity.ExecWorkflowCol.Status:        gconv.Int8(*flow.FlowExecutionStatusRunning.Code()),
+			entity.ExecWorkflowCol.LastHeartbeat: time.Now().UnixMilli(),
+			entity.ExecWorkflowCol.NodeGroupId:   nodeGroupId,
+		}).
+		Update()
+	return err
+}
+
+// TouchHeartbeat 更新执行心跳（毫秒时间戳），供后台心跳 goroutine 每 30s 调用一次
+func (d *execWorkflowDao) TouchHeartbeat(ctx context.Context, id int64) error {
+	_, err := gfdb.DB(ctx, public.DbNameBlackDeacon).Model(ctx, public.TableNameExecWorkflow).
+		Where(entity.ExecWorkflowCol.Id, id).
+		Data(map[string]any{entity.ExecWorkflowCol.LastHeartbeat: time.Now().UnixMilli()}).
+		Update()
+	return err
+}
+
+// UpdateRetry 更新重试标记与已重试次数（map 更新，retryable=0 也需写入）
+func (d *execWorkflowDao) UpdateRetry(ctx context.Context, id int64, retryable int, retryCount int) error {
+	_, err := gfdb.DB(ctx, public.DbNameBlackDeacon).Model(ctx, public.TableNameExecWorkflow).
+		Where(entity.ExecWorkflowCol.Id, id).
+		Data(map[string]any{
+			entity.ExecWorkflowCol.Retryable:  retryable,
+			entity.ExecWorkflowCol.RetryCount: retryCount,
+		}).
+		Update()
+	return err
+}
+
+// ListRecoverable 返回可恢复执行：僵尸运行中（status=1 且心跳陈旧）或可重试失败（status=3 且 retryable=1 且未耗尽）
+func (d *execWorkflowDao) ListRecoverable(ctx context.Context, now int64, staleBefore int64, maxRetry int) (res []*entity.ExecWorkflow, err error) {
+	r, err := gfdb.DB(ctx, public.DbNameBlackDeacon).Model(ctx, public.TableNameExecWorkflow).
+		Where(fmt.Sprintf("(%s = ? AND %s < ?) OR (%s = ? AND %s = 1 AND %s < ?)",
+			entity.ExecWorkflowCol.Status, entity.ExecWorkflowCol.LastHeartbeat,
+			entity.ExecWorkflowCol.Status, entity.ExecWorkflowCol.Retryable, entity.ExecWorkflowCol.RetryCount),
+			gconv.Int8(*flow.FlowExecutionStatusRunning.Code()), staleBefore,
+			gconv.Int8(*flow.FlowExecutionStatusFailed.Code()), maxRetry).
+		All()
+	if err != nil {
+		return
+	}
+	// 用 All+Structs 而非 Scan：Scan 生成的列清单会丢嵌入 SQLBaseDO 的 id 等基础列，恢复例程需要 id
+	err = r.Structs(&res)
+	return
+}
@@ -0,0 +1,106 @@
+package session
+
+import (
+	"context"
+	"os"
+	"testing"
+	"time"
+
+	"ai-agent/workflow/consts/flow"
+	sessionDto "ai-agent/workflow/model/dto/session"
+	"ai-agent/workflow/model/entity"
+
+	"gitea.redpowerfuture.com/red-future/common/beans"
+	_ "github.com/gogf/gf/contrib/drivers/pgsql/v2"
+)
+
+// TestExecWorkflowRetryColsRoundTrip 验证三新列读写 + ListRecoverable 三类判定（需 AI_AGENT_TEST_DB=1 且本地 PG 可用）
+func TestExecWorkflowRetryColsRoundTrip(t *testing.T) {
+	if !isTestDBEnabled(t) {
+		t.Skip("skip: AI_AGENT_TEST_DB not set")
+	}
+	// gfdb 的 Hook 依赖 ctx 中注入用户信息（租户/创建人），否则 Insert 报 "token 数据为空"
+	ctx := context.WithValue(context.Background(), "user", &beans.User{
+		UserName: "test-retry", TenantId: 1,
+	})
+	id, err := ExecWorkflowDao.Insert(ctx, &sessionDto.CreateWorkflowReq{
+		SessionId: "test-retry", FlowId: 1, NodeGroupId: "ng",
+		Status:        flow.FlowExecutionStatusRunning.Code(),
+		LastHeartbeat: time.Now().UnixMilli(),
+	})
+	if err != nil {
+		t.Fatalf("insert: %v", err)
+	}
+	defer ExecWorkflowDao.Delete(ctx, &sessionDto.DeleteExecWorkflowReq{Id: []int64{id}})
+
+	if err := ExecWorkflowDao.TouchHeartbeat(ctx, id); err != nil {
+		t.Fatalf("touch heartbeat: %v", err)
+	}
+	if err := ExecWorkflowDao.UpdateRetry(ctx, id, 1, 1); err != nil {
+		t.Fatalf("update retry: %v", err)
+	}
+	got, err := ExecWorkflowDao.GetById(ctx, id)
+	if err != nil || got == nil {
+		t.Fatalf("get: %v", err)
+	}
+	if got.Retryable != 1 || got.RetryCount != 1 || got.LastHeartbeat == 0 {
+		t.Fatalf("cols not persisted: retryable=%d retryCount=%d heartbeat=%d", got.Retryable, got.RetryCount, got.LastHeartbeat)
+	}
+
+	now := time.Now().UnixMilli()
+	stale := now - int64(120*time.Second/time.Millisecond)
+	if err := ExecWorkflowDao.UpdateRetry(ctx, id, 0, 0); err != nil {
+		t.Fatalf("update retry reset: %v", err)
+	}
+	// 情况1：status=3 retryable=1 retry_count<N → 可恢复
+	_, err = ExecWorkflowDao.Update(ctx, &sessionDto.UpdateWorkflowReq{Id: id, Status: flow.FlowExecutionStatusFailed.Code()})
+	if err != nil {
+		t.Fatalf("update failed: %v", err)
+	}
+	if err := ExecWorkflowDao.UpdateRetry(ctx, id, 1, 0); err != nil {
+		t.Fatalf("update retry: %v", err)
+	}
+	rows, err := ExecWorkflowDao.ListRecoverable(ctx, now, stale, 2)
+	if err != nil {
+		t.Fatalf("list recoverable: %v", err)
+	}
+	if !containsID(rows, id) {
+		t.Fatalf("status=3 retryable=1 应可恢复，未命中 id=%d", id)
+	}
+	// 情况2：retry_count 耗尽 → 不可恢复
+	if err := ExecWorkflowDao.UpdateRetry(ctx, id, 1, 2); err != nil {
+		t.Fatalf("update retry exhaust: %v", err)
+	}
+	rows, err = ExecWorkflowDao.ListRecoverable(ctx, now, stale, 2)
+	if err != nil {
+		t.Fatalf("list recoverable: %v", err)
+	}
+	if containsID(rows, id) {
+		t.Fatalf("retry 耗尽不应可恢复")
+	}
+	// 情况3：status=3 retryable=0（用户取消）→ 不可恢复
+	if err := ExecWorkflowDao.UpdateRetry(ctx, id, 0, 0); err != nil {
+		t.Fatalf("update retry: %v", err)
+	}
+	rows, err = ExecWorkflowDao.ListRecoverable(ctx, now, stale, 2)
+	if err != nil {
+		t.Fatalf("list recoverable: %v", err)
+	}
+	if containsID(rows, id) {
+		t.Fatalf("用户取消不可恢复")
+	}
+}
+
+func containsID(rows []*entity.ExecWorkflow, id int64) bool {
+	for _, r := range rows {
+		if r.Id == id {
+			return true
+		}
+	}
+	return false
+}
+
+// isTestDBEnabled 测试辅助：AI_AGENT_TEST_DB=1 时连接真实 PG
+func isTestDBEnabled(t *testing.T) bool {
+	return os.Getenv("AI_AGENT_TEST_DB") == "1"
+}