381 lines
13 KiB
Go
381 lines
13 KiB
Go
package store
|
||
|
||
import (
|
||
"database/sql"
|
||
"errors"
|
||
"fmt"
|
||
"time"
|
||
|
||
core "agent-system-core"
|
||
)
|
||
|
||
// DuplicateRequestError 带上已落库的那条消息。
|
||
type DuplicateRequestError struct {
|
||
Existing core.Message
|
||
}
|
||
|
||
func (e *DuplicateRequestError) Error() string { return ErrDuplicateRequest.Error() }
|
||
|
||
func (e *DuplicateRequestError) Unwrap() error { return ErrDuplicateRequest }
|
||
|
||
// MessageByRequest 按会话和 request_id 查找已落库消息。
|
||
func (s *Store) MessageByRequest(conversationID, requestID string) (core.Message, error) {
|
||
if requestID == "" {
|
||
return core.Message{}, ErrNotFound
|
||
}
|
||
var m core.Message
|
||
var parts, ca string
|
||
err := s.db.QueryRow(s.q(`SELECT id,conversation_id,run_id,role,parts,created_at FROM messages WHERE conversation_id=? AND request_id=?`),
|
||
conversationID, requestID).Scan(&m.ID, &m.ConversationID, &m.RunID, &m.Role, &parts, &ca)
|
||
if errors.Is(err, sql.ErrNoRows) {
|
||
return core.Message{}, ErrNotFound
|
||
}
|
||
if err != nil {
|
||
return core.Message{}, err
|
||
}
|
||
m.Role = core.Role(m.Role)
|
||
m.Parts, _ = unmarshalJSON[[]core.Part](parts)
|
||
m.CreatedAt, _ = parseTime(ca)
|
||
return m, nil
|
||
}
|
||
|
||
// SetMessageRunID 把用户消息挂到随后启动的 run 上，供同一 request_id 重试时带回。
|
||
func (s *Store) SetMessageRunID(messageID, runID string) error {
|
||
_, err := s.db.Exec(s.q(`UPDATE messages SET run_id=? WHERE id=?`), runID, messageID)
|
||
return err
|
||
}
|
||
|
||
// LatestOpenRun 返回会话里最新的非终态、非子代理 run。
|
||
// 幂等重试在用户消息还没写上 run_id 时，用它避免绑到更早的已完成 run。
|
||
func (s *Store) LatestOpenRun(conversationID string) (core.Run, error) {
|
||
r, err := scanRun(func(dst ...any) error {
|
||
return s.db.QueryRow(s.q(`SELECT `+runColumns+` FROM runs WHERE conversation_id=? AND kind<>? AND status IN ('pending','running','waiting_hitl') ORDER BY created_at DESC LIMIT 1`),
|
||
conversationID, string(core.RunSubagent)).Scan(dst...)
|
||
})
|
||
if errors.Is(err, sql.ErrNoRows) {
|
||
return r, ErrNotFound
|
||
}
|
||
return r, err
|
||
}
|
||
|
||
// ListRecoverableRuns 列出仍停在非终态的 run。
|
||
// all 为真时扫全部（SQLite 单进程）。否则只扫本实例，以及升级前 instance_id 为空的行。
|
||
func (s *Store) ListRecoverableRuns(instanceID string, all bool) ([]core.Run, error) {
|
||
q := `SELECT ` + runColumns + ` FROM runs WHERE status IN ('pending','running','waiting_hitl')`
|
||
args := []any{}
|
||
if !all {
|
||
q += ` AND (instance_id=? OR instance_id='')`
|
||
args = append(args, instanceID)
|
||
}
|
||
q += ` ORDER BY created_at`
|
||
rows, err := s.db.Query(s.q(q), args...)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
defer rows.Close()
|
||
out := []core.Run{}
|
||
for rows.Next() {
|
||
r, err := scanRun(rows.Scan)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
out = append(out, r)
|
||
}
|
||
return out, rows.Err()
|
||
}
|
||
|
||
// SealInterrupted 把仍在进行的 run 收成 interrupted。
|
||
// 条件更新保证并发对账时只有一个赢家；已经是终态的行不受影响。
|
||
// 赢家并在同一事务中为该 run 写入一条助手中断错误消息（若尚无助手消息）。
|
||
func (s *Store) SealInterrupted(id string) (bool, error) {
|
||
tx, err := s.db.Begin()
|
||
if err != nil {
|
||
return false, err
|
||
}
|
||
defer tx.Rollback()
|
||
|
||
var convID string
|
||
err = tx.QueryRow(s.q(`SELECT conversation_id FROM runs WHERE id=?`), id).Scan(&convID)
|
||
if errors.Is(err, sql.ErrNoRows) {
|
||
return false, nil
|
||
} else if err != nil {
|
||
return false, err
|
||
}
|
||
if convID != "" {
|
||
if err := s.dialect.lockConversation(tx, convID); err != nil {
|
||
return false, err
|
||
}
|
||
}
|
||
|
||
errPartObj := &core.ErrorPart{Code: "run.interrupted", Message: "进程重启或停机，运行已中断"}
|
||
errPartJSON := marshalJSON(errPartObj)
|
||
res, err := tx.Exec(s.q(`UPDATE runs SET status=?, phase=?, error=?, ended_at=? WHERE id=? AND status IN ('pending','running','waiting_hitl')`),
|
||
string(core.RunInterrupted), string(core.PhaseDone), errPartJSON, nowISO(), id)
|
||
if err != nil {
|
||
return false, err
|
||
}
|
||
n, err := res.RowsAffected()
|
||
if err != nil {
|
||
return false, err
|
||
}
|
||
if n != 1 {
|
||
return false, nil
|
||
}
|
||
|
||
var count int
|
||
if err := tx.QueryRow(s.q(`SELECT COUNT(*) FROM messages WHERE run_id=? AND role=?`), id, string(core.RoleAssistant)).Scan(&count); err != nil {
|
||
return false, err
|
||
}
|
||
if count == 0 && convID != "" {
|
||
msg := &core.Message{
|
||
ID: newID("msg"),
|
||
ConversationID: convID,
|
||
RunID: id,
|
||
Role: core.RoleAssistant,
|
||
Parts: []core.Part{{Type: core.PartError, Err: errPartObj}},
|
||
CreatedAt: time.Now().UTC(),
|
||
}
|
||
var variantGroup string
|
||
var variantIndex int
|
||
if err := tx.QueryRow(s.q(`SELECT variant_group, variant_index FROM regen_pending WHERE conversation_id=?`), convID).Scan(&variantGroup, &variantIndex); err == nil && variantGroup != "" {
|
||
msg.VariantGroup = variantGroup
|
||
msg.VariantIndex = variantIndex
|
||
}
|
||
if err := s.insertSettledMessageTx(tx, msg); err != nil {
|
||
return false, err
|
||
}
|
||
}
|
||
|
||
if err := tx.Commit(); err != nil {
|
||
return false, err
|
||
}
|
||
return true, nil
|
||
}
|
||
|
||
// ActiveHITLCardsByRun 返回某个 run 上尚未关闭的 HITL 卡片。
|
||
func (s *Store) ActiveHITLCardsByRun(runID string) ([]core.HITLCard, error) {
|
||
rows, err := s.db.Query(s.q(`SELECT payload FROM hitl_cards WHERE run_id=? AND resolved=0 ORDER BY created_at`), runID)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
defer rows.Close()
|
||
out := []core.HITLCard{}
|
||
for rows.Next() {
|
||
var payload string
|
||
if err := rows.Scan(&payload); err != nil {
|
||
return nil, err
|
||
}
|
||
c, err := unmarshalJSON[core.HITLCard](payload)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
out = append(out, c)
|
||
}
|
||
return out, rows.Err()
|
||
}
|
||
|
||
// SettleRun 在一个事务里更新 run，并可选地追加助手消息。
|
||
// updated 为假表示条件更新没有改到行（已被收成终态）；消息仍会写入。
|
||
// 序号冲突时整段事务重试。其它失败回滚，run 保持原状态。
|
||
// SettleRunFenced（协作存储 U5/KTD4）：终态结算带执行围栏——事务内校验
|
||
// (instance,epoch) 仍是该会话的当前执行权，被接管则拒写（fenced=true，run
|
||
// 不变）。无租约行（未启用执行所有权）恒放行，语义与 SettleRun 一致。
|
||
func (s *Store) SettleRunFenced(r *core.Run, msg *core.Message, instanceID string, epoch int64) (updated, fenced bool, err error) {
|
||
return s.settleRunRetried(r, msg, &fenceRef{instance: instanceID, epoch: epoch, conversation: string(r.ConversationID)}, nil)
|
||
}
|
||
|
||
// SettleRunFencedAccountChecked adds a transaction-time account check to the
|
||
// execution fence. A disable committing first prevents a later business
|
||
// message from reaching the authoritative message tree.
|
||
func (s *Store) SettleRunFencedAccountChecked(r *core.Run, msg *core.Message, instanceID string, epoch int64, allowImplicitLocal bool) (bool, bool, error) {
|
||
return s.settleRunRetried(r, msg, &fenceRef{instance: instanceID, epoch: epoch, conversation: string(r.ConversationID)},
|
||
&accountSettleCheck{allowImplicitLocal: allowImplicitLocal})
|
||
}
|
||
|
||
// fenceRef 执行围栏引用（nil = 不校验）。
|
||
type fenceRef struct {
|
||
instance, conversation string
|
||
epoch int64
|
||
}
|
||
|
||
type accountSettleCheck struct{ allowImplicitLocal bool }
|
||
|
||
func (s *Store) SettleRun(r *core.Run, msg *core.Message) (updated bool, err error) {
|
||
updated, _, err = s.settleRunRetried(r, msg, nil, nil)
|
||
return updated, err
|
||
}
|
||
|
||
func (s *Store) SettleRunAccountChecked(r *core.Run, msg *core.Message, allowImplicitLocal bool) (bool, error) {
|
||
updated, _, err := s.settleRunRetried(r, msg, nil, &accountSettleCheck{allowImplicitLocal: allowImplicitLocal})
|
||
return updated, err
|
||
}
|
||
|
||
func (s *Store) settleRunRetried(r *core.Run, msg *core.Message, fence *fenceRef, account *accountSettleCheck) (updated, fenced bool, err error) {
|
||
s.settleMu.Lock()
|
||
fault := s.settleFault
|
||
s.settleMu.Unlock()
|
||
if fault != nil {
|
||
return false, false, fault
|
||
}
|
||
if msg != nil && msg.ID == "" {
|
||
msg.ID = newID("msg")
|
||
}
|
||
var last error
|
||
for attempt := 0; attempt < 8; attempt++ {
|
||
updated, fenced, last = s.settleRunOnce(r, msg, fence, account)
|
||
if last == nil {
|
||
return updated, fenced, nil
|
||
}
|
||
if !isUniqueViolation(last) {
|
||
return false, false, last
|
||
}
|
||
time.Sleep(time.Duration(attempt) * time.Millisecond)
|
||
}
|
||
if last == nil {
|
||
last = fmt.Errorf("store: settle retries exhausted (run %s)", r.ID)
|
||
}
|
||
return false, false, last
|
||
}
|
||
|
||
func (s *Store) settleRunOnce(r *core.Run, msg *core.Message, fence *fenceRef, account *accountSettleCheck) (bool, bool, error) {
|
||
tx, err := s.db.Begin()
|
||
if err != nil {
|
||
return false, false, err
|
||
}
|
||
defer tx.Rollback()
|
||
if account != nil {
|
||
query := s.q(`SELECT enabled FROM users WHERE owner=?`)
|
||
if s.Dialect() == "postgres" {
|
||
query += ` FOR UPDATE`
|
||
}
|
||
var enabled bool
|
||
err := tx.QueryRow(query, r.Owner).Scan(&enabled)
|
||
if errors.Is(err, sql.ErrNoRows) {
|
||
if account.allowImplicitLocal && r.Owner == OwnerDefault {
|
||
enabled, err = true, nil
|
||
} else {
|
||
return false, false, ErrAccountForbidden
|
||
}
|
||
}
|
||
if err != nil {
|
||
return false, false, err
|
||
}
|
||
if !enabled {
|
||
// An already identified account-denial failure may persist a minimal
|
||
// error state; it contains no assistant content or business output.
|
||
allowedFailure := r.Status == core.RunFailed && r.Error != nil && r.Error.Code == "run.account_inactive"
|
||
if msg != nil {
|
||
for _, part := range msg.Parts {
|
||
if part.Type != core.PartError || part.Err == nil || part.Err.Code != "run.account_inactive" {
|
||
allowedFailure = false
|
||
break
|
||
}
|
||
}
|
||
}
|
||
if !allowedFailure {
|
||
return false, false, ErrAccountForbidden
|
||
}
|
||
}
|
||
}
|
||
if err := s.dialect.lockConversation(tx, string(r.ConversationID)); err != nil {
|
||
return false, false, err
|
||
}
|
||
// 执行围栏（KTD4）：会话存在租约行且 (instance,epoch) 不匹配 → 拒写。
|
||
// 无行 = 该会话未启用执行所有权（本地档/旧路径），放行。
|
||
if fence != nil {
|
||
var n int
|
||
if err := tx.QueryRow(s.q(`SELECT COUNT(*) FROM conversation_executions WHERE conversation_id=?`), fence.conversation).Scan(&n); err != nil {
|
||
return false, false, err
|
||
}
|
||
if n > 0 {
|
||
var m int
|
||
if err := tx.QueryRow(s.q(`SELECT COUNT(*) FROM conversation_executions WHERE conversation_id=? AND instance_id=? AND epoch=?`), fence.conversation, fence.instance, fence.epoch).Scan(&m); err != nil {
|
||
return false, false, err
|
||
}
|
||
if m == 0 {
|
||
return false, true, nil // fenced：旧执行代次，run 不变
|
||
}
|
||
}
|
||
}
|
||
var ended any
|
||
if r.EndedAt != nil {
|
||
ended = nowISO()
|
||
}
|
||
res, err := tx.Exec(s.q(`UPDATE runs SET status=?, phase=?, usage=?, error=?, ended_at=? WHERE id=? AND status IN ('pending','running','waiting_hitl')`),
|
||
string(r.Status), string(r.Phase), marshalJSON(r.Usage), marshalJSON(r.Error), ended, r.ID)
|
||
if err != nil {
|
||
return false, false, err
|
||
}
|
||
n, err := res.RowsAffected()
|
||
if err != nil {
|
||
return false, false, err
|
||
}
|
||
if msg != nil {
|
||
if err := s.insertSettledMessageTx(tx, msg); err != nil {
|
||
return false, false, err
|
||
}
|
||
}
|
||
if err := tx.Commit(); err != nil {
|
||
return false, false, err
|
||
}
|
||
return n == 1, false, nil
|
||
}
|
||
|
||
func (s *Store) insertSettledMessageTx(tx *sql.Tx, msg *core.Message) error {
|
||
if msg == nil {
|
||
return nil
|
||
}
|
||
if msg.ID == "" {
|
||
msg.ID = newID("msg")
|
||
}
|
||
if msg.CreatedAt.IsZero() {
|
||
msg.CreatedAt = time.Now().UTC()
|
||
}
|
||
insert := `INSERT INTO messages(id,conversation_id,run_id,role,parts,seq,created_at,request_id)
|
||
SELECT ?,?,?,?,?,COALESCE(MAX(seq),0)+1,?,? FROM messages WHERE conversation_id=?`
|
||
if _, err := tx.Exec(s.q(insert), msg.ID, msg.ConversationID, msg.RunID, string(msg.Role), marshalJSON(msg.Parts), msg.CreatedAt.UTC().Format(time.RFC3339Nano), "", msg.ConversationID); err != nil {
|
||
return err
|
||
}
|
||
if err := s.appendProjectMessageObservationsTx(tx, msg); err != nil {
|
||
return err
|
||
}
|
||
if msg.VariantGroup != "" || msg.Excluded {
|
||
ex := 0
|
||
if msg.Excluded {
|
||
ex = 1
|
||
}
|
||
if _, err := tx.Exec(s.q(`UPDATE messages SET excluded=?, variant_group=?, variant_index=? WHERE id=?`), ex, msg.VariantGroup, msg.VariantIndex, msg.ID); err != nil {
|
||
return err
|
||
}
|
||
}
|
||
if msg.VariantGroup != "" {
|
||
if _, err := tx.Exec(s.q(`DELETE FROM regen_pending WHERE conversation_id=? AND variant_group=? AND variant_index=?`), msg.ConversationID, msg.VariantGroup, msg.VariantIndex); err != nil {
|
||
return err
|
||
}
|
||
}
|
||
if _, err := tx.Exec(s.q(`UPDATE conversations SET updated_at=? WHERE id=?`), nowISO(), msg.ConversationID); err != nil {
|
||
return err
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// ListInterruptedSubagentRuns 列出被收成 interrupted 的子代理 run
|
||
// （N2 WP5 重启续跑的候选集；幂等过滤由转录目录的 done/resumed 标记完成）。
|
||
func (s *Store) ListInterruptedSubagentRuns() ([]core.Run, error) {
|
||
rows, err := s.db.Query(s.q(`SELECT `+runColumns+` FROM runs WHERE kind=? AND status=? ORDER BY created_at`),
|
||
string(core.RunSubagent), string(core.RunInterrupted))
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
defer rows.Close()
|
||
out := []core.Run{}
|
||
for rows.Next() {
|
||
r, err := scanRun(rows.Scan)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
out = append(out, r)
|
||
}
|
||
return out, rows.Err()
|
||
}
|