Agent-System/internal/projectinstance/store_pg_activate.go

package projectinstance

import (
	"context"
	"errors"
	"fmt"
	"time"

	"github.com/jackc/pgx/v5"

	"chat-agent-system/internal/workflowdef"
)

func (s *pgStore) ActivateInstance(ctx context.Context, projectID int64, inst *WorkflowInstance, version *workflowdef.Version, expectedVersion int, activatedBy int64) (*WorkflowInstance, error) {
	return s.activateInstanceTx(ctx, s.pool, "", projectID, inst, version, expectedVersion, activatedBy, "")
}

func (s *pgStore) ActivateInstanceWithToken(ctx context.Context, token string, projectID int64, inst *WorkflowInstance, version *workflowdef.Version, expectedVersion int, activatedBy int64, planHash string) (*WorkflowInstance, error) {
	return s.activateInstanceTx(ctx, s.pool, token, projectID, inst, version, expectedVersion, activatedBy, planHash)
}

func (s *pgStore) activateInstanceTx(ctx context.Context, db pgxDB, token string, projectID int64, inst *WorkflowInstance, version *workflowdef.Version, expectedVersion int, activatedBy int64, planHash string) (*WorkflowInstance, error) {
	workflowKey := ""
	if version != nil {
		workflowKey = version.WorkflowKey
	}
	tx, err := db.Begin(ctx)
	if err != nil {
		return nil, err
	}
	defer tx.Rollback(ctx)

	usesAuthorOnly := planUsesAuthorOnly(inst.ConfigSnapshot)
	var currentPolicyRev int64
	if usesAuthorOnly {
		snap, lockErr, ok := runActivateAuthorOnlyLock(ctx, tx)
		if !ok {
			return nil, ErrAuthorOnlyPolicyUnavailable
		}
		if lockErr != nil {
			return nil, fmt.Errorf("%w: %v", ErrAuthorOnlyPolicyUnavailable, lockErr)
		}
		if !snap.Allowed {
			return nil, ErrAuthorOnlyPolicyDisabled
		}
		currentPolicyRev = snap.Revision
	}

	if err := runActivateMemberRecheck(ctx); err != nil {
		return nil, err
	}
	if token != "" {
		tokenPolicyRev, err := consumeActivationTokenTx(ctx, tx, token, projectID, planHash)
		if err != nil {
			return nil, err
		}
		if usesAuthorOnly {
			if tokenPolicyRev == nil || *tokenPolicyRev != currentPolicyRev {
				return nil, ErrAuthorOnlyPolicyChanged
			}
		}
	}

	tag, err := tx.Exec(ctx, `
		UPDATE projects SET lifecycle_status = 'active', active_workflow_key = $2,
		       workflow_selection_version = workflow_selection_version + 1,
		       active_workflow_changed_at = now(), active_workflow_changed_by = $4,
		       updated_at = now()
		WHERE id = $1 AND workflow_selection_version = $3
	`, projectID, workflowKey, expectedVersion, activatedBy)
	if err != nil {
		return nil, err
	}
	if tag.RowsAffected() == 0 {
		return nil, ErrConcurrentConflict
	}

	now := time.Now()
	row := tx.QueryRow(ctx, `
		UPDATE project_workflow_instances SET status = 'active', activated_by = $2, activated_at = $3, updated_at = $3
		WHERE id = $1 AND status = 'planning'
		RETURNING id, project_id, workflow_version_id, instance_no, status, config_revision, plan_hash,
		          config_snapshot, activated_by, activated_at, created_by, created_at, updated_at
	`, inst.ID, activatedBy, now)
	activated, err := scanInstance(row)
	if err != nil {
		return nil, err
	}

	_, err = tx.Exec(ctx, `
		UPDATE project_stage_instances SET status = CASE WHEN enabled THEN 'active' ELSE 'skipped' END, updated_at = now()
		WHERE workflow_instance_id = $1 AND stage_key = (
			SELECT stage_key FROM project_stage_instances
			WHERE workflow_instance_id = $1 AND enabled = true
			ORDER BY id ASC LIMIT 1
		)
	`, inst.ID)
	if err != nil {
		return nil, err
	}
	_, err = tx.Exec(ctx, `
		UPDATE project_stage_instances SET status = 'pending', updated_at = now()
		WHERE workflow_instance_id = $1 AND status <> 'active' AND enabled = true
	`, inst.ID)
	if err != nil {
		return nil, err
	}

	if version != nil {
		if err := s.applyRuntimeMetadata(ctx, tx, inst.ID, version.Definition); err != nil {
			return nil, err
		}
	}

	if err := tx.Commit(ctx); err != nil {
		return nil, err
	}
	return activated, nil
}

type pgxDB interface {
	Begin(ctx context.Context) (pgx.Tx, error)
}

func consumeActivationTokenTx(ctx context.Context, tx pgx.Tx, token string, projectID int64, planHash string) (*int64, error) {
	tokenHash := hashActivationToken(token)
	var storedHash string
	var usedAt *time.Time
	var expiresAt time.Time
	var policyRev *int64
	err := tx.QueryRow(ctx, `
		SELECT plan_hash, used_at, expires_at, author_only_policy_revision
		FROM workflow_activation_tokens WHERE token = $1 AND project_id = $2
	`, tokenHash, projectID).Scan(&storedHash, &usedAt, &expiresAt, &policyRev)
	if errors.Is(err, pgx.ErrNoRows) {
		return nil, ErrTokenInvalid
	}
	if err != nil {
		return nil, err
	}
	if usedAt != nil || time.Now().After(expiresAt) || storedHash != planHash {
		return nil, ErrTokenInvalid
	}
	tag, err := tx.Exec(ctx, `
		UPDATE workflow_activation_tokens SET used_at = now() WHERE token = $1 AND used_at IS NULL
	`, tokenHash)
	if err != nil {
		return nil, err
	}
	if tag.RowsAffected() == 0 {
		return nil, ErrTokenInvalid
	}
	return policyRev, nil
}

func (s *pgStore) CountActiveTaskRuns(ctx context.Context, projectID int64) (int, error) {
	if projectID <= 0 {
		return 0, nil
	}
	var count int
	err := s.pool.QueryRow(ctx, `
		SELECT COUNT(*) FROM task_runs
		WHERE project_id = $1 AND status IN ('pending', 'running')
	`, projectID).Scan(&count)
	if err != nil {
		return 0, nil
	}
	return count, nil
}

func (s *pgStore) applyRuntimeMetadata(ctx context.Context, tx pgx.Tx, instanceID int64, def workflowdef.VersionDefinition) error {
	taskDefs := map[string]workflowdef.DocTaskDef{}
	for _, st := range def.Stages {
		for _, dt := range st.DocTasks {
			taskDefs[dt.Key] = dt
		}
	}
	// 先完整读取并关闭结果集，再执行 UPDATE：
	// pgx 在同一事务连接上存在未读完的 rows 时执行 Exec 会报 "conn busy"。
	rows, err := tx.Query(ctx, `
		SELECT dti.id, dti.task_key
		FROM project_doc_task_instances dti
		JOIN project_stage_instances si ON si.id = dti.stage_instance_id
		WHERE si.workflow_instance_id = $1
	`, instanceID)
	if err != nil {
		return err
	}
	type taskInstanceRow struct {
		id      int64
		taskKey string
	}
	taskRows := make([]taskInstanceRow, 0)
	for rows.Next() {
		var r taskInstanceRow
		if err := rows.Scan(&r.id, &r.taskKey); err != nil {
			rows.Close()
			return err
		}
		taskRows = append(taskRows, r)
	}
	iterErr := rows.Err()
	rows.Close()
	if iterErr != nil {
		return iterErr
	}
	for _, r := range taskRows {
		td, ok := taskDefs[r.taskKey]
		if !ok {
			continue
		}
		tplKey := td.TemplateAssetKey
		tplVer := td.TemplateVersion
		if td.Template != nil {
			if td.Template.AssetKey != "" {
				tplKey = td.Template.AssetKey
			}
			if td.Template.Version > 0 {
				tplVer = td.Template.Version
			}
		}
		if _, err := tx.Exec(ctx, `
			UPDATE project_doc_task_instances SET
				executor_key = $2, quality_gate_key = $3,
				template_asset_key = $4, template_version = $5,
				output_doc_type = $6, output_doc_title = $7, updated_at = now()
			WHERE id = $1
		`, r.id, td.ExecutorKey, td.QualityGateKey, tplKey, tplVer, td.OutputDocType, td.OutputDocTitle); err != nil {
			return err
		}
	}
	return nil
}