package agent import ( "context" "errors" "fmt" "os" "path/filepath" "sort" "sync" "time" "github.com/docker/docker/api/types/container" "github.com/docker/docker/api/types/filters" "github.com/docker/docker/client" ) // CoopTask 记录一次后台运行的 pi 协作任务。 // 任务由 run_coop 工具异步启动,完成后通过 CoopManager.Complete 通知主 agent。 type CoopTask struct { ID string `json:"id"` SessionID string `json:"session_id"` ContainerID string `json:"container_id,omitempty"` // ProcessID 是本地子进程模式的 worker 进程 PID(docker 模式为 0)。 ProcessID int `json:"process_id,omitempty"` Status string `json:"status"` // running | done ExitCode int `json:"exit_code,omitempty"` Output string `json:"output,omitempty"` Blackboard string `json:"blackboard,omitempty"` Error string `json:"error,omitempty"` RoundMax int `json:"round_max"` CreatedAt time.Time `json:"created_at"` FinishedAt time.Time `json:"finished_at,omitempty"` // Result 是 supervisor 写入黑板 result.json 的结构化结果(存在时优先于 Output)。 Result *CoopResult `json:"result,omitempty"` // ChallengeCode 任务描述的题目编号(如 c-06),用于通知控制层关闭/切换靶机 ChallengeCode string `json:"challenge_code,omitempty"` } // CoopResult 是 pi 协作容器 supervisor 生成的标准结果($BB/result.json), // 完成通知与主 agent 汇报轮直接使用这些结构化字段,不再依赖截断的容器 stdout。 type CoopResult struct { Status string `json:"status"` // solved | unsolved | error | timeout ExitCode int `json:"exit_code"` Summary string `json:"summary"` Flag string `json:"flag"` Artifacts []string `json:"artifacts"` } // Solved 返回任务是否成功解出/交付完成。 func (r *CoopResult) Solved() bool { return r != nil && r.Status == "solved" } // CoopManager 管理后台协作任务的生命周期,并在任务完成时回调 Agent, // 由 Agent 决定是否唤醒主 agent 汇报结果。 type CoopManager struct { mu sync.Mutex tasks map[string]*CoopTask notify func(task *CoopTask) } func NewCoopManager() *CoopManager { return &CoopManager{tasks: make(map[string]*CoopTask)} } // SetNotify 注册任务完成回调。 func (m *CoopManager) SetNotify(fn func(task *CoopTask)) { m.mu.Lock() m.notify = fn m.mu.Unlock() } // Register 登记一个正在运行的后台任务。 func (m *CoopManager) Register(task *CoopTask) { m.mu.Lock() m.tasks[task.ID] = task m.mu.Unlock() } // Get 按任务 ID 查询任务。 func (m *CoopManager) Get(id string) (*CoopTask, bool) { m.mu.Lock() defer m.mu.Unlock() task, ok := m.tasks[id] return task, ok } // List 返回全部任务,按创建时间倒序;sessionID 非空时只返回该会话的任务。 func (m *CoopManager) List(sessionID string) []*CoopTask { m.mu.Lock() defer m.mu.Unlock() out := make([]*CoopTask, 0, len(m.tasks)) for _, task := range m.tasks { if sessionID == "" || task.SessionID == sessionID { out = append(out, task) } } sort.Slice(out, func(i, j int) bool { return out[i].CreatedAt.After(out[j].CreatedAt) }) return out } // Complete 标记任务完成并触发完成回调(回调在锁外执行,避免死锁)。 func (m *CoopManager) Complete(task *CoopTask) { m.mu.Lock() task.Status = "done" task.FinishedAt = time.Now() m.tasks[task.ID] = task notify := m.notify m.mu.Unlock() if notify != nil { notify(task) } } // CurrentCoopRound 统计黑板 logs 目录中已出现的最大协作轮次。 // 单 agent 版日志命名为 round-N.log;兼容旧版 a-round-N.log / b-round-N.log。 func CurrentCoopRound(blackboardDir string) int { if blackboardDir == "" { return 0 } entries, err := os.ReadDir(filepath.Join(blackboardDir, "logs")) if err != nil { return 0 } maxRound := 0 for _, entry := range entries { var round int if _, err := fmt.Sscanf(entry.Name(), "round-%d.log", &round); err == nil { if round > maxRound { maxRound = round } continue } if _, err := fmt.Sscanf(entry.Name(), "a-round-%d.log", &round); err == nil { if round > maxRound { maxRound = round } } if _, err := fmt.Sscanf(entry.Name(), "b-round-%d.log", &round); err == nil { if round > maxRound { maxRound = round } } } return maxRound } // ContainerInspectStatus 查询容器当前的运行状态与人类可读状态描述 // (如 "running" + "Up 2 minutes" / "exited" + "Exited (0)"), // 供 Web 页面实时展示协作容器状态。 func ContainerInspectStatus(socket, containerID string) (state, status string, err error) { if containerID == "" { return "", "", errors.New("缺少容器 ID") } cli, err := client.NewClientWithOpts( client.WithHost(socket), client.WithAPIVersionNegotiation(), ) if err != nil { return "", "", err } defer cli.Close() ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() list, err := cli.ContainerList(ctx, container.ListOptions{ All: true, Filters: filters.NewArgs(filters.Arg("id", containerID)), }) if err != nil { return "", "", err } if len(list) == 0 { return "", "", errors.New("容器不存在") } return list[0].State, list[0].Status, nil }