685 lines
26 KiB
Go
685 lines
26 KiB
Go
package agent
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"log"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
"unicode"
|
|
)
|
|
|
|
type Agent struct {
|
|
cfg Config
|
|
llm *LLMClient
|
|
store *SessionStore
|
|
tools *Toolset
|
|
docker *DockerConfigStore
|
|
apiCfg *APIConfigStore
|
|
coop *CoopManager
|
|
locks map[string]*sync.Mutex
|
|
locksMu sync.Mutex
|
|
liveMu sync.Mutex
|
|
live map[string]*liveReg
|
|
}
|
|
|
|
// liveReg 记录某会话当前"在线"的 SSE 连接,用于推送后台任务完成通知。
|
|
type liveReg struct {
|
|
emit func(Event)
|
|
}
|
|
|
|
type Event struct {
|
|
Type string `json:"type"`
|
|
Data any `json:"data,omitempty"`
|
|
}
|
|
|
|
func NewAgent(cfg Config, store *SessionStore, dockerCfg *DockerConfigStore, apiCfg *APIConfigStore) *Agent {
|
|
agent := &Agent{
|
|
cfg: cfg,
|
|
llm: NewLLMClient(cfg, apiCfg),
|
|
store: store,
|
|
docker: dockerCfg,
|
|
apiCfg: apiCfg,
|
|
locks: make(map[string]*sync.Mutex),
|
|
live: make(map[string]*liveReg),
|
|
}
|
|
agent.tools = NewToolset(cfg, dockerCfg, apiCfg)
|
|
agent.coop = NewCoopManager()
|
|
agent.tools.coop = agent.coop
|
|
agent.coop.SetNotify(agent.notifyCoopDone)
|
|
return agent
|
|
}
|
|
|
|
func (a *Agent) Run(ctx context.Context, sessionID, userContent string, emit func(Event)) error {
|
|
if strings.TrimSpace(userContent) == "" {
|
|
emit(Event{Type: "error", Data: map[string]any{"message": "消息内容不能为空"}})
|
|
return errors.New("消息内容不能为空")
|
|
}
|
|
|
|
lock := a.sessionLock(sessionID)
|
|
lock.Lock()
|
|
defer lock.Unlock()
|
|
|
|
session, ok := a.store.Get(sessionID)
|
|
if !ok {
|
|
emit(Event{Type: "error", Data: map[string]any{"message": "会话不存在"}})
|
|
return errors.New("会话不存在")
|
|
}
|
|
// 在私有副本上操作,避免修改持久化对象与其他会话的序列化产生数据竞争
|
|
session = cloneSession(session)
|
|
|
|
if session.Title == "" || session.Title == "新对话" {
|
|
session.Title = firstRunes(userContent, 30)
|
|
}
|
|
session.Messages = append(session.Messages, Message{
|
|
Role: "user",
|
|
Content: stringPointer(userContent),
|
|
})
|
|
a.store.Save(session)
|
|
|
|
emit(Event{Type: "meta", Data: map[string]any{
|
|
"session_id": session.ID,
|
|
"title": session.Title,
|
|
}})
|
|
|
|
history := session.Messages
|
|
compressed := false
|
|
iterations := 0
|
|
emptyRetries := 0
|
|
|
|
for {
|
|
iterations++
|
|
systemMessage := Message{
|
|
Role: "system",
|
|
Content: stringPointer(a.systemPrompt(sessionID)),
|
|
}
|
|
apiMessages := append([]Message{systemMessage}, history...)
|
|
|
|
if estimateMessages(apiMessages) > a.cfg.MaxContextTokens {
|
|
if summary, err := a.compressHistory(ctx, sessionID, history); err != nil {
|
|
history = truncateHistory(history, a.cfg.KeepRecentMessages)
|
|
compressed = true
|
|
emit(Event{Type: "notice", Data: map[string]any{
|
|
"text": "上下文较长,自动压缩失败,已降级为截断早期对话: " + err.Error(),
|
|
}})
|
|
} else {
|
|
history = summary
|
|
compressed = true
|
|
emit(Event{Type: "notice", Data: map[string]any{
|
|
"text": "上下文较长,已自动压缩早期对话(分块摘要),关键信息会保留。",
|
|
}})
|
|
}
|
|
}
|
|
|
|
apiMessages = append([]Message{systemMessage}, history...)
|
|
assistant, err := a.streamAssistantReply(ctx, apiMessages, a.tools.Definitions(), emit)
|
|
if err != nil {
|
|
// 模型返回空内容:强制压缩上下文后重试,而非直接中断会话。
|
|
// 长会话(如跑分循环累积大量工具输出)时模型容易因输入过大而返回空,
|
|
// 压缩后通常可恢复;多次仍空则优雅退出,保留会话状态供用户继续。
|
|
if errors.Is(err, errEmptyResponse) {
|
|
emptyRetries++
|
|
if emptyRetries <= 2 {
|
|
emit(Event{Type: "notice", Data: map[string]any{
|
|
"text": "模型返回空内容,正在压缩上下文后重试…",
|
|
}})
|
|
if summary, serr := a.compressHistory(ctx, sessionID, history); serr == nil {
|
|
history = summary
|
|
} else {
|
|
history = truncateHistory(history, a.cfg.KeepRecentMessages)
|
|
}
|
|
session.Messages = history
|
|
a.store.Save(session)
|
|
continue
|
|
}
|
|
session.Messages = history
|
|
a.store.Save(session)
|
|
emit(Event{Type: "notice", Data: map[string]any{
|
|
"text": "模型多次返回空内容,本轮已停止。你可以继续发消息让我接着处理(程序不会退出)。",
|
|
}})
|
|
emit(Event{Type: "done", Data: map[string]any{
|
|
"compressed": compressed,
|
|
"iterations": iterations,
|
|
}})
|
|
return nil
|
|
}
|
|
session.Messages = history
|
|
a.store.Save(session)
|
|
if errors.Is(err, context.Canceled) {
|
|
// 用户主动停止(前端 stop 按钮 / Esc):是正常中断,不按错误展示
|
|
emit(Event{Type: "notice", Data: map[string]any{
|
|
"text": "已停止本次生成,已完成的对话进度已保存,可继续发送新消息。",
|
|
}})
|
|
} else {
|
|
emit(Event{Type: "error", Data: map[string]any{"message": err.Error()}})
|
|
}
|
|
return err
|
|
}
|
|
emptyRetries = 0
|
|
|
|
history = append(history, assistant)
|
|
if len(assistant.ToolCalls) == 0 {
|
|
break
|
|
}
|
|
|
|
// 本轮是否成功下发了协作任务:分发完成后本轮直接收尾,
|
|
// 避免控制层在 run_coop 之后继续空转(sleep 等待 / 反复轮询),
|
|
// coop 完成后会以 user 身份推送完成通知并启动新一轮处理。
|
|
dispatched := false
|
|
for _, call := range assistant.ToolCalls {
|
|
emit(Event{Type: "tool", Data: map[string]any{
|
|
"call_id": call.ID,
|
|
"name": call.Function.Name,
|
|
"input": call.Function.Arguments,
|
|
}})
|
|
|
|
result := a.tools.Execute(ctx, sessionID, call.Function.Name, call.Function.Arguments)
|
|
result.Output = truncateRunes(result.Output, a.cfg.MaxToolResultChars)
|
|
history = append(history, Message{
|
|
Role: "tool",
|
|
ToolCallID: call.ID,
|
|
Content: stringPointer(result.Output),
|
|
})
|
|
if call.Function.Name == "run_coop" && result.Success {
|
|
dispatched = true
|
|
}
|
|
|
|
emit(Event{Type: "tool_result", Data: map[string]any{
|
|
"call_id": call.ID,
|
|
"name": call.Function.Name,
|
|
"output": result.Output,
|
|
"success": result.Success,
|
|
}})
|
|
}
|
|
|
|
// 每轮迭代后落盘:长任务(如跑分循环)期间用户刷新页面也能看到
|
|
// 已完成轮次的消息与工具结果,而不是只能等整轮 Run 结束。
|
|
session.Messages = history
|
|
a.store.Save(session)
|
|
|
|
if dispatched {
|
|
emit(Event{Type: "notice", Data: map[string]any{
|
|
"text": "协作任务已全部下发,本轮收尾结束。任务完成后会自动以用户身份推送完成通知并启动新一轮处理,无需在等待上消耗本轮回合。",
|
|
}})
|
|
break
|
|
}
|
|
|
|
if iterations >= a.cfg.MaxIterations {
|
|
emit(Event{Type: "notice", Data: map[string]any{
|
|
"text": fmt.Sprintf("已达到 %d 次工具调用上限,本轮先基于已有结果收尾。你可以继续发消息让我接着处理(程序不会退出)。", a.cfg.MaxIterations),
|
|
}})
|
|
// 达到上限并不中断程序:用一次不带工具的收尾回答结束本回合,
|
|
// 避免在工具链中途戛然而止、用户得不到任何总结。
|
|
closingSystem := Message{
|
|
Role: "system",
|
|
Content: stringPointer(a.systemPrompt(sessionID) +
|
|
"\n\n注意:本轮已达到工具调用次数上限。现在必须直接输出最终回答:总结目前已确定的结果、未完成事项与建议的下一步,不要再调用任何工具。"),
|
|
}
|
|
if closing, err := a.streamAssistantReply(ctx, append([]Message{closingSystem}, history...), nil, emit); err == nil {
|
|
history = append(history, closing)
|
|
}
|
|
break
|
|
}
|
|
}
|
|
|
|
session.Messages = history
|
|
a.store.Save(session)
|
|
emit(Event{Type: "done", Data: map[string]any{
|
|
"compressed": compressed,
|
|
"iterations": iterations,
|
|
}})
|
|
return nil
|
|
}
|
|
|
|
// errEmptyResponse 表示模型流式返回既无文本也无工具调用。
|
|
// 区别于普通错误:调用方应压缩上下文后重试,而非直接中断会话。
|
|
var errEmptyResponse = errors.New("模型没有返回任何内容")
|
|
|
|
func (a *Agent) streamAssistantReply(ctx context.Context, messages []Message, tools []ToolDefinition, emit func(Event)) (Message, error) {
|
|
const maxRetries = 2
|
|
for attempt := 0; attempt <= maxRetries; attempt++ {
|
|
ch, err := a.llm.ChatStream(ctx, messages, tools)
|
|
if err != nil {
|
|
return Message{}, err
|
|
}
|
|
|
|
var content strings.Builder
|
|
var calls []ToolCall
|
|
for chunk := range ch {
|
|
if chunk.Error != nil {
|
|
return Message{}, chunk.Error
|
|
}
|
|
if chunk.Content != "" {
|
|
content.WriteString(chunk.Content)
|
|
emit(Event{Type: "message", Data: map[string]any{"delta": chunk.Content}})
|
|
}
|
|
for _, delta := range chunk.ToolCalls {
|
|
calls = mergeToolCall(calls, delta)
|
|
}
|
|
}
|
|
|
|
for i := range calls {
|
|
if calls[i].ID == "" {
|
|
calls[i].ID = "call_" + newID()
|
|
}
|
|
}
|
|
|
|
assistant := Message{
|
|
Role: "assistant",
|
|
Content: stringPointer(content.String()),
|
|
}
|
|
if len(calls) > 0 {
|
|
assistant.ToolCalls = calls
|
|
if content.Len() == 0 {
|
|
assistant.Content = nil
|
|
}
|
|
}
|
|
if content.Len() == 0 && len(calls) == 0 {
|
|
if attempt < maxRetries {
|
|
emit(Event{Type: "notice", Data: map[string]any{
|
|
"text": fmt.Sprintf("模型返回空内容,正在重试(第 %d/%d 次)…", attempt+1, maxRetries),
|
|
}})
|
|
continue
|
|
}
|
|
return Message{}, errEmptyResponse
|
|
}
|
|
return assistant, nil
|
|
}
|
|
return Message{}, errEmptyResponse
|
|
}
|
|
|
|
// compressHistory 压缩 history 中除最近 keep 条之外的全部早期消息。
|
|
// 早期消息可能远超单次模型输入窗口,因此按压缩预算分块、逐块摘要后再合并;
|
|
// LLM 压缩失败时用 fallbackSummary 保留最近消息要点,避免直接整体截断。
|
|
// 合并后若仍超窗口,对合并结果再压缩一层(递归收敛),极端情况才截断。
|
|
func (a *Agent) compressHistory(ctx context.Context, sessionID string, history []Message) ([]Message, error) {
|
|
keep := a.cfg.KeepRecentMessages
|
|
if keep <= 0 {
|
|
keep = 12
|
|
}
|
|
if len(history) <= keep {
|
|
return history, nil
|
|
}
|
|
old := history[:len(history)-keep]
|
|
recent := history[len(history)-keep:]
|
|
|
|
summary, err := a.summarizeBlocks(ctx, old)
|
|
if err != nil {
|
|
// LLM 压缩失败:降级为文本要点摘要(纯文本,不产生孤儿 tool 消息)
|
|
summary = fallbackSummary(old)
|
|
}
|
|
merged := append([]Message{{
|
|
Role: "system",
|
|
Content: stringPointer("[早期对话摘要]\n" + summary),
|
|
}}, recent...)
|
|
|
|
// 压缩后校验:摘要 + 最近消息仍超窗口时,对合并结果再压缩一层(递归收敛)
|
|
if estimateMessages(merged) > a.cfg.MaxContextTokens {
|
|
if inner, inerr := a.compressHistory(ctx, sessionID, merged); inerr == nil {
|
|
return inner, nil
|
|
}
|
|
merged = truncateHistory(merged, keep)
|
|
}
|
|
return merged, nil
|
|
}
|
|
|
|
// summarizeBlocks 把早期消息按单次压缩输入预算切成若干块,逐块交给 LLM 摘要,
|
|
// 最后把各块摘要按顺序合并返回。任一块失败即返回错误(由调用方降级)。
|
|
func (a *Agent) summarizeBlocks(ctx context.Context, history []Message) (string, error) {
|
|
budget := a.compactionInputBudget()
|
|
var summaries []string
|
|
block := make([]Message, 0, 32)
|
|
blockTokens := 0
|
|
flush := func() error {
|
|
if len(block) == 0 {
|
|
return nil
|
|
}
|
|
sum, cerr := a.llm.Compress(ctx, block)
|
|
block = block[:0]
|
|
blockTokens = 0
|
|
if cerr != nil {
|
|
return cerr
|
|
}
|
|
summaries = append(summaries, sum)
|
|
return nil
|
|
}
|
|
for _, msg := range history {
|
|
tokens := estimateMessageTokens(msg)
|
|
if len(block) > 0 && blockTokens+tokens > budget {
|
|
if err := flush(); err != nil {
|
|
return "", err
|
|
}
|
|
}
|
|
block = append(block, msg)
|
|
blockTokens += tokens
|
|
}
|
|
if err := flush(); err != nil {
|
|
return "", err
|
|
}
|
|
return strings.Join(summaries, "\n\n"), nil
|
|
}
|
|
|
|
// compactionInputBudget 返回单次 LLM 压缩请求的输入 token 预算:
|
|
// 总窗口减去压缩输出(CompactionTokens)与常规输出(StreamOutputTokens)的预留,
|
|
// 再留出安全余量,避免把输入塞满窗口导致请求被拒。
|
|
func (a *Agent) compactionInputBudget() int {
|
|
window := a.cfg.MaxContextTokens
|
|
output := a.cfg.CompactionTokens + a.cfg.StreamOutputTokens
|
|
budget := window - output - window/8
|
|
if budget < 4096 {
|
|
budget = 4096
|
|
}
|
|
return budget
|
|
}
|
|
|
|
// fallbackSummary 当 LLM 压缩失败时的兜底摘要:按时间顺序保留早期消息中
|
|
// 最近若干条的用户目标/助手结论/工具结果要点,输出纯文本(不产生孤立的
|
|
// tool 消息),保证后续请求协议合法、信息尽量不丢。
|
|
func fallbackSummary(history []Message) string {
|
|
const maxLines = 10
|
|
start := len(history) - maxLines
|
|
if start < 0 {
|
|
start = 0
|
|
}
|
|
labels := map[string]string{
|
|
"user": "用户",
|
|
"assistant": "助手",
|
|
"tool": "工具结果",
|
|
}
|
|
var lines []string
|
|
for i := start; i < len(history); i++ {
|
|
msg := history[i]
|
|
label, ok := labels[msg.Role]
|
|
if !ok {
|
|
label = msg.Role
|
|
}
|
|
text := ""
|
|
if msg.Content != nil {
|
|
text = strings.TrimSpace(*msg.Content)
|
|
}
|
|
if text == "" {
|
|
if len(msg.ToolCalls) > 0 {
|
|
text = "调用了 " + msg.ToolCalls[0].Function.Name + " 等工具"
|
|
} else {
|
|
continue
|
|
}
|
|
}
|
|
if countRunes(text) > 200 {
|
|
text = truncateRunes(text, 200)
|
|
}
|
|
lines = append(lines, fmt.Sprintf("[%s] %s", label, text))
|
|
}
|
|
if len(lines) == 0 {
|
|
return "(早期对话内容较多,摘要失败后仅保留最近消息。)"
|
|
}
|
|
return "(早期对话压缩失败,以下为最近消息要点,更早内容已截断)\n" + strings.Join(lines, "\n")
|
|
}
|
|
|
|
// truncateHistory 在模型压缩失败时兜底:仅保留最近消息,并丢弃开头孤立的
|
|
// tool 结果消息,避免出现没有对应 assistant 调用声明的 tool 消息导致协议错乱。
|
|
func truncateHistory(history []Message, keep int) []Message {
|
|
if len(history) <= keep {
|
|
return history
|
|
}
|
|
kept := append([]Message(nil), history[len(history)-keep:]...)
|
|
for len(kept) > 0 && kept[0].Role == "tool" {
|
|
kept = kept[1:]
|
|
}
|
|
return kept
|
|
}
|
|
|
|
func (a *Agent) systemPrompt(sessionID string) string {
|
|
return fmt.Sprintf(`你是 BlackBean,运行在用户本机上的 AI Agent,正在协助用户完成真实任务。
|
|
当前日期:%s
|
|
本会话工作目录(与其他会话隔离,相对路径与命令默认在此执行):%s
|
|
项目根目录(可读取,含 agent 自身文件、协作黑板产物等共享内容):%s
|
|
|
|
工作方式:
|
|
1. 先理解用户目标,必要时先用工具查看文件、目录和运行结果,再给出准确答案。
|
|
2. 你可以运行 Bash、读写文件、运行 Python。所有文件和命令默认都在本会话工作目录内进行,不要访问项目根目录之外的文件。
|
|
3. 使用工具时必须根据真实输出继续推理,绝不能编造命令结果、文件内容或错误信息。
|
|
4. 回答使用简洁的中文,重要代码、命令和路径用 Markdown 代码块呈现。
|
|
5. 需要执行可能造成不可逆影响的操作前,先说明风险和影响,再谨慎执行。
|
|
6. 如果一次工具调用没有解决问题,可以多次调用工具排查;不要在没有依据时下结论。
|
|
|
|
|
|
【任务调度规则】
|
|
- 解题类任务(CTF 题目、渗透测试、漏洞挖掘、Web 攻防、逆向、密码学、取证分析等)一律交由 coop 执行,你不负责具体解题,只做调度。
|
|
- 你的调度职责:
|
|
1. 从用户表述中提取足够信息:目标地址/文件、约束条件、平台提交规则(如 TSec Benchmark 的 submit API、BENCHMARK_TOKEN、unique_code 等);
|
|
2. 调用 run_coop 工具,把任务描述与提交规则写清楚,交给 pi 协作单 Agent 执行;
|
|
3. 任务完成后按完成通知中的结构化结果向用户汇报,并按约定规则提交 flag、回报得分。
|
|
- 主动并发调度(默认行为,无需用户提醒):解题跑分类任务要最大化并发,始终把在跑的靶机/协作任务维持到平台并发上限(TSec Benchmark 通常为 3 个并发容器),不要一次只开一个靶机慢慢等:
|
|
1. 开始解题时:先查看题目列表,选定多道未完成题目,用 start 接口一次性启动多个靶机,并逐一调用 run_coop 下发对应协作任务,把并发拉满;
|
|
2. 每有任务完成/失败释放名额时:立即主动 close 已结束题目 → start 下一道未完成题目 → run_coop 下发,始终补满并发,不要空出名额等待;
|
|
3. 只有平台已无可启动题目、或用户明确要求停止时才停止扩并发。
|
|
- 主动检测状态(默认行为,无需用户提醒):在收到协作完成通知、下发新任务、或推进到关键节点时,主动用短命令(≤60s)查询平台进度(GET challenges)与运行中容器(docker_ps / run_bash),并向用户简要汇报「已通关题数/总分、运行中靶机、下一步计划」,不要等用户说「检测状态」才去查。
|
|
- 调用 run_coop 成功下发目标后,本轮会话立即收尾结束:直接输出一句简短总结即可,不要原地等待、不要 sleep、不要反复轮询黑板。coop 完成后系统会自动以「用户」身份在会话中追加完成通知并启动新一轮处理,届时你再读取黑板产物、提交 flag、规划下一题。
|
|
- coop 拿到 flag 后,必须按任务中约定的平台提交规则(如 TSec Benchmark 的 submit API、BENCHMARK_TOKEN、unique_code)用 curl 提交 flag,并把提交响应与得分回报给用户;提交失败要重试并说明原因,不能只汇报不提交。
|
|
- 不要用 sleep 或超时超过 60 秒的长等待命令原地等待 coop 任务完成:coop 是异步运行的,任务完成后系统会自动推送「协作任务完成通知」并触发你继续处理,无需 sleep 阻塞。等待期间可以用短命令(≤60s)检查黑板进度,或直接推进其他不依赖该任务的工作(如规划下一题)。
|
|
- 非解题类任务(文档撰写、代码开发、日常问答等)按常规方式由你自己完成,不必交给 coop。
|
|
`,
|
|
time.Now().Format("2006-01-02"), a.tools.workspaceFor(sessionID), a.cfg.Workspace)
|
|
}
|
|
|
|
func (a *Agent) sessionLock(sessionID string) *sync.Mutex {
|
|
a.locksMu.Lock()
|
|
defer a.locksMu.Unlock()
|
|
lock := a.locks[sessionID]
|
|
if lock == nil {
|
|
lock = &sync.Mutex{}
|
|
a.locks[sessionID] = lock
|
|
}
|
|
return lock
|
|
}
|
|
|
|
// ForgetSession 释放会话持有的运行时资源(会话锁与 live 注册)。
|
|
// 在删除会话时调用,避免 locks / live map 随会话累积只增不减。
|
|
func (a *Agent) ForgetSession(sessionID string) {
|
|
a.locksMu.Lock()
|
|
delete(a.locks, sessionID)
|
|
a.locksMu.Unlock()
|
|
|
|
a.liveMu.Lock()
|
|
delete(a.live, sessionID)
|
|
a.liveMu.Unlock()
|
|
}
|
|
|
|
// RegisterLive 把某个 SSE 连接注册为会话的"在线接收者",供后台协作任务完成
|
|
func (a *Agent) RegisterLive(sessionID string, emit func(Event)) func() {
|
|
a.liveMu.Lock()
|
|
reg := &liveReg{emit: emit}
|
|
a.live[sessionID] = reg
|
|
a.liveMu.Unlock()
|
|
return func() {
|
|
a.liveMu.Lock()
|
|
if a.live[sessionID] == reg {
|
|
delete(a.live, sessionID)
|
|
}
|
|
a.liveMu.Unlock()
|
|
}
|
|
}
|
|
|
|
// CoopTasks 返回协作任务列表;sessionID 非空时只返回该会话的任务,
|
|
// 供 Web 页面实时展示容器运行状态。
|
|
func (a *Agent) CoopTasks(sessionID string) []*CoopTask {
|
|
return a.coop.List(sessionID)
|
|
}
|
|
|
|
// liveEmitter 返回动态事件发射器:每次发送时实时查询该会话当前在线的连接。
|
|
// 与启动时一次性捕获 emit 相比,即使汇报轮在页面连接之前启动,页面打开后
|
|
// 也能实时收到流式事件(修复"汇报轮事件前端收不到"的问题)。
|
|
func (a *Agent) liveEmitter(sessionID string) func(Event) {
|
|
return func(event Event) {
|
|
a.liveMu.Lock()
|
|
reg := a.live[sessionID]
|
|
a.liveMu.Unlock()
|
|
if reg != nil {
|
|
reg.emit(event)
|
|
}
|
|
}
|
|
}
|
|
|
|
// notifyCoopDone 是后台协作任务完成后的回调:给会话注入一条
|
|
// "【协作任务完成通知】"用户消息,并触发主 agent 新一轮处理
|
|
// (有在线 SSE 连接则流式推送,否则在后台静默完成并持久化结果)。
|
|
// 通知基于 supervisor 的结构化结果(CoopResult)渲染,成功时携带 flag/摘要/
|
|
// 产物清单,失败时携带明确的续跑指令(关闭靶机 → 开下一题 → 重新下发 coop)。
|
|
// 汇报轮失败会在会话中持久化错误信息,避免"通知已到但流程静默停摆"。
|
|
func (a *Agent) notifyCoopDone(task *CoopTask) {
|
|
notif := buildCoopNotification(task)
|
|
|
|
emit := a.liveEmitter(task.SessionID)
|
|
// 汇报轮加超时兜底:后台自动触发的 Run 不能永久持锁阻塞后续所有通知
|
|
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Minute)
|
|
emit(Event{Type: "turn_start", Data: map[string]any{"user": notif}})
|
|
// 使用独立后台上下文:即使前端已关闭页面,汇报轮仍会完成并持久化。
|
|
// 注意:cancel 必须在 Run 结束后调用(defer 到 goroutine 内),
|
|
// 否则本函数一返回 ctx 即被取消,汇报轮的首个 LLM 请求会立即报
|
|
// "Post .../v1/messages: context canceled"(历史 bug:会话停在完成通知处)。
|
|
go func() {
|
|
defer cancel()
|
|
if err := a.Run(ctx, task.SessionID, notif, emit); err != nil {
|
|
// 汇报轮失败兜底:把错误持久化到会话,避免静默断流
|
|
if session, ok := a.store.Get(task.SessionID); ok {
|
|
fallback := "[系统] 协作完成通知的自动处理失败(" + err.Error() +
|
|
")。请先读取黑板产物(" + task.Blackboard + ")向用户汇报,再按 TSec 流程继续推进:close 已通关题目释放名额 → start 下一题 → run_coop 下发协作。"
|
|
session.Messages = append(session.Messages, Message{Role: "user", Content: stringPointer(fallback)})
|
|
a.store.Save(session)
|
|
}
|
|
log.Printf("[coop] 汇报轮失败 session=%s task=%s err=%v", task.SessionID, task.ID, err)
|
|
}
|
|
}()
|
|
}
|
|
|
|
// buildCoopNotification 根据任务的结构化结果渲染完成通知。
|
|
// 有 result.json 时按 status 判定成功/失败并附带 flag/产物清单;
|
|
// 缺失时回退到容器 stdout(老版本兼容)。
|
|
func buildCoopNotification(task *CoopTask) string {
|
|
success := false
|
|
status := "完成"
|
|
var summary, flag string
|
|
var artifacts []string
|
|
|
|
if task.Result != nil {
|
|
res := task.Result
|
|
switch res.Status {
|
|
case "solved":
|
|
success = true
|
|
status = "完成"
|
|
case "unsolved":
|
|
status = "未解出"
|
|
case "timeout":
|
|
status = "超时"
|
|
default:
|
|
status = "失败"
|
|
}
|
|
summary, flag, artifacts = res.Summary, res.Flag, res.Artifacts
|
|
} else {
|
|
// 老版本 / result.json 缺失时回退到容器日志
|
|
detail := task.Output
|
|
if detail == "" {
|
|
detail = task.Error
|
|
}
|
|
failed := task.Error != "" || task.ExitCode != 0 ||
|
|
strings.Contains(detail, "未解出") || strings.Contains(detail, "未完成")
|
|
if failed {
|
|
status = "未解出/失败"
|
|
} else {
|
|
success = true
|
|
}
|
|
summary = detail
|
|
}
|
|
|
|
var b strings.Builder
|
|
fmt.Fprintf(&b, "【协作任务完成通知】你之前发起的 pi 协作任务(ID: %s)已%s。\n", task.ID, status)
|
|
if success && flag != "" {
|
|
fmt.Fprintf(&b, "\n【flag】%s\n", flag)
|
|
}
|
|
if strings.TrimSpace(summary) != "" {
|
|
b.WriteString("\n" + truncateRunes(strings.TrimSpace(summary), 1200) + "\n")
|
|
}
|
|
if task.Blackboard != "" {
|
|
fmt.Fprintf(&b, "\n黑板产物目录:%s", task.Blackboard)
|
|
if len(artifacts) > 0 {
|
|
b.WriteString("\n关键产物:")
|
|
for _, artifact := range artifacts {
|
|
if name := strings.TrimSpace(artifact); name != "" {
|
|
b.WriteString("\n- " + name)
|
|
}
|
|
}
|
|
}
|
|
b.WriteString("\n可选用 read_file / list_directory 查看黑板产物,向用户汇报结果与关键文件。")
|
|
}
|
|
if success {
|
|
b.WriteString("\n若产物中包含 flag 且尚未提交:请读取黑板下 task.md(内含平台提交规则 / BENCHMARK_TOKEN / unique_code),立即用 curl 向平台提交 flag 并给出提交响应与得分;不要只汇报而不提交。")
|
|
b.WriteString("\n随后按 TSec 流程持续推进以拿更高总分:close 已通关题目释放名额 → 主动补满并发(一次性 start 多道未完成题目并逐一 run_coop,把在跑靶机数拉满到平台上限)→ 不要空出名额等待;若所有题目已完成或平台任务超时,则停止并向用户汇报总分。")
|
|
} else {
|
|
b.WriteString("\n该任务未解出/失败,请按 TSec 标准流程继续推进,不要空等:")
|
|
b.WriteString("\n1) 用 POST {BENCHMARK_BASE_URL}/openapi/v1/challenges/close?unique_code=<该题编号> 关闭当前靶机容器释放名额(该题编号为 " + task.ChallengeCode + ",可核对黑板 task.md);")
|
|
b.WriteString("\n2) 用 GET {BENCHMARK_BASE_URL}/openapi/v1/challenges 查看剩余未完成题目,一次性选定多道未完成题目补满并发;")
|
|
b.WriteString("\n3) 用 POST {BENCHMARK_BASE_URL}/openapi/v1/challenges/start?unique_code=<新题> 逐个启动新靶机,把在跑靶机数拉满到平台并发上限;")
|
|
b.WriteString("\n4) 对每道新题调用 run_coop 工具下发协作解题,不要只开一个靶机空等。")
|
|
b.WriteString("\n若所有题目已完成或平台任务已结束(接口持续 invalid_state),则停止并向用户汇报总分。")
|
|
}
|
|
return b.String()
|
|
}
|
|
|
|
func mergeToolCall(calls []ToolCall, delta ToolCallDelta) []ToolCall {
|
|
index := delta.Index
|
|
if index < 0 {
|
|
index = len(calls)
|
|
}
|
|
for len(calls) <= index {
|
|
calls = append(calls, ToolCall{Type: "function"})
|
|
}
|
|
call := &calls[index]
|
|
if delta.ID != "" {
|
|
call.ID = delta.ID
|
|
}
|
|
if delta.Name != "" {
|
|
call.Function.Name = delta.Name
|
|
}
|
|
call.Function.Arguments += delta.ArgumentsDelta
|
|
return calls
|
|
}
|
|
|
|
func estimateMessages(messages []Message) int {
|
|
total := 0
|
|
for _, message := range messages {
|
|
total += estimateMessageTokens(message)
|
|
}
|
|
return total
|
|
}
|
|
|
|
// estimateMessageTokens 估算单条消息的 token 数(含 tool_calls 与消息开销)。
|
|
func estimateMessageTokens(message Message) int {
|
|
total := 0
|
|
if message.Content != nil {
|
|
total += estimateTokens(*message.Content)
|
|
}
|
|
for _, call := range message.ToolCalls {
|
|
total += estimateTokens(call.Function.Name + call.Function.Arguments)
|
|
}
|
|
total += 8
|
|
return total
|
|
}
|
|
|
|
func estimateTokens(value string) int {
|
|
runes := []rune(value)
|
|
hanCount := 0
|
|
for _, r := range runes {
|
|
if unicode.Is(unicode.Han, r) {
|
|
hanCount++
|
|
}
|
|
}
|
|
other := len(runes) - hanCount
|
|
return int(float64(other)/4.0+float64(hanCount)*0.8) + 4
|
|
}
|
|
|
|
func stringPointer(value string) *string {
|
|
return &value
|
|
}
|