package workbench import ( "context" "errors" "fmt" "github.com/Tencent/WeKnora/internal/accesspolicy" "runtime/debug" "time" "github.com/Tencent/WeKnora/internal/logger" "github.com/Tencent/WeKnora/internal/types" "github.com/google/uuid" "gorm.io/gorm" "gorm.io/gorm/clause" ) func requireIdle(tx *gorm.DB, tenant uint64, id string) error { var count int64 e := tx.Model(&Job{}).Where("tenant_id = ? AND resource_id = ? AND status IN ?", tenant, id, []string{"queued", "running"}).Count(&count).Error if e != nil { return e } if count > 0 { return ErrConflict } return nil } func allowedAction(kind, action string) bool { actions := map[string][]string{"training-maps": {"generate", "revise", "generate-day", "confirm-day", "confirm"}, "content-tasks": {"decompose", "import-image", "generate", "revise", "accept", "skip", "review", "review-settings", "compose"}} for _, a := range actions[kind] { if a == action { return true } } return false } func (s *Service) submit(ctx context.Context, kind, id, action string, request ActionRequest) (*Job, error) { if !allowedAction(kind, action) { return nil, ErrInvalid } if !accesspolicy.Allowed(ctx, accesspolicy.WorkbenchActionPermission(kind, action)) { return nil, ErrForbidden } if len([]rune(request.Advice)) > 6000 || len([]rune(request.Quote)) > 6000 { return nil, ErrInvalid } if action == "revise" && request.Advice == "" { return nil, fmt.Errorf("%w: 请填写修改建议", ErrInvalid) } if action == "review-settings" { if request.QualitySettings == nil { return nil, fmt.Errorf("%w: 请填写质检设置", ErrInvalid) } if err := request.QualitySettings.validate(); err != nil { return nil, err } } t, u, _, e := identity(ctx) if e != nil { return nil, e } j := &Job{ID: uuid.NewString(), TenantID: t, OwnerID: u, ResourceID: id, Kind: kind, Action: action, Status: "queued", Stage: "等待执行", Request: asJSON(request), CreatedAt: time.Now(), UpdatedAt: time.Now()} e = s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { current, err := lockResource(ctx, tx, kind, id) if err != nil { return err } // Generation is billed to its owner; administrator inspection and manual // edits must not silently enqueue billable work for somebody else. if current.OwnerID != u { return ErrForbidden } if current.Version != request.Version { return ErrConflict } if err = requireActionExecutable(current, action); err != nil { return err } if kind == "content-tasks" { if err = s.validateTaskAction(ctx, current, action); err != nil { return err } } else if err = s.validate(ctx, current); err != nil { return err } if err = requireIdle(tx, t, id); err != nil { return err } return tx.Create(j).Error }) if e != nil { return nil, e } _ = s.event(ctx, id, action, "提交任务:"+actionName(action), "info") return j, nil } // Start uses database leases so multiple service instances cannot claim the same // job. Expired jobs require an explicit retry: a provider request may have succeeded. func (s *Service) Start(ctx context.Context) { go func() { ticker := time.NewTicker(2 * time.Second) defer ticker.Stop() for { select { case <-ctx.Done(): return case <-ticker.C: s.tick(ctx) } } }() } func (s *Service) tick(ctx context.Context) { if err := s.releaseExpiredEvidence(ctx); err != nil { logJobFailure(ctx, nil, "evidence_recovery", err) } if err := s.releaseExpiredBalances(ctx); err != nil { logJobFailure(ctx, nil, "balance_recovery", err) } now := time.Now() if err := s.recoverExpiredJobs(ctx, now); err != nil { logJobFailure(ctx, nil, "lease_recovery", err) return } j, err := s.claim(ctx) if err != nil { logJobFailure(ctx, nil, "claim", err) return } if j == nil { return } s.run(ctx, j) } func (s *Service) claim(ctx context.Context) (*Job, error) { var j Job err := s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { q := tx.Clauses(clause.Locking{Strength: "UPDATE", Options: "SKIP LOCKED"}).Where("status = ?", "queued").Order("created_at ASC, id ASC").Limit(1).Find(&j) if q.Error != nil || q.RowsAffected == 0 { return q.Error } now := time.Now() until := now.Add(10 * time.Minute) j.Status = "running" j.Progress = 5 j.Stage = "检查任务配置" j.LeaseToken = uuid.NewString() j.LeaseUntil = &until q = tx.Model(&Job{}).Where("id = ? AND status = ?", j.ID, "queued").Updates(map[string]any{"status": j.Status, "progress": j.Progress, "stage": j.Stage, "lease_token": j.LeaseToken, "lease_until": until, "updated_at": now}) if q.Error != nil { return q.Error } if q.RowsAffected != 1 { return ErrConflict } return nil }) if err != nil { return nil, err } if j.ID == "" { return nil, nil } return &j, nil } type jobLease struct{ ID, Token string } type leaseContextKey struct{} func (s *Service) progress(ctx context.Context, id string, percent int, stage string) error { lease, ok := ctx.Value(leaseContextKey{}).(jobLease) if !ok || lease.ID != id { return errors.New("missing job lease") } q := s.db.WithContext(ctx).Model(&Job{}).Where("id = ? AND lease_token = ? AND status = ? AND lease_until > ?", id, lease.Token, "running", time.Now()).Updates(map[string]any{"progress": percent, "stage": stage, "updated_at": time.Now()}) if q.Error != nil { return q.Error } if q.RowsAffected != 1 { return ErrConflict } return nil } func (s *Service) run(parent context.Context, j *Job) { ctx, cancel := context.WithTimeout(parent, 8*time.Minute) defer cancel() ctx = context.WithValue(ctx, types.TenantIDContextKey, j.TenantID) ctx = context.WithValue(ctx, types.UserIDContextKey, j.OwnerID) ctx = context.WithValue(ctx, leaseContextKey{}, jobLease{j.ID, j.LeaseToken}) ctx = types.WithBackgroundTask(ctx) go s.watchJobLease(ctx, cancel, j) var runErr error defer func() { if panicValue := recover(); panicValue != nil { runErr = errors.New("后台任务发生异常,请重试") logger.GetLogger(ctx).WithFields(logger.Fields{"job_id": j.ID, "panic_type": fmt.Sprintf("%T", panicValue), "stack": string(debug.Stack())}).Error("工作台任务发生异常") } finishCtx, stop := context.WithTimeout(context.WithoutCancel(ctx), 10*time.Second) defer stop() status, stage, message := "succeeded", "已完成", "" updates := map[string]any{"status": status, "stage": stage, "progress": 100, "error": message, "updated_at": time.Now()} if runErr != nil { if err := s.releaseJobEvidence(finishCtx, j); err != nil { logJobFailure(finishCtx, j, "evidence_release", err) } status = "failed" stage = "执行失败" message = publicJobError(runErr) logJobFailure(ctx, j, "execute", runErr) updates["status"], updates["stage"], updates["error"] = status, stage, message // Preserve the latest persisted stage progress instead of resetting to 5%. delete(updates, "progress") } q := s.db.WithContext(finishCtx).Model(&Job{}).Where("id = ? AND lease_token = ? AND status = ? AND lease_until > ?", j.ID, j.LeaseToken, "running", time.Now()).Updates(updates) if q.Error != nil { logJobFailure(finishCtx, j, "finalize", q.Error) s.recordJobFailure(finishCtx, j, "finalize", q.Error) return } if q.RowsAffected != 1 { var current Job if err := s.db.WithContext(finishCtx).Where("id=?", j.ID).Take(¤t).Error; err == nil && (current.Status == "paused" || current.Status == "cancelled") { return } logJobFailure(finishCtx, j, "finalize_lease", ErrConflict) return } level := "info" if runErr != nil { level = "error" s.recordJobFailure(finishCtx, j, "execute", runErr) } if err := s.event(finishCtx, j.ResourceID, j.Action, stage+":"+message, level); err != nil { logJobFailure(finishCtx, j, "event", err) } }() // Re-read membership on execution; queued work does not preserve revoked roles. var member types.TenantMember runErr = s.db.WithContext(ctx).Where("user_id = ? AND tenant_id = ? AND status = ?", j.OwnerID, j.TenantID, types.TenantMemberStatusActive).First(&member).Error if errors.Is(runErr, gorm.ErrRecordNotFound) || (runErr == nil && !member.Role.HasPermission(types.TenantRoleContributor)) { runErr = ErrForbidden return } if runErr != nil { return } ctx = context.WithValue(ctx, types.TenantRoleContextKey, member.Role) ctx, runErr = s.restoreJobTenant(ctx, j.TenantID) if runErr != nil { return } if s.accessSnapshot != nil { var snapshot *accesspolicy.Snapshot snapshot, runErr = s.accessSnapshot(ctx) if runErr != nil { return } ctx = accesspolicy.WithSnapshot(ctx, snapshot) if !accesspolicy.Allowed(ctx, accesspolicy.WorkbenchActionPermission(j.Kind, j.Action)) { runErr = ErrForbidden return } } r, e := s.get(ctx, j.Kind, j.ResourceID) if e != nil { runErr = e return } request, e := decode[ActionRequest](j.Request) if e != nil { runErr = e return } if r.Version != request.Version { runErr = ErrConflict return } if runErr = requireActionExecutable(r, j.Action); runErr != nil { return } if r.Kind == "content-tasks" { if runErr = s.validateTaskAction(ctx, r, j.Action); runErr != nil { return } } else if runErr = s.validate(ctx, r); runErr != nil { return } switch r.Kind { case "training-maps": runErr = s.executeMap(ctx, j, r, request) case "content-tasks": runErr = s.executeTask(ctx, j, r, request) default: runErr = ErrInvalid } if runErr == nil { runErr = s.save(ctx, r, request.Version) } } func publicJobError(e error) string { _, message := jobErrorDetails(e) return message }