first commit
This commit is contained in:
@@ -0,0 +1,684 @@
|
||||
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
|
||||
}
|
||||
@@ -0,0 +1,258 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// 支持的接口类型。
|
||||
const (
|
||||
ProviderOpenAI = "openai"
|
||||
ProviderAnthropic = "anthropic"
|
||||
)
|
||||
|
||||
// 协作 worker 管理模式。
|
||||
const (
|
||||
// CoopModeDocker 通过 Docker 容器管理 worker（需本机/远程 Docker daemon）。
|
||||
CoopModeDocker = "docker"
|
||||
// CoopModeLocal 通过本地子进程管理 worker（无需 Docker，托管沙箱等场景）。
|
||||
CoopModeLocal = "local"
|
||||
)
|
||||
|
||||
// LLM 配置环境变量名。环境变量优先级最高，覆盖 api-config.json 用户配置与内置默认值。
|
||||
// 适用于托管沙箱等场景：平台注入环境变量，Agent 启动即生效，无需通过 Web 设置页配置。
|
||||
const (
|
||||
EnvLLMAPIKey = "LLM_API_KEY"
|
||||
EnvLLMBaseURL = "LLM_BASE_URL"
|
||||
EnvLLMModel = "LLM_MODEL"
|
||||
EnvLLMProvider = "LLM_PROVIDER"
|
||||
EnvLLMEngine = "LLM_ENGINE"
|
||||
EnvCoopMode = "COOP_MODE"
|
||||
)
|
||||
|
||||
// APIConfig 是用户可配置的 LLM API 连接信息。
|
||||
type APIConfig struct {
|
||||
// APIKey 是模型服务商 API Key。
|
||||
APIKey string `json:"api_key,omitempty"`
|
||||
// BaseURL 是模型接口地址。
|
||||
BaseURL string `json:"base_url,omitempty"`
|
||||
// Model 是模型名称。
|
||||
Model string `json:"model,omitempty"`
|
||||
// Provider 是接口类型：openai（OpenAI 兼容）或 anthropic（Anthropic Messages API）。
|
||||
Provider string `json:"provider,omitempty"`
|
||||
// Engine 是默认协作引擎：pi（默认）/ pigo / claude。
|
||||
// 仅影响 run_coop 未显式指定 engine 时的默认值，见 tools_coop.go。
|
||||
Engine string `json:"engine,omitempty"`
|
||||
// CoopMode 是协作 worker 管理模式：local（默认，本地子进程）/ docker（Docker 容器）。
|
||||
// local 模式无需 Docker，适用于托管沙箱等无 Docker 环境；docker 模式需本机/远程 Docker daemon。
|
||||
CoopMode string `json:"coop_mode,omitempty"`
|
||||
}
|
||||
|
||||
// APIConfigStore 持久化用户的 LLM API 配置，未配置的字段回退到默认值（环境变量）。
|
||||
type APIConfigStore struct {
|
||||
mu sync.RWMutex
|
||||
path string
|
||||
config APIConfig
|
||||
defaults APIConfig
|
||||
}
|
||||
|
||||
func NewAPIConfigStore(dataDir string, defaults APIConfig) (*APIConfigStore, error) {
|
||||
store := &APIConfigStore{
|
||||
path: filepath.Join(dataDir, "api-config.json"),
|
||||
defaults: defaults,
|
||||
}
|
||||
store.load()
|
||||
return store, nil
|
||||
}
|
||||
|
||||
func (s *APIConfigStore) load() {
|
||||
data, err := os.ReadFile(s.path)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
if len(data) == 0 {
|
||||
return
|
||||
}
|
||||
var config APIConfig
|
||||
if err := json.Unmarshal(data, &config); err != nil {
|
||||
return
|
||||
}
|
||||
s.config = config
|
||||
}
|
||||
|
||||
func (s *APIConfigStore) persist() error {
|
||||
data, err := json.MarshalIndent(s.config, "", " ")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
tmp := s.path + ".tmp"
|
||||
if err := os.WriteFile(tmp, data, 0o600); err != nil {
|
||||
return err
|
||||
}
|
||||
return os.Rename(tmp, s.path)
|
||||
}
|
||||
|
||||
// APIKey 返回生效的 API Key。
|
||||
// 优先级：环境变量 LLM_API_KEY > api-config.json > 内置默认值。
|
||||
func (s *APIConfigStore) APIKey() string {
|
||||
if v := strings.TrimSpace(os.Getenv(EnvLLMAPIKey)); v != "" {
|
||||
return v
|
||||
}
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
if s.config.APIKey != "" {
|
||||
return s.config.APIKey
|
||||
}
|
||||
return s.defaults.APIKey
|
||||
}
|
||||
|
||||
// BaseURL 返回生效的接口地址。
|
||||
// 优先级：环境变量 LLM_BASE_URL > api-config.json > 内置默认值。
|
||||
func (s *APIConfigStore) BaseURL() string {
|
||||
if v := strings.TrimSpace(os.Getenv(EnvLLMBaseURL)); v != "" {
|
||||
return v
|
||||
}
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
if s.config.BaseURL != "" {
|
||||
return s.config.BaseURL
|
||||
}
|
||||
return s.defaults.BaseURL
|
||||
}
|
||||
|
||||
// Model 返回生效的模型名称。
|
||||
// 优先级：环境变量 LLM_MODEL > api-config.json > 内置默认值。
|
||||
func (s *APIConfigStore) Model() string {
|
||||
if v := strings.TrimSpace(os.Getenv(EnvLLMModel)); v != "" {
|
||||
return v
|
||||
}
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
if s.config.Model != "" {
|
||||
return s.config.Model
|
||||
}
|
||||
return s.defaults.Model
|
||||
}
|
||||
|
||||
// IsAPIKeyConfigured 报告是否已配置 API Key（含环境变量、用户配置与内置默认值）。
|
||||
// LLM 客户端实际会按 APIKey() 的优先级（环境变量 > api-config.json > 默认值）取用，
|
||||
// 因此只要最终能拿到非空 Key 就视为已配置，避免误判导致无法对话。
|
||||
func (s *APIConfigStore) IsAPIKeyConfigured() bool {
|
||||
if v := strings.TrimSpace(os.Getenv(EnvLLMAPIKey)); v != "" {
|
||||
return true
|
||||
}
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
return s.config.APIKey != "" || s.defaults.APIKey != ""
|
||||
}
|
||||
|
||||
// Provider 返回生效的接口类型。
|
||||
// 优先级：环境变量 LLM_PROVIDER > api-config.json > 内置默认值 > openai。
|
||||
func (s *APIConfigStore) Provider() string {
|
||||
if v := strings.TrimSpace(os.Getenv(EnvLLMProvider)); v != "" {
|
||||
if v == ProviderOpenAI || v == ProviderAnthropic {
|
||||
return v
|
||||
}
|
||||
}
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
if s.config.Provider != "" {
|
||||
return s.config.Provider
|
||||
}
|
||||
if s.defaults.Provider != "" {
|
||||
return s.defaults.Provider
|
||||
}
|
||||
return ProviderOpenAI
|
||||
}
|
||||
|
||||
// Engine 返回生效的默认协作引擎：pi（默认）/ pigo / claude。
|
||||
// 优先级：环境变量 LLM_ENGINE > api-config.json > 内置默认值 > pi。
|
||||
// 仅在 run_coop 未显式指定 engine 时使用。
|
||||
func (s *APIConfigStore) Engine() string {
|
||||
if v := strings.TrimSpace(os.Getenv(EnvLLMEngine)); v != "" {
|
||||
switch v {
|
||||
case "pi", "pigo", "claude":
|
||||
return v
|
||||
}
|
||||
}
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
switch s.config.Engine {
|
||||
case "pi", "pigo", "claude":
|
||||
return s.config.Engine
|
||||
}
|
||||
if s.defaults.Engine != "" {
|
||||
return s.defaults.Engine
|
||||
}
|
||||
return "pi"
|
||||
}
|
||||
|
||||
// CoopMode 返回生效的协作 worker 管理模式：local（默认）/ docker。
|
||||
// 优先级：环境变量 COOP_MODE > api-config.json > 内置默认值 > local。
|
||||
// local 模式通过本地子进程运行 worker（无需 Docker），docker 模式通过 Docker 容器运行。
|
||||
func (s *APIConfigStore) CoopMode() string {
|
||||
if v := strings.TrimSpace(os.Getenv(EnvCoopMode)); v != "" {
|
||||
switch v {
|
||||
case CoopModeDocker, CoopModeLocal:
|
||||
return v
|
||||
}
|
||||
}
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
switch s.config.CoopMode {
|
||||
case CoopModeDocker, CoopModeLocal:
|
||||
return s.config.CoopMode
|
||||
}
|
||||
if s.defaults.CoopMode != "" {
|
||||
return s.defaults.CoopMode
|
||||
}
|
||||
return CoopModeLocal
|
||||
}
|
||||
|
||||
// Update 更新用户配置。每个参数为 nil 表示不修改该字段；
|
||||
// 非 nil（含空字符串）表示设置或清除该字段（空串=清除，回退默认）。
|
||||
func (s *APIConfigStore) Update(apiKey, baseURL, model, provider, engine, coopMode *string) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
if apiKey != nil {
|
||||
s.config.APIKey = *apiKey
|
||||
}
|
||||
if baseURL != nil {
|
||||
s.config.BaseURL = *baseURL
|
||||
}
|
||||
if model != nil {
|
||||
s.config.Model = *model
|
||||
}
|
||||
if provider != nil {
|
||||
value := strings.TrimSpace(*provider)
|
||||
if value != ProviderOpenAI && value != ProviderAnthropic {
|
||||
value = ""
|
||||
}
|
||||
s.config.Provider = value
|
||||
}
|
||||
if engine != nil {
|
||||
value := strings.TrimSpace(*engine)
|
||||
switch value {
|
||||
case "pi", "pigo", "claude":
|
||||
// 合法值
|
||||
default:
|
||||
value = ""
|
||||
}
|
||||
s.config.Engine = value
|
||||
}
|
||||
if coopMode != nil {
|
||||
value := strings.TrimSpace(*coopMode)
|
||||
switch value {
|
||||
case CoopModeDocker, CoopModeLocal:
|
||||
// 合法值
|
||||
default:
|
||||
value = ""
|
||||
}
|
||||
s.config.CoopMode = value
|
||||
}
|
||||
return s.persist()
|
||||
}
|
||||
@@ -0,0 +1,71 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// TestFallbackSummary 验证 LLM 压缩失败时的兜底摘要：
|
||||
// 纯文本、保留最近消息要点、无孤立 tool 消息。
|
||||
func TestFallbackSummary(t *testing.T) {
|
||||
history := []Message{
|
||||
{Role: "user", Content: stringPointer("第一轮任务：分析日志")},
|
||||
{Role: "assistant", Content: stringPointer("已确认方案，开始执行")},
|
||||
{Role: "assistant", ToolCalls: []ToolCall{{ID: "c1", Type: "function", Function: ToolFunctionCall{Name: "run_bash", Arguments: `{"command":"ls"}`}}}},
|
||||
{Role: "tool", ToolCallID: "c1", Content: stringPointer("file.txt")},
|
||||
{Role: "user", Content: stringPointer("继续，报告结果")},
|
||||
}
|
||||
summary := fallbackSummary(history)
|
||||
if summary == "" {
|
||||
t.Fatal("fallbackSummary 不应返回空")
|
||||
}
|
||||
if strings.Contains(summary, "tool_calls") || strings.Contains(summary, "\"id\"") {
|
||||
t.Fatalf("fallbackSummary 必须为纯文本，不得包含结构化 tool_calls: %s", summary)
|
||||
}
|
||||
if !strings.Contains(summary, "用户") || !strings.Contains(summary, "助手") || !strings.Contains(summary, "工具结果") {
|
||||
t.Fatalf("fallbackSummary 应包含用户/助手/工具结果要点: %s", summary)
|
||||
}
|
||||
if !strings.Contains(summary, "调用了 run_bash") {
|
||||
t.Fatalf("fallbackSummary 应记录工具调用名: %s", summary)
|
||||
}
|
||||
}
|
||||
|
||||
// TestCompactionInputBudget 验证单次压缩输入预算随窗口放大而放大，
|
||||
// 且始终留出输出与安全余量。
|
||||
func TestCompactionInputBudget(t *testing.T) {
|
||||
a := &Agent{cfg: Config{MaxContextTokens: 96000, CompactionTokens: 4096, StreamOutputTokens: 4096}}
|
||||
budget := a.compactionInputBudget()
|
||||
if budget <= 0 || budget >= 96000 {
|
||||
t.Fatalf("预算应介于 (0, 96000)，got %d", budget)
|
||||
}
|
||||
// 窗口放大后预算也应放大
|
||||
a2 := &Agent{cfg: Config{MaxContextTokens: 48000, CompactionTokens: 4096, StreamOutputTokens: 4096}}
|
||||
if a2.compactionInputBudget() >= budget {
|
||||
t.Fatalf("窗口更大的预算应更大，got %d vs %d", a2.compactionInputBudget(), budget)
|
||||
}
|
||||
// 极端小窗口时兜底下限 4096
|
||||
a3 := &Agent{cfg: Config{MaxContextTokens: 2000, CompactionTokens: 512, StreamOutputTokens: 512}}
|
||||
if got := a3.compactionInputBudget(); got < 4096 {
|
||||
t.Fatalf("小窗口预算应有下限，got %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestEstimateMessageTokens 验证单条消息估算包含 content、tool_calls 与开销。
|
||||
func TestEstimateMessageTokens(t *testing.T) {
|
||||
m := Message{
|
||||
Role: "assistant",
|
||||
Content: stringPointer("你好"),
|
||||
ToolCalls: []ToolCall{{
|
||||
ID: "c1", Type: "function",
|
||||
Function: ToolFunctionCall{Name: "read_file", Arguments: `{"path":"a.txt"}`},
|
||||
}},
|
||||
}
|
||||
n := estimateMessageTokens(m)
|
||||
if n <= 0 {
|
||||
t.Fatalf("估算应大于 0，got %d", n)
|
||||
}
|
||||
plain := estimateMessageTokens(Message{Role: "user", Content: stringPointer("你好")})
|
||||
if n <= plain {
|
||||
t.Fatalf("含 tool_calls 的消息估算应大于纯文本，got %d vs %d", n, plain)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,137 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
type Config struct {
|
||||
APIKey string
|
||||
BaseURL string
|
||||
Model string
|
||||
Port string
|
||||
Workspace string
|
||||
DataDir string
|
||||
// AuthCode 是 Web/API/WS 访问授权码。为空（默认）表示不启用授权；
|
||||
// 非空（Docker 部署时通过环境变量 AGENT_AUTH_CODE 注入）表示启用：
|
||||
// 访问者必须输入正确授权码才能使用本实例。
|
||||
AuthCode string
|
||||
MaxContextTokens int
|
||||
KeepRecentMessages int
|
||||
MaxToolResultChars int
|
||||
MaxIterations int
|
||||
ToolTimeout time.Duration
|
||||
RequestTimeout time.Duration
|
||||
StreamOutputTokens int
|
||||
CompactionTokens int
|
||||
Temperature float64
|
||||
}
|
||||
|
||||
func LoadConfig() Config {
|
||||
loadDotEnv(".env")
|
||||
|
||||
workspace := getenv("AGENT_WORKSPACE", "")
|
||||
if workspace == "" {
|
||||
workspace, _ = os.Getwd()
|
||||
}
|
||||
absWorkspace, err := filepath.Abs(workspace)
|
||||
if err != nil {
|
||||
absWorkspace = workspace
|
||||
}
|
||||
dataDir := getenv("AGENT_DATA_DIR", filepath.Join(absWorkspace, "data"))
|
||||
if !filepath.IsAbs(dataDir) {
|
||||
dataDir = filepath.Join(absWorkspace, dataDir)
|
||||
}
|
||||
|
||||
cfg := Config{
|
||||
APIKey: os.Getenv("SILICONFLOW_API_KEY"),
|
||||
BaseURL: getenv("SILICONFLOW_BASE_URL", "https://api.siliconflow.cn/v1/chat/completions"),
|
||||
Model: getenv("AGENT_MODEL", "Qwen/Qwen3-32B"),
|
||||
Port: getenv("AGENT_PORT", "8080"),
|
||||
Workspace: absWorkspace,
|
||||
DataDir: dataDir,
|
||||
// 全面放开限制以适配 DeepSeek（128K 上下文、大输出）跑复杂任务：
|
||||
// - MaxContextTokens 给到 1M，确保 compactionInputBudget 公式（window - Compaction - Stream - window/8）仍有充足余量
|
||||
// - StreamOutputTokens 提升到 65536：直接作为 max_tokens 传给 LLM（见 llm.go），是放开主 Agent 输出能力的关键
|
||||
// - CompactionTokens 提升到 16384：压缩摘要允许足够长度，避免历史信息丢失
|
||||
// - MaxToolResultChars 提升到 100000：复杂任务的脚本/日志能完整传回 LLM，不被截断
|
||||
// - KeepRecentMessages 提升到 50：压缩后保留更多历史，复杂多步任务上下文不丢
|
||||
// - ToolTimeout/RequestTimeout 提升到 600s：编译/长命令/深度推理有充足时间
|
||||
MaxContextTokens: getenvInt("AGENT_MAX_CONTEXT_TOKENS", 1000000),
|
||||
KeepRecentMessages: getenvInt("AGENT_KEEP_RECENT_MESSAGES", 50),
|
||||
MaxToolResultChars: getenvInt("AGENT_MAX_TOOL_RESULT_CHARS", 100000),
|
||||
MaxIterations: getenvInt("AGENT_MAX_ITERATIONS", 50),
|
||||
ToolTimeout: time.Duration(getenvInt("AGENT_TOOL_TIMEOUT_SECONDS", 600)) * time.Second,
|
||||
RequestTimeout: time.Duration(getenvInt("AGENT_REQUEST_TIMEOUT_SECONDS", 600)) * time.Second,
|
||||
StreamOutputTokens: getenvInt("AGENT_STREAM_OUTPUT_TOKENS", 65536),
|
||||
CompactionTokens: getenvInt("AGENT_COMPACTION_TOKENS", 16384),
|
||||
Temperature: getenvFloat("AGENT_TEMPERATURE", 0.3),
|
||||
// 授权码默认关闭：仅当 Docker 部署时显式注入 AGENT_AUTH_CODE 才启用访问授权。
|
||||
AuthCode: getenv("AGENT_AUTH_CODE", ""),
|
||||
}
|
||||
return cfg
|
||||
}
|
||||
|
||||
func loadDotEnv(path string) {
|
||||
file, err := os.Open(path)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer file.Close()
|
||||
|
||||
scanner := bufio.NewScanner(file)
|
||||
for scanner.Scan() {
|
||||
line := strings.TrimSpace(scanner.Text())
|
||||
if line == "" || strings.HasPrefix(line, "#") {
|
||||
continue
|
||||
}
|
||||
key, value, ok := strings.Cut(line, "=")
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
key = strings.TrimSpace(key)
|
||||
value = strings.TrimSpace(value)
|
||||
value = strings.Trim(value, `"'`)
|
||||
if key == "" {
|
||||
continue
|
||||
}
|
||||
if _, exists := os.LookupEnv(key); !exists {
|
||||
_ = os.Setenv(key, value)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func getenv(key, fallback string) string {
|
||||
if value := os.Getenv(key); value != "" {
|
||||
return value
|
||||
}
|
||||
return fallback
|
||||
}
|
||||
|
||||
func getenvInt(key string, fallback int) int {
|
||||
value := os.Getenv(key)
|
||||
if value == "" {
|
||||
return fallback
|
||||
}
|
||||
parsed, err := strconv.Atoi(value)
|
||||
if err != nil {
|
||||
return fallback
|
||||
}
|
||||
return parsed
|
||||
}
|
||||
|
||||
func getenvFloat(key string, fallback float64) float64 {
|
||||
value := os.Getenv(key)
|
||||
if value == "" {
|
||||
return fallback
|
||||
}
|
||||
parsed, err := strconv.ParseFloat(value, 64)
|
||||
if err != nil {
|
||||
return fallback
|
||||
}
|
||||
return parsed
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -0,0 +1,85 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestCoopBaseURL(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
provider string
|
||||
engine string
|
||||
baseURL string
|
||||
want string
|
||||
}{
|
||||
// ---- pi / claude（用 @anthropic-ai/sdk，自带 /v1/messages，不能补 /v1）----
|
||||
// Anthropic：剥离 /v1，避免双重 /v1/v1/messages
|
||||
{"pi/anthropic bare", "anthropic", "pi", "https://api.deepseek.com/anthropic", "https://api.deepseek.com/anthropic"},
|
||||
{"pi/anthropic trailing slash", "anthropic", "pi", "https://api.deepseek.com/anthropic/", "https://api.deepseek.com/anthropic"},
|
||||
{"pi/anthropic with /v1", "anthropic", "pi", "https://api.deepseek.com/anthropic/v1", "https://api.deepseek.com/anthropic"},
|
||||
{"pi/anthropic with /v1/messages", "anthropic", "pi", "https://api.deepseek.com/anthropic/v1/messages", "https://api.deepseek.com/anthropic"},
|
||||
{"pi/anthropic with /messages", "anthropic", "pi", "https://api.deepseek.com/anthropic/messages", "https://api.deepseek.com/anthropic"},
|
||||
{"pi/anthropic anthropic.com", "anthropic", "pi", "https://api.anthropic.com", "https://api.anthropic.com"},
|
||||
// claude 行为与 pi 一致
|
||||
{"claude/anthropic bare", "anthropic", "claude", "https://api.deepseek.com/anthropic", "https://api.deepseek.com/anthropic"},
|
||||
{"claude/anthropic with /v1", "anthropic", "claude", "https://api.deepseek.com/anthropic/v1", "https://api.deepseek.com/anthropic"},
|
||||
|
||||
// ---- pigo（anthropicCompatDriver 只追加 /messages，需补 /v1）----
|
||||
{"pigo/anthropic bare", "anthropic", "pigo", "https://api.deepseek.com/anthropic", "https://api.deepseek.com/anthropic/v1"},
|
||||
{"pigo/anthropic trailing slash", "anthropic", "pigo", "https://api.deepseek.com/anthropic/", "https://api.deepseek.com/anthropic/v1"},
|
||||
{"pigo/anthropic with /v1", "anthropic", "pigo", "https://api.deepseek.com/anthropic/v1", "https://api.deepseek.com/anthropic/v1"},
|
||||
{"pigo/anthropic with /v1/messages", "anthropic", "pigo", "https://api.deepseek.com/anthropic/v1/messages", "https://api.deepseek.com/anthropic/v1"},
|
||||
{"pigo/anthropic with /messages", "anthropic", "pigo", "https://api.deepseek.com/anthropic/messages", "https://api.deepseek.com/anthropic/v1"},
|
||||
{"pigo/anthropic anthropic.com", "anthropic", "pigo", "https://api.anthropic.com", "https://api.anthropic.com/v1"},
|
||||
|
||||
// ---- OpenAI 协议（三种 engine 行为一致：保留 /v1，剥离 /chat/completions）----
|
||||
{"pi/openai full url", "openai", "pi", "https://api.siliconflow.cn/v1/chat/completions", "https://api.siliconflow.cn/v1"},
|
||||
{"pi/openai with /v1", "openai", "pi", "https://api.siliconflow.cn/v1", "https://api.siliconflow.cn/v1"},
|
||||
{"pi/openai bare", "openai", "pi", "https://api.siliconflow.cn", "https://api.siliconflow.cn"},
|
||||
{"pigo/openai full url", "openai", "pigo", "https://api.siliconflow.cn/v1/chat/completions", "https://api.siliconflow.cn/v1"},
|
||||
{"pigo/openai bare", "openai", "pigo", "https://api.siliconflow.cn", "https://api.siliconflow.cn"},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := coopBaseURL(tt.provider, tt.baseURL, tt.engine)
|
||||
if got != tt.want {
|
||||
t.Errorf("coopBaseURL(%q, %q, %q) = %q, want %q", tt.provider, tt.baseURL, tt.engine, got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunCoopErrorPaths(t *testing.T) {
|
||||
apiCfg := &APIConfigStore{}
|
||||
apiCfg.config = APIConfig{APIKey: "test-key", BaseURL: "https://api.deepseek.com/anthropic", Model: "deepseek-chat"}
|
||||
|
||||
toolset := &Toolset{
|
||||
Workspace: ".",
|
||||
Timeout: 30 * time.Second,
|
||||
MaxOutput: 4000,
|
||||
apiCfg: apiCfg,
|
||||
dockerSocket: defaultDockerSocket,
|
||||
coop: NewCoopManager(),
|
||||
sessionRoot: t.TempDir(),
|
||||
}
|
||||
|
||||
// 1. 缺少 task
|
||||
r := toolset.runCoop(context.Background(), "test-session", map[string]any{})
|
||||
if r.Success || r.Output == "" {
|
||||
t.Fatalf("empty task should fail, got %#v", r)
|
||||
}
|
||||
|
||||
// 2. 正常参数：应进入 Docker 检查阶段（本机无 Docker 时返回连接/镜像错误而非 panic）
|
||||
r = toolset.runCoop(context.Background(), "test-session", map[string]any{
|
||||
"task": "测试任务",
|
||||
"round_max": 2,
|
||||
"timeout": 30,
|
||||
})
|
||||
t.Logf("runCoop output: %s", r.Output)
|
||||
if r.Success {
|
||||
// 若真的跑成功了（有 Docker 且镜像存在），也无妨
|
||||
t.Logf("coop unexpectedly succeeded")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,99 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// DockerConfig 是用户可配置的 Docker 连接信息。
|
||||
type DockerConfig struct {
|
||||
// Socket 是用户配置的 Docker 连接地址，支持
|
||||
// unix:///var/run/docker.sock、npipe:////./pipe/docker_engine、tcp://host:2375 等。
|
||||
// 为空时使用默认的本地 Docker。
|
||||
Socket string `json:"socket,omitempty"`
|
||||
}
|
||||
|
||||
// DockerConfigStore 持久化用户的 Docker socket 配置，并给出默认值。
|
||||
type DockerConfigStore struct {
|
||||
mu sync.RWMutex
|
||||
path string
|
||||
config DockerConfig
|
||||
defaultSocket string
|
||||
}
|
||||
|
||||
func NewDockerConfigStore(dataDir string) (*DockerConfigStore, error) {
|
||||
defaultSocket := getenv("AGENT_DOCKER_SOCKET", "")
|
||||
if defaultSocket == "" {
|
||||
defaultSocket = defaultDockerSocket()
|
||||
}
|
||||
store := &DockerConfigStore{
|
||||
path: filepath.Join(dataDir, "docker-config.json"),
|
||||
defaultSocket: defaultSocket,
|
||||
}
|
||||
store.load()
|
||||
return store, nil
|
||||
}
|
||||
|
||||
func defaultDockerSocket() string {
|
||||
if runtime.GOOS == "windows" {
|
||||
return "npipe:////./pipe/docker_engine"
|
||||
}
|
||||
return "unix:///var/run/docker.sock"
|
||||
}
|
||||
|
||||
func (s *DockerConfigStore) load() {
|
||||
data, err := os.ReadFile(s.path)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
if len(data) == 0 {
|
||||
return
|
||||
}
|
||||
var config DockerConfig
|
||||
if err := json.Unmarshal(data, &config); err != nil {
|
||||
return
|
||||
}
|
||||
s.config = config
|
||||
}
|
||||
|
||||
// Socket 返回当前生效的 docker 连接地址：用户配置优先，否则使用默认本地 Docker。
|
||||
func (s *DockerConfigStore) Socket() string {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
if s.config.Socket != "" {
|
||||
return s.config.Socket
|
||||
}
|
||||
return s.defaultSocket
|
||||
}
|
||||
|
||||
// DefaultSocket 返回未配置时使用的默认地址。
|
||||
func (s *DockerConfigStore) DefaultSocket() string {
|
||||
return s.defaultSocket
|
||||
}
|
||||
|
||||
// IsConfigured 报告用户是否手动配置过 socket。
|
||||
func (s *DockerConfigStore) IsConfigured() bool {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
return s.config.Socket != ""
|
||||
}
|
||||
|
||||
// SetSocket 保存用户的 socket 配置；传入空字符串表示清除配置、恢复默认。
|
||||
func (s *DockerConfigStore) SetSocket(socket string) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
s.config.Socket = socket
|
||||
data, err := json.MarshalIndent(s.config, "", " ")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
tmp := s.path + ".tmp"
|
||||
if err := os.WriteFile(tmp, data, 0o600); err != nil {
|
||||
return err
|
||||
}
|
||||
return os.Rename(tmp, s.path)
|
||||
}
|
||||
@@ -0,0 +1,470 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
type LLMClient struct {
|
||||
cfg Config
|
||||
apiConfig *APIConfigStore
|
||||
client *http.Client
|
||||
}
|
||||
|
||||
type StreamChunk struct {
|
||||
Content string
|
||||
ToolCalls []ToolCallDelta
|
||||
FinishReason string
|
||||
Usage *Usage
|
||||
Error error
|
||||
}
|
||||
|
||||
type ToolCallDelta struct {
|
||||
Index int
|
||||
ID string
|
||||
Name string
|
||||
ArgumentsDelta string
|
||||
}
|
||||
|
||||
func NewLLMClient(cfg Config, apiCfg *APIConfigStore) *LLMClient {
|
||||
return &LLMClient{
|
||||
cfg: cfg,
|
||||
apiConfig: apiCfg,
|
||||
client: &http.Client{
|
||||
Timeout: cfg.RequestTimeout + 30*time.Second,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// isAnthropic 报告当前是否使用 Anthropic Messages API。
|
||||
func (c *LLMClient) isAnthropic() bool {
|
||||
return c.apiConfig.Provider() == ProviderAnthropic
|
||||
}
|
||||
|
||||
func (c *LLMClient) ChatStream(ctx context.Context, messages []Message, tools []ToolDefinition) (<-chan StreamChunk, error) {
|
||||
if c.isAnthropic() {
|
||||
body := map[string]any{
|
||||
"model": c.apiConfig.Model(),
|
||||
"max_tokens": c.cfg.StreamOutputTokens,
|
||||
"stream": true,
|
||||
"temperature": c.cfg.Temperature,
|
||||
"messages": toAnthropicMessages(messages),
|
||||
}
|
||||
if system := extractSystemPrompt(messages); system != "" {
|
||||
body["system"] = system
|
||||
}
|
||||
if len(tools) > 0 {
|
||||
body["tools"] = toAnthropicTools(tools)
|
||||
}
|
||||
return c.stream(ctx, body)
|
||||
}
|
||||
|
||||
body := map[string]any{
|
||||
"model": c.apiConfig.Model(),
|
||||
"messages": toAPIMessages(messages),
|
||||
"stream": true,
|
||||
"temperature": c.cfg.Temperature,
|
||||
"max_tokens": c.cfg.StreamOutputTokens,
|
||||
}
|
||||
if len(tools) > 0 {
|
||||
body["tools"] = tools
|
||||
body["tool_choice"] = "auto"
|
||||
}
|
||||
return c.stream(ctx, body)
|
||||
}
|
||||
|
||||
func (c *LLMClient) Compress(ctx context.Context, history []Message) (string, error) {
|
||||
systemPrompt := `你是一个高信息密度的对话压缩器。把下面的历史对话压缩成一份中文摘要，供后续继续执行同一任务。
|
||||
必须保留（按重要程度排序）：
|
||||
1. 用户的核心目标与最新要求、尚未完成的任务与下一步计划；
|
||||
2. 已确认的决策与结论、得分/进度（如 TSec 累计分数与已通关/进行中的题目编号）；
|
||||
3. 关键文件路径、执行过的命令与重要输出要点、代码要点；
|
||||
4. 错误信息与解决方案、需要继续跟进的问题；
|
||||
5. 平台提交规则（如 BENCHMARK_TOKEN、unique_code、提交接口与提交方式）。
|
||||
规则：
|
||||
- 按时间顺序组织，最新信息优先，可适当合并同类项；
|
||||
- 若历史包含工具调用（run_bash/read_file/run_coop 等），只保留"做了什么、结果如何"的要点，不要逐字复制命令或输出；
|
||||
- 不要添加历史中不存在的信息，不要臆测；
|
||||
- 直接输出摘要正文，不要输出任何解释、标题或多余格式。`
|
||||
messages := append([]Message{{Role: "system", Content: &systemPrompt}}, history...)
|
||||
|
||||
if c.isAnthropic() {
|
||||
return c.compressAnthropic(ctx, messages)
|
||||
}
|
||||
|
||||
body := map[string]any{
|
||||
"model": c.apiConfig.Model(),
|
||||
"messages": toAPIMessages(messages),
|
||||
"stream": false,
|
||||
"temperature": 0.2,
|
||||
"max_tokens": c.cfg.CompactionTokens,
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(ctx, c.cfg.RequestTimeout)
|
||||
defer cancel()
|
||||
|
||||
payload, err := json.Marshal(body)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
respBody, err := c.doSyncRequestWithRetry(ctx, false, c.apiConfig.BaseURL(), payload)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
var result struct {
|
||||
Choices []struct {
|
||||
Message struct {
|
||||
Content *string `json:"content"`
|
||||
ReasoningContent *string `json:"reasoning_content"`
|
||||
} `json:"message"`
|
||||
} `json:"choices"`
|
||||
}
|
||||
if err := json.Unmarshal(respBody, &result); err != nil {
|
||||
return "", err
|
||||
}
|
||||
if len(result.Choices) == 0 {
|
||||
return "", errors.New("模型没有返回压缩摘要")
|
||||
}
|
||||
msg := result.Choices[0].Message
|
||||
if msg.Content != nil && strings.TrimSpace(*msg.Content) != "" {
|
||||
return strings.TrimSpace(*msg.Content), nil
|
||||
}
|
||||
// 推理模型（deepseek-reasoner 系）正文在 reasoning_content，content 可能为空
|
||||
if msg.ReasoningContent != nil && strings.TrimSpace(*msg.ReasoningContent) != "" {
|
||||
return strings.TrimSpace(*msg.ReasoningContent), nil
|
||||
}
|
||||
return "", errors.New("模型没有返回压缩摘要")
|
||||
}
|
||||
|
||||
// stream 发起流式 LLM 请求，对连接失败、限流(429/5xx/551)、首 chunk 前流断开
|
||||
// 做指数退避重试（3 次：1s→2s→4s）。保留 (<-chan, error) 签名，上层无需改动。
|
||||
func (c *LLMClient) stream(ctx context.Context, body map[string]any) (<-chan StreamChunk, error) {
|
||||
anthropic := c.isAnthropic()
|
||||
url := c.apiConfig.BaseURL()
|
||||
if anthropic {
|
||||
url = normalizeAnthropicURL(url)
|
||||
}
|
||||
payload, err := json.Marshal(body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
const maxRetries = 3
|
||||
var lastErr error
|
||||
|
||||
for attempt := 0; attempt <= maxRetries; attempt++ {
|
||||
if attempt > 0 {
|
||||
backoff := time.Duration(1<<(attempt-1)) * time.Second
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
case <-time.After(backoff):
|
||||
}
|
||||
}
|
||||
|
||||
ch, reqErr := c.doStreamRequest(ctx, anthropic, url, payload)
|
||||
if reqErr != nil {
|
||||
lastErr = reqErr
|
||||
if attempt < maxRetries && isRetriableErr(reqErr) {
|
||||
log.Printf("[llm] 流式请求失败（第 %d 次），将重试: %v", attempt+1, reqErr)
|
||||
continue
|
||||
}
|
||||
return nil, reqErr
|
||||
}
|
||||
|
||||
// 请求成功建立，但流可能在首个 chunk 前就断开（EOF）。
|
||||
// peek 第一个 chunk：如果是可重试错误，排空旧 channel 后重试。
|
||||
first, ok := <-ch
|
||||
if !ok {
|
||||
lastErr = errors.New("stream: 空响应")
|
||||
if attempt < maxRetries {
|
||||
log.Printf("[llm] 流式响应为空（第 %d 次），重试", attempt+1)
|
||||
continue
|
||||
}
|
||||
return nil, lastErr
|
||||
}
|
||||
if first.Error != nil && attempt < maxRetries && isRetriableErr(first.Error) {
|
||||
for range ch {
|
||||
}
|
||||
lastErr = first.Error
|
||||
log.Printf("[llm] 流式响应首 chunk 前断开（第 %d 次），重试: %v", attempt+1, first.Error)
|
||||
continue
|
||||
}
|
||||
|
||||
// 正常：转发 first + 后续 chunk
|
||||
out := make(chan StreamChunk, 64)
|
||||
go func() {
|
||||
defer close(out)
|
||||
out <- first
|
||||
for chunk := range ch {
|
||||
out <- chunk
|
||||
}
|
||||
}()
|
||||
return out, nil
|
||||
}
|
||||
|
||||
return nil, lastErr
|
||||
}
|
||||
|
||||
// doStreamRequest 发起单次流式 LLM 请求（不含重试）。
|
||||
func (c *LLMClient) doStreamRequest(ctx context.Context, anthropic bool, url string, payload []byte) (<-chan StreamChunk, error) {
|
||||
ctx, cancel := context.WithTimeout(ctx, c.cfg.RequestTimeout)
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(payload))
|
||||
if err != nil {
|
||||
cancel()
|
||||
return nil, err
|
||||
}
|
||||
c.setRequestHeaders(req, anthropic)
|
||||
|
||||
resp, err := c.client.Do(req)
|
||||
if err != nil {
|
||||
cancel()
|
||||
return nil, err
|
||||
}
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
apiErr := c.readAPIError(resp)
|
||||
resp.Body.Close()
|
||||
cancel()
|
||||
return nil, apiErr
|
||||
}
|
||||
|
||||
ch := make(chan StreamChunk, 64)
|
||||
go func() {
|
||||
defer cancel()
|
||||
defer close(ch)
|
||||
defer resp.Body.Close()
|
||||
if anthropic {
|
||||
readAnthropicSSE(ctx, resp, ch)
|
||||
} else {
|
||||
c.readSSE(ctx, resp, ch)
|
||||
}
|
||||
}()
|
||||
return ch, nil
|
||||
}
|
||||
|
||||
// doSyncRequestWithRetry 发起同步（非流式）LLM 请求，对连接失败、限流(429/5xx/551)
|
||||
// 做指数退避重试（3 次：1s→2s→4s）。返回响应体字节，由调用方解析。
|
||||
func (c *LLMClient) doSyncRequestWithRetry(ctx context.Context, anthropic bool, url string, payload []byte) ([]byte, error) {
|
||||
const maxRetries = 3
|
||||
var lastErr error
|
||||
|
||||
for attempt := 0; attempt <= maxRetries; attempt++ {
|
||||
if attempt > 0 {
|
||||
backoff := time.Duration(1<<(attempt-1)) * time.Second
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
case <-time.After(backoff):
|
||||
}
|
||||
}
|
||||
|
||||
body, err := c.doSyncRequest(ctx, anthropic, url, payload)
|
||||
if err == nil {
|
||||
return body, nil
|
||||
}
|
||||
lastErr = err
|
||||
if attempt < maxRetries && isRetriableErr(err) {
|
||||
log.Printf("[llm] 同步请求失败（第 %d 次），将重试: %v", attempt+1, err)
|
||||
continue
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
return nil, lastErr
|
||||
}
|
||||
|
||||
// doSyncRequest 发起单次同步（非流式）LLM 请求（不含重试）。
|
||||
func (c *LLMClient) doSyncRequest(ctx context.Context, anthropic bool, url string, payload []byte) ([]byte, error) {
|
||||
ctx, cancel := context.WithTimeout(ctx, c.cfg.RequestTimeout)
|
||||
defer cancel()
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(payload))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
c.setRequestHeaders(req, anthropic)
|
||||
|
||||
resp, err := c.client.Do(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
return nil, c.readAPIError(resp)
|
||||
}
|
||||
return io.ReadAll(resp.Body)
|
||||
}
|
||||
|
||||
// isRetriableErr 判断错误是否值得重试（网络错误、限流、服务端错误）。
|
||||
func isRetriableErr(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) {
|
||||
return true
|
||||
}
|
||||
msg := err.Error()
|
||||
if strings.Contains(msg, "EOF") ||
|
||||
strings.Contains(msg, "connection reset") ||
|
||||
strings.Contains(msg, "broken pipe") ||
|
||||
strings.Contains(msg, "timeout") ||
|
||||
strings.Contains(msg, "deadline exceeded") ||
|
||||
strings.Contains(msg, "connection refused") ||
|
||||
strings.Contains(msg, "no such host") {
|
||||
return true
|
||||
}
|
||||
// HTTP 状态码错误：5xx 服务端错误、429 限流、551 网关熔断
|
||||
if strings.HasPrefix(msg, "HTTP 5") ||
|
||||
strings.HasPrefix(msg, "HTTP 429") ||
|
||||
strings.HasPrefix(msg, "HTTP 551") {
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (c *LLMClient) readSSE(ctx context.Context, resp *http.Response, ch chan<- StreamChunk) {
|
||||
scanner := bufio.NewScanner(resp.Body)
|
||||
scanner.Buffer(make([]byte, 64*1024), 1024*1024)
|
||||
for scanner.Scan() {
|
||||
line := scanner.Text()
|
||||
if !strings.HasPrefix(line, "data:") {
|
||||
continue
|
||||
}
|
||||
data := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
|
||||
if data == "[DONE]" {
|
||||
break
|
||||
}
|
||||
|
||||
var chunk struct {
|
||||
Choices []struct {
|
||||
Delta struct {
|
||||
Content *string `json:"content"`
|
||||
ToolCalls []struct {
|
||||
Index *int `json:"index"`
|
||||
ID string `json:"id"`
|
||||
Function struct {
|
||||
Name string `json:"name"`
|
||||
Arguments string `json:"arguments"`
|
||||
} `json:"function"`
|
||||
} `json:"tool_calls"`
|
||||
} `json:"delta"`
|
||||
FinishReason *string `json:"finish_reason"`
|
||||
} `json:"choices"`
|
||||
Usage *Usage `json:"usage"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(data), &chunk); err != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
event := StreamChunk{}
|
||||
for _, choice := range chunk.Choices {
|
||||
if choice.Delta.Content != nil {
|
||||
event.Content += *choice.Delta.Content
|
||||
}
|
||||
for _, tc := range choice.Delta.ToolCalls {
|
||||
index := 0
|
||||
if tc.Index != nil {
|
||||
index = *tc.Index
|
||||
}
|
||||
event.ToolCalls = append(event.ToolCalls, ToolCallDelta{
|
||||
Index: index,
|
||||
ID: tc.ID,
|
||||
Name: tc.Function.Name,
|
||||
ArgumentsDelta: tc.Function.Arguments,
|
||||
})
|
||||
}
|
||||
if choice.FinishReason != nil {
|
||||
event.FinishReason = *choice.FinishReason
|
||||
}
|
||||
}
|
||||
if chunk.Usage != nil {
|
||||
event.Usage = chunk.Usage
|
||||
}
|
||||
if event.Content != "" || len(event.ToolCalls) > 0 || event.FinishReason != "" || event.Usage != nil {
|
||||
ch <- event
|
||||
}
|
||||
}
|
||||
if err := scanner.Err(); err != nil && ctx.Err() == nil {
|
||||
ch <- StreamChunk{Error: err}
|
||||
}
|
||||
}
|
||||
|
||||
func (c *LLMClient) setRequestHeaders(req *http.Request, anthropic bool) {
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Accept", "text/event-stream")
|
||||
if anthropic {
|
||||
req.Header.Set("x-api-key", c.apiConfig.APIKey())
|
||||
req.Header.Set("anthropic-version", "2023-06-01")
|
||||
} else {
|
||||
req.Header.Set("Authorization", "Bearer "+c.apiConfig.APIKey())
|
||||
}
|
||||
}
|
||||
|
||||
func (c *LLMClient) readAPIError(resp *http.Response) error {
|
||||
body, _ := io.ReadAll(io.LimitReader(resp.Body, 8*1024))
|
||||
var apiErr struct {
|
||||
Error struct {
|
||||
Message string `json:"message"`
|
||||
} `json:"error"`
|
||||
}
|
||||
_ = json.Unmarshal(body, &apiErr)
|
||||
message := strings.TrimSpace(apiErr.Error.Message)
|
||||
if message == "" {
|
||||
message = strings.TrimSpace(string(body))
|
||||
}
|
||||
if message == "" {
|
||||
message = fmt.Sprintf("HTTP %d", resp.StatusCode)
|
||||
}
|
||||
return errors.New(message)
|
||||
}
|
||||
|
||||
func toAPIMessages(messages []Message) []map[string]any {
|
||||
result := make([]map[string]any, 0, len(messages))
|
||||
for _, message := range messages {
|
||||
item := map[string]any{
|
||||
"role": message.Role,
|
||||
"content": contentValue(message.Content),
|
||||
}
|
||||
if len(message.ToolCalls) > 0 {
|
||||
item["tool_calls"] = toAPIToolCalls(message.ToolCalls)
|
||||
}
|
||||
if message.ToolCallID != "" {
|
||||
item["tool_call_id"] = message.ToolCallID
|
||||
}
|
||||
result = append(result, item)
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
func toAPIToolCalls(calls []ToolCall) []map[string]any {
|
||||
result := make([]map[string]any, 0, len(calls))
|
||||
for _, call := range calls {
|
||||
result = append(result, map[string]any{
|
||||
"id": call.ID,
|
||||
"type": "function",
|
||||
"function": map[string]any{
|
||||
"name": call.Function.Name,
|
||||
"arguments": call.Function.Arguments,
|
||||
},
|
||||
})
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
func contentValue(content *string) any {
|
||||
if content == nil {
|
||||
return nil
|
||||
}
|
||||
return *content
|
||||
}
|
||||
@@ -0,0 +1,298 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// normalizeAnthropicURL 把用户填写的 Base URL 规范化为 Anthropic Messages 接口地址。
|
||||
// 兼容多种填写方式：https://api.anthropic.com、.../v1、.../v1/messages、.../v1/chat/completions。
|
||||
func normalizeAnthropicURL(baseURL string) string {
|
||||
trimmed := strings.TrimRight(strings.TrimSpace(baseURL), "/")
|
||||
if trimmed == "" {
|
||||
return "https://api.anthropic.com/v1/messages"
|
||||
}
|
||||
switch {
|
||||
case strings.HasSuffix(trimmed, "/v1/chat/completions"):
|
||||
return strings.TrimSuffix(trimmed, "/chat/completions") + "/messages"
|
||||
case strings.HasSuffix(trimmed, "/v1/messages"):
|
||||
return trimmed
|
||||
case strings.HasSuffix(trimmed, "/v1"):
|
||||
return trimmed + "/messages"
|
||||
default:
|
||||
return trimmed + "/v1/messages"
|
||||
}
|
||||
}
|
||||
|
||||
// extractSystemPrompt 汇总消息中的 system 角色内容，Anthropic 要求 system 放在顶层字段。
|
||||
func extractSystemPrompt(messages []Message) string {
|
||||
var parts []string
|
||||
for _, message := range messages {
|
||||
if message.Role == "system" && message.Content != nil && strings.TrimSpace(*message.Content) != "" {
|
||||
parts = append(parts, *message.Content)
|
||||
}
|
||||
}
|
||||
return strings.Join(parts, "\n\n")
|
||||
}
|
||||
|
||||
// toAnthropicTools 把内部工具定义转换为 Anthropic 的 tools 数组（input_schema 替代 parameters）。
|
||||
func toAnthropicTools(tools []ToolDefinition) []map[string]any {
|
||||
result := make([]map[string]any, 0, len(tools))
|
||||
for _, tool := range tools {
|
||||
schema := tool.Function.Parameters
|
||||
if schema == nil {
|
||||
schema = map[string]any{"type": "object"}
|
||||
}
|
||||
result = append(result, map[string]any{
|
||||
"name": tool.Function.Name,
|
||||
"description": tool.Function.Description,
|
||||
"input_schema": schema,
|
||||
})
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// toAnthropicMessages 把内部 Message 列表转换为 Anthropic messages 数组。
|
||||
// system 消息被过滤（走顶层 system 字段）；assistant 的 tool_use 与 user 的
|
||||
// tool_result 都以 content block 形式表达。
|
||||
// 注意：Anthropic 要求上一条 assistant 消息中所有 tool_use 的 tool_result
|
||||
// 必须放在紧邻的同一条 user 消息里，因此连续的 tool 结果消息需要合并。
|
||||
func toAnthropicMessages(messages []Message) []map[string]any {
|
||||
result := make([]map[string]any, 0, len(messages))
|
||||
for i := 0; i < len(messages); i++ {
|
||||
message := messages[i]
|
||||
switch message.Role {
|
||||
case "system":
|
||||
continue
|
||||
case "assistant":
|
||||
blocks := make([]map[string]any, 0, 1+len(message.ToolCalls))
|
||||
if message.Content != nil && strings.TrimSpace(*message.Content) != "" {
|
||||
blocks = append(blocks, map[string]any{"type": "text", "text": *message.Content})
|
||||
}
|
||||
for _, call := range message.ToolCalls {
|
||||
blocks = append(blocks, map[string]any{
|
||||
"type": "tool_use",
|
||||
"id": call.ID,
|
||||
"name": call.Function.Name,
|
||||
"input": parseJSONValue(call.Function.Arguments),
|
||||
})
|
||||
}
|
||||
if len(blocks) == 0 {
|
||||
blocks = append(blocks, map[string]any{"type": "text", "text": ""})
|
||||
}
|
||||
result = append(result, map[string]any{"role": "assistant", "content": blocks})
|
||||
case "tool":
|
||||
// 合并连续的 tool 消息：同一条 user 消息包含所有 tool_result 块
|
||||
blocks := []map[string]any{{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": message.ToolCallID,
|
||||
"content": contentString(message.Content),
|
||||
}}
|
||||
for i+1 < len(messages) && messages[i+1].Role == "tool" {
|
||||
i++
|
||||
next := messages[i]
|
||||
blocks = append(blocks, map[string]any{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": next.ToolCallID,
|
||||
"content": contentString(next.Content),
|
||||
})
|
||||
}
|
||||
result = append(result, map[string]any{"role": "user", "content": blocks})
|
||||
default: // user
|
||||
blocks := make([]map[string]any, 0, 1)
|
||||
if message.Content != nil && strings.TrimSpace(*message.Content) != "" {
|
||||
blocks = append(blocks, map[string]any{"type": "text", "text": *message.Content})
|
||||
}
|
||||
if len(blocks) == 0 {
|
||||
blocks = append(blocks, map[string]any{"type": "text", "text": ""})
|
||||
}
|
||||
result = append(result, map[string]any{"role": "user", "content": blocks})
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// parseJSONValue 把工具参数 JSON 字符串解析为任意值；解析失败时回退为空对象。
|
||||
func parseJSONValue(raw string) any {
|
||||
var value any
|
||||
if err := json.Unmarshal([]byte(raw), &value); err != nil || value == nil {
|
||||
return map[string]any{}
|
||||
}
|
||||
return value
|
||||
}
|
||||
|
||||
func contentString(content *string) string {
|
||||
if content == nil {
|
||||
return ""
|
||||
}
|
||||
return *content
|
||||
}
|
||||
|
||||
// readAnthropicSSE 解析 Anthropic Messages API 的流式响应。
|
||||
// 事件格式为 `event: <type>` 与 `data: <json>` 两行一组。
|
||||
func readAnthropicSSE(ctx context.Context, resp *http.Response, ch chan<- StreamChunk) {
|
||||
scanner := bufio.NewScanner(resp.Body)
|
||||
scanner.Buffer(make([]byte, 64*1024), 1024*1024)
|
||||
|
||||
// content block index -> 正在累积的 tool_use 状态
|
||||
type toolState struct {
|
||||
callIndex int
|
||||
id string
|
||||
name string
|
||||
}
|
||||
tools := make(map[int]*toolState)
|
||||
nextCallIndex := 0
|
||||
|
||||
var eventType string
|
||||
for scanner.Scan() {
|
||||
line := scanner.Text()
|
||||
switch {
|
||||
case strings.HasPrefix(line, "event:"):
|
||||
eventType = strings.TrimSpace(strings.TrimPrefix(line, "event:"))
|
||||
continue
|
||||
case strings.HasPrefix(line, "data:"):
|
||||
default:
|
||||
continue
|
||||
}
|
||||
|
||||
data := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
|
||||
switch eventType {
|
||||
case "content_block_start":
|
||||
var ev struct {
|
||||
Index int `json:"index"`
|
||||
Block struct {
|
||||
Type string `json:"type"`
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
} `json:"content_block"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(data), &ev); err != nil {
|
||||
continue
|
||||
}
|
||||
if ev.Block.Type == "tool_use" {
|
||||
state := &toolState{callIndex: nextCallIndex, id: ev.Block.ID, name: ev.Block.Name}
|
||||
nextCallIndex++
|
||||
tools[ev.Index] = state
|
||||
ch <- StreamChunk{ToolCalls: []ToolCallDelta{
|
||||
{Index: state.callIndex, ID: state.id, Name: state.name},
|
||||
}}
|
||||
}
|
||||
case "content_block_delta":
|
||||
var ev struct {
|
||||
Index int `json:"index"`
|
||||
Delta struct {
|
||||
Type string `json:"type"`
|
||||
Text string `json:"text"`
|
||||
PartialJSON string `json:"partial_json"`
|
||||
} `json:"delta"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(data), &ev); err != nil {
|
||||
continue
|
||||
}
|
||||
switch ev.Delta.Type {
|
||||
case "text_delta":
|
||||
ch <- StreamChunk{Content: ev.Delta.Text}
|
||||
case "input_json_delta":
|
||||
if state, ok := tools[ev.Index]; ok {
|
||||
ch <- StreamChunk{ToolCalls: []ToolCallDelta{
|
||||
{Index: state.callIndex, ArgumentsDelta: ev.Delta.PartialJSON},
|
||||
}}
|
||||
}
|
||||
}
|
||||
case "message_delta":
|
||||
var ev struct {
|
||||
Delta struct {
|
||||
StopReason string `json:"stop_reason"`
|
||||
} `json:"delta"`
|
||||
Usage *struct {
|
||||
OutputTokens int `json:"output_tokens"`
|
||||
} `json:"usage"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(data), &ev); err != nil {
|
||||
continue
|
||||
}
|
||||
if ev.Delta.StopReason != "" {
|
||||
ch <- StreamChunk{FinishReason: ev.Delta.StopReason}
|
||||
}
|
||||
if ev.Usage != nil {
|
||||
ch <- StreamChunk{Usage: &Usage{CompletionTokens: ev.Usage.OutputTokens}}
|
||||
}
|
||||
case "error":
|
||||
var ev struct {
|
||||
Error struct {
|
||||
Type string `json:"type"`
|
||||
Message string `json:"message"`
|
||||
} `json:"error"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(data), &ev); err != nil {
|
||||
continue
|
||||
}
|
||||
msg := strings.TrimSpace(ev.Error.Message)
|
||||
if msg == "" {
|
||||
msg = "Anthropic API 错误"
|
||||
}
|
||||
ch <- StreamChunk{Error: errors.New(msg)}
|
||||
}
|
||||
}
|
||||
if err := scanner.Err(); err != nil && ctx.Err() == nil {
|
||||
ch <- StreamChunk{Error: err}
|
||||
}
|
||||
}
|
||||
|
||||
// compressAnthropic 使用 Anthropic 非流式接口执行对话压缩。
|
||||
func (c *LLMClient) compressAnthropic(ctx context.Context, messages []Message) (string, error) {
|
||||
body := map[string]any{
|
||||
"model": c.apiConfig.Model(),
|
||||
"max_tokens": c.cfg.CompactionTokens,
|
||||
"stream": false,
|
||||
"temperature": 0.2,
|
||||
// 关闭思考：推理模型（如 deepseek-v4-flash）默认先输出 thinking 块，
|
||||
// 会把 max_tokens 预算耗尽而拿不到 text 块，导致压缩被判为失败。
|
||||
"thinking": map[string]any{"type": "disabled"},
|
||||
"messages": toAnthropicMessages(messages),
|
||||
}
|
||||
if system := extractSystemPrompt(messages); system != "" {
|
||||
body["system"] = system
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(ctx, c.cfg.RequestTimeout)
|
||||
defer cancel()
|
||||
|
||||
payload, err := json.Marshal(body)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
respBody, err := c.doSyncRequestWithRetry(ctx, true, normalizeAnthropicURL(c.apiConfig.BaseURL()), payload)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
var result struct {
|
||||
Content []struct {
|
||||
Type string `json:"type"`
|
||||
Text string `json:"text"`
|
||||
Thinking string `json:"thinking"`
|
||||
} `json:"content"`
|
||||
}
|
||||
if err := json.Unmarshal(respBody, &result); err != nil {
|
||||
return "", err
|
||||
}
|
||||
fallback := ""
|
||||
for _, block := range result.Content {
|
||||
if block.Type == "text" && strings.TrimSpace(block.Text) != "" {
|
||||
return strings.TrimSpace(block.Text), nil
|
||||
}
|
||||
// 记录 thinking 作为兜底（仅当端点不支持 thinking:disabled 时才会出现）
|
||||
if block.Type == "thinking" && fallback == "" {
|
||||
fallback = strings.TrimSpace(block.Thinking)
|
||||
}
|
||||
}
|
||||
if fallback != "" {
|
||||
return fallback, nil
|
||||
}
|
||||
return "", errors.New("模型没有返回压缩摘要")
|
||||
}
|
||||
@@ -0,0 +1,338 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"crypto/rand"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
type SessionStore struct {
|
||||
mu sync.RWMutex
|
||||
path string
|
||||
sessions map[string]*Session
|
||||
order []string
|
||||
}
|
||||
|
||||
func NewSessionStore(dataDir string) (*SessionStore, error) {
|
||||
if err := os.MkdirAll(dataDir, 0o755); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
path := filepath.Join(dataDir, "sessions.json")
|
||||
store := &SessionStore{
|
||||
path: path,
|
||||
sessions: make(map[string]*Session),
|
||||
order: make([]string, 0),
|
||||
}
|
||||
_ = store.load()
|
||||
return store, nil
|
||||
}
|
||||
|
||||
func (s *SessionStore) load() error {
|
||||
data, err := os.ReadFile(s.path)
|
||||
if err != nil {
|
||||
if errors.Is(err, os.ErrNotExist) {
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
}
|
||||
if len(strings.TrimSpace(string(data))) == 0 {
|
||||
return nil
|
||||
}
|
||||
var sessions map[string]*Session
|
||||
if err := json.Unmarshal(data, &sessions); err != nil {
|
||||
return err
|
||||
}
|
||||
for id, session := range sessions {
|
||||
if session == nil || session.ID == "" {
|
||||
continue
|
||||
}
|
||||
s.sessions[id] = session
|
||||
s.order = append(s.order, id)
|
||||
}
|
||||
s.sortOrder()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *SessionStore) saveLocked() error {
|
||||
data, err := json.MarshalIndent(s.sessions, "", " ")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
tmp := s.path + ".tmp"
|
||||
if err := os.WriteFile(tmp, data, 0o600); err != nil {
|
||||
return err
|
||||
}
|
||||
return os.Rename(tmp, s.path)
|
||||
}
|
||||
|
||||
func (s *SessionStore) List() []SessionSummary {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
return s.listLocked("", 0, 0)
|
||||
}
|
||||
|
||||
// Search 按关键词过滤会话（标题或任意消息内容，忽略大小写）。
|
||||
func (s *SessionStore) Search(query string) []SessionSummary {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
return s.listLocked(query, 0, 0)
|
||||
}
|
||||
|
||||
// listLocked 返回会话摘要列表；query 非空时按关键词过滤，limit>0 时截断。
|
||||
// 必须持读锁调用。
|
||||
func (s *SessionStore) listLocked(query string, limit, offset int) []SessionSummary {
|
||||
query = strings.ToLower(strings.TrimSpace(query))
|
||||
out := make([]SessionSummary, 0, len(s.order))
|
||||
skip := 0
|
||||
for _, id := range s.order {
|
||||
session, ok := s.sessions[id]
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
if query != "" && !sessionMatches(session, query) {
|
||||
continue
|
||||
}
|
||||
if offset > 0 && skip < offset {
|
||||
skip++
|
||||
continue
|
||||
}
|
||||
if limit > 0 && len(out) >= limit {
|
||||
break
|
||||
}
|
||||
out = append(out, summarize(session))
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func sessionMatches(session *Session, query string) bool {
|
||||
if strings.Contains(strings.ToLower(session.Title), query) {
|
||||
return true
|
||||
}
|
||||
for _, message := range session.Messages {
|
||||
if message.Content != nil && strings.Contains(strings.ToLower(*message.Content), query) {
|
||||
return true
|
||||
}
|
||||
for _, call := range message.ToolCalls {
|
||||
if strings.Contains(strings.ToLower(call.Function.Name), query) ||
|
||||
strings.Contains(strings.ToLower(call.Function.Arguments), query) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// Messages 分页返回会话消息（按时间正序），offset 从最新一条往前数，
|
||||
// 即 offset=0 返回最新 limit 条；返回 has_more 表示还有更早的消息。
|
||||
func (s *SessionStore) Messages(id string, limit, offset int) ([]Message, bool, bool) {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
session, ok := s.sessions[id]
|
||||
if !ok {
|
||||
return nil, false, false
|
||||
}
|
||||
total := len(session.Messages)
|
||||
if limit <= 0 {
|
||||
limit = 50
|
||||
}
|
||||
if offset < 0 {
|
||||
offset = 0
|
||||
}
|
||||
start := total - offset - limit
|
||||
if start < 0 {
|
||||
start = 0
|
||||
}
|
||||
messages := make([]Message, 0, limit)
|
||||
for i := start; i < total-offset; i++ {
|
||||
messages = append(messages, cloneMessage(session.Messages[i]))
|
||||
}
|
||||
hasMore := start > 0
|
||||
return messages, true, hasMore
|
||||
}
|
||||
|
||||
// Clear 清空会话消息（保留会话本身）。
|
||||
func (s *SessionStore) Clear(id string) bool {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
session, ok := s.sessions[id]
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
session.Messages = make([]Message, 0)
|
||||
session.UpdatedAt = time.Now()
|
||||
_ = s.saveLocked()
|
||||
return true
|
||||
}
|
||||
|
||||
// TotalMessages 返回会话消息总数（用于分页计算）。
|
||||
func (s *SessionStore) TotalMessages(id string) (int, bool) {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
session, ok := s.sessions[id]
|
||||
if !ok {
|
||||
return 0, false
|
||||
}
|
||||
return len(session.Messages), true
|
||||
}
|
||||
|
||||
func cloneMessage(message Message) Message {
|
||||
copy := message
|
||||
if message.Content != nil {
|
||||
content := *message.Content
|
||||
copy.Content = &content
|
||||
}
|
||||
copy.ToolCalls = append([]ToolCall(nil), message.ToolCalls...)
|
||||
return copy
|
||||
}
|
||||
|
||||
func (s *SessionStore) Get(id string) (*Session, bool) {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
session, ok := s.sessions[id]
|
||||
return session, ok
|
||||
}
|
||||
|
||||
func (s *SessionStore) Create() *Session {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
session := &Session{
|
||||
ID: newID(),
|
||||
Title: "新对话",
|
||||
CreatedAt: time.Now(),
|
||||
UpdatedAt: time.Now(),
|
||||
Messages: make([]Message, 0),
|
||||
}
|
||||
s.sessions[session.ID] = session
|
||||
s.order = append([]string{session.ID}, s.order...)
|
||||
_ = s.saveLocked()
|
||||
return session
|
||||
}
|
||||
|
||||
func (s *SessionStore) Delete(id string) bool {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if _, ok := s.sessions[id]; !ok {
|
||||
return false
|
||||
}
|
||||
delete(s.sessions, id)
|
||||
for i, item := range s.order {
|
||||
if item == id {
|
||||
s.order = append(s.order[:i], s.order[i+1:]...)
|
||||
break
|
||||
}
|
||||
}
|
||||
_ = s.saveLocked()
|
||||
return true
|
||||
}
|
||||
|
||||
func (s *SessionStore) Save(session *Session) {
|
||||
session.UpdatedAt = time.Now()
|
||||
s.mu.Lock()
|
||||
// 存储克隆而非原对象：会话对象可能仍在 Agent 循环中被修改，
|
||||
// 直接存指针会导致 saveLocked 序列化时与其他会话的写入产生数据竞争。
|
||||
s.sessions[session.ID] = cloneSession(session)
|
||||
s.order = append([]string{session.ID}, removeString(s.order, session.ID)...)
|
||||
err := s.saveLocked()
|
||||
s.mu.Unlock()
|
||||
if err != nil {
|
||||
// A failed save should not break an in-memory conversation.
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
func (s *SessionStore) sortOrder() {
|
||||
sort.SliceStable(s.order, func(i, j int) bool {
|
||||
a, aOK := s.sessions[s.order[i]]
|
||||
b, bOK := s.sessions[s.order[j]]
|
||||
if !aOK || !bOK {
|
||||
return aOK && !bOK
|
||||
}
|
||||
return a.UpdatedAt.After(b.UpdatedAt)
|
||||
})
|
||||
}
|
||||
|
||||
func summarize(session *Session) SessionSummary {
|
||||
summary := SessionSummary{
|
||||
ID: session.ID,
|
||||
Title: session.Title,
|
||||
CreatedAt: session.CreatedAt,
|
||||
UpdatedAt: session.UpdatedAt,
|
||||
MessageCount: len(session.Messages),
|
||||
}
|
||||
for i := len(session.Messages) - 1; i >= 0; i-- {
|
||||
message := session.Messages[i]
|
||||
if message.Role != "user" && message.Role != "assistant" {
|
||||
continue
|
||||
}
|
||||
if message.Content != nil {
|
||||
summary.Preview = truncateRunes(*message.Content, 120)
|
||||
break
|
||||
}
|
||||
}
|
||||
return summary
|
||||
}
|
||||
|
||||
func newID() string {
|
||||
buf := make([]byte, 12)
|
||||
_, _ = rand.Read(buf)
|
||||
return hex.EncodeToString(buf)
|
||||
}
|
||||
|
||||
// cloneSession 深拷贝会话对象，保证持久化对象在存储后不再被外部修改，
|
||||
// 从而消除 Agent 循环写入与 saveLocked 序列化之间的数据竞争。
|
||||
func cloneSession(session *Session) *Session {
|
||||
if session == nil {
|
||||
return nil
|
||||
}
|
||||
clone := *session
|
||||
clone.Messages = make([]Message, len(session.Messages))
|
||||
for i, message := range session.Messages {
|
||||
copy := message
|
||||
if message.Content != nil {
|
||||
content := *message.Content
|
||||
copy.Content = &content
|
||||
}
|
||||
copy.ToolCalls = append([]ToolCall(nil), message.ToolCalls...)
|
||||
clone.Messages[i] = copy
|
||||
}
|
||||
return &clone
|
||||
}
|
||||
|
||||
func removeString(items []string, target string) []string {
|
||||
result := make([]string, 0, len(items))
|
||||
for _, item := range items {
|
||||
if item != target {
|
||||
result = append(result, item)
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
func truncateRunes(value string, max int) string {
|
||||
runes := []rune(value)
|
||||
if len(runes) <= max {
|
||||
return value
|
||||
}
|
||||
return string(runes[:max]) + "..."
|
||||
}
|
||||
|
||||
func firstRunes(value string, max int) string {
|
||||
runes := []rune(strings.TrimSpace(value))
|
||||
if len(runes) <= max {
|
||||
return string(runes)
|
||||
}
|
||||
return string(runes[:max])
|
||||
}
|
||||
|
||||
func countRunes(value string) int {
|
||||
return utf8.RuneCountInString(value)
|
||||
}
|
||||
@@ -0,0 +1,559 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
type Toolset struct {
|
||||
Workspace string
|
||||
Timeout time.Duration
|
||||
MaxOutput int
|
||||
shell string
|
||||
shellArgs []string
|
||||
python string
|
||||
dockerSocket func() string
|
||||
apiCfg *APIConfigStore
|
||||
coop *CoopManager
|
||||
// sessionRoot 是各会话独立工作目录的根（data/workspaces），
|
||||
// 每个会话在其下拥有 <sessionRoot>/<sessionID> 目录，互不干扰。
|
||||
sessionRoot string
|
||||
}
|
||||
|
||||
func NewToolset(cfg Config, dockerCfg *DockerConfigStore, apiCfg *APIConfigStore) *Toolset {
|
||||
toolset := &Toolset{
|
||||
Workspace: cfg.Workspace,
|
||||
Timeout: cfg.ToolTimeout,
|
||||
MaxOutput: cfg.MaxToolResultChars,
|
||||
apiCfg: apiCfg,
|
||||
sessionRoot: filepath.Join(cfg.DataDir, "workspaces"),
|
||||
}
|
||||
|
||||
if dockerCfg != nil {
|
||||
toolset.dockerSocket = dockerCfg.Socket
|
||||
} else {
|
||||
toolset.dockerSocket = defaultDockerSocket
|
||||
}
|
||||
|
||||
if shell, err := exec.LookPath("bash"); err == nil {
|
||||
toolset.shell = shell
|
||||
toolset.shellArgs = []string{"-lc"}
|
||||
} else if shell, err := exec.LookPath("pwsh"); err == nil {
|
||||
toolset.shell = shell
|
||||
toolset.shellArgs = []string{"-NoProfile", "-NonInteractive", "-Command"}
|
||||
} else if shell, err := exec.LookPath("powershell"); err == nil {
|
||||
toolset.shell = shell
|
||||
toolset.shellArgs = []string{"-NoProfile", "-NonInteractive", "-Command"}
|
||||
} else {
|
||||
toolset.shell = "cmd.exe"
|
||||
toolset.shellArgs = []string{"/C"}
|
||||
}
|
||||
|
||||
if python, err := exec.LookPath("python"); err == nil {
|
||||
toolset.python = python
|
||||
} else if python, err := exec.LookPath("python3"); err == nil {
|
||||
toolset.python = python
|
||||
} else {
|
||||
toolset.python = "python"
|
||||
}
|
||||
return toolset
|
||||
}
|
||||
|
||||
func (t *Toolset) Definitions() []ToolDefinition {
|
||||
return []ToolDefinition{
|
||||
{
|
||||
Type: "function",
|
||||
Function: ToolDefinitionFunction{
|
||||
Name: "run_bash",
|
||||
Description: "在用户工作区目录中执行一条 Bash 命令。Windows 上没有 Bash 时自动回退到 PowerShell 或 cmd。适合查看文件、运行构建、安装依赖、搜索代码、启动程序等。",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"command": map[string]any{
|
||||
"type": "string",
|
||||
"description": "要执行的完整命令，例如 `ls -la` 或 `go test ./...`。",
|
||||
},
|
||||
"timeout_seconds": map[string]any{
|
||||
"type": "integer",
|
||||
"description": "命令超时秒数，默认 120 秒。",
|
||||
},
|
||||
},
|
||||
"required": []string{"command"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Type: "function",
|
||||
Function: ToolDefinitionFunction{
|
||||
Name: "read_file",
|
||||
Description: "读取工作区内的文本文件并返回内容。超长文件会被截断。",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"path": map[string]any{
|
||||
"type": "string",
|
||||
"description": "文件路径，可以是相对路径或绝对路径。必须位于工作区内。",
|
||||
},
|
||||
},
|
||||
"required": []string{"path"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Type: "function",
|
||||
Function: ToolDefinitionFunction{
|
||||
Name: "write_file",
|
||||
Description: "写入或覆盖工作区内的文件。父目录不存在时自动创建。",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"path": map[string]any{
|
||||
"type": "string",
|
||||
"description": "目标文件路径，可以是相对路径或绝对路径。必须位于工作区内。",
|
||||
},
|
||||
"content": map[string]any{
|
||||
"type": "string",
|
||||
"description": "要写入的完整文件内容。",
|
||||
},
|
||||
},
|
||||
"required": []string{"path", "content"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Type: "function",
|
||||
Function: ToolDefinitionFunction{
|
||||
Name: "run_python",
|
||||
Description: "在工作区目录中执行一段 Python 代码。适合数据处理、批量修改、生成脚本和自动化任务。",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"code": map[string]any{
|
||||
"type": "string",
|
||||
"description": "要执行的完整 Python 代码。",
|
||||
},
|
||||
},
|
||||
"required": []string{"code"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Type: "function",
|
||||
Function: ToolDefinitionFunction{
|
||||
Name: "list_directory",
|
||||
Description: "列出工作区内目录的内容，包括文件大小和修改时间。",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"path": map[string]any{
|
||||
"type": "string",
|
||||
"description": "目录路径，默认为工作区根目录。",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Type: "function",
|
||||
Function: ToolDefinitionFunction{
|
||||
Name: "docker_ps",
|
||||
Description: "列出 Docker 容器。默认连接本地 Docker，可通过 Web 设置中的 Docker Socket 切换到远程 Docker 服务器。",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"all": map[string]any{
|
||||
"type": "boolean",
|
||||
"description": "是否列出所有容器（包括已停止的），默认 false 只显示运行中。",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Type: "function",
|
||||
Function: ToolDefinitionFunction{
|
||||
Name: "docker_images",
|
||||
Description: "列出 Docker 镜像列表。",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Type: "function",
|
||||
Function: ToolDefinitionFunction{
|
||||
Name: "docker_logs",
|
||||
Description: "查看 Docker 容器的最近日志。",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"container": map[string]any{
|
||||
"type": "string",
|
||||
"description": "容器名称或 ID。",
|
||||
},
|
||||
"tail": map[string]any{
|
||||
"type": "integer",
|
||||
"description": "返回最近的日志行数，默认 100。",
|
||||
},
|
||||
},
|
||||
"required": []string{"container"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Type: "function",
|
||||
Function: ToolDefinitionFunction{
|
||||
Name: "docker_inspect",
|
||||
Description: "查看 Docker 容器的详细配置和状态信息。",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"container": map[string]any{
|
||||
"type": "string",
|
||||
"description": "容器名称或 ID。",
|
||||
},
|
||||
},
|
||||
"required": []string{"container"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Type: "function",
|
||||
Function: ToolDefinitionFunction{
|
||||
Name: "docker_exec",
|
||||
Description: "在 Docker 容器内执行一条命令（使用 sh -c），返回执行输出。",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"container": map[string]any{
|
||||
"type": "string",
|
||||
"description": "容器名称或 ID。",
|
||||
},
|
||||
"command": map[string]any{
|
||||
"type": "string",
|
||||
"description": "要在容器内执行的命令，例如 `ls -la` 或 `cat /etc/os-release`。",
|
||||
},
|
||||
},
|
||||
"required": []string{"container", "command"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Type: "function",
|
||||
Function: ToolDefinitionFunction{
|
||||
Name: "docker_start",
|
||||
Description: "启动一个 Docker 容器。",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"container": map[string]any{
|
||||
"type": "string",
|
||||
"description": "容器名称或 ID。",
|
||||
},
|
||||
},
|
||||
"required": []string{"container"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Type: "function",
|
||||
Function: ToolDefinitionFunction{
|
||||
Name: "docker_stop",
|
||||
Description: "停止一个 Docker 容器。",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"container": map[string]any{
|
||||
"type": "string",
|
||||
"description": "容器名称或 ID。",
|
||||
},
|
||||
"timeout": map[string]any{
|
||||
"type": "integer",
|
||||
"description": "等待优雅停止的秒数，默认由 Docker 决定。",
|
||||
},
|
||||
},
|
||||
"required": []string{"container"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Type: "function",
|
||||
Function: ToolDefinitionFunction{
|
||||
Name: "run_coop",
|
||||
Description: "调用协作容器异步完成一个任务（容器内单个 agent 独立执行，自动注入当前 Agent 已配置的模型、Base URL、API Key 与接口协议，无需重复填写）。支持三种引擎：engine=pi（默认，pi-coding-agent，支持 openai+anthropic 协议）、engine=pigo（pigo，支持 openai+anthropic 协议）或 engine=claude（Claude Code，仅支持 anthropic 协议）。需先构建对应镜像（pi-coop / pigo-coop / claude-coop）。该工具是异步的：调用后立即返回任务 ID，容器在后台运行；任务完成（成功 / 失败 / 超时）后系统会自动注入一条【协作任务完成通知】消息（含结构化结果摘要），由你确认结果、必要时提交 flag 并向用户汇报，无需在此等待。",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"task": map[string]any{
|
||||
"type": "string",
|
||||
"description": "要交给协作容器完成的任务描述，例如“请为 XX 编写设计文档并实现原型”。解题类任务需写明目标地址、平台提交规则（BENCHMARK_TOKEN / unique_code / submit API）。",
|
||||
},
|
||||
"engine": map[string]any{
|
||||
"type": "string",
|
||||
"description": "协作引擎：pi（默认，pi-coding-agent）、pigo（pigo）或 claude（Claude Code，仅支持 anthropic 协议）。pi 与 pigo 支持 openai+anthropic 协议；claude 仅支持 anthropic 协议。",
|
||||
"enum": []string{"pi", "pigo", "claude"},
|
||||
},
|
||||
"round_max": map[string]any{
|
||||
"type": "integer",
|
||||
"description": "最大轮次，默认 1（单 agent 一次运行完成，未完成则失败并重新下发）。",
|
||||
},
|
||||
"timeout": map[string]any{
|
||||
"type": "integer",
|
||||
"description": "单轮超时秒数，默认 1800（30 分钟）。单 agent 需在此时限内完成全部工作并创建 DONE。",
|
||||
},
|
||||
"blackboard": map[string]any{
|
||||
"type": "string",
|
||||
"description": "可选：主机上的目录路径，挂载到容器 /blackboard 保留黑板产物与会话；Windows 路径会自动转换为 WSL 挂载路径（/mnt/盘符/...）。留空则自动使用本会话工作区下的独立子目录 blackboard/<任务ID>（同一会话多次协作互不污染）。",
|
||||
},
|
||||
},
|
||||
"required": []string{"task"},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (t *Toolset) Execute(ctx context.Context, sessionID, name, arguments string) ToolResult {
|
||||
var params map[string]any
|
||||
if err := json.Unmarshal([]byte(arguments), ¶ms); err != nil {
|
||||
return ToolResult{Success: false, Output: "无法解析工具参数: " + err.Error()}
|
||||
}
|
||||
|
||||
switch name {
|
||||
case "run_bash":
|
||||
command, _ := params["command"].(string)
|
||||
if strings.TrimSpace(command) == "" {
|
||||
return ToolResult{Success: false, Output: "command 不能为空"}
|
||||
}
|
||||
return t.runShell(ctx, sessionID, command, intParam(params, "timeout_seconds"))
|
||||
case "read_file":
|
||||
path, _ := params["path"].(string)
|
||||
return t.readFile(sessionID, path)
|
||||
case "write_file":
|
||||
path, _ := params["path"].(string)
|
||||
content, _ := params["content"].(string)
|
||||
return t.writeFile(sessionID, path, content)
|
||||
case "run_python":
|
||||
code, _ := params["code"].(string)
|
||||
if strings.TrimSpace(code) == "" {
|
||||
return ToolResult{Success: false, Output: "code 不能为空"}
|
||||
}
|
||||
return t.runPython(ctx, sessionID, code)
|
||||
case "list_directory":
|
||||
path, _ := params["path"].(string)
|
||||
return t.listDirectory(sessionID, path)
|
||||
case "docker_ps":
|
||||
return t.dockerPS(ctx, params)
|
||||
case "docker_images":
|
||||
return t.dockerImages(ctx, params)
|
||||
case "docker_logs":
|
||||
return t.dockerLogs(ctx, params)
|
||||
case "docker_inspect":
|
||||
return t.dockerInspect(ctx, params)
|
||||
case "docker_exec":
|
||||
return t.dockerExec(ctx, params)
|
||||
case "docker_start":
|
||||
return t.dockerStart(ctx, params)
|
||||
case "docker_stop":
|
||||
return t.dockerStop(ctx, params)
|
||||
case "run_coop":
|
||||
return t.runCoop(ctx, sessionID, params)
|
||||
default:
|
||||
return ToolResult{Success: false, Output: "未知工具: " + name}
|
||||
}
|
||||
}
|
||||
|
||||
func (t *Toolset) runShell(ctx context.Context, sessionID, command string, timeoutSeconds int) ToolResult {
|
||||
// 不信任 LLM 传入的超长 timeout：超过全局工具超时（默认 120s）一律截断为全局值。
|
||||
// 否则 sleep / 轮询类命令会长时间阻塞主 Agent（持有会话锁），期间前端
|
||||
// 长时间收不到事件会被判"任务中断"，且 coop 完成通知也拿不到锁无法及时处理。
|
||||
if timeoutSeconds <= 0 || timeoutSeconds > int(t.Timeout.Seconds()) {
|
||||
timeoutSeconds = int(t.Timeout.Seconds())
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(ctx, time.Duration(timeoutSeconds)*time.Second)
|
||||
defer cancel()
|
||||
|
||||
var stdout, stderr bytes.Buffer
|
||||
cmd := exec.CommandContext(ctx, t.shell, t.shellArgs...)
|
||||
cmd.Args = append(cmd.Args, command)
|
||||
cmd.Dir = t.workspaceFor(sessionID)
|
||||
cmd.Env = append(os.Environ(), "AGENT_WORKSPACE="+t.workspaceFor(sessionID))
|
||||
cmd.Stdout = &stdout
|
||||
cmd.Stderr = &stderr
|
||||
|
||||
err := cmd.Run()
|
||||
var out strings.Builder
|
||||
out.WriteString(strings.TrimSpace(stdout.String()))
|
||||
if stderrText := strings.TrimSpace(stderr.String()); stderrText != "" {
|
||||
if out.Len() > 0 {
|
||||
out.WriteString("\n")
|
||||
}
|
||||
out.WriteString("[stderr]\n")
|
||||
out.WriteString(stderrText)
|
||||
}
|
||||
if err != nil {
|
||||
if out.Len() > 0 {
|
||||
out.WriteString("\n")
|
||||
}
|
||||
out.WriteString("[命令失败] " + err.Error())
|
||||
}
|
||||
return ToolResult{Success: err == nil, Output: t.truncate(out.String())}
|
||||
}
|
||||
|
||||
func (t *Toolset) runPython(ctx context.Context, sessionID, code string) ToolResult {
|
||||
ctx, cancel := context.WithTimeout(ctx, t.Timeout)
|
||||
defer cancel()
|
||||
|
||||
var stdout, stderr bytes.Buffer
|
||||
cmd := exec.CommandContext(ctx, t.python, "-c", code)
|
||||
cmd.Dir = t.workspaceFor(sessionID)
|
||||
cmd.Env = append(os.Environ(), "AGENT_WORKSPACE="+t.workspaceFor(sessionID))
|
||||
cmd.Stdout = &stdout
|
||||
cmd.Stderr = &stderr
|
||||
|
||||
err := cmd.Run()
|
||||
var out strings.Builder
|
||||
out.WriteString(strings.TrimSpace(stdout.String()))
|
||||
if stderrText := strings.TrimSpace(stderr.String()); stderrText != "" {
|
||||
if out.Len() > 0 {
|
||||
out.WriteString("\n")
|
||||
}
|
||||
out.WriteString("[stderr]\n")
|
||||
out.WriteString(stderrText)
|
||||
}
|
||||
if err != nil {
|
||||
if out.Len() > 0 {
|
||||
out.WriteString("\n")
|
||||
}
|
||||
out.WriteString("[运行失败] " + err.Error())
|
||||
}
|
||||
return ToolResult{Success: err == nil, Output: t.truncate(out.String())}
|
||||
}
|
||||
|
||||
func (t *Toolset) readFile(sessionID, path string) ToolResult {
|
||||
absPath, err := t.resolvePath(sessionID, path)
|
||||
if err != nil {
|
||||
return ToolResult{Success: false, Output: err.Error()}
|
||||
}
|
||||
data, err := os.ReadFile(absPath)
|
||||
if err != nil {
|
||||
return ToolResult{Success: false, Output: err.Error()}
|
||||
}
|
||||
return ToolResult{Success: true, Output: t.truncate(string(data))}
|
||||
}
|
||||
|
||||
func (t *Toolset) writeFile(sessionID, path, content string) ToolResult {
|
||||
absPath, err := t.resolvePath(sessionID, path)
|
||||
if err != nil {
|
||||
return ToolResult{Success: false, Output: err.Error()}
|
||||
}
|
||||
if err := os.MkdirAll(filepath.Dir(absPath), 0o755); err != nil {
|
||||
return ToolResult{Success: false, Output: err.Error()}
|
||||
}
|
||||
if err := os.WriteFile(absPath, []byte(content), 0o644); err != nil {
|
||||
return ToolResult{Success: false, Output: err.Error()}
|
||||
}
|
||||
return ToolResult{Success: true, Output: fmt.Sprintf("已写入 %d 字节到 %s", len(content), absPath)}
|
||||
}
|
||||
|
||||
func (t *Toolset) listDirectory(sessionID, path string) ToolResult {
|
||||
absPath, err := t.resolvePath(sessionID, path)
|
||||
if err != nil {
|
||||
return ToolResult{Success: false, Output: err.Error()}
|
||||
}
|
||||
entries, err := os.ReadDir(absPath)
|
||||
if err != nil {
|
||||
return ToolResult{Success: false, Output: err.Error()}
|
||||
}
|
||||
sort.Slice(entries, func(i, j int) bool {
|
||||
return entries[i].Name() < entries[j].Name()
|
||||
})
|
||||
|
||||
var out strings.Builder
|
||||
for _, entry := range entries {
|
||||
info, infoErr := entry.Info()
|
||||
if infoErr != nil {
|
||||
continue
|
||||
}
|
||||
name := entry.Name()
|
||||
if entry.IsDir() {
|
||||
name += "/"
|
||||
}
|
||||
fmt.Fprintf(&out, "%-40s %10d %s\n", name, info.Size(), info.ModTime().Format("2006-01-02 15:04"))
|
||||
}
|
||||
if out.Len() == 0 {
|
||||
return ToolResult{Success: true, Output: "目录为空"}
|
||||
}
|
||||
return ToolResult{Success: true, Output: t.truncate(out.String())}
|
||||
}
|
||||
|
||||
// resolvePath 把工具传入的路径解析为绝对路径：
|
||||
// - 相对路径基于该会话的独立工作目录（data/workspaces/<sessionID>）；
|
||||
// - 绝对路径允许在项目根目录内（协作黑板产物、agent 自身文件等共享内容）；
|
||||
// - 项目根目录之外一律拒绝。
|
||||
func (t *Toolset) resolvePath(sessionID, raw string) (string, error) {
|
||||
if strings.TrimSpace(raw) == "" {
|
||||
raw = "."
|
||||
}
|
||||
path := filepath.FromSlash(raw)
|
||||
if !filepath.IsAbs(path) {
|
||||
path = filepath.Join(t.workspaceFor(sessionID), path)
|
||||
}
|
||||
absPath, err := filepath.Abs(path)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if !pathWithin(absPath, t.Workspace) {
|
||||
return "", errors.New("路径超出项目根目录范围，已拒绝: " + absPath)
|
||||
}
|
||||
return absPath, nil
|
||||
}
|
||||
|
||||
// workspaceFor 返回某会话的独立工作目录并确保其存在。
|
||||
func (t *Toolset) workspaceFor(sessionID string) string {
|
||||
root := t.sessionRoot
|
||||
if root == "" {
|
||||
root = t.Workspace
|
||||
}
|
||||
dir := filepath.Join(root, sessionID)
|
||||
_ = os.MkdirAll(dir, 0o755)
|
||||
return dir
|
||||
}
|
||||
|
||||
// pathWithin 判断 child 是否位于 parent 目录内（含自身）。
|
||||
func pathWithin(child, parent string) bool {
|
||||
rel, err := filepath.Rel(parent, child)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
return rel == "." || (rel != ".." && !strings.HasPrefix(rel, ".."+string(filepath.Separator)))
|
||||
}
|
||||
|
||||
func (t *Toolset) truncate(value string) string {
|
||||
runes := []rune(value)
|
||||
if len(runes) <= t.MaxOutput {
|
||||
return value
|
||||
}
|
||||
return string(runes[:t.MaxOutput]) + "\n...[输出过长，已截断]"
|
||||
}
|
||||
|
||||
func intParam(params map[string]any, key string) int {
|
||||
switch value := params[key].(type) {
|
||||
case float64:
|
||||
return int(value)
|
||||
case int:
|
||||
return value
|
||||
default:
|
||||
return 0
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,556 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/docker/docker/api/types/container"
|
||||
"github.com/docker/docker/api/types/mount"
|
||||
)
|
||||
|
||||
// 协作镜像名（构建方式见各自目录的 RUN.md）。
|
||||
const (
|
||||
coopImagePigo = "pigo-coop" // pigo 协作镜像（支持 openai + anthropic 协议）
|
||||
coopImage = "pi-coop" // pi 协作镜像（支持 openai + anthropic 协议）
|
||||
coopImageClaude = "claude-coop" // Claude Code 协作镜像（仅支持 anthropic 协议）
|
||||
)
|
||||
|
||||
// 协作引擎到镜像 / 容器名前缀 / 展示标签 / 运行时目录的映射。
|
||||
// 三种引擎共享同一套环境变量契约（MODEL/BASE_URL/API_KEY/PROTOCOL/TASK 等），
|
||||
// supervisor 各自适配，控制层只需选择镜像并做协议约束 + baseURL 规范化。
|
||||
// runtimeDir 是 local 模式下 supervisor.sh / prompts / extensions 所在的相对目录。
|
||||
var coopEngines = map[string]struct {
|
||||
image, namePrefix, label, runtimeDir string
|
||||
}{
|
||||
"pigo": {coopImagePigo, "pigo-coop-", "pigo", "pigo/coop"},
|
||||
"pi": {coopImage, "pi-coop-", "pi", "pi-coop"},
|
||||
"claude": {coopImageClaude, "claude-coop-", "claude code", "claude code"},
|
||||
}
|
||||
|
||||
// codePattern 从任务描述中提取题目编号（如 "c-06" / "a-05" / "e1-01"）。
|
||||
var codePattern = regexp.MustCompile(`[A-Za-z0-9]{1,8}-\d{1,4}`)
|
||||
|
||||
// runCoop 异步启动 pi 协作任务：
|
||||
// 校验参数与镜像后立即返回，容器创建 / 运行 / 等待 / 清理全部放到后台 goroutine。
|
||||
// 任务结束后由 CoopManager 通知 Agent，主 agent 会自动收到一条
|
||||
// "【协作任务完成通知】"消息并汇报结果，无需在前端长时间等待。
|
||||
func (t *Toolset) runCoop(ctx context.Context, sessionID string, params map[string]any) ToolResult {
|
||||
if t.coop == nil {
|
||||
return ToolResult{Success: false, Output: "协作任务管理器未初始化"}
|
||||
}
|
||||
task, _ := params["task"].(string)
|
||||
task = strings.TrimSpace(task)
|
||||
if task == "" {
|
||||
return ToolResult{Success: false, Output: "task 不能为空"}
|
||||
}
|
||||
|
||||
// 完成协议：约束协作 agent 主动收尾，避免"已提交成功却没写 DONE 被强杀"、
|
||||
// "解不出却空耗到超时"两类问题。supervisor 只在黑板根目录出现 DONE 标记时正常退出。
|
||||
completionProtocol := `
|
||||
|
||||
【完成协议（务必严格遵守，决定协作能否正常收尾）】
|
||||
- 找到 flag 并提交成功（平台响应 correct=true）后：用 blackboard 工具 action=done，在黑板根目录创建 DONE 标记，内容写入完成摘要（含 flag 值、提交响应、解题路径）。supervisor 检测到 DONE 即正常结束（exit 0）；不创建 DONE 会空耗到超时被强杀，协作被视为失败。
|
||||
- 若经充分尝试后确认本轮无法解出（目标不可达 / 无漏洞 / 试错过多）：同样用 action=done 创建 DONE，内容开头写明「未解出」与已尝试内容，让调度方及时关闭靶机并切换下一题，不要空耗到超时。
|
||||
- 已通关题目不要重复提交：平台对已通关题目的后续提交统一返回 correct:false（而非 duplicate），属正常现象、不影响已得分数，不要误判为失败。`
|
||||
task += completionProtocol
|
||||
|
||||
// 从任务描述解析题目编号（如 "c-06"），供完成通知携带，便于控制层关靶机/切题
|
||||
challengeCode := ""
|
||||
if m := codePattern.FindString(task); m != "" {
|
||||
challengeCode = m
|
||||
}
|
||||
roundMax := intParam(params, "round_max")
|
||||
if roundMax <= 0 {
|
||||
// 单 agent 默认 1 轮：一次运行完成全部工作，未完成则失败并重新下发
|
||||
roundMax = 1
|
||||
}
|
||||
if roundMax > 30 {
|
||||
roundMax = 30
|
||||
}
|
||||
timeoutSec := intParam(params, "timeout")
|
||||
if timeoutSec <= 0 {
|
||||
// 默认 1800s：600s/900s 对需要写脚本+多步探测的渗透/解题任务偏紧，
|
||||
// 实测多因单轮超时（exit 143）导致协作失败。
|
||||
timeoutSec = 1800
|
||||
}
|
||||
|
||||
if t.apiCfg == nil {
|
||||
return ToolResult{Success: false, Output: "缺少 LLM API 配置，无法注入模型配置"}
|
||||
}
|
||||
|
||||
// 解析协作引擎：engine=pi（默认）/ pigo / claude。
|
||||
// 未显式指定时使用设置页配置的默认引擎（apiCfg.Engine()）。
|
||||
// 三种镜像共享同一套环境变量契约（MODEL/BASE_URL/API_KEY/PROTOCOL/TASK 等），
|
||||
// supervisor 各自适配，控制层只需选择镜像并做协议约束 + baseURL 规范化。
|
||||
engine, _ := params["engine"].(string)
|
||||
engine = strings.TrimSpace(strings.ToLower(engine))
|
||||
if engine == "" {
|
||||
engine = t.apiCfg.Engine()
|
||||
}
|
||||
eng, ok := coopEngines[engine]
|
||||
if !ok {
|
||||
return ToolResult{Success: false, Output: "不支持的 engine \"" + engine + "\"，可选值：pigo | pi（默认）| claude"}
|
||||
}
|
||||
|
||||
model := t.apiCfg.Model()
|
||||
apiKey := t.apiCfg.APIKey()
|
||||
provider := t.apiCfg.Provider()
|
||||
// coopBaseURL 根据 engine 规范化 baseURL：
|
||||
// pigo 的 anthropicCompatDriver 只追加 /messages，需补 /v1；
|
||||
// pi/claude 用官方 SDK（自带 /v1/messages），需剥离 /v1 避免双重路径。
|
||||
baseURL := coopBaseURL(provider, t.apiCfg.BaseURL(), engine)
|
||||
if model == "" || apiKey == "" || baseURL == "" {
|
||||
return ToolResult{Success: false, Output: "LLM API 配置不完整（model / base_url / api_key 缺一不可），请先在设置页配置"}
|
||||
}
|
||||
|
||||
// Claude Code 仅支持 Anthropic 协议端点（ANTHROPIC_BASE_URL），
|
||||
// 若配置为 openai 协议则直接拒绝，避免容器启动后才报错。
|
||||
if engine == "claude" && provider != ProviderAnthropic {
|
||||
return ToolResult{Success: false, Output: "claude-coop 仅支持 Anthropic 协议端点，当前 provider 为 " + provider +
|
||||
"。请改用 engine=pi 或 engine=pigo，或将 API 配置切换为 anthropic 协议（如 DeepSeek 的 /anthropic 端点）。"}
|
||||
}
|
||||
|
||||
env := []string{
|
||||
"MODEL=" + model,
|
||||
"BASE_URL=" + baseURL,
|
||||
"API_KEY=" + apiKey,
|
||||
"PROTOCOL=" + provider,
|
||||
"TASK=" + task,
|
||||
fmt.Sprintf("ROUND_MAX=%d", roundMax),
|
||||
fmt.Sprintf("TIMEOUT=%d", timeoutSec),
|
||||
}
|
||||
|
||||
// 生成任务 ID（也用作容器名与默认黑板子目录）
|
||||
taskID := newID()
|
||||
|
||||
// 黑板目录：用户显式指定 blackboard 时使用指定路径；
|
||||
// 未指定时为该任务分配独立的会话工作区子目录（data/workspaces/<sessionID>/blackboard/<taskID>），
|
||||
// 保证同一会话发起的多个 coop 任务互相隔离、互不污染。
|
||||
blackboardDir := ""
|
||||
if dirParam, _ := params["blackboard"].(string); strings.TrimSpace(dirParam) != "" {
|
||||
dirParam = strings.TrimSpace(dirParam)
|
||||
abs, err := filepath.Abs(dirParam)
|
||||
if err != nil {
|
||||
return ToolResult{Success: false, Output: "blackboard 路径无效: " + err.Error()}
|
||||
}
|
||||
// 安全约束：blackboard 会被 worker 进程/容器读写，必须限制在项目根目录内，
|
||||
// 否则 LLM 可通过指定任意主机目录（如 ~/.ssh）让 worker 读写敏感文件。
|
||||
if !pathWithin(abs, t.Workspace) {
|
||||
return ToolResult{Success: false, Output: "blackboard 路径超出项目根目录范围，已拒绝: " + abs}
|
||||
}
|
||||
if err := os.MkdirAll(abs, 0o755); err != nil {
|
||||
return ToolResult{Success: false, Output: "创建 blackboard 目录失败: " + err.Error()}
|
||||
}
|
||||
blackboardDir = abs
|
||||
} else {
|
||||
dir := filepath.Join(t.workspaceFor(sessionID), "blackboard", taskID)
|
||||
if err := os.MkdirAll(dir, 0o755); err != nil {
|
||||
return ToolResult{Success: false, Output: "创建 blackboard 目录失败: " + err.Error()}
|
||||
}
|
||||
blackboardDir = dir
|
||||
}
|
||||
|
||||
record := &CoopTask{
|
||||
ID: taskID,
|
||||
SessionID: sessionID,
|
||||
Status: "running",
|
||||
Blackboard: blackboardDir,
|
||||
RoundMax: roundMax,
|
||||
CreatedAt: time.Now(),
|
||||
ChallengeCode: challengeCode,
|
||||
}
|
||||
t.coop.Register(record)
|
||||
|
||||
mode := t.apiCfg.CoopMode()
|
||||
log.Printf("[coop] 协作任务启动 task=%s session=%s mode=%s engine=%s challenge=%s round_max=%d timeout=%ds blackboard=%s",
|
||||
record.ID, sessionID, mode, engine, challengeCode, roundMax, timeoutSec, blackboardDir)
|
||||
|
||||
overallSec := timeoutSec*roundMax + 300
|
||||
if mode == CoopModeLocal {
|
||||
// 本地子进程模式：无需 Docker，直接 exec supervisor.sh
|
||||
supervisor, promptsDir, extensionsDir, err := findCoopRuntime(engine)
|
||||
if err != nil {
|
||||
t.coop.Complete(t.failedTask(record, err))
|
||||
return ToolResult{Success: false, Output: err.Error()}
|
||||
}
|
||||
go t.runCoopLocalAsync(record, env, blackboardDir, supervisor, promptsDir, extensionsDir, roundMax, timeoutSec, eng.label)
|
||||
} else {
|
||||
// Docker 容器模式：现有逻辑
|
||||
var mounts []mount.Mount
|
||||
cli, err := t.dockerClient()
|
||||
if err != nil {
|
||||
t.coop.Complete(t.failedTask(record, err))
|
||||
return ToolResult{Success: false, Output: err.Error()}
|
||||
}
|
||||
defer cli.Close()
|
||||
|
||||
// 判断目标 daemon 系统类型：远程 Linux daemon（如 WSL）的 bind 挂载
|
||||
// 只认容器侧路径，Windows 路径需转换为 /mnt/<盘符>/... 形式。
|
||||
info, infoErr := cli.Info(ctx)
|
||||
linuxDaemon := infoErr == nil && info.OSType == "linux"
|
||||
|
||||
// 检查镜像是否存在（同步快速失败，避免后台任务因镜像缺失白跑）
|
||||
checkCtx, cancel := context.WithTimeout(ctx, t.Timeout)
|
||||
defer cancel()
|
||||
if _, err := cli.ImageInspect(checkCtx, eng.image); err != nil {
|
||||
t.coop.Complete(t.failedTask(record, fmt.Errorf("未找到镜像 %s", eng.image)))
|
||||
return ToolResult{Success: false, Output: "未找到镜像 " + eng.image + "。请先在仓库根构建：\n" +
|
||||
coopBuildHint(engine) + "\n（详见对应目录的 RUN.md）"}
|
||||
}
|
||||
|
||||
// 容器内以非 root 用户（uid=1000 agent）运行，主机目录必须对任何用户可写，
|
||||
// 否则 supervisor 初始化（mkdir/写 task.md/复制 AGENTS.md）会失败。
|
||||
_ = os.Chmod(blackboardDir, 0o777)
|
||||
source := blackboardDir
|
||||
if linuxDaemon {
|
||||
source = wslBindPath(blackboardDir)
|
||||
}
|
||||
mounts = append(mounts, mount.Mount{
|
||||
Type: mount.TypeBind,
|
||||
Source: source,
|
||||
Target: "/blackboard",
|
||||
})
|
||||
|
||||
// 容器创建 / 启动 / 等待 / 清理放到后台，并使用独立上下文，
|
||||
// 避免阻塞当前 SSE 流（此前同步等待最长可达 timeout×round_max+300 秒）。
|
||||
go t.runCoopAsync(record, env, mounts, roundMax, timeoutSec, eng.image, eng.namePrefix)
|
||||
}
|
||||
|
||||
return ToolResult{Success: true, Output: fmt.Sprintf(
|
||||
"协作任务已在后台启动，任务 ID：%s。\n运行方式：%s 协作单 Agent（%s 模式），最多 %d 轮，整体上限约 %d 分钟。\n"+
|
||||
"你无需在此等待，可以继续处理其他请求；任务完成后系统会自动通知你并汇报结果。",
|
||||
record.ID, eng.label, mode, roundMax, overallSec/60)}
|
||||
}
|
||||
|
||||
// runCoopAsync 在后台完成协作容器的创建、启动、等待、日志收集与清理。
|
||||
// 结束（成功 / 失败 / 超时）后通过 coop.Complete 通知 Agent 唤醒主 agent。
|
||||
// image / namePrefix 由 runCoop 根据 engine 选择（pi-coop 或 claude-coop）。
|
||||
func (t *Toolset) runCoopAsync(record *CoopTask, env []string, mounts []mount.Mount, roundMax, timeoutSec int, image, namePrefix string) {
|
||||
// 使用独立后台上下文，避免随 SSE 请求断开而中断协作任务
|
||||
ctx := context.Background()
|
||||
|
||||
cli, err := t.dockerClient()
|
||||
if err != nil {
|
||||
t.coop.Complete(t.failedTask(record, err))
|
||||
return
|
||||
}
|
||||
defer cli.Close()
|
||||
|
||||
name := namePrefix + record.ID[:8]
|
||||
created, err := cli.ContainerCreate(ctx, &container.Config{
|
||||
Image: image,
|
||||
Env: env,
|
||||
}, &container.HostConfig{
|
||||
Mounts: mounts,
|
||||
}, nil, nil, name)
|
||||
if err != nil {
|
||||
t.coop.Complete(t.failedTask(record, fmt.Errorf("创建容器失败: %w", err)))
|
||||
return
|
||||
}
|
||||
record.ContainerID = created.ID
|
||||
|
||||
// 运行结束后无论如何清理容器（等价 docker run --rm）
|
||||
cleanup := func() {
|
||||
cleanupCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
||||
defer cancel()
|
||||
_ = cli.ContainerRemove(cleanupCtx, created.ID, container.RemoveOptions{Force: true})
|
||||
}
|
||||
|
||||
if err := cli.ContainerStart(ctx, created.ID, container.StartOptions{}); err != nil {
|
||||
cleanup()
|
||||
t.coop.Complete(t.failedTask(record, fmt.Errorf("启动容器失败: %w", err)))
|
||||
return
|
||||
}
|
||||
|
||||
// 等待容器退出；整体超时上限 = 轮次 × 单轮超时 + 缓冲
|
||||
overall := time.Duration(timeoutSec*roundMax+300) * time.Second
|
||||
waitCtx, waitCancel := context.WithTimeout(ctx, overall)
|
||||
defer waitCancel()
|
||||
waitCh, errCh := cli.ContainerWait(waitCtx, created.ID, container.WaitConditionNotRunning)
|
||||
|
||||
exitCode := -1
|
||||
select {
|
||||
case res := <-waitCh:
|
||||
exitCode = int(res.StatusCode)
|
||||
case err := <-errCh:
|
||||
cleanup()
|
||||
t.coop.Complete(t.failedTask(record, fmt.Errorf("等待容器退出失败: %w", err)))
|
||||
return
|
||||
case <-waitCtx.Done():
|
||||
cleanup()
|
||||
t.coop.Complete(t.failedTask(record, fmt.Errorf(
|
||||
"协作运行超时（整体上限 %d 秒，约 %.0f 分钟），已强制清理容器。可增大 round_max / timeout 或缩小任务规模后重试。",
|
||||
int(overall.Seconds()), overall.Minutes())))
|
||||
return
|
||||
}
|
||||
|
||||
// 读取容器日志（supervisor 输出 + DONE 总结）
|
||||
logCtx, logCancel := context.WithTimeout(context.Background(), 30*time.Second)
|
||||
defer logCancel()
|
||||
result := ""
|
||||
if logs, err := cli.ContainerLogs(logCtx, created.ID, container.LogsOptions{
|
||||
ShowStdout: true,
|
||||
ShowStderr: true,
|
||||
}); err == nil {
|
||||
if raw, readErr := io.ReadAll(logs); readErr == nil {
|
||||
result = demuxDockerLogs(raw)
|
||||
_ = logs.Close()
|
||||
}
|
||||
}
|
||||
|
||||
result = strings.TrimSpace(result)
|
||||
switch exitCode {
|
||||
case 0:
|
||||
result += "\n[协作完成，退出码 0]"
|
||||
case 1:
|
||||
result += "\n[达到最大轮次仍未完成，退出码 1；可增大 round_max 后重试]"
|
||||
default:
|
||||
result += fmt.Sprintf("\n[容器退出码 %d，请检查上方日志]", exitCode)
|
||||
}
|
||||
if len(mounts) > 0 {
|
||||
result += "\n黑板产物已保留在: " + mounts[0].Source
|
||||
}
|
||||
cleanup()
|
||||
|
||||
record.ExitCode = exitCode
|
||||
record.Output = t.truncate(result)
|
||||
|
||||
// 解析 supervisor 生成的结构化结果 result.json（如有）：
|
||||
// 完成通知与汇报轮优先使用 Result 字段，容器 stdout 仅作兜底。
|
||||
status := "unknown"
|
||||
if data, rerr := os.ReadFile(filepath.Join(record.Blackboard, "result.json")); rerr == nil {
|
||||
var res CoopResult
|
||||
if json.Unmarshal(data, &res) == nil {
|
||||
record.Result = &res
|
||||
status = res.Status
|
||||
}
|
||||
}
|
||||
log.Printf("[coop] 协作任务结束 task=%s session=%s exit_code=%d status=%s", record.ID, record.SessionID, exitCode, status)
|
||||
t.coop.Complete(record)
|
||||
}
|
||||
|
||||
// failedTask 生成一个失败任务记录，供 runCoopAsync 失败路径统一上报。
|
||||
func (t *Toolset) failedTask(record *CoopTask, err error) *CoopTask {
|
||||
record.Error = err.Error()
|
||||
record.Output = "协作任务失败: " + err.Error()
|
||||
return record
|
||||
}
|
||||
|
||||
// coopBuildHint 返回指定 engine 对应镜像的构建命令提示。
|
||||
func coopBuildHint(engine string) string {
|
||||
switch engine {
|
||||
case "pigo":
|
||||
// pigo-coop 构建上下文是 pigo/ 目录，且需先交叉编译 pigo 二进制
|
||||
return "cd pigo && GOOS=linux CGO_ENABLED=0 go build -trimpath -ldflags=\"-s -w\" -o coop/tmp/pigo-linux-amd64 ./cmd/pigo && docker build -f coop/Dockerfile -t pigo-coop ."
|
||||
case "claude":
|
||||
return "docker build -f \"claude code/Dockerfile\" -t claude-coop \"claude code/\""
|
||||
default:
|
||||
return "docker build -f pi-coop/Dockerfile -t pi-coop pi-coop/"
|
||||
}
|
||||
}
|
||||
|
||||
// coopBaseURL 把当前 Agent 配置的 Base URL 规范化为协作容器期望的"基础地址"。
|
||||
// 不同 engine 的模型接入层对 baseURL 的路径处理不同，需按 engine 区分：
|
||||
// - pigo：anthropicCompatDriver 只追加 /messages（不含 /v1），因此 anthropic
|
||||
// 端点需补 /v1（最终 /v1/messages）；openai 端点需保留 /v1。
|
||||
// - pi：provider 使用官方 SDK（@anthropic-ai/sdk / openai），SDK 自行追加完整
|
||||
// 路径（anthropic: /v1/messages；openai: /chat/completions），不能再补 /v1，
|
||||
// 否则产生 /v1/v1/messages 双重路径导致 404。
|
||||
// - claude：同 pi，用 @anthropic-ai/sdk，不能补 /v1。
|
||||
//
|
||||
// 对用户误填的完整路径后缀（/v1/messages、/chat/completions 等）统一剥离，
|
||||
// 再按 engine + provider 决定是否补 /v1。
|
||||
func coopBaseURL(provider, baseURL, engine string) string {
|
||||
base := strings.TrimRight(strings.TrimSpace(baseURL), "/")
|
||||
// 先剥离用户误填的完整路径后缀，统一回退到"基础地址"
|
||||
switch {
|
||||
case strings.HasSuffix(base, "/v1/chat/completions"):
|
||||
base = strings.TrimSuffix(base, "/chat/completions")
|
||||
case strings.HasSuffix(base, "/chat/completions"):
|
||||
base = strings.TrimSuffix(base, "/chat/completions")
|
||||
case strings.HasSuffix(base, "/v1/messages"):
|
||||
base = strings.TrimSuffix(base, "/v1/messages")
|
||||
case strings.HasSuffix(base, "/messages"):
|
||||
base = strings.TrimSuffix(base, "/messages")
|
||||
}
|
||||
|
||||
if provider != ProviderAnthropic {
|
||||
return base
|
||||
}
|
||||
|
||||
// Anthropic 协议端点处理
|
||||
switch engine {
|
||||
case "pigo":
|
||||
// pigo 的 anthropicCompatDriver 只追加 /messages，需补 /v1。
|
||||
// 若用户已填 /v1 结尾则保持，否则补上。
|
||||
if strings.HasSuffix(base, "/v1") {
|
||||
return base
|
||||
}
|
||||
return base + "/v1"
|
||||
default:
|
||||
// pi / claude 用 @anthropic-ai/sdk，SDK 自带 /v1/messages，
|
||||
// 不能补 /v1；若用户已填 /v1 则剥离（SDK 会补回完整的 /v1/messages）。
|
||||
if strings.HasSuffix(base, "/v1") {
|
||||
return strings.TrimSuffix(base, "/v1")
|
||||
}
|
||||
return base
|
||||
}
|
||||
}
|
||||
|
||||
// wslBindPath 把 Windows 绝对路径转换为 WSL 挂载路径（E:\path → /mnt/e/path），
|
||||
// 供运行在 WSL 内的 Linux Docker daemon 做 bind 挂载。
|
||||
func wslBindPath(path string) string {
|
||||
if len(path) < 2 || path[1] != ':' {
|
||||
return path
|
||||
}
|
||||
drive := strings.ToLower(path[:1])
|
||||
rest := strings.ReplaceAll(path[2:], "\\", "/")
|
||||
return "/mnt/" + drive + rest
|
||||
}
|
||||
|
||||
// findCoopRuntime 查找指定 engine 的 supervisor.sh / prompts / extensions 目录。
|
||||
// 查找顺序：环境变量 COOP_DIR > 相对于可执行文件 > 相对于工作目录。
|
||||
// 返回的路径均为绝对路径，供本地子进程模式（runCoopLocalAsync）使用。
|
||||
func findCoopRuntime(engine string) (supervisor, prompts, extensions string, err error) {
|
||||
eng, ok := coopEngines[engine]
|
||||
if !ok {
|
||||
return "", "", "", fmt.Errorf("不支持的 engine: %s", engine)
|
||||
}
|
||||
dirName := eng.runtimeDir
|
||||
|
||||
// 候选基目录列表：COOP_DIR 环境变量 > 可执行文件同级/上级 > 当前工作目录
|
||||
var candidates []string
|
||||
if envDir := os.Getenv("COOP_DIR"); envDir != "" {
|
||||
candidates = append(candidates, filepath.Join(envDir, dirName))
|
||||
}
|
||||
if exe, exeErr := os.Executable(); exeErr == nil {
|
||||
exeDir := filepath.Dir(exe)
|
||||
candidates = append(candidates, filepath.Join(exeDir, dirName))
|
||||
candidates = append(candidates, filepath.Join(exeDir, "..", dirName))
|
||||
}
|
||||
if wd, wdErr := os.Getwd(); wdErr == nil {
|
||||
candidates = append(candidates, filepath.Join(wd, dirName))
|
||||
}
|
||||
|
||||
for _, base := range candidates {
|
||||
sp := filepath.Join(base, "supervisor.sh")
|
||||
if st, statErr := os.Stat(sp); statErr == nil && !st.IsDir() {
|
||||
pp := filepath.Join(base, "prompts")
|
||||
ep := filepath.Join(base, "extensions")
|
||||
// prompts / extensions 可选：缺失时传空串，supervisor 用内置默认
|
||||
return sp, pp, ep, nil
|
||||
}
|
||||
}
|
||||
return "", "", "", fmt.Errorf(
|
||||
"本地模式未找到 %s 的 supervisor.sh，已查找目录: %v\n"+
|
||||
"请确保 %s 目录存在且包含 supervisor.sh，或设置 COOP_DIR 环境变量指向包含该目录的父目录",
|
||||
engine, candidates, dirName)
|
||||
}
|
||||
|
||||
// runCoopLocalAsync 在后台以本地子进程方式运行 supervisor.sh 完成 worker 任务。
|
||||
// 与 runCoopAsync（Docker 模式）对应：无需 Docker daemon，直接 exec supervisor.sh，
|
||||
// 通过 BLACKBOARD/PROMPTS/EXTENSIONS 环境变量指向本地路径。
|
||||
// 结束（成功 / 失败 / 超时）后通过 coop.Complete 通知 Agent。
|
||||
func (t *Toolset) runCoopLocalAsync(record *CoopTask, env []string, blackboardDir, supervisor, promptsDir, extensionsDir string, roundMax, timeoutSec int, label string) {
|
||||
// 使用独立后台上下文，避免随 SSE 请求断开而中断协作任务
|
||||
ctx := context.Background()
|
||||
overall := time.Duration(timeoutSec*roundMax+300) * time.Second
|
||||
runCtx, runCancel := context.WithTimeout(ctx, overall)
|
||||
defer runCancel()
|
||||
|
||||
// Windows 上 bash 通常是 WSL bash，不认反斜杠路径（E:\foo → E:foo 被吞）。
|
||||
// 需把传给 bash 的路径转为 /mnt/<盘符>/... 格式；Go 侧文件操作仍用原始路径。
|
||||
toBashPath := func(p string) string {
|
||||
if len(p) >= 2 && p[1] == ':' {
|
||||
return wslBindPath(p)
|
||||
}
|
||||
return p
|
||||
}
|
||||
bashSupervisor := toBashPath(supervisor)
|
||||
bashBlackboard := toBashPath(blackboardDir)
|
||||
bashPrompts := toBashPath(promptsDir)
|
||||
bashExtensions := toBashPath(extensionsDir)
|
||||
|
||||
// 构建子进程环境：继承父进程环境（PATH 等）+ 注入协作环境变量
|
||||
procEnv := os.Environ()
|
||||
procEnv = append(procEnv, env...)
|
||||
procEnv = append(procEnv, "BLACKBOARD="+bashBlackboard)
|
||||
if bashPrompts != "" {
|
||||
procEnv = append(procEnv, "PROMPTS="+bashPrompts)
|
||||
}
|
||||
if bashExtensions != "" {
|
||||
procEnv = append(procEnv, "EXTENSIONS="+bashExtensions)
|
||||
}
|
||||
|
||||
cmd := exec.CommandContext(runCtx, "bash", bashSupervisor)
|
||||
cmd.Env = procEnv
|
||||
// stdout+stderr 合并捕获（supervisor 的日志输出）
|
||||
var buf bytes.Buffer
|
||||
cmd.Stdout = &buf
|
||||
cmd.Stderr = &buf
|
||||
|
||||
log.Printf("[coop] 本地协作进程启动 task=%s pid=pending blackboard=%s supervisor=%s",
|
||||
record.ID, bashBlackboard, bashSupervisor)
|
||||
|
||||
if err := cmd.Start(); err != nil {
|
||||
t.coop.Complete(t.failedTask(record, fmt.Errorf("启动 supervisor 失败: %w", err)))
|
||||
return
|
||||
}
|
||||
record.ProcessID = cmd.Process.Pid
|
||||
log.Printf("[coop] 本地协作进程已启动 task=%s pid=%d", record.ID, record.ProcessID)
|
||||
|
||||
// 等待进程退出（exec.CommandContext 在 runCtx 超时时自动发送 SIGKILL）
|
||||
waitErr := cmd.Wait()
|
||||
exitCode := 0
|
||||
if waitErr != nil {
|
||||
if exitErr, ok := waitErr.(*exec.ExitError); ok {
|
||||
exitCode = exitErr.ExitCode()
|
||||
} else {
|
||||
exitCode = -1
|
||||
}
|
||||
}
|
||||
|
||||
// 判断是否因超时被杀
|
||||
timedOut := runCtx.Err() == context.DeadlineExceeded
|
||||
|
||||
result := strings.TrimSpace(buf.String())
|
||||
switch {
|
||||
case timedOut:
|
||||
result += fmt.Sprintf("\n[协作运行超时（整体上限 %d 秒，约 %.0f 分钟），已强制终止进程]",
|
||||
int(overall.Seconds()), overall.Minutes())
|
||||
exitCode = 124
|
||||
case exitCode == 0:
|
||||
result += "\n[协作完成，退出码 0]"
|
||||
case exitCode == 1:
|
||||
result += "\n[达到最大轮次仍未完成，退出码 1；可增大 round_max 后重试]"
|
||||
default:
|
||||
result += fmt.Sprintf("\n[进程退出码 %d，请检查上方日志]", exitCode)
|
||||
}
|
||||
result += "\n黑板产物已保留在: " + blackboardDir
|
||||
|
||||
record.ExitCode = exitCode
|
||||
record.Output = t.truncate(result)
|
||||
|
||||
// 解析 supervisor 生成的结构化结果 result.json（与 docker 模式一致）
|
||||
status := "unknown"
|
||||
if data, rerr := os.ReadFile(filepath.Join(blackboardDir, "result.json")); rerr == nil {
|
||||
var res CoopResult
|
||||
if json.Unmarshal(data, &res) == nil {
|
||||
record.Result = &res
|
||||
status = res.Status
|
||||
}
|
||||
}
|
||||
log.Printf("[coop] 本地协作任务结束 task=%s pid=%d exit_code=%d status=%s",
|
||||
record.ID, record.ProcessID, exitCode, status)
|
||||
t.coop.Complete(record)
|
||||
}
|
||||
@@ -0,0 +1,320 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/binary"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/docker/docker/api/types/container"
|
||||
"github.com/docker/docker/api/types/image"
|
||||
"github.com/docker/docker/client"
|
||||
)
|
||||
|
||||
// Docker 工具：通过 Docker API 操作本地或远程 Docker。
|
||||
// socket 地址通过 DockerConfigStore 动态读取，用户在 Web 页面配置后即时生效。
|
||||
|
||||
func (t *Toolset) dockerClient() (*client.Client, error) {
|
||||
socket := t.dockerSocket()
|
||||
cli, err := client.NewClientWithOpts(
|
||||
client.WithHost(socket),
|
||||
client.WithAPIVersionNegotiation(),
|
||||
)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("创建 Docker 客户端失败: %w", err)
|
||||
}
|
||||
return cli, nil
|
||||
}
|
||||
|
||||
func (t *Toolset) dockerPS(ctx context.Context, params map[string]any) ToolResult {
|
||||
all, _ := params["all"].(bool)
|
||||
cli, err := t.dockerClient()
|
||||
if err != nil {
|
||||
return ToolResult{Success: false, Output: err.Error()}
|
||||
}
|
||||
defer cli.Close()
|
||||
|
||||
ctx, cancel := context.WithTimeout(ctx, t.Timeout)
|
||||
defer cancel()
|
||||
list, err := cli.ContainerList(ctx, container.ListOptions{All: all})
|
||||
if err != nil {
|
||||
return ToolResult{Success: false, Output: "Docker 操作失败: " + err.Error()}
|
||||
}
|
||||
if len(list) == 0 {
|
||||
return ToolResult{Success: true, Output: "没有找到容器"}
|
||||
}
|
||||
|
||||
sort.Slice(list, func(i, j int) bool {
|
||||
return list[i].Created > list[j].Created
|
||||
})
|
||||
var out strings.Builder
|
||||
fmt.Fprintf(&out, "%-14s %-26s %-20s %-12s %s\n", "CONTAINER ID", "NAME", "IMAGE", "STATUS", "PORTS")
|
||||
for _, c := range list {
|
||||
name := strings.TrimPrefix(strings.Join(c.Names, ","), "/")
|
||||
fmt.Fprintf(&out, "%-14s %-26s %-20s %-12s %s\n",
|
||||
c.ID[:min(12, len(c.ID))], name, c.Image, c.Status, dockerPorts(c.Ports))
|
||||
}
|
||||
return ToolResult{Success: true, Output: t.truncate(out.String())}
|
||||
}
|
||||
|
||||
func (t *Toolset) dockerImages(ctx context.Context, params map[string]any) ToolResult {
|
||||
cli, err := t.dockerClient()
|
||||
if err != nil {
|
||||
return ToolResult{Success: false, Output: err.Error()}
|
||||
}
|
||||
defer cli.Close()
|
||||
|
||||
ctx, cancel := context.WithTimeout(ctx, t.Timeout)
|
||||
defer cancel()
|
||||
list, err := cli.ImageList(ctx, image.ListOptions{All: true})
|
||||
if err != nil {
|
||||
return ToolResult{Success: false, Output: "Docker 操作失败: " + err.Error()}
|
||||
}
|
||||
if len(list) == 0 {
|
||||
return ToolResult{Success: true, Output: "没有找到镜像"}
|
||||
}
|
||||
|
||||
sort.Slice(list, func(i, j int) bool {
|
||||
return list[i].Created > list[j].Created
|
||||
})
|
||||
var out strings.Builder
|
||||
fmt.Fprintf(&out, "%-22s %-28s %-14s %s\n", "REPOSITORY", "TAG", "IMAGE ID", "SIZE")
|
||||
for _, img := range list {
|
||||
if len(img.RepoTags) == 0 {
|
||||
continue
|
||||
}
|
||||
for _, tag := range img.RepoTags {
|
||||
repo, imageTag, _ := strings.Cut(tag, ":")
|
||||
id := img.ID
|
||||
if len(id) > 12 {
|
||||
id = id[7:19]
|
||||
}
|
||||
fmt.Fprintf(&out, "%-22s %-28s %-14s %s\n", repo, imageTag, id, formatBytes(img.Size))
|
||||
}
|
||||
}
|
||||
return ToolResult{Success: true, Output: t.truncate(out.String())}
|
||||
}
|
||||
|
||||
func (t *Toolset) dockerLogs(ctx context.Context, params map[string]any) ToolResult {
|
||||
containerID, _ := params["container"].(string)
|
||||
containerID = strings.TrimSpace(containerID)
|
||||
if containerID == "" {
|
||||
return ToolResult{Success: false, Output: "container 不能为空"}
|
||||
}
|
||||
tail := intParam(params, "tail")
|
||||
if tail <= 0 {
|
||||
tail = 100
|
||||
}
|
||||
|
||||
cli, err := t.dockerClient()
|
||||
if err != nil {
|
||||
return ToolResult{Success: false, Output: err.Error()}
|
||||
}
|
||||
defer cli.Close()
|
||||
|
||||
ctx, cancel := context.WithTimeout(ctx, t.Timeout)
|
||||
defer cancel()
|
||||
resp, err := cli.ContainerLogs(ctx, containerID, container.LogsOptions{
|
||||
ShowStdout: true,
|
||||
ShowStderr: true,
|
||||
Tail: strconv.Itoa(tail),
|
||||
})
|
||||
if err != nil {
|
||||
return ToolResult{Success: false, Output: "Docker 操作失败: " + err.Error()}
|
||||
}
|
||||
defer resp.Close()
|
||||
|
||||
raw, err := io.ReadAll(resp)
|
||||
if err != nil {
|
||||
return ToolResult{Success: false, Output: "读取 Docker 日志失败: " + err.Error()}
|
||||
}
|
||||
return ToolResult{Success: true, Output: t.truncate(demuxDockerLogs(raw))}
|
||||
}
|
||||
|
||||
// demuxDockerLogs 解析 Docker 日志流：多路复用流（stdout/stderr 分帧）按帧解出，
|
||||
// 普通文本流（TTY 容器）直接返回。
|
||||
func demuxDockerLogs(data []byte) string {
|
||||
if len(data) >= 8 && isMultiplexedFrame(data[0:8]) {
|
||||
var out strings.Builder
|
||||
for i := 0; i+8 <= len(data); {
|
||||
frameSize := int(binary.BigEndian.Uint32(data[i+4 : i+8]))
|
||||
i += 8
|
||||
if frameSize < 0 || i+frameSize > len(data) {
|
||||
break
|
||||
}
|
||||
out.Write(data[i : i+frameSize])
|
||||
i += frameSize
|
||||
}
|
||||
return out.String()
|
||||
}
|
||||
return string(data)
|
||||
}
|
||||
|
||||
func isMultiplexedFrame(header []byte) bool {
|
||||
streamType := header[0]
|
||||
if streamType != 0 && streamType != 1 && streamType != 2 {
|
||||
return false
|
||||
}
|
||||
return header[1] == 0 && header[2] == 0 && header[3] == 0
|
||||
}
|
||||
|
||||
func (t *Toolset) dockerInspect(ctx context.Context, params map[string]any) ToolResult {
|
||||
containerID, _ := params["container"].(string)
|
||||
containerID = strings.TrimSpace(containerID)
|
||||
if containerID == "" {
|
||||
return ToolResult{Success: false, Output: "container 不能为空"}
|
||||
}
|
||||
|
||||
cli, err := t.dockerClient()
|
||||
if err != nil {
|
||||
return ToolResult{Success: false, Output: err.Error()}
|
||||
}
|
||||
defer cli.Close()
|
||||
|
||||
ctx, cancel := context.WithTimeout(ctx, t.Timeout)
|
||||
defer cancel()
|
||||
inspect, err := cli.ContainerInspect(ctx, containerID)
|
||||
if err != nil {
|
||||
return ToolResult{Success: false, Output: "Docker 操作失败: " + err.Error()}
|
||||
}
|
||||
|
||||
state := inspect.State
|
||||
if state == nil {
|
||||
state = &container.State{}
|
||||
}
|
||||
summary := map[string]any{
|
||||
"id": inspect.ID,
|
||||
"name": strings.TrimPrefix(inspect.Name, "/"),
|
||||
"image": inspect.Config.Image,
|
||||
"status": state.Status,
|
||||
"running": state.Running,
|
||||
"exit_code": state.ExitCode,
|
||||
"restart_count": inspect.RestartCount,
|
||||
"created": inspect.Created,
|
||||
"command": inspect.Config.Cmd,
|
||||
"env_count": len(inspect.Config.Env),
|
||||
"ports": inspect.NetworkSettings.Ports,
|
||||
"network_mode": string(inspect.HostConfig.NetworkMode),
|
||||
"ip_address": inspect.NetworkSettings.IPAddress,
|
||||
}
|
||||
data, _ := json.MarshalIndent(summary, "", " ")
|
||||
return ToolResult{Success: true, Output: t.truncate(string(data))}
|
||||
}
|
||||
|
||||
func (t *Toolset) dockerExec(ctx context.Context, params map[string]any) ToolResult {
|
||||
containerID, _ := params["container"].(string)
|
||||
containerID = strings.TrimSpace(containerID)
|
||||
command, _ := params["command"].(string)
|
||||
command = strings.TrimSpace(command)
|
||||
if containerID == "" || command == "" {
|
||||
return ToolResult{Success: false, Output: "container 和 command 不能为空"}
|
||||
}
|
||||
|
||||
cli, err := t.dockerClient()
|
||||
if err != nil {
|
||||
return ToolResult{Success: false, Output: err.Error()}
|
||||
}
|
||||
defer cli.Close()
|
||||
|
||||
ctx, cancel := context.WithTimeout(ctx, t.Timeout)
|
||||
defer cancel()
|
||||
execCfg := container.ExecOptions{
|
||||
AttachStdout: true,
|
||||
AttachStderr: true,
|
||||
Cmd: []string{"sh", "-c", command},
|
||||
}
|
||||
created, err := cli.ContainerExecCreate(ctx, containerID, execCfg)
|
||||
if err != nil {
|
||||
return ToolResult{Success: false, Output: "Docker 操作失败: " + err.Error()}
|
||||
}
|
||||
attach, err := cli.ContainerExecAttach(ctx, created.ID, container.ExecAttachOptions{})
|
||||
if err != nil {
|
||||
return ToolResult{Success: false, Output: "Docker 操作失败: " + err.Error()}
|
||||
}
|
||||
defer attach.Close()
|
||||
|
||||
raw, err := io.ReadAll(attach.Reader)
|
||||
if err != nil {
|
||||
return ToolResult{Success: false, Output: "读取命令输出失败: " + err.Error()}
|
||||
}
|
||||
return ToolResult{Success: true, Output: t.truncate(demuxDockerLogs(raw))}
|
||||
}
|
||||
|
||||
func (t *Toolset) dockerStart(ctx context.Context, params map[string]any) ToolResult {
|
||||
containerID, _ := params["container"].(string)
|
||||
containerID = strings.TrimSpace(containerID)
|
||||
if containerID == "" {
|
||||
return ToolResult{Success: false, Output: "container 不能为空"}
|
||||
}
|
||||
|
||||
cli, err := t.dockerClient()
|
||||
if err != nil {
|
||||
return ToolResult{Success: false, Output: err.Error()}
|
||||
}
|
||||
defer cli.Close()
|
||||
|
||||
ctx, cancel := context.WithTimeout(ctx, t.Timeout)
|
||||
defer cancel()
|
||||
if err := cli.ContainerStart(ctx, containerID, container.StartOptions{}); err != nil {
|
||||
return ToolResult{Success: false, Output: "Docker 操作失败: " + err.Error()}
|
||||
}
|
||||
return ToolResult{Success: true, Output: "容器 " + containerID + " 已启动"}
|
||||
}
|
||||
|
||||
func (t *Toolset) dockerStop(ctx context.Context, params map[string]any) ToolResult {
|
||||
containerID, _ := params["container"].(string)
|
||||
containerID = strings.TrimSpace(containerID)
|
||||
if containerID == "" {
|
||||
return ToolResult{Success: false, Output: "container 不能为空"}
|
||||
}
|
||||
timeoutSeconds := intParam(params, "timeout")
|
||||
|
||||
cli, err := t.dockerClient()
|
||||
if err != nil {
|
||||
return ToolResult{Success: false, Output: err.Error()}
|
||||
}
|
||||
defer cli.Close()
|
||||
|
||||
ctx, cancel := context.WithTimeout(ctx, t.Timeout)
|
||||
defer cancel()
|
||||
var timeout *int
|
||||
if timeoutSeconds > 0 {
|
||||
timeout = &timeoutSeconds
|
||||
}
|
||||
if err := cli.ContainerStop(ctx, containerID, container.StopOptions{Timeout: timeout}); err != nil {
|
||||
return ToolResult{Success: false, Output: "Docker 操作失败: " + err.Error()}
|
||||
}
|
||||
return ToolResult{Success: true, Output: "容器 " + containerID + " 已停止"}
|
||||
}
|
||||
|
||||
func dockerPorts(ports []container.Port) string {
|
||||
if len(ports) == 0 {
|
||||
return ""
|
||||
}
|
||||
var parts []string
|
||||
for _, p := range ports {
|
||||
host := ""
|
||||
if p.PublicPort > 0 {
|
||||
host = fmt.Sprintf("%d->", p.PublicPort)
|
||||
}
|
||||
parts = append(parts, fmt.Sprintf("%s%d/%s", host, p.PrivatePort, p.Type))
|
||||
}
|
||||
return strings.Join(parts, ", ")
|
||||
}
|
||||
|
||||
func formatBytes(size int64) string {
|
||||
const unit = 1024
|
||||
if size < unit {
|
||||
return fmt.Sprintf("%dB", size)
|
||||
}
|
||||
div, exp := int64(unit), 0
|
||||
for n := size / unit; n >= unit; n /= unit {
|
||||
div *= unit
|
||||
exp++
|
||||
}
|
||||
return fmt.Sprintf("%.1f%cB", float64(size)/float64(div), "KMGTPE"[exp])
|
||||
}
|
||||
@@ -0,0 +1,76 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// TestPathWithin 验证路径越界校验：read_file / write_file / run_coop 的 blackboard
|
||||
// 都依赖该函数拦截指向项目根目录之外的路径，是核心安全边界。
|
||||
func TestPathWithin(t *testing.T) {
|
||||
var parent string
|
||||
if runtime.GOOS == "windows" {
|
||||
parent = `e:\proj`
|
||||
} else {
|
||||
parent = "/proj"
|
||||
}
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
child string
|
||||
expect bool
|
||||
}{
|
||||
{"自身", parent, true},
|
||||
{"直接子文件", filepath.Join(parent, "a.txt"), true},
|
||||
{"嵌套子目录", filepath.Join(parent, "sub", "deep", "f.txt"), true},
|
||||
{"上级目录", filepath.Join(parent, "..", "secret"), false},
|
||||
{"同级兄弟目录", filepath.Join(filepath.Dir(parent), "other"), false},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
if got := pathWithin(tc.child, parent); got != tc.expect {
|
||||
t.Fatalf("pathWithin(%q, %q) = %v, want %v", tc.child, parent, got, tc.expect)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestMergeToolCall 验证 SSE 流式 tool_call 增量合并：
|
||||
// 跨多个 chunk 按 index 聚合 id / name / arguments，是 LLM 工具调用协议正确性的关键。
|
||||
func TestMergeToolCall(t *testing.T) {
|
||||
var calls []ToolCall
|
||||
// 第一个 chunk：声明 index 0 的调用（id + name）
|
||||
calls = mergeToolCall(calls, ToolCallDelta{Index: 0, ID: "call_1", Name: "run_bash"})
|
||||
// 同一调用的参数分片到达
|
||||
calls = mergeToolCall(calls, ToolCallDelta{Index: 0, ArgumentsDelta: `{"comm`})
|
||||
calls = mergeToolCall(calls, ToolCallDelta{Index: 0, ArgumentsDelta: `and":"ls"}`})
|
||||
|
||||
if len(calls) != 1 {
|
||||
t.Fatalf("应合并为 1 个调用，got %d", len(calls))
|
||||
}
|
||||
c := calls[0]
|
||||
if c.ID != "call_1" || c.Function.Name != "run_bash" {
|
||||
t.Fatalf("id/name 不匹配: %+v", c)
|
||||
}
|
||||
if c.Function.Arguments != `{"command":"ls"}` {
|
||||
t.Fatalf("arguments 拼接错误: %q", c.Function.Arguments)
|
||||
}
|
||||
|
||||
// 第二个调用在 index 1
|
||||
calls = mergeToolCall(calls, ToolCallDelta{Index: 1, ID: "call_2", Name: "read_file"})
|
||||
calls = mergeToolCall(calls, ToolCallDelta{Index: 1, ArgumentsDelta: `{"path":"a"}`})
|
||||
if len(calls) != 2 {
|
||||
t.Fatalf("应有 2 个调用，got %d", len(calls))
|
||||
}
|
||||
if calls[1].Function.Name != "read_file" {
|
||||
t.Fatalf("第二个调用名错误: %+v", calls[1])
|
||||
}
|
||||
|
||||
// index 为负时回退为追加新调用
|
||||
prev := len(calls)
|
||||
calls = mergeToolCall(calls, ToolCallDelta{Index: -1, ID: "call_3", Name: "list_directory"})
|
||||
if len(calls) != prev+1 {
|
||||
t.Fatalf("负 index 应追加新调用，got %d", len(calls))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,60 @@
|
||||
package agent
|
||||
|
||||
import "time"
|
||||
|
||||
type Message struct {
|
||||
Role string `json:"role"`
|
||||
Content *string `json:"content"`
|
||||
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
|
||||
ToolCallID string `json:"tool_call_id,omitempty"`
|
||||
}
|
||||
|
||||
type ToolCall struct {
|
||||
ID string `json:"id"`
|
||||
Type string `json:"type"`
|
||||
Function ToolFunctionCall `json:"function"`
|
||||
}
|
||||
|
||||
type ToolFunctionCall struct {
|
||||
Name string `json:"name"`
|
||||
Arguments string `json:"arguments"`
|
||||
}
|
||||
|
||||
type Session struct {
|
||||
ID string `json:"id"`
|
||||
Title string `json:"title"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
Messages []Message `json:"messages"`
|
||||
}
|
||||
|
||||
type SessionSummary struct {
|
||||
ID string `json:"id"`
|
||||
Title string `json:"title"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
MessageCount int `json:"message_count"`
|
||||
Preview string `json:"preview"`
|
||||
}
|
||||
|
||||
type ToolDefinition struct {
|
||||
Type string `json:"type"`
|
||||
Function ToolDefinitionFunction `json:"function"`
|
||||
}
|
||||
|
||||
type ToolDefinitionFunction struct {
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
Parameters map[string]any `json:"parameters"`
|
||||
}
|
||||
|
||||
type ToolResult struct {
|
||||
Success bool `json:"success"`
|
||||
Output string `json:"output"`
|
||||
}
|
||||
|
||||
type Usage struct {
|
||||
PromptTokens int `json:"prompt_tokens,omitempty"`
|
||||
CompletionTokens int `json:"completion_tokens,omitempty"`
|
||||
TotalTokens int `json:"total_tokens,omitempty"`
|
||||
}
|
||||
@@ -0,0 +1,80 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"crypto/sha256"
|
||||
"crypto/subtle"
|
||||
"encoding/hex"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// authCookieName 保存已通过授权校验的凭证。值为授权码的 SHA-256 十六进制摘要，
|
||||
// 不落明文；HttpOnly 使前端 JS 无法读取，降低泄露面。
|
||||
const authCookieName = "blackbean_auth"
|
||||
|
||||
// authCookieValue 计算授权码的稳定凭证值（SHA-256 摘要）。
|
||||
func authCookieValue(code string) string {
|
||||
sum := sha256.Sum256([]byte(code))
|
||||
return hex.EncodeToString(sum[:])
|
||||
}
|
||||
|
||||
// authRequired 是访问授权中间件：cfg.AuthCode 为空（未启用授权）时直接放行；
|
||||
// 否则要求请求携带与授权码匹配的 Cookie，不匹配则 401 并中止后续处理。
|
||||
func (s *Server) authRequired(c *gin.Context) {
|
||||
if s.cfg.AuthCode == "" {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
if cookie, err := c.Cookie(authCookieName); err == nil && cookie == authCookieValue(s.cfg.AuthCode) {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
fail(c, http.StatusUnauthorized, "需要授权码才能访问，请先在首页输入授权码")
|
||||
c.Abort()
|
||||
}
|
||||
|
||||
// authStatus 返回授权状态：required=是否启用了授权，authorized=当前请求是否已通过授权。
|
||||
// 该接口始终放行（不挂授权中间件），前端据此决定是否展示授权界面。
|
||||
func (s *Server) authStatus(c *gin.Context) {
|
||||
required := s.cfg.AuthCode != ""
|
||||
authorized := false
|
||||
if required {
|
||||
if cookie, err := c.Cookie(authCookieName); err == nil {
|
||||
authorized = cookie == authCookieValue(s.cfg.AuthCode)
|
||||
}
|
||||
}
|
||||
ok(c, gin.H{"required": required, "authorized": authorized})
|
||||
}
|
||||
|
||||
// authLogin 校验授权码并签发 Cookie。
|
||||
// 该接口始终放行（不挂授权中间件）；未启用授权时返回 400。
|
||||
func (s *Server) authLogin(c *gin.Context) {
|
||||
if s.cfg.AuthCode == "" {
|
||||
fail(c, http.StatusBadRequest, "本实例未启用访问授权")
|
||||
return
|
||||
}
|
||||
var req struct {
|
||||
Code string `json:"code"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
fail(c, http.StatusBadRequest, "请求体格式错误")
|
||||
return
|
||||
}
|
||||
// 常量时间比较，避免通过响应时差探测授权码
|
||||
want := []byte(s.cfg.AuthCode)
|
||||
got := []byte(strings.TrimSpace(req.Code))
|
||||
if len(want) != len(got) || subtle.ConstantTimeCompare(want, got) != 1 {
|
||||
fail(c, http.StatusUnauthorized, "授权码错误")
|
||||
return
|
||||
}
|
||||
http.SetCookie(c.Writer, &http.Cookie{
|
||||
Name: authCookieName,
|
||||
Value: authCookieValue(s.cfg.AuthCode),
|
||||
Path: "/",
|
||||
HttpOnly: true,
|
||||
SameSite: http.SameSiteLaxMode,
|
||||
})
|
||||
ok(c, gin.H{"authorized": true})
|
||||
}
|
||||
@@ -0,0 +1,45 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log"
|
||||
|
||||
"agent/internal/agent"
|
||||
)
|
||||
|
||||
// benchmarkInitialTask 是托管模式启动时发给 Agent 的初始任务消息。
|
||||
// 系统提示中已包含完整的 Tsecbench 调度规则（并发管理、flag 提交、靶机开关），
|
||||
// 此处只需一条简短触发指令让 Agent 开始拉取题目并解题。
|
||||
const benchmarkInitialTask = `Tsecbench 托管评测已启动。平台已通过环境变量注入 BENCHMARK_TOKEN 和 BENCHMARK_BASE_URL。
|
||||
|
||||
请立即开始评测：
|
||||
1. 用 shell 工具读取环境变量：echo $BENCHMARK_BASE_URL && echo $BENCHMARK_TOKEN
|
||||
2. 用 curl GET $BENCHMARK_BASE_URL/openapi/v1/challenges 获取题目列表
|
||||
3. 选定多道未完成题目，逐一 POST $BENCHMARK_BASE_URL/openapi/v1/challenges/start?unique_code=<题号> 启动靶机，把并发拉满（默认上限 3）
|
||||
4. 对每道题调用 run_coop 下发解题任务（task 中写明目标地址、unique_code、BENCHMARK_TOKEN、提交规则）
|
||||
5. worker 返回 flag 后用 curl 向平台 submit API 提交，关闭已通关靶机，继续补满并发
|
||||
6. 直到所有题目完成或时限到达
|
||||
|
||||
目标：在时限内尽可能多地完成题目、拿到更高总分。`
|
||||
|
||||
// StartBenchmark 在托管模式下自动创建会话并触发 Agent 开始解题。
|
||||
// 由 main.go 在检测到 BENCHMARK_TOKEN 环境变量时以 goroutine 调用。
|
||||
// 不依赖 HTTP 服务就绪 —— agent.Run 直接操作 Agent 实例，不经 HTTP 路由。
|
||||
//
|
||||
// 关键：必须注册持久化的 live emitter。coop 完成后 notifyCoopDone 通过
|
||||
// liveEmitter 获取 emit 回调来触发汇报轮，若未注册则汇报轮无日志输出、
|
||||
// 事件丢失，表现为"评测流程结束"后无任何后续日志（实际仍在运行）。
|
||||
// 注册后不注销：整个评测期间（含所有后续汇报轮）保持日志通道畅通。
|
||||
func (s *Server) StartBenchmark(ctx context.Context) error {
|
||||
session := s.store.Create()
|
||||
sessionID := session.ID
|
||||
|
||||
emit := func(event agent.Event) {
|
||||
log.Printf("[benchmark] [%s] %v", event.Type, event.Data)
|
||||
}
|
||||
// 注册 live emitter，让 notifyCoopDone 触发的汇报轮也能输出日志
|
||||
s.agent.RegisterLive(sessionID, emit)
|
||||
|
||||
log.Printf("[benchmark] 自动启动评测会话 %s", sessionID)
|
||||
return s.agent.Run(ctx, sessionID, benchmarkInitialTask, emit)
|
||||
}
|
||||
@@ -0,0 +1,111 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func (s *Server) health(c *gin.Context) {
|
||||
ok(c, gin.H{"status": "ok"})
|
||||
}
|
||||
|
||||
func (s *Server) config(c *gin.Context) {
|
||||
ok(c, gin.H{
|
||||
"model": s.cfg.Model,
|
||||
"base_url": s.cfg.BaseURL,
|
||||
"workspace": s.cfg.Workspace,
|
||||
"max_context_tokens": s.cfg.MaxContextTokens,
|
||||
"max_tool_result_chars": s.cfg.MaxToolResultChars,
|
||||
"max_iterations": s.cfg.MaxIterations,
|
||||
"keep_recent_messages": s.cfg.KeepRecentMessages,
|
||||
"api_key_configured": s.cfg.APIKey != "",
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) dockerConfig(c *gin.Context) {
|
||||
ok(c, gin.H{
|
||||
"docker_socket": s.dockerCfg.Socket(),
|
||||
"default_socket": s.dockerCfg.DefaultSocket(),
|
||||
"configured": s.dockerCfg.IsConfigured(),
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) setDockerConfig(c *gin.Context) {
|
||||
var req struct {
|
||||
Socket string `json:"socket"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
fail(c, http.StatusBadRequest, "请求体格式错误")
|
||||
return
|
||||
}
|
||||
req.Socket = strings.TrimSpace(req.Socket)
|
||||
// 空字符串表示清除配置、恢复默认本地 Docker
|
||||
if err := s.dockerCfg.SetSocket(req.Socket); err != nil {
|
||||
fail(c, http.StatusInternalServerError, "保存 Docker 配置失败: "+err.Error())
|
||||
return
|
||||
}
|
||||
ok(c, gin.H{
|
||||
"docker_socket": s.dockerCfg.Socket(),
|
||||
"default_socket": s.dockerCfg.DefaultSocket(),
|
||||
"configured": s.dockerCfg.IsConfigured(),
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) llmConfig(c *gin.Context) {
|
||||
apiKey := s.apiCfg.APIKey()
|
||||
ok(c, gin.H{
|
||||
"provider": s.apiCfg.Provider(),
|
||||
"api_key_configured": apiKey != "",
|
||||
"api_key_masked": maskSecret(apiKey),
|
||||
"base_url": s.apiCfg.BaseURL(),
|
||||
"model": s.apiCfg.Model(),
|
||||
"engine": s.apiCfg.Engine(),
|
||||
"coop_mode": s.apiCfg.CoopMode(),
|
||||
"default_base_url": s.cfg.BaseURL,
|
||||
"default_model": s.cfg.Model,
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) setLLMConfig(c *gin.Context) {
|
||||
var req struct {
|
||||
APIKey *string `json:"api_key"`
|
||||
BaseURL *string `json:"base_url"`
|
||||
Model *string `json:"model"`
|
||||
Provider *string `json:"provider"`
|
||||
Engine *string `json:"engine"`
|
||||
CoopMode *string `json:"coop_mode"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
fail(c, http.StatusBadRequest, "请求体格式错误")
|
||||
return
|
||||
}
|
||||
// 字段为 nil 表示不修改；空字符串表示清除该字段、回退默认
|
||||
if err := s.apiCfg.Update(req.APIKey, req.BaseURL, req.Model, req.Provider, req.Engine, req.CoopMode); err != nil {
|
||||
fail(c, http.StatusInternalServerError, "保存 LLM 配置失败: "+err.Error())
|
||||
return
|
||||
}
|
||||
apiKey := s.apiCfg.APIKey()
|
||||
ok(c, gin.H{
|
||||
"provider": s.apiCfg.Provider(),
|
||||
"api_key_configured": apiKey != "",
|
||||
"api_key_masked": maskSecret(apiKey),
|
||||
"base_url": s.apiCfg.BaseURL(),
|
||||
"model": s.apiCfg.Model(),
|
||||
"engine": s.apiCfg.Engine(),
|
||||
"coop_mode": s.apiCfg.CoopMode(),
|
||||
})
|
||||
}
|
||||
|
||||
// maskSecret 对密钥做掩码展示，避免在 Web 页面泄露完整值。
|
||||
func maskSecret(value string) string {
|
||||
runes := []rune(value)
|
||||
if len(runes) <= 6 {
|
||||
if len(runes) == 0 {
|
||||
return ""
|
||||
}
|
||||
return "****"
|
||||
}
|
||||
return string(runes[:3]) + "****" + string(runes[len(runes)-3:])
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"agent/internal/agent"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// coopTaskView 是返回给前端的协作任务视图，附带容器的实时状态与协作轮次进度。
|
||||
type coopTaskView struct {
|
||||
ID string `json:"id"`
|
||||
SessionID string `json:"session_id"`
|
||||
Status string `json:"status"` // running | done
|
||||
Error string `json:"error,omitempty"`
|
||||
ExitCode int `json:"exit_code,omitempty"`
|
||||
Blackboard string `json:"blackboard,omitempty"`
|
||||
RoundMax int `json:"round_max"`
|
||||
CurrentRound int `json:"current_round"`
|
||||
ContainerState string `json:"container_state,omitempty"`
|
||||
ContainerStatus string `json:"container_status,omitempty"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
FinishedAt time.Time `json:"finished_at,omitempty"`
|
||||
}
|
||||
|
||||
// coopTasks 返回协作任务列表，支持 ?session=<id> 过滤到某个会话。
|
||||
func (s *Server) coopTasks(c *gin.Context) {
|
||||
tasks := s.agent.CoopTasks(strings.TrimSpace(c.Query("session")))
|
||||
views := make([]coopTaskView, 0, len(tasks))
|
||||
for _, task := range tasks {
|
||||
views = append(views, s.coopTaskView(task))
|
||||
}
|
||||
ok(c, gin.H{"tasks": views})
|
||||
}
|
||||
|
||||
func (s *Server) coopTaskView(task *agent.CoopTask) coopTaskView {
|
||||
view := coopTaskView{
|
||||
ID: task.ID,
|
||||
SessionID: task.SessionID,
|
||||
Status: task.Status,
|
||||
Error: task.Error,
|
||||
ExitCode: task.ExitCode,
|
||||
Blackboard: task.Blackboard,
|
||||
RoundMax: task.RoundMax,
|
||||
CreatedAt: task.CreatedAt,
|
||||
FinishedAt: task.FinishedAt,
|
||||
}
|
||||
// 轮次进度以黑板 logs 目录实际出现的轮次为准（容器清理后日志仍在）
|
||||
view.CurrentRound = agent.CurrentCoopRound(task.Blackboard)
|
||||
if view.CurrentRound == 0 {
|
||||
view.CurrentRound = view.RoundMax
|
||||
}
|
||||
// 运行中的任务实时查询容器状态；已完成任务容器已被清理，直接读记录即可
|
||||
if task.Status == "running" && task.ContainerID != "" {
|
||||
if state, status, err := agent.ContainerInspectStatus(s.dockerCfg.Socket(), task.ContainerID); err == nil {
|
||||
view.ContainerState = state
|
||||
view.ContainerStatus = status
|
||||
}
|
||||
}
|
||||
return view
|
||||
}
|
||||
@@ -0,0 +1,89 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"agent/internal/agent"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// listSessions 会话列表，支持 ?q=<关键词> 搜索、?limit=&offset= 分页。
|
||||
func (s *Server) listSessions(c *gin.Context) {
|
||||
query := strings.TrimSpace(c.Query("q"))
|
||||
var sessions []agent.SessionSummary
|
||||
if query != "" {
|
||||
sessions = s.store.Search(query)
|
||||
} else {
|
||||
sessions = s.store.List()
|
||||
}
|
||||
limit, offset := parsePagination(c)
|
||||
if offset > 0 || limit > 0 {
|
||||
start := offset
|
||||
if start > len(sessions) {
|
||||
start = len(sessions)
|
||||
}
|
||||
end := len(sessions)
|
||||
if limit > 0 && start+limit < end {
|
||||
end = start + limit
|
||||
}
|
||||
sessions = sessions[start:end]
|
||||
}
|
||||
ok(c, gin.H{"sessions": sessions})
|
||||
}
|
||||
|
||||
func (s *Server) createSession(c *gin.Context) {
|
||||
session := s.store.Create()
|
||||
c.JSON(http.StatusCreated, session)
|
||||
}
|
||||
|
||||
func (s *Server) getSession(c *gin.Context) {
|
||||
session, found := s.store.Get(c.Param("id"))
|
||||
if !found {
|
||||
fail(c, http.StatusNotFound, "会话不存在")
|
||||
return
|
||||
}
|
||||
ok(c, gin.H{"session": session})
|
||||
}
|
||||
|
||||
// sessionMessages 分页返回会话消息（offset 从最新一条往前数），
|
||||
// 响应含 has_more 与 total 供前端"加载更早"。
|
||||
func (s *Server) sessionMessages(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
total, exists := s.store.TotalMessages(id)
|
||||
if !exists {
|
||||
fail(c, http.StatusNotFound, "会话不存在")
|
||||
return
|
||||
}
|
||||
limit, offset := parsePagination(c)
|
||||
if limit <= 0 || limit > 200 {
|
||||
limit = 50
|
||||
}
|
||||
messages, _, hasMore := s.store.Messages(id, limit, offset)
|
||||
ok(c, gin.H{
|
||||
"messages": messages,
|
||||
"has_more": hasMore,
|
||||
"total": total,
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) deleteSession(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
if !s.store.Delete(id) {
|
||||
fail(c, http.StatusNotFound, "会话不存在")
|
||||
return
|
||||
}
|
||||
// 释放会话持有的运行时资源（会话锁与 live 注册），避免内存泄漏
|
||||
s.agent.ForgetSession(id)
|
||||
c.Status(http.StatusNoContent)
|
||||
}
|
||||
|
||||
// clearSession 清空会话历史消息（保留会话）。
|
||||
func (s *Server) clearSession(c *gin.Context) {
|
||||
if !s.store.Clear(c.Param("id")) {
|
||||
fail(c, http.StatusNotFound, "会话不存在")
|
||||
return
|
||||
}
|
||||
c.Status(http.StatusNoContent)
|
||||
}
|
||||
@@ -0,0 +1,31 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strconv"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// fail 统一错误响应：{ "error": "..." }
|
||||
func fail(c *gin.Context, status int, message string) {
|
||||
c.JSON(status, gin.H{"error": message})
|
||||
}
|
||||
|
||||
// ok 统一成功响应。
|
||||
func ok(c *gin.Context, data gin.H) {
|
||||
c.JSON(http.StatusOK, data)
|
||||
}
|
||||
|
||||
// parsePagination 解析 limit / offset 查询参数（非法值归零）。
|
||||
func parsePagination(c *gin.Context) (limit, offset int) {
|
||||
limit, _ = strconv.Atoi(c.Query("limit"))
|
||||
offset, _ = strconv.Atoi(c.Query("offset"))
|
||||
if limit < 0 {
|
||||
limit = 0
|
||||
}
|
||||
if offset < 0 {
|
||||
offset = 0
|
||||
}
|
||||
return
|
||||
}
|
||||
@@ -0,0 +1,86 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"path/filepath"
|
||||
"time"
|
||||
|
||||
"agent/internal/agent"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// NewRouter 组装全部路由：静态资源、会话直达、/api/v1 REST、/ws WebSocket。
|
||||
// 返回 (gin.Engine, *Server)：*Server 供 main.go 在托管模式下调用 StartBenchmark 自动触发评测。
|
||||
func NewRouter(cfg agent.Config, store *agent.SessionStore, dockerCfg *agent.DockerConfigStore, apiCfg *agent.APIConfigStore) (*gin.Engine, *Server) {
|
||||
s := &Server{
|
||||
cfg: cfg,
|
||||
store: store,
|
||||
agent: agent.NewAgent(cfg, store, dockerCfg, apiCfg),
|
||||
dockerCfg: dockerCfg,
|
||||
apiCfg: apiCfg,
|
||||
}
|
||||
|
||||
gin.SetMode(gin.ReleaseMode)
|
||||
r := gin.New()
|
||||
r.Use(gin.Recovery(), requestLogger())
|
||||
|
||||
staticDir := filepath.Join(cfg.Workspace, "web", "static")
|
||||
indexFile := filepath.Join(staticDir, "index.html")
|
||||
r.Static("/static", staticDir)
|
||||
r.StaticFile("/", indexFile)
|
||||
|
||||
// 会话直达：/<sessionId> 返回前端页面，由前端 JS 从路径解析会话 ID
|
||||
r.GET("/:sessionId", func(c *gin.Context) {
|
||||
if !sessionIDPattern.MatchString(c.Param("sessionId")) {
|
||||
c.Status(http.StatusNotFound)
|
||||
return
|
||||
}
|
||||
c.File(indexFile)
|
||||
})
|
||||
|
||||
v1 := r.Group("/api/v1")
|
||||
v1.GET("/health", s.health)
|
||||
// 授权相关接口始终放行：前端据此判断是否需要授权 / 提交授权码。
|
||||
// 其余 API 与 WebSocket 统一挂 authRequired（未启用授权时自动放行）。
|
||||
auth := v1.Group("/auth")
|
||||
auth.GET("/status", s.authStatus)
|
||||
auth.POST("/login", s.authLogin)
|
||||
|
||||
authed := v1.Group("")
|
||||
authed.Use(s.authRequired)
|
||||
authed.GET("/config", s.config)
|
||||
authed.GET("/docker/config", s.dockerConfig)
|
||||
authed.PUT("/docker/config", s.setDockerConfig)
|
||||
authed.GET("/llm/config", s.llmConfig)
|
||||
authed.PUT("/llm/config", s.setLLMConfig)
|
||||
authed.GET("/sessions", s.listSessions)
|
||||
authed.POST("/sessions", s.createSession)
|
||||
authed.GET("/sessions/:id", s.getSession)
|
||||
authed.GET("/sessions/:id/messages", s.sessionMessages)
|
||||
authed.DELETE("/sessions/:id", s.deleteSession)
|
||||
authed.DELETE("/sessions/:id/messages", s.clearSession)
|
||||
authed.GET("/coop/tasks", s.coopTasks)
|
||||
|
||||
// WebSocket 实时通道：Agent 流式输出 / 状态 / 协作通知（同样受授权保护）
|
||||
r.GET("/ws", s.authRequired, s.ws)
|
||||
|
||||
return r, s
|
||||
}
|
||||
|
||||
// requestLogger 轻量请求日志中间件：统一 [http] 前缀，便于与 [coop] 等日志区分。
|
||||
func requestLogger() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
start := time.Now()
|
||||
c.Next()
|
||||
gin.DefaultWriter.Write([]byte(fmt.Sprintf(
|
||||
"[http] %s | %s %s | %s | %d\n",
|
||||
time.Now().Format("2006/01/02 15:04:05"),
|
||||
c.Request.Method,
|
||||
c.Request.URL.Path,
|
||||
time.Since(start).Round(time.Millisecond).String(),
|
||||
c.Writer.Status(),
|
||||
)))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,19 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"regexp"
|
||||
|
||||
"agent/internal/agent"
|
||||
)
|
||||
|
||||
// Server 持有全部依赖，各 handler 文件通过它访问共享状态。
|
||||
type Server struct {
|
||||
cfg agent.Config
|
||||
store *agent.SessionStore
|
||||
agent *agent.Agent
|
||||
dockerCfg *agent.DockerConfigStore
|
||||
apiCfg *agent.APIConfigStore
|
||||
}
|
||||
|
||||
// sessionIDPattern 匹配会话直达路径 / WebSocket 绑定的会话 ID 格式（24 位十六进制）。
|
||||
var sessionIDPattern = regexp.MustCompile(`^[0-9a-f]{24}$`)
|
||||
@@ -0,0 +1,177 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"log"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"agent/internal/agent"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
var upgrader = websocket.Upgrader{
|
||||
ReadBufferSize: 1024,
|
||||
WriteBufferSize: 4096,
|
||||
// 本地单用户 Web，允许任意 Origin
|
||||
CheckOrigin: func(r *http.Request) bool { return true },
|
||||
}
|
||||
|
||||
// wsInbound 客户端 → 服务端消息。
|
||||
// chat：在绑定会话上运行 Agent；switch：切换绑定会话；stop：停止当前生成；ping：保活。
|
||||
type wsInbound struct {
|
||||
Type string `json:"type"`
|
||||
Content string `json:"content"`
|
||||
SessionID string `json:"session_id"`
|
||||
}
|
||||
|
||||
// wsConn 封装一个 WebSocket 连接：绑定一个会话、串行处理消息。
|
||||
type wsConn struct {
|
||||
server *Server
|
||||
conn *websocket.Conn
|
||||
sendMu sync.Mutex // 串行化 WriteJSON
|
||||
|
||||
sessionID string
|
||||
unreg func() // 当前会话的 live 注销函数
|
||||
|
||||
runMu sync.Mutex
|
||||
runCancel context.CancelFunc
|
||||
}
|
||||
|
||||
func (s *Server) ws(c *gin.Context) {
|
||||
sessionID := strings.TrimSpace(c.Query("session"))
|
||||
if sessionID == "" || !sessionIDPattern.MatchString(sessionID) {
|
||||
log.Printf("ws: 连接被拒（session 参数无效）session=%q", sessionID)
|
||||
fail(c, http.StatusBadRequest, "缺少有效的 session 参数")
|
||||
return
|
||||
}
|
||||
if _, ok := s.store.Get(sessionID); !ok {
|
||||
log.Printf("ws: 连接被拒（会话不存在）session=%q", sessionID)
|
||||
fail(c, http.StatusNotFound, "会话不存在")
|
||||
return
|
||||
}
|
||||
|
||||
conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
client := &wsConn{server: s, conn: conn}
|
||||
client.bind(sessionID)
|
||||
client.readPump()
|
||||
}
|
||||
|
||||
// sendJSON 安全地推送一个事件帧 {type, data}。
|
||||
func (cl *wsConn) sendJSON(eventType string, data any) {
|
||||
cl.sendMu.Lock()
|
||||
defer cl.sendMu.Unlock()
|
||||
_ = cl.conn.WriteJSON(map[string]any{"type": eventType, "data": data})
|
||||
}
|
||||
|
||||
// bind 切换绑定会话：注销旧会话的 live 注册，注册新会话。
|
||||
// 传入空串表示仅注销。
|
||||
func (cl *wsConn) bind(sessionID string) {
|
||||
if cl.unreg != nil {
|
||||
cl.unreg()
|
||||
cl.unreg = nil
|
||||
}
|
||||
cl.sessionID = sessionID
|
||||
if sessionID == "" {
|
||||
return
|
||||
}
|
||||
cl.unreg = cl.server.agent.RegisterLive(sessionID, func(event agent.Event) {
|
||||
cl.sendJSON(event.Type, event.Data)
|
||||
})
|
||||
}
|
||||
|
||||
func (cl *wsConn) handleChat(content string) {
|
||||
content = strings.TrimSpace(content)
|
||||
if content == "" {
|
||||
cl.sendJSON("error", map[string]any{"message": "消息内容不能为空"})
|
||||
return
|
||||
}
|
||||
// 未配置 LLM API Key 时直接提示去配置，避免进入无效的模型调用。
|
||||
// 环境变量 LLM_API_KEY（托管模式）与设置页 api-config.json 都计入已配置。
|
||||
if !cl.server.apiCfg.IsAPIKeyConfigured() {
|
||||
cl.sendJSON("error", map[string]any{"message": "尚未配置 LLM API，请先点击左下角设置按钮配置 API Key 与模型"})
|
||||
return
|
||||
}
|
||||
cl.runMu.Lock()
|
||||
if cl.runCancel != nil {
|
||||
cl.runMu.Unlock()
|
||||
cl.sendJSON("error", map[string]any{"message": "Agent 正在执行，请等待完成或先停止"})
|
||||
return
|
||||
}
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cl.runCancel = cancel
|
||||
cl.runMu.Unlock()
|
||||
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
defer close(done)
|
||||
_ = cl.server.agent.Run(ctx, cl.sessionID, content, func(event agent.Event) {
|
||||
cl.sendJSON(event.Type, event.Data)
|
||||
})
|
||||
}()
|
||||
go func() {
|
||||
<-done
|
||||
cl.runMu.Lock()
|
||||
cl.runCancel = nil
|
||||
cl.runMu.Unlock()
|
||||
}()
|
||||
}
|
||||
|
||||
func (cl *wsConn) handleStop() {
|
||||
cl.runMu.Lock()
|
||||
if cl.runCancel != nil {
|
||||
cl.runCancel()
|
||||
}
|
||||
cl.runMu.Unlock()
|
||||
}
|
||||
|
||||
func (cl *wsConn) readPump() {
|
||||
defer func() {
|
||||
cl.bind("") // 注销 live
|
||||
_ = cl.conn.Close()
|
||||
}()
|
||||
cl.conn.SetReadLimit(1 << 20) // 1MB
|
||||
_ = cl.conn.SetReadDeadline(time.Now().Add(90 * time.Second))
|
||||
cl.conn.SetPongHandler(func(string) error {
|
||||
_ = cl.conn.SetReadDeadline(time.Now().Add(90 * time.Second))
|
||||
return nil
|
||||
})
|
||||
|
||||
for {
|
||||
// 任何收到消息（含客户端心跳 {type:"ping"}）都刷新 90s 读超时，
|
||||
// 否则浏览器只能发协议层 Pong，应用层 ping 不会触发 SetPongHandler，连接必死在第 90 秒
|
||||
_ = cl.conn.SetReadDeadline(time.Now().Add(90 * time.Second))
|
||||
_, raw, err := cl.conn.ReadMessage()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
var msg wsInbound
|
||||
if err := json.Unmarshal(raw, &msg); err != nil {
|
||||
continue
|
||||
}
|
||||
switch msg.Type {
|
||||
case "chat":
|
||||
cl.handleChat(msg.Content)
|
||||
case "stop":
|
||||
cl.handleStop()
|
||||
case "switch":
|
||||
if msg.SessionID != "" && sessionIDPattern.MatchString(msg.SessionID) {
|
||||
if _, ok := cl.server.store.Get(msg.SessionID); ok {
|
||||
cl.bind(msg.SessionID)
|
||||
}
|
||||
}
|
||||
case "ping":
|
||||
cl.sendJSON("pong", map[string]any{"time": time.Now().UnixMilli()})
|
||||
default:
|
||||
log.Printf("ws: unknown message type %q", msg.Type)
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user