487 lines
13 KiB
Go
487 lines
13 KiB
Go
package squad
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"chat-agent-system/internal/agent"
|
|
)
|
|
|
|
// StepEvent 表示小队单步执行生命周期的事件
|
|
type StepEvent struct {
|
|
StepID string `json:"step_id"`
|
|
Role string `json:"role"`
|
|
Title string `json:"title"`
|
|
Mode string `json:"mode"`
|
|
Status string `json:"status"` // started, completed, failed, skipped
|
|
DurationMs int64 `json:"duration_ms,omitempty"`
|
|
Error string `json:"error,omitempty"`
|
|
}
|
|
|
|
// ExecuteOptions 执行调度选项
|
|
type ExecuteOptions struct {
|
|
Timeout time.Duration
|
|
StepNotify func(StepEvent)
|
|
}
|
|
|
|
// WorkflowExecutionResult 工作流执行结果
|
|
type WorkflowExecutionResult struct {
|
|
Success bool `json:"success"`
|
|
StepsRun int `json:"steps_run"`
|
|
Deliverable string `json:"deliverable"`
|
|
FinalContent []byte `json:"-"`
|
|
Duration time.Duration `json:"duration"`
|
|
Fields map[string]string `json:"fields,omitempty"`
|
|
ErrorReason string `json:"error_reason,omitempty"`
|
|
}
|
|
|
|
// WorkflowDriver 负责动态工作流调度与沙箱隔离执行
|
|
type WorkflowDriver struct {
|
|
engine *WorkspaceEngine
|
|
}
|
|
|
|
func NewWorkflowDriver(engine *WorkspaceEngine) *WorkflowDriver {
|
|
if engine == nil {
|
|
engine = NewWorkspaceEngine()
|
|
}
|
|
return &WorkflowDriver{engine: engine}
|
|
}
|
|
|
|
// LoadWorkflow 读取工作区 workflow.json
|
|
func (d *WorkflowDriver) LoadWorkflow(wsDir string) (*WorkflowPlan, error) {
|
|
p := filepath.Join(wsDir, "workflow.json")
|
|
data, err := os.ReadFile(p)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("read workflow.json: %w", err)
|
|
}
|
|
var wf WorkflowPlan
|
|
if err := json.Unmarshal(data, &wf); err != nil {
|
|
return nil, fmt.Errorf("unmarshal workflow.json: %w", err)
|
|
}
|
|
return &wf, nil
|
|
}
|
|
|
|
// SaveWorkflow 持久化 workflow.json
|
|
func (d *WorkflowDriver) SaveWorkflow(wsDir string, wf *WorkflowPlan) error {
|
|
p := filepath.Join(wsDir, "workflow.json")
|
|
return writeJSON(p, wf)
|
|
}
|
|
|
|
// LoadBlackboard 读取工作区 blackboard.json
|
|
func (d *WorkflowDriver) LoadBlackboard(wsDir string) (*BlackboardState, error) {
|
|
p := filepath.Join(wsDir, "blackboard.json")
|
|
data, err := os.ReadFile(p)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("read blackboard.json: %w", err)
|
|
}
|
|
var bb BlackboardState
|
|
if err := json.Unmarshal(data, &bb); err != nil {
|
|
return nil, fmt.Errorf("unmarshal blackboard.json: %w", err)
|
|
}
|
|
return &bb, nil
|
|
}
|
|
|
|
// SaveBlackboard 保存 blackboard.json
|
|
func (d *WorkflowDriver) SaveBlackboard(wsDir string, bb *BlackboardState) error {
|
|
bb.UpdatedAt = time.Now()
|
|
p := filepath.Join(wsDir, "blackboard.json")
|
|
return writeJSON(p, bb)
|
|
}
|
|
|
|
// AddStep 动态向工作流追加任务步骤并自动构建其子工作区
|
|
func (d *WorkflowDriver) AddStep(wsDir string, step StepSpec) error {
|
|
wf, err := d.LoadWorkflow(wsDir)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
if step.SubWorkspace == "" {
|
|
step.SubWorkspace = fmt.Sprintf("tasks/%02d_%s", len(wf.Steps)+1, step.Role)
|
|
}
|
|
if step.Status == "" {
|
|
step.Status = "pending"
|
|
}
|
|
if step.Mode == "" {
|
|
step.Mode = "execute"
|
|
}
|
|
|
|
// 准备子任务沙箱
|
|
squadCfg, _ := d.loadSquadConfig(wsDir)
|
|
tmpl, _ := os.ReadFile(filepath.Join(wsDir, "template.md"))
|
|
singleStepPlan := WorkflowPlan{Steps: []StepSpec{step}}
|
|
if err := d.engine.buildSubWorkspaces(wsDir, singleStepPlan, squadCfg, string(tmpl)); err != nil {
|
|
return fmt.Errorf("build added sub workspace: %w", err)
|
|
}
|
|
|
|
wf.Steps = append(wf.Steps, step)
|
|
return d.SaveWorkflow(wsDir, wf)
|
|
}
|
|
|
|
// RemoveStep 动态删除未执行的任务步骤
|
|
func (d *WorkflowDriver) RemoveStep(wsDir string, stepID string) error {
|
|
wf, err := d.LoadWorkflow(wsDir)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
idx := -1
|
|
for i, s := range wf.Steps {
|
|
if s.ID == stepID {
|
|
idx = i
|
|
break
|
|
}
|
|
}
|
|
if idx == -1 {
|
|
return fmt.Errorf("step %s not found in workflow", stepID)
|
|
}
|
|
|
|
wf.Steps = append(wf.Steps[:idx], wf.Steps[idx+1:]...)
|
|
return d.SaveWorkflow(wsDir, wf)
|
|
}
|
|
|
|
// UpdateStep 动态修改现有任务步骤（如重写验收标准或切换模式）
|
|
func (d *WorkflowDriver) UpdateStep(wsDir string, step StepSpec) error {
|
|
wf, err := d.LoadWorkflow(wsDir)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
found := false
|
|
for i, s := range wf.Steps {
|
|
if s.ID == step.ID {
|
|
wf.Steps[i] = step
|
|
found = true
|
|
break
|
|
}
|
|
}
|
|
if !found {
|
|
return fmt.Errorf("step %s not found to update", step.ID)
|
|
}
|
|
|
|
// 同步更新子工作区 TASK.md
|
|
squadCfg, _ := d.loadSquadConfig(wsDir)
|
|
tmpl, _ := os.ReadFile(filepath.Join(wsDir, "template.md"))
|
|
_ = d.engine.buildSubWorkspaces(wsDir, WorkflowPlan{Steps: []StepSpec{step}}, squadCfg, string(tmpl))
|
|
|
|
return d.SaveWorkflow(wsDir, wf)
|
|
}
|
|
|
|
// Execute 调度并执行工作区内的全套步骤
|
|
func (d *WorkflowDriver) Execute(ctx context.Context, wsDir string, backend agent.Backend, opts ExecuteOptions) (*WorkflowExecutionResult, error) {
|
|
start := time.Now()
|
|
|
|
wf, err := d.LoadWorkflow(wsDir)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
wf.Status = "in_progress"
|
|
_ = d.SaveWorkflow(wsDir, wf)
|
|
|
|
_ = d.updateWorkspaceStatus(wsDir, "running")
|
|
|
|
timeout := opts.Timeout
|
|
if timeout <= 0 {
|
|
timeout = 30 * time.Minute
|
|
}
|
|
runCtx, cancel := context.WithTimeout(ctx, timeout)
|
|
defer cancel()
|
|
|
|
var mu sync.Mutex
|
|
stepsRun := 0
|
|
var producedFiles []string
|
|
|
|
i := 0
|
|
for i < len(wf.Steps) {
|
|
step := wf.Steps[i]
|
|
if step.Parallel {
|
|
var parallelGroup []StepSpec
|
|
for i < len(wf.Steps) && wf.Steps[i].Parallel {
|
|
parallelGroup = append(parallelGroup, wf.Steps[i])
|
|
i++
|
|
}
|
|
|
|
var wg sync.WaitGroup
|
|
var groupErr error
|
|
for _, s := range parallelGroup {
|
|
sCopy := s
|
|
wg.Add(1)
|
|
go func() {
|
|
defer wg.Done()
|
|
if err := d.runOneStep(runCtx, wsDir, sCopy, backend, opts.StepNotify, &mu, &producedFiles); err != nil {
|
|
mu.Lock()
|
|
if groupErr == nil {
|
|
groupErr = err
|
|
}
|
|
mu.Unlock()
|
|
}
|
|
}()
|
|
}
|
|
wg.Wait()
|
|
if groupErr != nil {
|
|
wf.Status = "failed"
|
|
_ = d.SaveWorkflow(wsDir, wf)
|
|
_ = d.updateWorkspaceStatus(wsDir, "failed")
|
|
return &WorkflowExecutionResult{
|
|
Success: false,
|
|
ErrorReason: groupErr.Error(),
|
|
Duration: time.Since(start),
|
|
}, groupErr
|
|
}
|
|
stepsRun += len(parallelGroup)
|
|
} else {
|
|
i++
|
|
if err := d.runOneStep(runCtx, wsDir, step, backend, opts.StepNotify, &mu, &producedFiles); err != nil {
|
|
wf.Status = "failed"
|
|
_ = d.SaveWorkflow(wsDir, wf)
|
|
_ = d.updateWorkspaceStatus(wsDir, "failed")
|
|
return &WorkflowExecutionResult{
|
|
Success: false,
|
|
ErrorReason: err.Error(),
|
|
Duration: time.Since(start),
|
|
}, err
|
|
}
|
|
stepsRun++
|
|
}
|
|
}
|
|
|
|
// 收集最终交付物并汇总到工作区根目录 output.md
|
|
finalContent, _ := d.resolveFinalDeliverable(wsDir, wf)
|
|
if len(finalContent) == 0 {
|
|
// 容错：若未提取到内容，读取 template.md
|
|
finalContent, _ = os.ReadFile(filepath.Join(wsDir, "template.md"))
|
|
}
|
|
rootOut := filepath.Join(wsDir, "output.md")
|
|
_ = os.WriteFile(rootOut, finalContent, 0o644)
|
|
|
|
wf.Status = "completed"
|
|
_ = d.SaveWorkflow(wsDir, wf)
|
|
_ = d.updateWorkspaceStatus(wsDir, "completed")
|
|
|
|
bb, _ := d.LoadBlackboard(wsDir)
|
|
var finalFields map[string]string
|
|
if bb != nil {
|
|
finalFields = bb.Fields
|
|
}
|
|
|
|
return &WorkflowExecutionResult{
|
|
Success: true,
|
|
StepsRun: stepsRun,
|
|
Deliverable: rootOut,
|
|
FinalContent: finalContent,
|
|
Duration: time.Since(start),
|
|
Fields: finalFields,
|
|
}, nil
|
|
}
|
|
|
|
func (d *WorkflowDriver) runOneStep(
|
|
ctx context.Context,
|
|
wsDir string,
|
|
step StepSpec,
|
|
backend agent.Backend,
|
|
notify func(StepEvent),
|
|
mu *sync.Mutex,
|
|
producedFiles *[]string,
|
|
) error {
|
|
if mu == nil {
|
|
var localMu sync.Mutex
|
|
mu = &localMu
|
|
}
|
|
if producedFiles == nil {
|
|
var localFiles []string
|
|
producedFiles = &localFiles
|
|
}
|
|
|
|
subDir := filepath.Join(wsDir, filepath.FromSlash(step.SubWorkspace))
|
|
_ = os.MkdirAll(subDir, 0o755)
|
|
|
|
deliverableName := "output.md"
|
|
if step.Mode == "plan" || step.Mode == "replan" {
|
|
deliverableName = "output_plan.md"
|
|
} else if step.Mode == "review_plan" {
|
|
deliverableName = "review_notes.md"
|
|
}
|
|
|
|
if notify != nil {
|
|
notify(StepEvent{
|
|
StepID: step.ID,
|
|
Role: step.Role,
|
|
Title: step.Title,
|
|
Mode: step.Mode,
|
|
Status: "started",
|
|
})
|
|
}
|
|
stepStart := time.Now()
|
|
|
|
// 1. 读取 TASK.md 与 AGENTS.md
|
|
taskData, _ := os.ReadFile(filepath.Join(subDir, "TASK.md"))
|
|
agentsData, _ := os.ReadFile(filepath.Join(subDir, "AGENTS.md"))
|
|
|
|
targets := []string{deliverableName}
|
|
if step.Role == "organizer" {
|
|
targets = append(targets, "fields.json")
|
|
}
|
|
|
|
req := agent.Request{
|
|
WorkDir: subDir,
|
|
Prompt: string(taskData),
|
|
Instructions: string(agentsData),
|
|
Targets: targets,
|
|
Timeout: 30 * time.Minute,
|
|
}
|
|
|
|
res, err := backend.Execute(ctx, req)
|
|
duration := time.Since(stepStart).Milliseconds()
|
|
|
|
if err != nil {
|
|
if notify != nil {
|
|
notify(StepEvent{
|
|
StepID: step.ID,
|
|
Role: step.Role,
|
|
Title: step.Title,
|
|
Mode: step.Mode,
|
|
Status: "failed",
|
|
DurationMs: duration,
|
|
Error: err.Error(),
|
|
})
|
|
}
|
|
return fmt.Errorf("step %s (%s) backend execute: %w", step.ID, step.Role, err)
|
|
}
|
|
|
|
if res != nil && !res.Success {
|
|
errMsg := res.ErrorReason
|
|
if notify != nil {
|
|
notify(StepEvent{
|
|
StepID: step.ID,
|
|
Role: step.Role,
|
|
Title: step.Title,
|
|
Mode: step.Mode,
|
|
Status: "failed",
|
|
DurationMs: duration,
|
|
Error: errMsg,
|
|
})
|
|
}
|
|
return fmt.Errorf("step %s failed: %s", step.ID, errMsg)
|
|
}
|
|
|
|
// 2. 确认 deliverable 文件落盘
|
|
targetPath := filepath.Join(subDir, deliverableName)
|
|
if _, err := os.Stat(targetPath); err != nil {
|
|
// 若后端补丁在内存中返回，写入磁盘
|
|
if res != nil && len(res.Patches[deliverableName]) > 0 {
|
|
_ = os.WriteFile(targetPath, res.Patches[deliverableName], 0o644)
|
|
} else {
|
|
return fmt.Errorf("step %s missing target deliverable %s", step.ID, deliverableName)
|
|
}
|
|
}
|
|
|
|
// 3. 检查是否有 fields.json 并合并入 blackboard.json
|
|
fieldsPath := filepath.Join(subDir, "fields.json")
|
|
if fData, err := os.ReadFile(fieldsPath); err == nil {
|
|
var extractedFields map[string]string
|
|
if err := json.Unmarshal(fData, &extractedFields); err == nil {
|
|
mu.Lock()
|
|
if bb, err := d.LoadBlackboard(wsDir); err == nil {
|
|
if bb.Fields == nil {
|
|
bb.Fields = make(map[string]string)
|
|
}
|
|
for k, v := range extractedFields {
|
|
bb.Fields[k] = v
|
|
}
|
|
_ = d.SaveBlackboard(wsDir, bb)
|
|
}
|
|
mu.Unlock()
|
|
}
|
|
}
|
|
|
|
// 4. 记录日志
|
|
logEntry := fmt.Sprintf("[%s] 步骤 %s (%s) 执行完成，耗时 %dms\n",
|
|
time.Now().Format(time.RFC3339), step.ID, step.Role, duration)
|
|
_ = os.WriteFile(filepath.Join(subDir, "execution.log"), []byte(logEntry), 0o644)
|
|
|
|
mu.Lock()
|
|
*producedFiles = append(*producedFiles, filepath.ToSlash(filepath.Join(step.SubWorkspace, deliverableName)))
|
|
mu.Unlock()
|
|
|
|
if notify != nil {
|
|
notify(StepEvent{
|
|
StepID: step.ID,
|
|
Role: step.Role,
|
|
Title: step.Title,
|
|
Mode: step.Mode,
|
|
Status: "completed",
|
|
DurationMs: duration,
|
|
})
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (d *WorkflowDriver) resolveFinalDeliverable(wsDir string, wf *WorkflowPlan) ([]byte, string) {
|
|
// 查找优先级：leader 的 output.md > writer 的 output.md > 最后一个步骤的 output.md
|
|
for i := len(wf.Steps) - 1; i >= 0; i-- {
|
|
step := wf.Steps[i]
|
|
if step.Role == "leader" && (step.Mode == "execute" || step.Mode == "") {
|
|
p := filepath.Join(wsDir, filepath.FromSlash(step.SubWorkspace), "output.md")
|
|
if data, err := os.ReadFile(p); err == nil && len(strings.TrimSpace(string(data))) > 0 {
|
|
return data, p
|
|
}
|
|
}
|
|
}
|
|
|
|
for i := len(wf.Steps) - 1; i >= 0; i-- {
|
|
step := wf.Steps[i]
|
|
if step.Role == "writer" && (step.Mode == "execute" || step.Mode == "") {
|
|
p := filepath.Join(wsDir, filepath.FromSlash(step.SubWorkspace), "output.md")
|
|
if data, err := os.ReadFile(p); err == nil && len(strings.TrimSpace(string(data))) > 0 {
|
|
return data, p
|
|
}
|
|
}
|
|
}
|
|
|
|
for i := len(wf.Steps) - 1; i >= 0; i-- {
|
|
step := wf.Steps[i]
|
|
p := filepath.Join(wsDir, filepath.FromSlash(step.SubWorkspace), "output.md")
|
|
if data, err := os.ReadFile(p); err == nil && len(strings.TrimSpace(string(data))) > 0 {
|
|
return data, p
|
|
}
|
|
}
|
|
|
|
return nil, ""
|
|
}
|
|
|
|
func (d *WorkflowDriver) loadSquadConfig(wsDir string) (SquadConfig, error) {
|
|
p := filepath.Join(wsDir, "squad.json")
|
|
var sc SquadConfig
|
|
data, err := os.ReadFile(p)
|
|
if err != nil {
|
|
return sc, err
|
|
}
|
|
err = json.Unmarshal(data, &sc)
|
|
return sc, err
|
|
}
|
|
|
|
func (d *WorkflowDriver) updateWorkspaceStatus(wsDir string, status string) error {
|
|
p := filepath.Join(wsDir, "workspace.json")
|
|
data, err := os.ReadFile(p)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
var meta WorkspaceMetadata
|
|
if err := json.Unmarshal(data, &meta); err != nil {
|
|
return err
|
|
}
|
|
meta.Status = status
|
|
meta.UpdatedAt = time.Now()
|
|
if status == "completed" {
|
|
now := time.Now()
|
|
meta.CompletedAt = &now
|
|
}
|
|
return writeJSON(p, meta)
|
|
}
|