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"`
	// Clarifications 已澄清问答（追加进 Goal 上下文）。
	Clarifications []string `json:"clarifications,omitempty"`
	// Steering 运行中人工补充指令（在下一步骤边界注入后清空）。
	Steering []string `json:"steering,omitempty"`
	// Usage 按角色累计的 token 用量。
	Usage map[string]*Usage `json:"usage,omitempty"`
	// Leader leader-delegation 模式的循环状态（workspace 内 leader.json 持久化）。
	Leader *LeaderState `json:"-"`
	// mu 保护并行步骤对共享状态（黑板/save/steering/usage）的变更。
	mu sync.Mutex `json:"-"`
}

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()
	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
		}
	}
	return writeJSON(filepath.Join(st.Meta.WorkDir, "runstate.json"), diskRunState{
		Approved:              st.Approved,
		RepairBudgetRemaining: st.RepairBudgetRemaining,
		Clarifications:        st.Clarifications,
	})
}

type diskRunState struct {
	Approved              map[string]bool `json:"approved,omitempty"`
	RepairBudgetRemaining int             `json:"repair_budget_remaining,omitempty"`
	Clarifications        []string        `json:"clarifications,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.Clarifications = disk.Clarifications
	} 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()
}

// 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
}