Agent-System/internal/taskrun/finalize.go

package taskrun

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"sync"
	"time"

	"chat-agent-system/internal/app/uow"
	"chat-agent-system/internal/artifactvault"
	"chat-agent-system/internal/bizdev"
	"chat-agent-system/internal/ctxmgr"
	"chat-agent-system/internal/events"
	"chat-agent-system/internal/projectworkflow"
	"chat-agent-system/internal/workspace"

	"github.com/jackc/pgx/v5"
)

const (
	StatusPending                   = "pending"
	StatusRunning                   = "running"
	StatusFinalizing                = "finalizing"
	StatusCompleted                 = "completed"
	StatusFailed                    = "failed"
	StatusCancelled                 = "cancelled"
	StatusFinalizingFailed          = "finalizing_failed"
	StatusQualityFailed             = "quality_failed"
	StatusInterrupted               = "interrupted"
	StatusPendingAuthorConfirmation = "pending_author_confirmation"
	StatusBlockedPendingGovernance  = "blocked_pending_governance"

	EventCompleted = "taskrun.completed"
	EventCancelled = "taskrun.cancelled"
)

var (
	ErrTerminalConflict = errors.New("taskrun: terminal status already set")
	ErrCleanupBlocked   = errors.New("taskrun: cleanup blocked until verified receipts match the frozen manifest")
)

var openStatuses = map[string]struct{}{
	StatusPending:                  {},
	StatusRunning:                  {},
	StatusFinalizing:               {},
	StatusInterrupted:              {},
	StatusBlockedPendingGovernance: {},
}

var terminalStatuses = map[string]struct{}{
	StatusCompleted:                 {},
	StatusFailed:                    {},
	StatusCancelled:                 {},
	StatusQualityFailed:             {},
	StatusPendingAuthorConfirmation: {},
	StatusFinalizingFailed:          {},
}

// Commands is the unique finalize/cancel command. Server is a thin adapter.
type Commands struct {
	UoW             uow.UnitOfWork
	Ctx             *ctxmgr.ContextManager
	PW              *projectworkflow.Service
	Fav             *artifactvault.Ingest
	Env             string
	StorageIdentity string

	mu  sync.Mutex
	mem map[string]*FinalizeResult
}

// FinalizeResult is the durable outcome of a terminal command.
type FinalizeResult struct {
	Status   string
	Manifest *RunManifest
	Applied  bool
}

type VerifiedReceipt struct {
	LogicalID       string
	Digest          string
	StorageIdentity string
	ArtifactUUID    string
	Verified        bool
}

func (c *Commands) Finalize(ctx context.Context, record *RunRecord, result *bizdev.TaskRunResult) (string, error) {
	if record == nil || result == nil {
		return StatusFailed, fmt.Errorf("finalize: missing record or result")
	}
	status := result.TerminalStatus()
	if _, err := ResolveBackend(result.BackendUsed, ""); err != nil && strings.TrimSpace(result.BackendUsed) != "" {
		return StatusFailed, err
	}
	out, err := c.apply(ctx, record, result, status)
	if err != nil {
		if errors.Is(err, ErrTerminalConflict) && out != nil && out.Status != "" {
			return out.Status, nil
		}
		return StatusFinalizingFailed, err
	}
	return out.Status, nil
}

func (c *Commands) Cancel(ctx context.Context, record *RunRecord) (string, error) {
	if record == nil {
		return StatusCancelled, fmt.Errorf("cancel: missing record")
	}
	result := &bizdev.TaskRunResult{Success: false, ErrorReason: "cancelled", TaskKey: record.TaskKey, WorkDir: record.WorkspaceID}
	out, err := c.apply(ctx, record, result, StatusCancelled)
	if err != nil {
		if errors.Is(err, ErrTerminalConflict) && out != nil {
			return out.Status, nil
		}
		return StatusCancelled, err
	}
	return out.Status, nil
}

func (c *Commands) apply(ctx context.Context, record *RunRecord, result *bizdev.TaskRunResult, want string) (*FinalizeResult, error) {
	if c != nil && c.UoW != nil {
		var out *FinalizeResult
		err := c.UoW.Within(ctx, func(scope uow.TxScope) error {
			ctx = uow.WithScope(ctx, scope)
			got, err := c.applyLocked(ctx, record, result, want)
			out = got
			return err
		})
		return out, err
	}
	return c.applyLocked(ctx, record, result, want)
}

func (c *Commands) applyLocked(ctx context.Context, record *RunRecord, result *bizdev.TaskRunResult, want string) (*FinalizeResult, error) {
	if err := rejectSensitiveText(result.MarkdownBody, result.ErrorReason); err != nil {
		return nil, err
	}
	workspaceID := workspace.OpaqueID(result.WorkDir, record.WorkspaceID)
	manifest, err := buildRunManifest(record, result, workspaceID, c.storageID(), want)
	if err != nil {
		return nil, err
	}

	existing, applied, err := c.casTerminal(ctx, record.RunID, want)
	if err != nil {
		return existing, err
	}
	if !applied {
		if existing != nil {
			return existing, nil
		}
		return &FinalizeResult{Status: want, Manifest: manifest, Applied: false}, nil
	}

	if err := c.persistManifest(ctx, record.RunID, manifest); err != nil {
		return nil, err
	}
	if err := c.projectWorkflow(ctx, record, result, want); err != nil {
		return nil, err
	}
	if err := c.finalizeAssistant(ctx, record, result, want, workspaceID, manifest); err != nil {
		return nil, err
	}
	if err := c.appendTerminalEvent(ctx, record, result, want, manifest); err != nil {
		return nil, err
	}
	if err := c.stageControlledDoc(ctx, record, result, want, manifest); err != nil {
		return nil, err
	}

	out := &FinalizeResult{Status: want, Manifest: manifest, Applied: true}
	c.remember(record.RunID, out)
	return out, nil
}

func (c *Commands) storageID() string {
	if c != nil && strings.TrimSpace(c.StorageIdentity) != "" {
		return c.StorageIdentity
	}
	env := "dev"
	if c != nil && c.Env != "" {
		env = c.Env
	}
	return "env:" + env
}

func (c *Commands) casTerminal(ctx context.Context, runID, want string) (*FinalizeResult, bool, error) {
	if scope, ok := uow.ScopeFrom(ctx); ok {
		tx, err := uow.AsPG(scope)
		if err != nil {
			return nil, false, err
		}
		return casPG(ctx, tx, runID, want)
	}
	return c.casMemory(runID, want)
}

func casPG(ctx context.Context, tx pgx.Tx, runID, want string) (*FinalizeResult, bool, error) {
	var current string
	err := tx.QueryRow(ctx, `SELECT status FROM task_runs WHERE run_id=$1 FOR UPDATE`, runID).Scan(&current)
	if errors.Is(err, pgx.ErrNoRows) {
		return nil, true, nil
	}
	if err != nil {
		return nil, false, err
	}
	if current == want {
		man, _ := loadManifestTX(ctx, tx, runID)
		return &FinalizeResult{Status: current, Manifest: man, Applied: false}, false, nil
	}
	if _, open := openStatuses[current]; !open {
		man, _ := loadManifestTX(ctx, tx, runID)
		return &FinalizeResult{Status: current, Manifest: man, Applied: false}, false, ErrTerminalConflict
	}
	tag, err := tx.Exec(ctx, `
		UPDATE task_runs SET status=$2, updated_at=now()
		WHERE run_id=$1 AND status=$3
	`, runID, want, current)
	if err != nil {
		return nil, false, err
	}
	if tag.RowsAffected() != 1 {
		return nil, false, ErrTerminalConflict
	}
	return nil, true, nil
}

func (c *Commands) casMemory(runID, want string) (*FinalizeResult, bool, error) {
	if c == nil {
		return nil, true, nil
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.mem == nil {
		c.mem = map[string]*FinalizeResult{}
	}
	if snap, ok := c.mem[runID]; ok {
		if snap.Status == want {
			cpy := *snap
			cpy.Applied = false
			return &cpy, false, nil
		}
		cpy := *snap
		cpy.Applied = false
		return &cpy, false, ErrTerminalConflict
	}
	c.mem[runID] = &FinalizeResult{Status: want, Applied: true}
	return nil, true, nil
}

func (c *Commands) remember(runID string, out *FinalizeResult) {
	if c == nil || out == nil || runID == "" {
		return
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.mem == nil {
		c.mem = map[string]*FinalizeResult{}
	}
	cpy := *out
	c.mem[runID] = &cpy
}

func (c *Commands) persistManifest(ctx context.Context, runID string, man *RunManifest) error {
	scope, ok := uow.ScopeFrom(ctx)
	if !ok {
		return nil
	}
	tx, err := uow.AsPG(scope)
	if err != nil {
		return err
	}
	members, err := json.Marshal(man.Members)
	if err != nil {
		return err
	}
	var existingSHA string
	err = tx.QueryRow(ctx, `SELECT sha256 FROM task_run_manifests WHERE run_id=$1`, runID).Scan(&existingSHA)
	if err == nil {
		if strings.TrimSpace(existingSHA) != man.SHA256 {
			return uow.ErrDigestConflict
		}
		return nil
	}
	if !errors.Is(err, pgx.ErrNoRows) {
		return err
	}
	_, err = tx.Exec(ctx, `
		INSERT INTO task_run_manifests (run_id, canonical_bytes, sha256, workspace_id, storage_identity, members)
		VALUES ($1,$2,$3,$4,$5,$6::jsonb)
		ON CONFLICT (run_id) DO NOTHING
	`, runID, man.CanonicalBytes, man.SHA256, man.WorkspaceID, man.StorageIdentity, members)
	if err != nil {
		// Unit tests without a durable task_runs row skip FK persistence.
		if strings.Contains(strings.ToLower(err.Error()), "foreign key") {
			return nil
		}
		return err
	}
	return nil
}

func loadManifestTX(ctx context.Context, tx pgx.Tx, runID string) (*RunManifest, error) {
	var man RunManifest
	var members []byte
	err := tx.QueryRow(ctx, `
		SELECT run_id, canonical_bytes, sha256, workspace_id, storage_identity, members
		FROM task_run_manifests WHERE run_id=$1
	`, runID).Scan(&man.RunID, &man.CanonicalBytes, &man.SHA256, &man.WorkspaceID, &man.StorageIdentity, &members)
	if err != nil {
		return nil, err
	}
	_ = json.Unmarshal(members, &man.Members)
	man.SchemaVersion = ManifestSchemaVersion
	return &man, nil
}

func (c *Commands) projectWorkflow(ctx context.Context, record *RunRecord, result *bizdev.TaskRunResult, status string) error {
	if c == nil || c.PW == nil || record.ProjectID <= 0 || record.RunID == "" {
		return nil
	}
	docID := ""
	title := result.ArtifactTitle
	if result.DocRecord != nil {
		docID = result.DocRecord.DocID
	}
	if docID == "" && record.TaskKey != "" {
		docID = fmt.Sprintf("doc_proj%d_%s", record.ProjectID, record.TaskKey)
	}
	if title == "" && record.TaskKey != "" {
		title = record.TaskKey
	}
	if err := c.PW.UpdateWorkflowRun(ctx, record.RunID, status, docID, title, ""); err != nil {
		if strings.Contains(strings.ToLower(err.Error()), "not found") {
			return nil
		}
		return err
	}
	return nil
}

func (c *Commands) finalizeAssistant(ctx context.Context, record *RunRecord, result *bizdev.TaskRunResult, status, workspaceID string, man *RunManifest) error {
	if c == nil || c.Ctx == nil || record.AssistantMsgID <= 0 {
		return nil
	}
	summary := result.CompletionSummary()
	if status == StatusCancelled {
		summary = "任务已取消。"
		if result.ErrorReason != "" && result.ErrorReason != "cancelled" {
			summary += "\n\n" + result.ErrorReason
		}
	}
	meta := StripSecrets(map[string]any{
		"run_id":       record.RunID,
		"workspace_id": workspaceID,
		"task_status":  status,
		"task_key":     record.TaskKey,
	})
	if record.ProjectID > 0 {
		meta["project_id"] = record.ProjectID
	}
	if man != nil {
		meta["run_manifest_sha256"] = man.SHA256
		meta["storage_identity"] = man.StorageIdentity
	}
	if art := result.ArtifactMetadata(); art != nil {
		if workspaceID != "" {
			art["workspace_id"] = workspaceID
		}
		if result.DocRecord != nil && result.DocRecord.ContentSHA256 != "" {
			art["sha256"] = result.DocRecord.ContentSHA256
		}
		meta["artifact"] = StripSecrets(art)
	} else if result.DocRecord != nil && result.DocRecord.DocID != "" {
		artifact := map[string]any{
			"doc_id":       result.DocRecord.DocID,
			"title":        result.ArtifactTitle,
			"task_key":     record.TaskKey,
			"workspace_id": workspaceID,
		}
		if result.DocRecord.ContentSHA256 != "" {
			artifact["sha256"] = result.DocRecord.ContentSHA256
		}
		if result.QualityReport != nil {
			artifact["quality_passed"] = result.QualityReport.Passed
			artifact["quality_score"] = result.QualityReport.Score
			if len(result.QualityReport.Violations) > 0 {
				artifact["quality_violations"] = result.QualityReport.Violations
			}
		}
		if result.AwaitsAgentConfirmation() {
			artifact["pending_agent_candidate_id"] = result.PendingAgentCandidateID
			if result.PendingChangeReportID > 0 {
				artifact["pending_change_report_id"] = result.PendingChangeReportID
			}
			if result.PendingRiskLevel != "" {
				artifact["pending_risk_level"] = result.PendingRiskLevel
			}
			artifact["awaiting_author_confirmation"] = true
		}
		meta["artifact"] = artifact
	}
	return c.Ctx.FinalizeAssistantTaskMessage(ctx, record.AssistantMsgID, summary, meta, status)
}

func (c *Commands) appendTerminalEvent(ctx context.Context, record *RunRecord, result *bizdev.TaskRunResult, status string, man *RunManifest) error {
	if _, ok := uow.ScopeFrom(ctx); !ok {
		return nil
	}
	name := EventCompleted
	if status == StatusCancelled {
		name = EventCancelled
	}
	body, err := json.Marshal(map[string]any{
		"run_id":           record.RunID,
		"status":           status,
		"workspace_id":     man.WorkspaceID,
		"storage_identity": man.StorageIdentity,
		"manifest_sha256":  man.SHA256,
		"task_key":         record.TaskKey,
		"doc_id":           docIDOf(result),
	})
	if err != nil {
		return err
	}
	envName := "dev"
	if c != nil && c.Env != "" {
		envName = c.Env
	}
	env := events.Envelope{
		EventID:          fmt.Sprintf("taskrun:%s:%s", record.RunID, status),
		EventName:        name,
		SchemaVersion:    events.EnvelopeSchema,
		Environment:      envName,
		ProjectID:        record.ProjectID,
		AggregateType:    "task_run",
		AggregateID:      record.RunID,
		AggregateVersion: 1,
		ActorRef:         fmt.Sprintf("user:%d", record.UserID),
		IdempotencyKey:   fmt.Sprintf("finalize:%s:%s", record.RunID, status),
		PayloadDigest:    events.Digest(body),
		Producer:         "taskrun",
		OccurredAt:       time.Now().UTC(),
		Payload:          body,
	}
	return (events.PGSink{}).Append(ctx, env, []string{"projection", "fav"})
}

func (c *Commands) stageControlledDoc(ctx context.Context, record *RunRecord, result *bizdev.TaskRunResult, status string, man *RunManifest) error {
	if c == nil || c.Fav == nil {
		return nil
	}
	if status != StatusCompleted {
		return nil
	}
	if result.DocRecord == nil || result.DocRecord.DocID == "" {
		return nil
	}
	payload := []byte(result.MarkdownBody)
	if len(payload) == 0 {
		payload = man.CanonicalBytes
	}
	if len(payload) == 0 {
		return nil
	}
	snapID := int64(result.DocRecord.Version)
	if snapID <= 0 {
		snapID = 1
	}
	eventID := fmt.Sprintf("taskrun:%s:%s", record.RunID, status)
	return c.Fav.StageAgentCommit(ctx, record.ProjectID, result.DocRecord.DocID, eventID, snapID, payload)
}

func buildRunManifest(record *RunRecord, result *bizdev.TaskRunResult, workspaceID, storageID, status string) (*RunManifest, error) {
	man := &RunManifest{
		SchemaVersion:   ManifestSchemaVersion,
		RunID:           record.RunID,
		WorkspaceID:     workspaceID,
		StorageIdentity: storageID,
	}
	if result.DocRecord != nil && result.DocRecord.DocID != "" {
		digest := strings.ToLower(strings.TrimSpace(result.DocRecord.ContentSHA256))
		if digest == "" && result.MarkdownBody != "" {
			sum := sha256Hex([]byte(result.MarkdownBody))
			digest = sum
		}
		if digest != "" {
			man.Members = append(man.Members, ManifestMember{
				LogicalID: result.DocRecord.DocID,
				Digest:    digest,
				Kind:      MemberKindDocument,
			})
		}
	}
	if record.AssistantMsgID > 0 {
		digest := sha256Hex([]byte(result.CompletionSummary()))
		if status == StatusCancelled {
			digest = sha256Hex([]byte("任务已取消。"))
		}
		man.Members = append(man.Members, ManifestMember{
			LogicalID: fmt.Sprintf("assistant:%d", record.AssistantMsgID),
			Digest:    digest,
			Kind:      MemberKindAssistantTurn,
		})
	}
	if err := man.Freeze(); err != nil {
		return nil, err
	}
	return man, nil
}

func sha256Hex(b []byte) string {
	sum := sha256.Sum256(b)
	return hex.EncodeToString(sum[:])
}

func docIDOf(result *bizdev.TaskRunResult) string {
	if result != nil && result.DocRecord != nil {
		return result.DocRecord.DocID
	}
	return ""
}

func IsTerminalStatus(status string) bool {
	_, ok := terminalStatuses[status]
	return ok
}

func IsOpenStatus(status string) bool {
	_, ok := openStatuses[status]
	return ok
}

// AssertCleanupAllowed requires a frozen manifest and a verified receipt for every retained member.
func (m *RunManifest) AssertCleanupAllowed(receipts []VerifiedReceipt) error {
	if m == nil || m.SHA256 == "" || m.StorageIdentity == "" {
		return ErrCleanupBlocked
	}
	byLogical := map[string]VerifiedReceipt{}
	for _, r := range receipts {
		byLogical[r.LogicalID] = r
	}
	for _, mem := range m.Members {
		if !memberNeedsRetention(mem.Kind) {
			continue
		}
		r, ok := byLogical[mem.LogicalID]
		if !ok || !r.Verified {
			return ErrCleanupBlocked
		}
		if r.Digest != mem.Digest || r.StorageIdentity != m.StorageIdentity {
			return ErrCleanupBlocked
		}
	}
	return nil
}