Agent-Squad/squad/runstate.go

package squad

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

// RunState 引擎内部运行状态（磁盘即权威：workspace.json / workflow.json /
// squad.json / blackboard.json / pending.json 全部落盘，进程重启可恢复）。
type RunState struct {
	Meta     RunMetadata     `json:"meta"`
	Plan     SquadPlan       `json:"plan"`
	BB       BlackboardState `json:"bb"`
	Template string          `json:"template,omitempty"`
	// Approved 运行内一次性授权（permission/confirmation），键为 role:/step: 前缀。
	Approved map[string]bool `json:"approved,omitempty"`
	// RepairBudgetRemaining 质检失败后的自动重跑剩余次数。
	RepairBudgetRemaining int `json:"repair_budget_remaining,omitempty"`
	// RepairBudgetGranted 人工 retry 累计追加次数（Q26 上限判定用）。
	RepairBudgetGranted int `json:"repair_budget_granted,omitempty"`
	// Clarifications 已澄清问答（追加进 Goal 上下文）。
	Clarifications []string `json:"clarifications,omitempty"`
	// Steering 运行中人工补充指令（在下一步骤边界注入后清空）。
	Steering []string `json:"steering,omitempty"`
	// Mailbox 运行内点对点消息（Q2）：随 RunState 持久化；只在步骤启动时
	// 并入任务书，永不触碰 HITL 门。
	Mailbox []MailMessage `json:"mailbox,omitempty"`
	// Handoffs 所有权转移记录（Q2 handoff 原语）：随 RunState 持久化；
	// 目标角色的下一步骤一次性接收移交上下文，finalize 取链尾产出。
	Handoffs []HandoffRecord `json:"handoffs,omitempty"`
	// Board task_board 拓扑的队列状态（Q2）：claim/lease/complete 三原语，
	// 随 RunState 持久化；孤儿租约在驱动重入时再武装。
	Board []BoardTask `json:"board,omitempty"`
	// Usage 按角色累计的 token 用量。
	Usage map[string]*Usage `json:"usage,omitempty"`
	// Leader leader-delegation 模式的循环状态（workspace 内 leader.json 持久化）。
	Leader *LeaderState `json:"-"`
	// store 运行状态持久化端口（Q1 WP4）；nil 回落 FileStore（现状文件态）。
	// 替换实现（如 SQLiteStore）在引擎装配期注入。
	store Store `json:"-"`
	// mu 保护并行步骤对共享状态（黑板/save/steering/usage）的变更。
	mu sync.Mutex `json:"-"`
}

// storeOr 返回注入的 Store，缺省 FileStore。
func (st *RunState) storeOr() Store {
	if st.store != nil {
		return st.store
	}
	return FileStore{}
}

func (st *RunState) policy() HitlPolicy {
	return mergeHitlPolicy(HitlPolicy{}, st.Meta.HitlPolicyOf())
}

// HitlPolicyOf 从 Goal 推导策略（Goal 覆盖全局）。
func (m RunMetadata) HitlPolicyOf() *HitlPolicy {
	p := &HitlPolicy{}
	if m.Goal.HitlMode != "" {
		p.Mode = m.Goal.HitlMode
	}
	if m.Goal.HitlTimeout != 0 {
		p.TimeoutAfter = m.Goal.HitlTimeout
	}
	return p
}

func (st *RunState) save() error {
	st.Meta.UpdatedAt = time.Now()
	return st.storeOr().SaveRun(st)
}

// persist 是 FileStore 语义的低层落盘（全部写文件）。Store.SaveRun 调它；
// RunState.save 走 Store 端口，替换实现（SQLiteStore）在装配期注入后自然接管。
func (st *RunState) persist() error {
	if err := writeJSON(filepath.Join(st.Meta.WorkDir, "workspace.json"), &st.Meta); err != nil {
		return err
	}
	if err := writeJSON(filepath.Join(st.Meta.WorkDir, "workflow.json"), st.Plan.Workflow); err != nil {
		return err
	}
	if err := writeJSON(filepath.Join(st.Meta.WorkDir, "squad.json"), &st.Plan.Squad); err != nil {
		return err
	}
	if err := saveBlackboard(st.Meta.WorkDir, &st.BB); err != nil {
		return err
	}
	if len(st.Usage) > 0 {
		if err := writeJSON(filepath.Join(st.Meta.WorkDir, "usage.json"), st.Usage); err != nil {
			return err
		}
	}
	// steering.json 与其他 RunState 字段同权：SaveRun 后的 LoadRun 必须能读回
	// 未消费的补充指令（与 SQLiteStore 的 steering 列同一契约）。
	if err := writeJSON(filepath.Join(st.Meta.WorkDir, "steering.json"), st.Steering); err != nil {
		return err
	}
	return writeJSON(filepath.Join(st.Meta.WorkDir, "runstate.json"), diskRunState{
		Approved:              st.Approved,
		RepairBudgetRemaining: st.RepairBudgetRemaining,
		RepairBudgetGranted:   st.RepairBudgetGranted,
		Clarifications:        st.Clarifications,
		Mailbox:               st.Mailbox,
		Handoffs:              st.Handoffs,
		Board:                 st.Board,
	})
}

type diskRunState struct {
	Approved              map[string]bool `json:"approved,omitempty"`
	RepairBudgetRemaining int             `json:"repair_budget_remaining,omitempty"`
	RepairBudgetGranted   int             `json:"repair_budget_granted,omitempty"`
	Clarifications        []string        `json:"clarifications,omitempty"`
	Mailbox               []MailMessage   `json:"mailbox,omitempty"`
	Handoffs              []HandoffRecord `json:"handoffs,omitempty"`
	Board                 []BoardTask     `json:"board,omitempty"`
}

// Workflow 返回计划的编排清单（简洁访问器）。
func (st *RunState) Workflow() WorkflowPlan { return st.Plan.Workflow }

// setWorkflow 写回编排清单。
func (st *RunState) setWorkflow(wf WorkflowPlan) { st.Plan.Workflow = wf }

// saveGate 落盘 HITL 挂起上下文（pending.json）。
func (st *RunState) saveGate(gate *HitlGate) error {
	if gate == nil {
		err := os.Remove(filepath.Join(st.Meta.WorkDir, "pending.json"))
		if os.IsNotExist(err) {
			return nil
		}
		return err
	}
	return writeJSON(filepath.Join(st.Meta.WorkDir, "pending.json"), gate)
}

// loadRunState 从工作区恢复运行状态。
func loadRunState(wsDir string) (*RunState, error) {
	st := &RunState{}
	if err := readJSON(filepath.Join(wsDir, "workspace.json"), &st.Meta); err != nil {
		return nil, fmt.Errorf("load workspace.json: %w", err)
	}
	if err := readJSON(filepath.Join(wsDir, "workflow.json"), &st.Plan.Workflow); err != nil {
		return nil, fmt.Errorf("load workflow.json: %w", err)
	}
	if err := readJSON(filepath.Join(wsDir, "squad.json"), &st.Plan.Squad); err != nil {
		return nil, fmt.Errorf("load squad.json: %w", err)
	}
	bb, err := loadBlackboard(wsDir)
	if err != nil {
		return nil, err
	}
	st.BB = *bb
	if data, err := os.ReadFile(filepath.Join(wsDir, "template.md")); err == nil {
		st.Template = string(data)
	}
	if st.Approved == nil {
		st.Approved = map[string]bool{}
	}
	if data, err := os.ReadFile(filepath.Join(wsDir, "runstate.json")); err == nil {
		var disk diskRunState
		if err := json.Unmarshal(data, &disk); err != nil {
			return nil, fmt.Errorf("parse runstate.json: %w", err)
		}
		if disk.Approved != nil {
			st.Approved = disk.Approved
		}
		st.RepairBudgetRemaining = disk.RepairBudgetRemaining
		st.RepairBudgetGranted = disk.RepairBudgetGranted
		st.Clarifications = disk.Clarifications
		st.Mailbox = disk.Mailbox
		st.Handoffs = disk.Handoffs
		st.Board = disk.Board
	} else if !os.IsNotExist(err) {
		return nil, err
	}
	st.loadSteering()
	if data, err := os.ReadFile(filepath.Join(wsDir, "usage.json")); err == nil {
		_ = json.Unmarshal(data, &st.Usage)
	}
	return st, nil
}

// loadGate 恢复挂起上下文；无挂起返回 nil。
func loadGate(wsDir string) (*HitlGate, error) {
	data, err := os.ReadFile(filepath.Join(wsDir, "pending.json"))
	if err != nil {
		if os.IsNotExist(err) {
			return nil, nil
		}
		return nil, err
	}
	gate := &HitlGate{}
	if err := json.Unmarshal(data, gate); err != nil {
		return nil, fmt.Errorf("parse pending.json: %w", err)
	}
	return gate, nil
}

func readJSON(path string, v any) error {
	data, err := os.ReadFile(path)
	if err != nil {
		return err
	}
	return json.Unmarshal(data, v)
}

// setStep 并行安全地回填步骤状态（与 save 的全量 marshal 互斥）。
func (st *RunState) setStep(idx int, s StepSpec) {
	st.mu.Lock()
	st.Plan.Workflow.Steps[idx] = s
	st.mu.Unlock()
}

// injectSteering 把挂起的运行中人工补充指令追加到 prompt（下一步骤边界注入），
// 注入后清空并落盘。注入发生在步骤开始前，不会污染正在执行的步骤。
func (st *RunState) injectSteering(prompt string) string {
	st.mu.Lock()
	defer st.mu.Unlock()
	if len(st.Steering) == 0 {
		return prompt
	}
	var b strings.Builder
	b.WriteString(prompt)
	b.WriteString("\n\n【运行中人工补充指令】以下补充来自人工，作用于当前进行中的运行（非新任务）。")
	b.WriteString("请把它当作当前任务的补充要求继续执行；仅当人工明确要求放弃原目标时才替换原目标，且不要单独回复确认。\n")
	for _, s := range st.Steering {
		b.WriteString("- " + s + "\n")
	}
	st.Steering = nil
	_ = writeJSON(filepath.Join(st.Meta.WorkDir, "steering.json"), st.Steering)
	return b.String()
}

// addSteering 追加一条尚未注入的人工补充，并立刻落盘。
func (st *RunState) addSteering(text string) error {
	text = strings.TrimSpace(text)
	if text == "" {
		return nil
	}
	st.mu.Lock()
	defer st.mu.Unlock()
	st.Steering = append(st.Steering, text)
	return writeJSON(filepath.Join(st.Meta.WorkDir, "steering.json"), st.Steering)
}

// loadSteering 恢复未消费的补充指令。
func (st *RunState) loadSteering() {
	data, err := os.ReadFile(filepath.Join(st.Meta.WorkDir, "steering.json"))
	if err != nil {
		return
	}
	var pending []string
	if json.Unmarshal(data, &pending) == nil && len(pending) > 0 {
		st.Steering = pending
	}
}

// addUsage 累计角色用量。
func (st *RunState) addUsage(role string, promptTokens, completionTokens int) {
	if promptTokens == 0 && completionTokens == 0 {
		return
	}
	st.mu.Lock()
	defer st.mu.Unlock()
	if st.Usage == nil {
		st.Usage = map[string]*Usage{}
	}
	u, ok := st.Usage[role]
	if !ok {
		u = &Usage{}
		st.Usage[role] = u
	}
	u.PromptTokens += promptTokens
	u.CompletionTokens += completionTokens
	u.Steps++
}

// usageSnapshot 输出用量快照（值拷贝）。
func (st *RunState) usageSnapshot() map[string]Usage {
	if len(st.Usage) == 0 {
		return nil
	}
	out := make(map[string]Usage, len(st.Usage))
	for k, v := range st.Usage {
		out[k] = *v
	}
	return out
}