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

package store

import (
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"time"
)

// 协作存储 U5（KTD4/KTD5）：会话执行租约 + 持久命令收件箱。
//
// 租约：conversation_executions 一行一会话。TryClaimExecution 空闲/租约到期
// 才可占用；接管（含同实例重新占用）递增 epoch——旧 owner 的续租与结算按
//（instance, epoch）双重校验被拒（fencing）。lease_until 用应用侧 UTC ISO
// （单库写入方一致即可；跨库时钟偏移属部署合同，见 runbook）。
//
// 命令：run_commands 收件箱承载跨实例取消/HITL（KTD5）。command_id 幂等；
// pending→consumed 由消费方原子认领（PG FOR UPDATE SKIP LOCKED；SQLite 单写
// 者条件更新）。命令不删除——审计与重放判定保留。

// ErrExecutionBusy 会话执行权被其他实例持有且租约未到期。
var ErrExecutionBusy = errors.New("store: execution owned by another instance")

// ErrFenced 旧执行代次的写入被围栏拒绝（epoch 已被接管递增）。
var ErrFenced = errors.New("store: fenced by newer execution epoch")

// ExecutionState 当前执行租约快照（诊断/测试）。
type ExecutionState struct {
	ConversationID string
	InstanceID     string
	RunID          string
	Epoch          int64
	LeaseUntil     string
}

// TryClaimExecution 占用（或租约到期后接管）会话执行权。
// 返回接管后的 epoch（接管必然 +1）。被他人持有且未到期 → ErrExecutionBusy。
func (s *Store) TryClaimExecution(conversationID, instanceID, runID string, leaseSec int) (int64, error) {
	if leaseSec <= 0 {
		leaseSec = 90
	}
	tx, err := s.db.Begin()
	if err != nil {
		return 0, err
	}
	defer tx.Rollback()
	if err := s.lockAdmission(tx); err != nil { // 复用全局账本锁串行化清点
		return 0, err
	}
	now := time.Now().UTC()
	until := now.Add(time.Duration(leaseSec) * time.Second).UTC().Format(time.RFC3339Nano)
	var inst, lease string
	var epoch int64
	err = tx.QueryRow(s.q(`SELECT instance_id, epoch, lease_until FROM conversation_executions WHERE conversation_id=?`), conversationID).Scan(&inst, &epoch, &lease)
	switch {
	case errors.Is(err, sql.ErrNoRows):
		if _, err := tx.Exec(s.q(`INSERT INTO conversation_executions (conversation_id, instance_id, run_id, epoch, lease_until, updated_at) VALUES (?,?,?,1,?,?)`),
			conversationID, instanceID, runID, until, nowISO()); err != nil {
			return 0, err
		}
		if err := tx.Commit(); err != nil {
			return 0, err
		}
		return 1, nil
	case err != nil:
		return 0, err
	}
	if inst != instanceID && lease > nowISO() {
		return epoch, fmt.Errorf("conv %s: %w", conversationID, ErrExecutionBusy)
	}
	// 接管（到期或同实例重占）：epoch+1，fencing 由此产生。
	if _, err := tx.Exec(s.q(`UPDATE conversation_executions SET instance_id=?, run_id=?, epoch=epoch+1, lease_until=?, updated_at=? WHERE conversation_id=?`),
		instanceID, runID, until, nowISO(), conversationID); err != nil {
		return 0, err
	}
	if err := tx.Commit(); err != nil {
		return 0, err
	}
	return epoch + 1, nil
}

// RenewExecution 续租（fencing：instance+epoch 双校验）。假 = 已被接管或释放。
func (s *Store) RenewExecution(conversationID, instanceID string, epoch int64, leaseSec int) bool {
	if leaseSec <= 0 {
		leaseSec = 90
	}
	until := time.Now().UTC().Add(time.Duration(leaseSec) * time.Second).UTC().Format(time.RFC3339Nano)
	res, err := s.db.Exec(s.q(`UPDATE conversation_executions SET lease_until=?, updated_at=? WHERE conversation_id=? AND instance_id=? AND epoch=?`),
		until, nowISO(), conversationID, instanceID, epoch)
	if err != nil {
		return false
	}
	n, _ := res.RowsAffected()
	return n > 0
}

// ReleaseExecution 主动释放（优雅停机）：行保留、租约置为已过期——epoch 保留，
// 下一次占用仍递增（fencing 单调）。
func (s *Store) ReleaseExecution(conversationID, instanceID string, epoch int64) {
	_, _ = s.db.Exec(s.q(`UPDATE conversation_executions SET lease_until=?, updated_at=? WHERE conversation_id=? AND instance_id=? AND epoch=?`),
		time.Now().UTC().Add(-time.Second).UTC().Format(time.RFC3339Nano), nowISO(), conversationID, instanceID, epoch)
}

// GetExecution 读取当前租约（无行 = 从未占用）。
func (s *Store) GetExecution(conversationID string) (ExecutionState, error) {
	var e ExecutionState
	err := s.db.QueryRow(s.q(`SELECT conversation_id, instance_id, COALESCE(run_id,''), epoch, lease_until FROM conversation_executions WHERE conversation_id=?`), conversationID).
		Scan(&e.ConversationID, &e.InstanceID, &e.RunID, &e.Epoch, &e.LeaseUntil)
	if errors.Is(err, sql.ErrNoRows) {
		return e, ErrNotFound
	}
	return e, err
}

// ExecutionFenceOK 校验 (instance, epoch) 仍是当前执行权（结算/事件写入的
// 前置检查；事务内校验见 SettleRunFenced）。
func (s *Store) ExecutionFenceOK(conversationID, instanceID string, epoch int64) bool {
	var n int
	if err := s.db.QueryRow(s.q(`SELECT COUNT(*) FROM conversation_executions WHERE conversation_id=? AND instance_id=? AND epoch=?`), conversationID, instanceID, epoch).Scan(&n); err != nil {
		return false
	}
	return n > 0
}

// SettleRunFenced 带执行围栏的终态结算（KTD4：旧 owner 恢复后不能提交）。
// fenced=true 表示 (instance,epoch) 已不是当前执行权——run 未被改动，调用方
// 应将本地 run 视为被接管（不写终态、事件标注接管）。无租约行（SQLite 本地
// 档/未启用执行权）恒放行。

// RunCommand 收件箱命令。
type RunCommand struct {
	CommandID      string
	Kind           string // cancel | hitl
	RunID          string
	ConversationID string
	Epoch          int64
	Payload        string
	CreatedAt      string
}

// EnqueueCommand 投递跨实例命令（幂等：同 command_id 重复投递不产生新行）。
func (s *Store) EnqueueCommand(commandID, kind, runID, conversationID string, epoch int64, payload any) error {
	raw, err := json.Marshal(payload)
	if err != nil {
		return err
	}
	_, err = s.db.Exec(s.q(`INSERT INTO run_commands (command_id, kind, run_id, conversation_id, epoch, payload, state, created_at)
		VALUES (?,?,?,?,?,?,'pending',?)
		ON CONFLICT(command_id) DO NOTHING`),
		commandID, kind, runID, conversationID, epoch, string(raw), nowISO())
	if err != nil {
		return fmt.Errorf("enqueue command %s: %w", commandID, err)
	}
	return nil
}

// ClaimPendingCommands 认领待处理命令（原子 pending→consumed）。
// PG：FOR UPDATE SKIP LOCKED（多实例并发认领不重复）；SQLite：单写者条件更新。
// 只认领「本实例相关」的命令由调用方过滤——收件箱按执行权路由：命令在
// 被执行实例消费（cancel/hitl 目标 run 的当前 owner）。
func (s *Store) ClaimPendingCommands(consumer string, limit int) ([]RunCommand, error) {
	if limit <= 0 {
		limit = 16
	}
	tx, err := s.db.Begin()
	if err != nil {
		return nil, err
	}
	defer tx.Rollback()
	var rows *sql.Rows
	if s.dialect.name() == "postgres" {
		rows, err = tx.Query(`SELECT command_id, kind, run_id, conversation_id, epoch, payload, created_at
			FROM run_commands WHERE state='pending' ORDER BY created_at
			LIMIT $1 FOR UPDATE SKIP LOCKED`, limit)
	} else {
		rows, err = tx.Query(s.q(`SELECT command_id, kind, run_id, conversation_id, epoch, payload, created_at
			FROM run_commands WHERE state='pending' ORDER BY created_at LIMIT ?`), limit)
	}
	if err != nil {
		return nil, err
	}
	var out []RunCommand
	for rows.Next() {
		var c RunCommand
		if err := rows.Scan(&c.CommandID, &c.Kind, &c.RunID, &c.ConversationID, &c.Epoch, &c.Payload, &c.CreatedAt); err != nil {
			rows.Close()
			return nil, err
		}
		out = append(out, c)
	}
	rows.Close()
	if err := rows.Err(); err != nil {
		return nil, err
	}
	for _, c := range out {
		res, err := tx.Exec(s.q(`UPDATE run_commands SET state='consumed', consumed_at=?, consumer=? WHERE command_id=? AND state='pending'`),
			nowISO(), consumer, c.CommandID)
		if err != nil {
			return nil, err
		}
		if n, _ := res.RowsAffected(); n == 0 {
			return nil, fmt.Errorf("command %s claimed concurrently", c.CommandID)
		}
	}
	if err := tx.Commit(); err != nil {
		return nil, err
	}
	return out, nil
}

// CommandState 查询单条命令状态（测试/诊断）。
func (s *Store) CommandState(commandID string) (state, consumer string, err error) {
	err = s.db.QueryRow(s.q(`SELECT state, COALESCE(consumer,'') FROM run_commands WHERE command_id=?`), commandID).Scan(&state, &consumer)
	if errors.Is(err, sql.ErrNoRows) {
		return "", "", ErrNotFound
	}
	return
}