Agent-System-Backend/internal/store/regeneration_account_write.go

package store

import (
	"database/sql"
	"errors"

	core "agent-system-core"
)

// PrepareRegenerationAccountChecked excludes the selected assistant run and
// records the next variant in one final actor-checked transaction. The target
// and conversation owner are read again after the account and conversation
// locks, so an earlier HTTP lookup cannot authorize a late write.
func (s *Store) PrepareRegenerationAccountChecked(actor, conversationID, messageID string, allowImplicitDefault bool) (string, int, error) {
	var group string
	var nextIndex int
	err := s.withEffectiveAccountWrite(actor, allowImplicitDefault, func(tx *sql.Tx, _ Account) error {
		if err := s.dialect.lockConversation(tx, conversationID); err != nil {
			return err
		}
		query := s.q(`SELECT COALESCE(m.run_id,''),m.variant_group,m.variant_index,m.role,c.owner
			FROM messages m JOIN conversations c ON c.id=m.conversation_id
			WHERE m.id=? AND m.conversation_id=?`)
		if s.Dialect() == "postgres" {
			query += ` FOR UPDATE OF m`
		}
		var runID, existingGroup, role, owner string
		var variantIndex int
		if err := tx.QueryRow(query, messageID, conversationID).Scan(&runID, &existingGroup, &variantIndex, &role, &owner); errors.Is(err, sql.ErrNoRows) {
			return ErrNotFound
		} else if err != nil {
			return err
		}
		if owner != actor || role != string(core.RoleAssistant) {
			return ErrNotFound
		}
		maxIndex := variantIndex
		if existingGroup != "" {
			if err := tx.QueryRow(s.q(`SELECT COALESCE(MAX(variant_index),0) FROM messages WHERE conversation_id=? AND variant_group=?`),
				conversationID, existingGroup).Scan(&maxIndex); err != nil {
				return err
			}
			if variantIndex > maxIndex {
				maxIndex = variantIndex
			}
		}
		group = existingGroup
		if group == "" {
			group = "var_" + messageID
		}
		var pendingGroup string
		var pendingIndex int
		err := tx.QueryRow(s.q(`SELECT variant_group,variant_index FROM regen_pending WHERE conversation_id=?`), conversationID).Scan(&pendingGroup, &pendingIndex)
		if err != nil && !errors.Is(err, sql.ErrNoRows) {
			return err
		}
		if pendingGroup == group && pendingIndex > maxIndex {
			maxIndex = pendingIndex
		}
		nextIndex = maxIndex + 1
		if _, err := tx.Exec(s.q(`UPDATE messages SET excluded=1,variant_group=?,variant_index=? WHERE id=? AND conversation_id=?`),
			group, variantIndex, messageID, conversationID); err != nil {
			return err
		}
		if runID != "" {
			if _, err := tx.Exec(s.q(`UPDATE messages SET excluded=1 WHERE run_id=? AND conversation_id=?`), runID, conversationID); err != nil {
				return err
			}
		}
		_, err = tx.Exec(s.q(`INSERT INTO regen_pending(conversation_id,variant_group,variant_index) VALUES(?,?,?)
			ON CONFLICT(conversation_id) DO UPDATE SET variant_group=excluded.variant_group,variant_index=excluded.variant_index`),
			conversationID, group, nextIndex)
		return err
	})
	if err != nil {
		return "", 0, err
	}
	return group, nextIndex, nil
}

// ClearRegenPendingAfterFailedStart is safety cleanup, not a new business
// write. It may remove this request's stale marker after account disable, but
// still locks the actor and verifies persisted conversation ownership. A later
// regeneration with another group or index is never cleared.
func (s *Store) ClearRegenPendingAfterFailedStart(actor, conversationID, group string, index int, allowImplicitDefault bool) error {
	if actor == "" || conversationID == "" || group == "" || index <= 0 {
		return ErrAccountForbidden
	}
	tx, err := s.db.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()
	_, err = s.lockAccountUsers(tx, actor)
	if errors.Is(err, ErrNotFound) && allowImplicitDefault && actor == OwnerDefault {
		err = nil
	}
	if errors.Is(err, ErrNotFound) {
		return ErrAccountForbidden
	}
	if err != nil {
		return err
	}
	if err := s.dialect.lockConversation(tx, conversationID); err != nil {
		return err
	}
	var owner string
	if err := tx.QueryRow(s.q(`SELECT owner FROM conversations WHERE id=?`), conversationID).Scan(&owner); errors.Is(err, sql.ErrNoRows) {
		return ErrNotFound
	} else if err != nil {
		return err
	}
	if owner != actor {
		return ErrNotFound
	}
	if _, err := tx.Exec(s.q(`DELETE FROM regen_pending WHERE conversation_id=? AND variant_group=? AND variant_index=?`), conversationID, group, index); err != nil {
		return err
	}
	return tx.Commit()
}