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

package store

import (
	"database/sql"
	"errors"

	core "agent-system-core"
)

// UpdateRunAccountChecked fences forward run state and completed output at the
// final run-row mutation. Terminal failure cleanup has a separate legacy path.
func (s *Store) UpdateRunAccountChecked(actor string, r *core.Run, allowImplicitDefault bool) error {
	if r == nil || actor == "" || string(r.Owner) != actor {
		return ErrAccountForbidden
	}
	return s.withEffectiveAccountWrite(actor, allowImplicitDefault, func(tx *sql.Tx, _ Account) error {
		var owner string
		if err := tx.QueryRow(s.q(`SELECT owner FROM runs WHERE id=?`), r.ID).Scan(&owner); errors.Is(err, sql.ErrNoRows) {
			return ErrNotFound
		} else if err != nil {
			return err
		}
		if owner != actor {
			return ErrNotFound
		}
		var ended any
		if r.EndedAt != nil {
			ended = nowISO()
		}
		_, err := tx.Exec(s.q(`UPDATE runs SET status=?,phase=?,usage=?,error=?,ended_at=? WHERE id=? AND owner=?`),
			string(r.Status), string(r.Phase), marshalJSON(r.Usage), marshalJSON(r.Error), ended, r.ID, actor)
		return err
	})
}

// CreateRunAccountChecked fences run admission against account disable. The
// conversation's stored owner is checked in the same final transaction.
func (s *Store) CreateRunAccountChecked(r *core.Run, allowImplicitDefault bool) error {
	if r == nil || r.Owner == "" {
		return ErrAccountForbidden
	}
	if r.ID == "" {
		r.ID = newID("run")
	}
	r.CreatedAt = r.CreatedAt.UTC()
	return s.withEffectiveAccountWrite(string(r.Owner), allowImplicitDefault, func(tx *sql.Tx, _ Account) error {
		var owner string
		if err := tx.QueryRow(s.q(`SELECT owner FROM conversations WHERE id=?`), r.ConversationID).Scan(&owner); errors.Is(err, sql.ErrNoRows) {
			return ErrNotFound
		} else if err != nil {
			return err
		}
		if owner != string(r.Owner) {
			return ErrNotFound
		}
		_, err := tx.Exec(s.q(`INSERT INTO runs(id,parent_run_id,conversation_id,project_id,owner,kind,profile_id,status,phase,agent_name,created_at,instance_id) VALUES(?,?,?,?,?,?,?,?,?,?,?,?)`),
			r.ID, r.ParentRunID, r.ConversationID, r.ProjectID, r.Owner, string(r.Kind), r.ProfileID, string(r.Status), string(r.Phase), r.AgentName, nowISO(), s.instanceID)
		return err
	})
}

// SetMessageRunIDAccountChecked links an admitted run only while its actor is
// still effective. Both rows must belong to the actor's conversation.
func (s *Store) SetMessageRunIDAccountChecked(actor, messageID, runID string, allowImplicitDefault bool) error {
	var routeConversation string
	if err := s.db.QueryRow(s.q(`SELECT conversation_id FROM messages WHERE id=?`), messageID).Scan(&routeConversation); errors.Is(err, sql.ErrNoRows) {
		return ErrNotFound
	} else if err != nil {
		return err
	}
	return s.withEffectiveAccountWrite(actor, allowImplicitDefault, func(tx *sql.Tx, _ Account) error {
		if err := s.dialect.lockConversation(tx, routeConversation); err != nil {
			return err
		}
		var messageConversation, messageOwner, runConversation, runOwner string
		if err := tx.QueryRow(s.q(`SELECT m.conversation_id,c.owner FROM messages m JOIN conversations c ON c.id=m.conversation_id WHERE m.id=?`), messageID).
			Scan(&messageConversation, &messageOwner); errors.Is(err, sql.ErrNoRows) {
			return ErrNotFound
		} else if err != nil {
			return err
		}
		if err := tx.QueryRow(s.q(`SELECT conversation_id,owner FROM runs WHERE id=?`), runID).Scan(&runConversation, &runOwner); errors.Is(err, sql.ErrNoRows) {
			return ErrNotFound
		} else if err != nil {
			return err
		}
		if messageConversation != routeConversation || runConversation != routeConversation || messageOwner != actor || runOwner != actor {
			return ErrNotFound
		}
		res, err := tx.Exec(s.q(`UPDATE messages SET run_id=? WHERE id=? AND conversation_id=?`), runID, messageID, routeConversation)
		if err != nil {
			return err
		}
		if n, err := res.RowsAffected(); err != nil {
			return err
		} else if n != 1 {
			return ErrNotFound
		}
		return nil
	})
}