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

package store

import (
	"database/sql"
	"errors"
	"fmt"
)

// 协作存储 U3 切片 2（R8/KTD「接纳与执行许可分别计数」）：执行许可的共享
// 权威落在数据库（admission_leases），多实例天然共享；接纳数（排队中的 run）
// 由 runmgr 的有界队列管理，不进本表。预算关闭（全零）时不产生任何租约行。

// ExecutionBudget 执行许可预算。零值 = 不限（本地档默认，行为与历史一致）。
type ExecutionBudget struct {
	Global  int // 全局并发执行上限（含子代理与索引任务）；0=不限
	PerUser int // 每租户并发执行上限；0=不限
}

// Off 表示预算完全关闭（不限流、不留租约）。
func (b ExecutionBudget) Off() bool { return b.Global <= 0 && b.PerUser <= 0 }

// ErrBudgetExceeded 预算已满（可排队）；ErrAdmissionQueueFull 是其排队也满的
// 上层形态（runmgr 定义，映射 429）。
var ErrBudgetExceeded = errors.New("store: execution budget exceeded")

// TryAcquireExecutionSlot 原子获取（或重新激活）run 的执行许可：在 admission
// 咨询锁下清点活跃租约（全局 + 租户），不超预算才登记。父 run 让出后重新
// 获取走同一路径（upsert 复活既有行）。
func (s *Store) TryAcquireExecutionSlot(runID, owner string, budget ExecutionBudget) error {
	if budget.Off() {
		return nil
	}
	tx, err := s.db.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()
	if err := s.lockAdmission(tx); err != nil {
		return err
	}
	if err := s.acquireExecutionSlotTx(tx, runID, owner, budget); err != nil {
		return err
	}
	return tx.Commit()
}

// TryAcquireExecutionSlotAccountChecked takes the actor account lock before
// the global admission lock. The run row, not a caller-supplied owner string,
// supplies the authoritative owner for a new or reactivated slot.
func (s *Store) TryAcquireExecutionSlotAccountChecked(actor, runID string, budget ExecutionBudget, allowImplicitDefault bool) error {
	if budget.Off() {
		return nil // no lease row or other DB mutation
	}
	return s.withEffectiveAccountWrite(actor, allowImplicitDefault, func(tx *sql.Tx, _ Account) error {
		if err := s.lockAdmission(tx); err != nil {
			return err
		}
		var owner string
		if err := tx.QueryRow(s.q(`SELECT owner FROM runs WHERE id=?`), runID).Scan(&owner); errors.Is(err, sql.ErrNoRows) {
			return ErrNotFound
		} else if err != nil {
			return err
		}
		if owner != actor {
			return ErrNotFound
		}
		return s.acquireExecutionSlotTx(tx, runID, actor, budget)
	})
}

func (s *Store) acquireExecutionSlotTx(tx *sql.Tx, runID, owner string, budget ExecutionBudget) error {
	var global, perUser int
	if budget.Global > 0 {
		if err := tx.QueryRow(s.q(`SELECT COUNT(*) FROM admission_leases WHERE released_at IS NULL AND run_id<>?`), runID).Scan(&global); err != nil {
			return err
		}
		if global >= budget.Global {
			return fmt.Errorf("exec slot %s: %w", runID, ErrBudgetExceeded)
		}
	}
	if budget.PerUser > 0 {
		if err := tx.QueryRow(s.q(`SELECT COUNT(*) FROM admission_leases WHERE released_at IS NULL AND owner=? AND run_id<>?`), owner, runID).Scan(&perUser); err != nil {
			return err
		}
		if perUser >= budget.PerUser {
			return fmt.Errorf("exec slot %s: %w", runID, ErrBudgetExceeded)
		}
	}
	if _, err := tx.Exec(s.q(`INSERT INTO admission_leases (run_id, owner, acquired_at, released_at) VALUES (?,?,?,NULL)
		ON CONFLICT(run_id) DO UPDATE SET owner=excluded.owner, acquired_at=excluded.acquired_at, released_at=NULL`),
		runID, owner, nowISO()); err != nil {
		return err
	}
	return nil
}

// ReleaseExecutionSlot 释放执行许可（幂等）：父 run 等待子代理/HITL 时让出，
// 恢复执行须重新获取（方案 KTD「让出执行许可、保留会话所有权」）。
func (s *Store) ReleaseExecutionSlot(runID string) error {
	_, err := s.db.Exec(s.q(`UPDATE admission_leases SET released_at=? WHERE run_id=? AND released_at IS NULL`), nowISO(), runID)
	return err
}

// ActiveExecutionSlots 返回活跃执行许可数（诊断/验收：全局与按租户）。
func (s *Store) ActiveExecutionSlots() (global int, byOwner map[string]int, err error) {
	byOwner = map[string]int{}
	rows, err := s.db.Query(s.q(`SELECT owner, COUNT(*) FROM admission_leases WHERE released_at IS NULL GROUP BY owner`))
	if err != nil {
		return 0, nil, err
	}
	defer rows.Close()
	for rows.Next() {
		var o string
		var n int
		if err := rows.Scan(&o, &n); err != nil {
			return 0, nil, err
		}
		byOwner[o] = n
		global += n
	}
	return global, byOwner, rows.Err()
}

// lockAdmission 在事务内锁定执行许可账本（PG 事务级咨询锁；SQLite 单写者
// 天然串行）。与 lockConversation/lockScope 同模式（dialect 收敛）。
func (s *Store) lockAdmission(tx *sql.Tx) error {
	return s.dialect.lockAdmission(tx)
}