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
@@ -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"
+}