story/internal/orchestration/recovery.go

package orchestration

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"strconv"
	"time"

	"data-terminal-cloud/internal/conversation"
	"data-terminal-cloud/internal/outbox"

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

type RecoveryOptions struct {
	ScanInterval  time.Duration
	StartingAfter time.Duration
	StoppingAfter time.Duration
	BatchSize     int
}

type RecoveryScanner struct {
	pool    *pgxpool.Pool
	outbox  *outbox.Store
	logger  *slog.Logger
	options RecoveryOptions
	now     func() time.Time
}

type recoveryCandidate struct {
	InternalID      string
	PublicID        string
	TenantID        string
	Status          conversation.Status
	ProviderTaskRef string
	StopReason      string
	Trigger         string
	RuntimeType     string
}

func NewRecoveryScanner(pool *pgxpool.Pool, outboxStore *outbox.Store, logger *slog.Logger, options RecoveryOptions) (*RecoveryScanner, error) {
	if pool == nil || outboxStore == nil || logger == nil || options.ScanInterval <= 0 ||
		options.StartingAfter <= 0 || options.StoppingAfter <= 0 || options.BatchSize < 1 {
		return nil, errors.New("valid recovery scanner dependencies and options are required")
	}
	return &RecoveryScanner{pool: pool, outbox: outboxStore, logger: logger, options: options, now: time.Now}, nil
}

func (s *RecoveryScanner) Run(ctx context.Context) {
	s.logger.Info("conversation recovery scanner started")
	ticker := time.NewTicker(s.options.ScanInterval)
	defer ticker.Stop()
	for {
		if count, err := s.ScanOnce(ctx); err != nil && !errors.Is(err, context.Canceled) {
			s.logger.Error("conversation recovery scan failed", "error", err)
		} else if count > 0 {
			s.logger.Info("conversation recovery actions scheduled", "count", count)
		}
		select {
		case <-ctx.Done():
			s.logger.Info("conversation recovery scanner stopped")
			return
		case <-ticker.C:
		}
	}
}

func (s *RecoveryScanner) ScanOnce(ctx context.Context) (int, error) {
	tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{})
	if err != nil {
		return 0, err
	}
	defer tx.Rollback(context.Background())
	now := s.now().UTC()
	rows, err := tx.Query(ctx, `
		SELECT c.id::text, c.public_id::text, c.tenant_id::text, c.status,
			c.provider_task_ref, COALESCE(c.stop_reason,''),
			CASE WHEN c.status='ACTIVE' AND cub.in_progress_action='STOP_IMMEDIATELY'
			          AND c.started_at IS NOT NULL
			          AND c.started_at+make_interval(secs=>cub.max_call_duration_sec)<=$1
			     THEN 'quota_duration_limit' ELSE '' END,c.runtime_type
		FROM conversation c LEFT JOIN rtc_task r ON r.conversation_id=c.id
		LEFT JOIN hardware_agent_session hs ON hs.conversation_id=c.id
		LEFT JOIN conversation_usage_budget cub ON cub.conversation_id=c.id AND cub.tenant_id=c.tenant_id
		WHERE (c.status='ROOM_ISSUED' AND c.allocation_expires_at <= $1)
		   OR (c.status='ACTIVE' AND c.lease_expires_at IS NOT NULL AND c.lease_expires_at <= $1)
		   OR (c.status='ACTIVE' AND cub.in_progress_action='STOP_IMMEDIATELY' AND c.started_at IS NOT NULL AND c.started_at+make_interval(secs=>cub.max_call_duration_sec)<=$1)
		   OR (c.status='STARTING_AGENT' AND c.updated_at <= $2)
		   OR (c.status='STOPPING' AND c.updated_at <= $3)
		ORDER BY c.updated_at, c.id
		LIMIT $4 FOR UPDATE OF c SKIP LOCKED`,
		now, now.Add(-s.options.StartingAfter), now.Add(-s.options.StoppingAfter), s.options.BatchSize)
	if err != nil {
		return 0, fmt.Errorf("select conversations for recovery: %w", err)
	}
	candidates := make([]recoveryCandidate, 0, s.options.BatchSize)
	for rows.Next() {
		var candidate recoveryCandidate
		if err := rows.Scan(&candidate.InternalID, &candidate.PublicID, &candidate.TenantID, &candidate.Status, &candidate.ProviderTaskRef, &candidate.StopReason, &candidate.Trigger, &candidate.RuntimeType); err != nil {
			rows.Close()
			return 0, err
		}
		candidates = append(candidates, candidate)
	}
	if err := rows.Err(); err != nil {
		rows.Close()
		return 0, err
	}
	rows.Close()
	for _, candidate := range candidates {
		if err := s.recoverCandidate(ctx, tx, candidate, now); err != nil {
			return 0, err
		}
	}
	if err := tx.Commit(ctx); err != nil {
		return 0, err
	}
	return len(candidates), nil
}

func (s *RecoveryScanner) recoverCandidate(ctx context.Context, tx pgx.Tx, candidate recoveryCandidate, now time.Time) error {
	action := ""
	switch candidate.Status {
	case conversation.StatusRoomIssued:
		action = "expire_unactivated_room"
		if _, err := tx.Exec(ctx, `UPDATE conversation SET status='EXPIRED', ended_at=$2,
			failure_code='activation_timeout', lease_expires_at=NULL, version=version+1, updated_at=$2
			WHERE id=$1::uuid AND status='ROOM_ISSUED'`, candidate.InternalID, now); err != nil {
			return err
		}
		if candidate.RuntimeType == RuntimeVolcHardwareAgent {
			if _, err := tx.Exec(ctx, `UPDATE hardware_agent_session SET status='EXPIRED',updated_at=$2 WHERE conversation_id=$1::uuid`, candidate.InternalID, now); err != nil {
				return err
			}
		} else {
			if _, err := tx.Exec(ctx, `UPDATE rtc_task SET status='COMPLETED', lease_expires_at=NULL,
				last_checked_at=$2, updated_at=$2 WHERE conversation_id=$1::uuid`, candidate.InternalID, now); err != nil {
				return err
			}
		}
		if _, err := tx.Exec(ctx, `INSERT INTO conversation_state_history (
			conversation_id, from_status, to_status, reason, source
		) VALUES ($1::uuid, 'ROOM_ISSUED', 'EXPIRED', 'activation deadline expired', 'worker')`, candidate.InternalID); err != nil {
			return err
		}
	case conversation.StatusActive:
		action = "stop_expired_lease"
		reason := "lease_timeout"
		historyReason := "conversation lease expired"
		if candidate.Trigger == "quota_duration_limit" {
			action = "stop_quota_duration_limit"
			reason = "quota_duration_limit"
			historyReason = "maximum call duration reached"
		}
		if _, err := tx.Exec(ctx, `UPDATE conversation SET status='STOPPING', stop_reason=$3,
			lease_expires_at=NULL, version=version+1, updated_at=$2
			WHERE id=$1::uuid AND status='ACTIVE'`, candidate.InternalID, now, reason); err != nil {
			return err
		}
		if candidate.RuntimeType == RuntimeVolcHardwareAgent {
			if _, err := tx.Exec(ctx, `UPDATE hardware_agent_session SET status='STOPPING',updated_at=$2 WHERE conversation_id=$1::uuid`, candidate.InternalID, now); err != nil {
				return err
			}
		} else {
			if _, err := tx.Exec(ctx, `UPDATE rtc_task SET status='STOPPING', lease_expires_at=NULL,
				updated_at=$2 WHERE conversation_id=$1::uuid`, candidate.InternalID, now); err != nil {
				return err
			}
		}
		if _, err := tx.Exec(ctx, `INSERT INTO conversation_state_history (
			conversation_id, from_status, to_status, reason, source
		) VALUES ($1::uuid, 'ACTIVE', 'STOPPING', $2, 'worker')`, candidate.InternalID, historyReason); err != nil {
			return err
		}
		if candidate.RuntimeType != RuntimeVolcHardwareAgent {
			if err := s.enqueueRecovery(ctx, tx, candidate, "conversation.agent_stop_requested", candidate.ProviderTaskRef, reason); err != nil {
				return err
			}
		}
	case conversation.StatusStartingAgent:
		action = "reconcile_starting_task"
		deduplicationKey := candidate.ProviderTaskRef + ":" + strconv.FormatInt(recoveryBucket(now, s.options.StartingAfter), 10)
		if err := s.enqueueRecovery(ctx, tx, candidate, "conversation.agent_start_reconcile_requested", deduplicationKey, "starting_timeout"); err != nil {
			return err
		}
		if _, err := tx.Exec(ctx, `UPDATE conversation SET updated_at=$2 WHERE id=$1::uuid`, candidate.InternalID, now); err != nil {
			return err
		}
	case conversation.StatusStopping:
		if candidate.RuntimeType == RuntimeVolcHardwareAgent {
			action = "fail_unconfirmed_hardware_stop"
			if _, err := tx.Exec(ctx, `UPDATE conversation SET status='FAILED',ended_at=$2,failure_code='hardware_agent_stop_timeout',version=version+1,updated_at=$2 WHERE id=$1::uuid AND status='STOPPING'`, candidate.InternalID, now); err != nil {
				return err
			}
			if _, err := tx.Exec(ctx, `UPDATE hardware_agent_session SET status='FAILED',provider_ended_at=$2,updated_at=$2 WHERE conversation_id=$1::uuid`, candidate.InternalID, now); err != nil {
				return err
			}
			if _, err := tx.Exec(ctx, `INSERT INTO conversation_state_history(conversation_id,from_status,to_status,reason,source) VALUES($1::uuid,'STOPPING','FAILED','hardware agent stop was not confirmed before deadline','worker')`, candidate.InternalID); err != nil {
				return err
			}
			if err := enqueueHardwareContentExtraction(ctx, tx, s.outbox, candidate.TenantID, candidate.InternalID); err != nil {
				return err
			}
			break
		}
		action = "reconcile_stopping_task"
		deduplicationKey := candidate.ProviderTaskRef + ":" + strconv.FormatInt(recoveryBucket(now, s.options.StoppingAfter), 10)
		reason := candidate.StopReason
		if reason == "" {
			reason = "recovery_stop"
		}
		if err := s.enqueueRecovery(ctx, tx, candidate, "conversation.agent_stop_reconcile_requested", deduplicationKey, reason); err != nil {
			return err
		}
		if _, err := tx.Exec(ctx, `UPDATE conversation SET updated_at=$2 WHERE id=$1::uuid`, candidate.InternalID, now); err != nil {
			return err
		}
	default:
		return nil
	}
	detail, _ := json.Marshal(map[string]any{"previous_status": candidate.Status, "action": action})
	_, err := tx.Exec(ctx, `INSERT INTO audit_event (
		tenant_id, actor_type, action, resource_type, resource_id, purpose, result, detail
	) VALUES ($1::uuid, 'system', 'conversation.recovery', 'conversation', $2::uuid, $3, 'success', $4::jsonb)`,
		candidate.TenantID, candidate.InternalID, action, detail)
	return err
}

func recoveryBucket(now time.Time, interval time.Duration) int64 {
	bucketWidth := int64(interval / time.Second)
	if bucketWidth < 1 {
		bucketWidth = 1
	}
	return now.Unix() / bucketWidth
}

func (s *RecoveryScanner) enqueueRecovery(
	ctx context.Context,
	tx pgx.Tx,
	candidate recoveryCandidate,
	eventType, deduplicationKey, reason string,
) error {
	payload, _ := json.Marshal(map[string]string{
		"conversation_id": candidate.PublicID, "provider_task_ref": candidate.ProviderTaskRef, "reason": reason,
	})
	_, err := s.outbox.Enqueue(ctx, tx, outbox.NewEvent{
		TenantID: &candidate.TenantID, AggregateType: "conversation", AggregateID: candidate.InternalID,
		EventType: eventType, DeduplicationKey: &deduplicationKey, Payload: payload,
	})
	return err
}