BlackBean/internal/agent/coop_async.go

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
}