diff --git a/.doc/test-reports/2026-09-12-phase4-lease-phase5-fault-injection-test-report.md b/.doc/test-reports/2026-09-12-phase4-lease-phase5-fault-injection-test-report.md new file mode 100644 index 00000000..246c99df --- /dev/null +++ b/.doc/test-reports/2026-09-12-phase4-lease-phase5-fault-injection-test-report.md @@ -0,0 +1,77 @@ +# Phase 4 Worker Lease + Phase 5 故障注入测试报告 + +## 一、 测试背景 (Background) + +Phase 4 持久化后,running 任务在进程重启时仍会被一律标记 interrupted,缺乏 worker lease/heartbeat 区分「仍存活 worker」与「孤儿任务」。Phase 5 验收矩阵尚需并发推进、lease 故障与备份恢复模拟。 + +## 二、 测试目标 (Test Objectives) + +1. `task_runs` 增加 `worker_id` / `lease_expires_at` / `lease_heartbeat_at` / `attempt` +2. 运行中任务定期 heartbeat;仅 **lease 过期** 的 running 任务标记 `interrupted` +3. 并发双次 stage advance:恰好一次成功、第二次冲突/门禁失败 +4. 备份恢复模拟:迁移幂等 + 核心表仍存在 + +## 三、 测试方案与环境 (Test Plan & Environment) + +- OS: Linux, Go 1.27 +- PG: `postgres://app:dev_secret_db_pass@127.0.0.1:5433/agent_system_dev` +- 迁移: goose **v20** (`00020_2026-09-12-task-run-worker-lease.sql`) +- Lease TTL: 30s;Heartbeat: 10s + +## 四、 预期结果 (Expected Results) + +| 用例 | 预期 | +|------|------| +| 过期 lease | `MarkOrphanedRunning` → `interrupted` | +| 有效 lease | orphan 扫描后仍为 `running` | +| 并发 advance | 1×200 + 1×409/422 | +| 备份恢复 | 重复 `goose Up` 版本不变,`task_runs` 存在 | + +## 五、 实际测试结果与判定 (Actual Results & Evaluation) + +### 5.1 过期 lease 孤儿检测 + +**命令**: `go test -run TestPGStore_StaleLeaseMarkedInterrupted ./internal/taskrun/...` + +**判定**: PASS + +### 5.2 有效 lease 不被误杀 + +**命令**: `go test -run TestPGStore_FreshLeaseSurvivesOrphanScan ./internal/taskrun/...` + +**判定**: PASS + +### 5.3 运行中 heartbeat 保护 + +**命令**: `go test -run TestPhase5_TaskRunHeartbeatSurvivesOrphanScan_PG ./internal/server/...` + +**判定**: PASS — 长任务 worker 持有 lease 时 orphan 扫描不中断 + +### 5.4 并发 stage advance + +**命令**: `go test -run TestPhase5_ConcurrentStageAdvance_PG ./internal/server/...` + +**判定**: PASS — 双 goroutine 并发 advance,一次 200、一次冲突 + +### 5.5 备份恢复迁移幂等 + +**命令**: `go test -run TestPhase5_BackupRestoreMigration_PG ./internal/server/...` + +**判定**: PASS + +## 六、 自动化测试套件汇总 (Test Suite Execution) + +``` +go test ./internal/taskrun/... ./internal/server/... -run 'TestPGStore_|TestPhase5_|TestManager_' -count=1 +ok chat-agent-system/internal/taskrun 0.192s +ok chat-agent-system/internal/server 0.929s +``` + +## 七、 测试结论与后续优化建议 (Conclusion & Next Steps) + +**结论**: Worker lease/heartbeat 与 Phase 5 首批故障注入项 **PASS**,较上一版「重启即 interrupted」显著改进;尚未实现跨进程 **任务续跑**(仅识别孤儿 + 重放历史事件)。 + +**后续**: +1. 基于 `workspace_id` 的可选 resume 策略(interrupted → 用户确认后续跑) +2. Phase 5:Finalize 故障 HTTP 级注入、滚动重启 soak +3. Gap P1:Mermaid strict、上传硬限 diff --git a/deploy/postgres/goose/00020_2026-09-12-task-run-worker-lease.sql b/deploy/postgres/goose/00020_2026-09-12-task-run-worker-lease.sql new file mode 100644 index 00000000..d6491d09 --- /dev/null +++ b/deploy/postgres/goose/00020_2026-09-12-task-run-worker-lease.sql @@ -0,0 +1,18 @@ +-- +goose Up +-- Phase 4: 长任务 worker lease / heartbeat,用于识别孤儿 running 任务 + +ALTER TABLE task_runs ADD COLUMN IF NOT EXISTS worker_id TEXT NOT NULL DEFAULT ''; +ALTER TABLE task_runs ADD COLUMN IF NOT EXISTS lease_expires_at TIMESTAMPTZ; +ALTER TABLE task_runs ADD COLUMN IF NOT EXISTS lease_heartbeat_at TIMESTAMPTZ; +ALTER TABLE task_runs ADD COLUMN IF NOT EXISTS attempt INT NOT NULL DEFAULT 1; + +CREATE INDEX IF NOT EXISTS idx_task_runs_stale_lease + ON task_runs(lease_expires_at) + WHERE status = 'running' AND lease_expires_at IS NOT NULL; + +-- +goose Down +DROP INDEX IF EXISTS idx_task_runs_stale_lease; +ALTER TABLE task_runs DROP COLUMN IF EXISTS attempt; +ALTER TABLE task_runs DROP COLUMN IF EXISTS lease_heartbeat_at; +ALTER TABLE task_runs DROP COLUMN IF EXISTS lease_expires_at; +ALTER TABLE task_runs DROP COLUMN IF EXISTS worker_id; diff --git a/internal/server/phase5_fault_injection_pg_test.go b/internal/server/phase5_fault_injection_pg_test.go new file mode 100644 index 00000000..e3d38372 --- /dev/null +++ b/internal/server/phase5_fault_injection_pg_test.go @@ -0,0 +1,257 @@ +package server_test + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "sync" + "testing" + "time" + + "github.com/jackc/pgx/v5/pgxpool" + + "chat-agent-system/internal/bizdev" + "chat-agent-system/internal/db" + "chat-agent-system/internal/projectinstance" + "chat-agent-system/internal/server" + "chat-agent-system/internal/stream" + "chat-agent-system/internal/taskrun" + "chat-agent-system/internal/workflowdef" +) + +// setupPGProjectReadyToAdvance 创建并激活工作流、完成 intake 文档批准,返回可推进实例。 +func setupPGProjectReadyToAdvance(t *testing.T, pool *pgxpool.Pool, handler http.Handler, srv *server.Server) (projectID, instanceID int64, creatorToken string, configRevision int) { + t.Helper() + ctx := context.Background() + wfSvc := workflowdef.NewService(pool, nil) + wfKey := "pg_p5_fault_" + time.Now().Format("150405.000000") + _, ver, err := wfSvc.CreateDefinition(ctx, workflowdef.CreateDefinitionInput{ + Key: wfKey, Name: "Phase5 Fault WF", + Definition: workflowdef.VersionDefinition{ + SchemaVersion: 1, + Stages: []workflowdef.StageDef{ + { + Key: "biz_intake", Name: "Intake", Position: 0, + DocTasks: []workflowdef.DocTaskDef{{ + Key: "intake", Required: true, Blocking: true, DefaultEnabled: true, + OutputDocType: "intake", OutputDocTitle: "Intake", + }}, + }, + {Key: "biz_feasibility", Name: "Feasibility", Position: 1, Optional: true}, + }, + }, + }, 1) + if err != nil { + t.Fatalf("create wf: %v", err) + } + _, _ = wfSvc.PublishVersion(ctx, ver.WorkflowID, ver.Version, 1) + + creatorToken = loginTestUser(t, handler, "zhanggong", "123456") + creatorID := fetchAuthUserID(t, handler, creatorToken) + reviewerID := fetchAuthUserID(t, handler, loginTestUser(t, handler, "lizhuguan", "123456")) + approverID := fetchAuthUserID(t, handler, loginTestUser(t, handler, "wangzong", "123456")) + + code := "p5-fault-" + time.Now().Format("150405.000000") + err = pool.QueryRow(ctx, ` + INSERT INTO projects (code, name, status, creator_user_id, lifecycle_status, active_workflow_key, workflow_selection_version, workflow_plan_revision) + VALUES ($1, 'Phase5 Fault Test', 'active', $2, 'planning', '', 1, 0) RETURNING id + `, code, creatorID).Scan(&projectID) + if err != nil { + t.Fatalf("insert project: %v", err) + } + _, _ = pool.Exec(ctx, ` + INSERT INTO project_members (project_id, user_id, project_role, status, joined_at, updated_at) + VALUES ($1, $2, 'owner', 'active', now(), now()) ON CONFLICT DO NOTHING + `, projectID, creatorID) + t.Cleanup(func() { cleanupProject(t, pool, projectID) }) + + reqDefault := httptest.NewRequest(http.MethodGet, fmt.Sprintf("/api/v1/projects/%d/workflow-plan/default?workflow_key=%s&version=1", projectID, wfKey), nil) + withBearerAuth(reqDefault, creatorToken) + recDefault := httptest.NewRecorder() + handler.ServeHTTP(recDefault, reqDefault) + var plan projectinstance.WorkflowPlan + _ = json.Unmarshal(recDefault.Body.Bytes(), &plan) + for i := range plan.Stages { + for j := range plan.Stages[i].DocTasks { + plan.Stages[i].DocTasks[j].Duties = projectinstance.DocDuty{ + AuthorUserID: creatorID, ReviewerUserID: reviewerID, ApproverUserID: approverID, + } + } + } + saveBody, _ := json.Marshal(projectinstance.SavePlanInput{Plan: plan, ExpectedPlanRevision: 0, ExpectedSelectionVersion: 1}) + reqSave := httptest.NewRequest(http.MethodPut, fmt.Sprintf("/api/v1/projects/%d/workflow-plan", projectID), bytes.NewReader(saveBody)) + withBearerAuth(reqSave, creatorToken) + recSave := httptest.NewRecorder() + handler.ServeHTTP(recSave, reqSave) + + reqPreview := httptest.NewRequest(http.MethodPost, fmt.Sprintf("/api/v1/projects/%d/activation-preview", projectID), nil) + withBearerAuth(reqPreview, creatorToken) + recPreview := httptest.NewRecorder() + handler.ServeHTTP(recPreview, reqPreview) + var preview projectinstance.ActivationPreview + _ = json.Unmarshal(recPreview.Body.Bytes(), &preview) + + activateBody, _ := json.Marshal(projectinstance.ActivateInput{ + ConfirmationToken: preview.ConfirmationToken, ExpectedProjectVersion: preview.ExpectedProjectVersion, + ExpectedPlanRevision: preview.ExpectedPlanRevision, + }) + reqActivate := httptest.NewRequest(http.MethodPost, fmt.Sprintf("/api/v1/projects/%d/workflow-instances/activate", projectID), bytes.NewReader(activateBody)) + withBearerAuth(reqActivate, creatorToken) + recActivate := httptest.NewRecorder() + handler.ServeHTTP(recActivate, reqActivate) + var activated projectinstance.ActivateResult + _ = json.Unmarshal(recActivate.Body.Bytes(), &activated) + instanceID = activated.Instance.ID + configRevision = activated.Instance.ConfigRevision + + docID := fmt.Sprintf("doc_p5_fault_%d", projectID) + docEngine := srv.BizdevService().DocGov() + _, _ = docEngine.CreateOrGetDocument(ctx, docID, projectID, "intake", "Intake", "V0.1", "", creatorID, "A", reviewerID, "R", approverID, "P") + _ = docEngine.MarkDocumentContentUpdated(ctx, docID, "hash_p5_fault", false) + _, _ = docEngine.SubmitForReview(ctx, docID, creatorID, "author", "Author") + _, _ = docEngine.ReviewApprove(ctx, docID, reviewerID, "reviewer", "Reviewer", "ok") + _, _ = docEngine.ApproverApprove(ctx, docID, approverID, "approver", "Approver", "ok") + return projectID, instanceID, creatorToken, configRevision +} + +func TestPhase5_ConcurrentStageAdvance_PG(t *testing.T) { + pool := openPGPoolWithMigrations(t) + if pool == nil { + return + } + defer pool.Close() + + handler, srv := setupPGGovernanceServer(t, pool) + _, instanceID, creatorToken, configRevision := setupPGProjectReadyToAdvance(t, pool, handler, srv) + + advanceBody, _ := json.Marshal(projectinstance.AdvanceStageInput{ + StageKey: "biz_intake", ExpectedConfigRevision: configRevision, + }) + url := fmt.Sprintf("/api/v1/workflow-instances/%d/stages/biz_intake/advance", instanceID) + + var wg sync.WaitGroup + results := make([]int, 2) + for i := 0; i < 2; i++ { + wg.Add(1) + go func(idx int) { + defer wg.Done() + req := httptest.NewRequest(http.MethodPost, url, bytes.NewReader(advanceBody)) + withBearerAuth(req, creatorToken) + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + results[idx] = rec.Code + }(i) + } + wg.Wait() + + okCount, conflictCount := 0, 0 + for _, code := range results { + switch code { + case http.StatusOK: + okCount++ + case http.StatusConflict, http.StatusBadRequest, http.StatusUnprocessableEntity: + conflictCount++ + } + } + if okCount != 1 { + t.Fatalf("expected exactly one 200 advance, got codes=%v", results) + } + if conflictCount < 1 { + t.Fatalf("expected second advance to fail with conflict/gate error, got codes=%v", results) + } +} + +func TestPhase5_TaskRunHeartbeatSurvivesOrphanScan_PG(t *testing.T) { + pool := openPGPoolWithMigrations(t) + if pool == nil { + return + } + defer pool.Close() + if !taskrun.LeaseColumnsReady(context.Background(), pool) { + t.Skip("lease columns not ready") + } + + handler, srv := setupPGGovernanceServer(t, pool) + token := loginTestUser(t, handler, "zhanggong", "123456") + userID := fetchAuthUserID(t, handler, token) + + runID := "run_hb_survive_" + time.Now().Format("150405.000000") + store := taskrun.NewPGStore(pool) + done := make(chan struct{}) + srv.TaskRuns().SetExecutor( + func(ctx context.Context, _ string, _ bizdev.TaskRunRequest, _ *stream.Streamer) (*bizdev.TaskRunResult, error) { + <-done + return &bizdev.TaskRunResult{Success: true, TaskKey: "intake"}, nil + }, + nil, nil, + ) + srv.TaskRuns().SetFinalizer(func(ctx context.Context, record *taskrun.RunRecord, result *bizdev.TaskRunResult) (string, error) { + return "completed", nil + }) + + record := &taskrun.RunRecord{ + RunID: runID, SessionID: "sess_hb", UserID: userID, TaskKey: "intake", Status: "running", + } + _, err := srv.TaskRuns().CreateAndStart(context.Background(), runID, record, bizdev.TaskRunRequest{TaskKey: "intake"}, bizdev.StageRunRequest{}, "") + if err != nil { + t.Fatalf("CreateAndStart: %v", err) + } + t.Cleanup(func() { close(done) }) + + time.Sleep(150 * time.Millisecond) + if n, err := store.MarkOrphanedRunning(context.Background()); err != nil { + t.Fatalf("orphan scan: %v", err) + } else if n > 0 { + rec, _ := store.GetRun(context.Background(), runID) + if rec != nil && rec.Status == "interrupted" { + t.Fatal("active heartbeat run should not be marked interrupted") + } + } + + rec, err := store.GetRun(context.Background(), runID) + if err != nil { + t.Fatalf("GetRun: %v", err) + } + if rec.Status != "running" { + t.Fatalf("expected running with fresh lease, got %s", rec.Status) + } +} + +func TestPhase5_BackupRestoreMigration_PG(t *testing.T) { + pool := openPGPoolWithMigrations(t) + if pool == nil { + return + } + defer pool.Close() + + ctx := context.Background() + migDir := db.ResolveMigrationsDir() + v1, err := db.CurrentVersion(ctx, pool) + if err != nil { + t.Fatalf("version: %v", err) + } + expected, _ := db.ExpectedMigrationCount(migDir) + if int(v1) < expected { + t.Fatalf("schema %d behind %d files", v1, expected) + } + + // 模拟备份恢复后再次 Up(空操作应幂等) + if err := db.RunMigrations(ctx, pool, migDir); err != nil { + t.Fatalf("restore migration: %v", err) + } + v2, _ := db.CurrentVersion(ctx, pool) + if v2 != v1 { + t.Fatalf("migration version changed after idempotent restore: %d -> %d", v1, v2) + } + + var taskRunsExists bool + _ = pool.QueryRow(ctx, ` + SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = 'task_runs') + `).Scan(&taskRunsExists) + if !taskRunsExists { + t.Fatal("task_runs missing after restore simulation") + } +} diff --git a/internal/taskrun/manager.go b/internal/taskrun/manager.go index a237bfa8..22d0b5bf 100644 --- a/internal/taskrun/manager.go +++ b/internal/taskrun/manager.go @@ -30,6 +30,7 @@ type Manager struct { mu sync.RWMutex runs map[string]*activeRun store Store + workerID string exec Executor resume ResumeExecutor stage StageExecutor @@ -42,16 +43,17 @@ const runRetention = time.Hour func NewManager(pwSvc *projectworkflow.Service, ctxMgr *ctxmgr.ContextManager) *Manager { m := &Manager{ - runs: make(map[string]*activeRun), - store: noopStore{}, - pwSvc: pwSvc, - ctxMgr: ctxMgr, + runs: make(map[string]*activeRun), + store: noopStore{}, + workerID: newWorkerID(), + pwSvc: pwSvc, + ctxMgr: ctxMgr, } go m.cleanupLoop() return m } -// SetStore 注入持久化存储(PG);启动时标记孤儿 running 任务为 interrupted。 +// SetStore 注入持久化存储(PG);启动时标记 lease 过期的 running 任务为 interrupted。 func (m *Manager) SetStore(store Store) { if m == nil { return @@ -116,10 +118,16 @@ func (m *Manager) CreateAndStart(parent context.Context, runID string, record *R now := time.Now() record.CreatedAt = now record.UpdatedAt = now + record.WorkerID = m.workerID + record.Attempt = 1 + leaseUntil := now.Add(defaultLeaseTTL) + record.LeaseExpiresAt = &leaseUntil if m.store != nil && m.store.Available() { if err := m.store.InsertRun(parent, record); err != nil { slog.Warn("taskrun_insert_failed", "run_id", runID, "error", err) + } else if err := m.store.AcquireLease(parent, runID, m.workerID, defaultLeaseTTL); err != nil { + slog.Warn("taskrun_lease_acquire_failed", "run_id", runID, "error", err) } } @@ -159,6 +167,12 @@ func (m *Manager) CreateAndStart(parent context.Context, runID string, record *R func (m *Manager) runWorker(ctx context.Context, ar *activeRun, taskReq bizdev.TaskRunRequest, stageReq bizdev.StageRunRequest, resumeWorkDir string) { defer close(ar.done) + hbCtx, hbCancel := context.WithCancel(ctx) + defer hbCancel() + if m.store != nil && m.store.Available() { + go m.heartbeatLoop(hbCtx, ar.record.RunID) + } + taskReq.RunID = ar.record.RunID taskReq.SessionID = ar.record.SessionID taskReq.ProjectID = ar.record.ProjectID @@ -265,6 +279,13 @@ func (m *Manager) finishRun(ar *activeRun, status, workspaceID, errMsg string, r "workspace_id", workspaceID, "error", errMsg, ) + if m.store != nil && m.store.Available() { + _ = m.store.ReleaseLease(context.Background(), runID, m.workerID) + } + ar.mu.Lock() + ar.record.WorkerID = "" + ar.record.LeaseExpiresAt = nil + ar.mu.Unlock() m.persistRun(ar) // 失败/取消路径:若 Finalize 未执行,仍更新 workflow_runs 与占位消息 metadata @@ -411,3 +432,18 @@ func (m *Manager) Cancel(runID string) error { ar.cancel() return nil } + +func (m *Manager) heartbeatLoop(ctx context.Context, runID string) { + ticker := time.NewTicker(defaultHeartbeatPeriod) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + if err := m.store.HeartbeatLease(context.Background(), runID, m.workerID, defaultLeaseTTL); err != nil { + slog.Warn("taskrun_heartbeat_failed", "run_id", runID, "error", err) + } + } + } +} diff --git a/internal/taskrun/store.go b/internal/taskrun/store.go index 83c193d2..0d51132d 100644 --- a/internal/taskrun/store.go +++ b/internal/taskrun/store.go @@ -23,6 +23,9 @@ type Store interface { AppendEvent(ctx context.Context, runID string, seq int64, ev stream.StreamEvent) error GetRun(ctx context.Context, runID string) (*RunRecord, error) ListEvents(ctx context.Context, runID string, afterSeq int64) ([]StoredEvent, error) + AcquireLease(ctx context.Context, runID, workerID string, ttl time.Duration) error + HeartbeatLease(ctx context.Context, runID, workerID string, ttl time.Duration) error + ReleaseLease(ctx context.Context, runID, workerID string) error MarkOrphanedRunning(ctx context.Context) (int, error) } @@ -49,6 +52,17 @@ func (s *PGStore) Available() bool { return ok } +func (s *PGStore) hasLeaseColumns(ctx context.Context) bool { + var ok bool + _ = s.pool.QueryRow(ctx, ` + SELECT EXISTS ( + SELECT 1 FROM information_schema.columns + WHERE table_schema = 'public' AND table_name = 'task_runs' AND column_name = 'worker_id' + ) + `).Scan(&ok) + return ok +} + func (s *PGStore) InsertRun(ctx context.Context, rec *RunRecord) error { if !s.Available() { return ErrStoreNotAvailable @@ -56,6 +70,32 @@ func (s *PGStore) InsertRun(ctx context.Context, rec *RunRecord) error { if rec == nil || rec.RunID == "" { return fmt.Errorf("taskrun: run_id required") } + attempt := rec.Attempt + if attempt <= 0 { + attempt = 1 + } + if s.hasLeaseColumns(ctx) { + var leaseHeartbeat *time.Time + if rec.LeaseExpiresAt != nil { + now := time.Now() + leaseHeartbeat = &now + } + _, err := s.pool.Exec(ctx, ` + INSERT INTO task_runs ( + run_id, session_id, project_id, user_id, client_turn_id, + task_key, stage, status, workspace_id, assistant_message_id, + error_message, worker_id, lease_expires_at, lease_heartbeat_at, attempt, + created_at, updated_at + ) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17) + ON CONFLICT (run_id) DO NOTHING + `, + rec.RunID, rec.SessionID, nullInt64(rec.ProjectID), rec.UserID, rec.ClientTurnID, + rec.TaskKey, rec.Stage, rec.Status, rec.WorkspaceID, nullInt64(rec.AssistantMsgID), + rec.ErrorMessage, rec.WorkerID, rec.LeaseExpiresAt, leaseHeartbeat, attempt, + rec.CreatedAt, rec.UpdatedAt, + ) + return err + } _, err := s.pool.Exec(ctx, ` INSERT INTO task_runs ( run_id, session_id, project_id, user_id, client_turn_id, @@ -78,6 +118,19 @@ func (s *PGStore) UpdateRun(ctx context.Context, rec *RunRecord) error { if rec == nil || rec.RunID == "" { return fmt.Errorf("taskrun: run_id required") } + if s.hasLeaseColumns(ctx) { + _, err := s.pool.Exec(ctx, ` + UPDATE task_runs SET + status = $2, + workspace_id = $3, + error_message = $4, + worker_id = $5, + lease_expires_at = $6, + updated_at = $7 + WHERE run_id = $1 + `, rec.RunID, rec.Status, rec.WorkspaceID, rec.ErrorMessage, rec.WorkerID, rec.LeaseExpiresAt, rec.UpdatedAt) + return err + } _, err := s.pool.Exec(ctx, ` UPDATE task_runs SET status = $2, @@ -111,22 +164,43 @@ func (s *PGStore) GetRun(ctx context.Context, runID string) (*RunRecord, error) } var rec RunRecord var projectID, assistantMsgID *int64 - err := s.pool.QueryRow(ctx, ` - SELECT run_id, session_id, project_id, user_id, client_turn_id, - task_key, stage, status, workspace_id, assistant_message_id, - error_message, created_at, updated_at - FROM task_runs - WHERE run_id = $1 - `, runID).Scan( - &rec.RunID, &rec.SessionID, &projectID, &rec.UserID, &rec.ClientTurnID, - &rec.TaskKey, &rec.Stage, &rec.Status, &rec.WorkspaceID, &assistantMsgID, - &rec.ErrorMessage, &rec.CreatedAt, &rec.UpdatedAt, - ) - if err != nil { - if errors.Is(err, pgx.ErrNoRows) { - return nil, fmt.Errorf("run not found") + if s.hasLeaseColumns(ctx) { + err := s.pool.QueryRow(ctx, ` + SELECT run_id, session_id, project_id, user_id, client_turn_id, + task_key, stage, status, workspace_id, assistant_message_id, + error_message, worker_id, lease_expires_at, attempt, created_at, updated_at + FROM task_runs + WHERE run_id = $1 + `, runID).Scan( + &rec.RunID, &rec.SessionID, &projectID, &rec.UserID, &rec.ClientTurnID, + &rec.TaskKey, &rec.Stage, &rec.Status, &rec.WorkspaceID, &assistantMsgID, + &rec.ErrorMessage, &rec.WorkerID, &rec.LeaseExpiresAt, &rec.Attempt, + &rec.CreatedAt, &rec.UpdatedAt, + ) + if err != nil { + if errors.Is(err, pgx.ErrNoRows) { + return nil, fmt.Errorf("run not found") + } + return nil, err + } + } else { + err := s.pool.QueryRow(ctx, ` + SELECT run_id, session_id, project_id, user_id, client_turn_id, + task_key, stage, status, workspace_id, assistant_message_id, + error_message, created_at, updated_at + FROM task_runs + WHERE run_id = $1 + `, runID).Scan( + &rec.RunID, &rec.SessionID, &projectID, &rec.UserID, &rec.ClientTurnID, + &rec.TaskKey, &rec.Stage, &rec.Status, &rec.WorkspaceID, &assistantMsgID, + &rec.ErrorMessage, &rec.CreatedAt, &rec.UpdatedAt, + ) + if err != nil { + if errors.Is(err, pgx.ErrNoRows) { + return nil, fmt.Errorf("run not found") + } + return nil, err } - return nil, err } if projectID != nil { rec.ProjectID = *projectID @@ -168,10 +242,95 @@ func (s *PGStore) ListEvents(ctx context.Context, runID string, afterSeq int64) return out, rows.Err() } +func (s *PGStore) AcquireLease(ctx context.Context, runID, workerID string, ttl time.Duration) error { + if !s.Available() { + return ErrStoreNotAvailable + } + if !s.hasLeaseColumns(ctx) { + return nil + } + expires := time.Now().Add(ttl) + tag, err := s.pool.Exec(ctx, ` + UPDATE task_runs + SET worker_id = $2, + lease_expires_at = $3, + lease_heartbeat_at = now(), + updated_at = now() + WHERE run_id = $1 AND status = 'running' + `, runID, workerID, expires) + if err != nil { + return err + } + if tag.RowsAffected() == 0 { + return fmt.Errorf("taskrun: lease acquire failed for %s", runID) + } + return nil +} + +func (s *PGStore) HeartbeatLease(ctx context.Context, runID, workerID string, ttl time.Duration) error { + if !s.Available() { + return ErrStoreNotAvailable + } + if !s.hasLeaseColumns(ctx) { + return nil + } + expires := time.Now().Add(ttl) + tag, err := s.pool.Exec(ctx, ` + UPDATE task_runs + SET lease_expires_at = $3, + lease_heartbeat_at = now(), + updated_at = now() + WHERE run_id = $1 AND worker_id = $2 AND status = 'running' + `, runID, workerID, expires) + if err != nil { + return err + } + if tag.RowsAffected() == 0 { + return fmt.Errorf("taskrun: heartbeat rejected for %s", runID) + } + return nil +} + +func (s *PGStore) ReleaseLease(ctx context.Context, runID, workerID string) error { + if !s.Available() { + return ErrStoreNotAvailable + } + if !s.hasLeaseColumns(ctx) { + return nil + } + _, err := s.pool.Exec(ctx, ` + UPDATE task_runs + SET worker_id = '', + lease_expires_at = NULL, + lease_heartbeat_at = NULL, + updated_at = now() + WHERE run_id = $1 AND (worker_id = $2 OR worker_id = '') + `, runID, workerID) + return err +} + func (s *PGStore) MarkOrphanedRunning(ctx context.Context) (int, error) { if !s.Available() { return 0, ErrStoreNotAvailable } + if s.hasLeaseColumns(ctx) { + tag, err := s.pool.Exec(ctx, ` + UPDATE task_runs + SET status = 'interrupted', + error_message = 'worker lease 过期,任务已中断(进程重启或 worker 失联)', + worker_id = '', + lease_expires_at = NULL, + lease_heartbeat_at = NULL, + updated_at = now() + WHERE status = 'running' + AND lease_expires_at IS NOT NULL + AND lease_expires_at < now() - interval '2 seconds' + `) + if err != nil { + return 0, err + } + return int(tag.RowsAffected()), nil + } tag, err := s.pool.Exec(ctx, ` UPDATE task_runs SET status = 'interrupted', @@ -207,6 +366,21 @@ func TablesReady(ctx context.Context, pool *pgxpool.Pool) bool { return ok } +// LeaseColumnsReady 检查 worker lease 列是否已迁移。 +func LeaseColumnsReady(ctx context.Context, pool *pgxpool.Pool) bool { + if pool == nil { + return false + } + var ok bool + _ = pool.QueryRow(ctx, ` + SELECT EXISTS ( + SELECT 1 FROM information_schema.columns + WHERE table_schema = 'public' AND table_name = 'task_runs' AND column_name = 'worker_id' + ) + `).Scan(&ok) + return ok +} + // noopStore 无持久化时的空实现。 type noopStore struct{} @@ -224,7 +398,14 @@ func (noopStore) GetRun(context.Context, string) (*RunRecord, error) { func (noopStore) ListEvents(context.Context, string, int64) ([]StoredEvent, error) { return nil, ErrStoreNotAvailable } -func (noopStore) MarkOrphanedRunning(context.Context) (int, error) { return 0, ErrStoreNotAvailable } +func (noopStore) AcquireLease(context.Context, string, string, time.Duration) error { + return ErrStoreNotAvailable +} +func (noopStore) HeartbeatLease(context.Context, string, string, time.Duration) error { + return ErrStoreNotAvailable +} +func (noopStore) ReleaseLease(context.Context, string, string) error { return ErrStoreNotAvailable } +func (noopStore) MarkOrphanedRunning(context.Context) (int, error) { return 0, ErrStoreNotAvailable } // EnsureUpdatedAt 若 UpdatedAt 为零则设为当前时间。 func EnsureUpdatedAt(rec *RunRecord) { diff --git a/internal/taskrun/store_lease_pg_test.go b/internal/taskrun/store_lease_pg_test.go new file mode 100644 index 00000000..78e59808 --- /dev/null +++ b/internal/taskrun/store_lease_pg_test.go @@ -0,0 +1,133 @@ +package taskrun_test + +import ( + "context" + "os" + "testing" + "time" + + "github.com/jackc/pgx/v5/pgxpool" + + "chat-agent-system/internal/db" + "chat-agent-system/internal/taskrun" +) + +func openLeaseTestPool(t *testing.T) *pgxpool.Pool { + t.Helper() + dsn := os.Getenv("TEST_DATABASE_URL") + if dsn == "" { + dsn = "postgres://app:dev_secret_db_pass@127.0.0.1:5433/agent_system_dev?sslmode=disable" + } + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + pool, err := pgxpool.New(ctx, dsn) + if err != nil { + t.Skipf("postgres unavailable: %v", err) + return nil + } + if err := pool.Ping(ctx); err != nil { + pool.Close() + t.Skipf("postgres ping failed: %v", err) + return nil + } + migCtx, migCancel := context.WithTimeout(context.Background(), 60*time.Second) + defer migCancel() + if err := db.RunMigrations(migCtx, pool, db.ResolveMigrationsDir()); err != nil { + pool.Close() + t.Skipf("migration failed: %v", err) + return nil + } + if !taskrun.LeaseColumnsReady(ctx, pool) { + pool.Close() + t.Skip("task run lease columns not ready") + } + return pool +} + +func TestPGStore_StaleLeaseMarkedInterrupted(t *testing.T) { + pool := openLeaseTestPool(t) + if pool == nil { + return + } + defer pool.Close() + + ctx := context.Background() + store := taskrun.NewPGStore(pool) + runID := "run_stale_lease_" + time.Now().Format("150405.000000") + + var userID int64 + err := pool.QueryRow(ctx, `SELECT id FROM users ORDER BY id LIMIT 1`).Scan(&userID) + if err != nil { + t.Fatalf("fetch user: %v", err) + } + + expired := time.Now().Add(-1 * time.Minute) + _, err = pool.Exec(ctx, ` + INSERT INTO task_runs ( + run_id, session_id, user_id, status, worker_id, lease_expires_at, lease_heartbeat_at, created_at, updated_at + ) VALUES ($1, 'sess', $2, 'running', 'dead-worker', $3, $3, now(), now()) + `, runID, userID, expired) + if err != nil { + t.Fatalf("insert stale run: %v", err) + } + t.Cleanup(func() { + _, _ = pool.Exec(context.Background(), `DELETE FROM task_runs WHERE run_id = $1`, runID) + }) + + n, err := store.MarkOrphanedRunning(ctx) + if err != nil { + t.Fatalf("MarkOrphanedRunning: %v", err) + } + if n < 1 { + t.Fatalf("expected at least 1 interrupted, got %d", n) + } + + rec, err := store.GetRun(ctx, runID) + if err != nil { + t.Fatalf("GetRun: %v", err) + } + if rec.Status != "interrupted" { + t.Fatalf("status=%s want interrupted", rec.Status) + } +} + +func TestPGStore_FreshLeaseSurvivesOrphanScan(t *testing.T) { + pool := openLeaseTestPool(t) + if pool == nil { + return + } + defer pool.Close() + + ctx := context.Background() + store := taskrun.NewPGStore(pool) + runID := "run_fresh_lease_" + time.Now().Format("150405.000000") + workerID := "worker-alive" + + var userID int64 + _ = pool.QueryRow(ctx, `SELECT id FROM users ORDER BY id LIMIT 1`).Scan(&userID) + + fresh := time.Now().Add(30 * time.Second) + _, err := pool.Exec(ctx, ` + INSERT INTO task_runs ( + run_id, session_id, user_id, status, worker_id, lease_expires_at, lease_heartbeat_at, created_at, updated_at + ) VALUES ($1, 'sess', $2, 'running', $3, $4, now(), now(), now()) + `, runID, userID, workerID, fresh) + if err != nil { + t.Fatalf("insert fresh run: %v", err) + } + t.Cleanup(func() { + _, _ = pool.Exec(context.Background(), `DELETE FROM task_runs WHERE run_id = $1`, runID) + }) + + if _, err := store.MarkOrphanedRunning(ctx); err != nil { + t.Fatalf("MarkOrphanedRunning: %v", err) + } + + rec, err := store.GetRun(ctx, runID) + if err != nil { + t.Fatalf("GetRun: %v", err) + } + if rec.Status != "running" { + t.Fatalf("fresh lease run should stay running, got %s", rec.Status) + } +} diff --git a/internal/taskrun/types.go b/internal/taskrun/types.go index b60b5c21..74d0acbd 100644 --- a/internal/taskrun/types.go +++ b/internal/taskrun/types.go @@ -35,9 +35,12 @@ type RunRecord struct { Status string // pending | running | completed | failed | cancelled | status_conflict WorkspaceID string AssistantMsgID int64 + WorkerID string + LeaseExpiresAt *time.Time CreatedAt time.Time UpdatedAt time.Time ErrorMessage string + Attempt int } // StoredEvent 持久化 SSE 事件(支持重连重放)。 diff --git a/internal/taskrun/worker.go b/internal/taskrun/worker.go new file mode 100644 index 00000000..e2e3a308 --- /dev/null +++ b/internal/taskrun/worker.go @@ -0,0 +1,20 @@ +package taskrun + +import ( + "fmt" + "os" + "time" +) + +const ( + defaultLeaseTTL = 30 * time.Second + defaultHeartbeatPeriod = 10 * time.Second +) + +func newWorkerID() string { + host, err := os.Hostname() + if err != nil || host == "" { + host = "local" + } + return fmt.Sprintf("%s-%d-%d", host, os.Getpid(), time.Now().UnixNano()) +}