first commit
This commit is contained in:
@@ -0,0 +1,180 @@
|
||||
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
|
||||
}
|
||||
Reference in New Issue
Block a user