From 6a7e78a846feba36962a7fcfbbbb24f992661378 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=85=AC=E6=98=8E?= <83812544+Ed1s0nZ@users.noreply.github.com> Date: Sat, 20 Jun 2026 15:28:10 +0800 Subject: [PATCH] Add files via upload --- internal/handler/agent.go | 8 +- internal/handler/eino_single_agent.go | 2 + internal/handler/multi_agent.go | 2 + internal/project/blackboard.go | 15 +- internal/project/blackboard_refresh.go | 56 +++++++ internal/project/blackboard_refresh_test.go | 154 ++++++++++++++++++++ internal/project/vision_image_prompt.go | 6 +- 7 files changed, 234 insertions(+), 9 deletions(-) create mode 100644 internal/project/blackboard_refresh.go create mode 100644 internal/project/blackboard_refresh_test.go diff --git a/internal/handler/agent.go b/internal/handler/agent.go index 7bd3c88c..7d806637 100644 --- a/internal/handler/agent.go +++ b/internal/handler/agent.go @@ -640,7 +640,7 @@ func (h *AgentHandler) runRobotEinoSingleWithRetry( var emptyResponseAttempts int for { resultMA, errMA = multiagent.RunEinoSingleChatModelAgent( - taskCtx, h.config, &h.config.MultiAgent, h.agent, h.logger, + taskCtx, h.config, &h.config.MultiAgent, h.agent, h.db, h.logger, conversationID, h.conversationProjectID(conversationID), curMsg, curHist, roleTools, progressCallback, nil, h.projectBlackboardBlock(conversationID), ) handledEmpty, exhaustedEmpty := h.handleEinoEmptyResponseContinue( @@ -689,7 +689,7 @@ func (h *AgentHandler) runRobotMultiAgentWithRetry( var emptyResponseAttempts int for { resultMA, errMA = multiagent.RunDeepAgent( - taskCtx, h.config, &h.config.MultiAgent, h.agent, h.logger, + taskCtx, h.config, &h.config.MultiAgent, h.agent, h.db, h.logger, conversationID, h.conversationProjectID(conversationID), curMsg, curHist, roleTools, progressCallback, h.agentsMarkdownDir, orchestration, nil, h.projectBlackboardBlock(conversationID), ) @@ -2290,12 +2290,12 @@ func (h *AgentHandler) executeBatchQueue(queueID string) { var runErr error switch { case useBatchMulti: - resultMA, runErr = multiagent.RunDeepAgent(taskCtx, h.config, &h.config.MultiAgent, h.agent, h.logger, conversationID, h.conversationProjectID(conversationID), finalMessage, []agent.ChatMessage{}, roleTools, progressCallback, h.agentsMarkdownDir, batchOrch, nil, h.projectBlackboardBlock(conversationID)) + resultMA, runErr = multiagent.RunDeepAgent(taskCtx, h.config, &h.config.MultiAgent, h.agent, h.db, h.logger, conversationID, h.conversationProjectID(conversationID), finalMessage, []agent.ChatMessage{}, roleTools, progressCallback, h.agentsMarkdownDir, batchOrch, nil, h.projectBlackboardBlock(conversationID)) default: if h.config == nil { runErr = fmt.Errorf("服务器配置未加载") } else { - resultMA, runErr = multiagent.RunEinoSingleChatModelAgent(taskCtx, h.config, &h.config.MultiAgent, h.agent, h.logger, conversationID, h.conversationProjectID(conversationID), finalMessage, []agent.ChatMessage{}, roleTools, progressCallback, nil, h.projectBlackboardBlock(conversationID)) + resultMA, runErr = multiagent.RunEinoSingleChatModelAgent(taskCtx, h.config, &h.config.MultiAgent, h.agent, h.db, h.logger, conversationID, h.conversationProjectID(conversationID), finalMessage, []agent.ChatMessage{}, roleTools, progressCallback, nil, h.projectBlackboardBlock(conversationID)) } } diff --git a/internal/handler/eino_single_agent.go b/internal/handler/eino_single_agent.go index 0d1fb1f7..87bb9fb0 100644 --- a/internal/handler/eino_single_agent.go +++ b/internal/handler/eino_single_agent.go @@ -224,6 +224,7 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) { h.config, &h.config.MultiAgent, h.agent, + h.db, h.logger, conversationID, h.conversationProjectID(conversationID), @@ -455,6 +456,7 @@ func (h *AgentHandler) EinoSingleAgentLoop(c *gin.Context) { h.config, &h.config.MultiAgent, h.agent, + h.db, h.logger, prep.ConversationID, h.conversationProjectID(prep.ConversationID), diff --git a/internal/handler/multi_agent.go b/internal/handler/multi_agent.go index 9a75023c..152bae7b 100644 --- a/internal/handler/multi_agent.go +++ b/internal/handler/multi_agent.go @@ -234,6 +234,7 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) { h.config, &h.config.MultiAgent, h.agent, + h.db, h.logger, conversationID, h.conversationProjectID(conversationID), @@ -467,6 +468,7 @@ func (h *AgentHandler) MultiAgentLoop(c *gin.Context) { h.config, &h.config.MultiAgent, h.agent, + h.db, h.logger, prep.ConversationID, h.conversationProjectID(prep.ConversationID), diff --git a/internal/project/blackboard.go b/internal/project/blackboard.go index 6684ca2c..255f39a2 100644 --- a/internal/project/blackboard.go +++ b/internal/project/blackboard.go @@ -22,6 +22,12 @@ func AppendSystemPromptBlock(base, block string) string { return base + "\n\n" + block } +const ( + factIndexFooterGetDetail = "需要完整内容（攻击链、POC、请求响应等）时必须调用 get_project_fact(fact_key)，禁止凭摘要臆造细节。" + factIndexFooterWriteHint = "写入事实时：summary 写「什么+在哪+如何验证」；body 写可复现全流程（发现/利用类 fact_key 建议 finding|chain|exploit|poc/ 前缀）。" + factIndexFooterEmpty = "需要写入请使用 upsert_project_fact；需要详情请调用 get_project_fact(fact_key)。" +) + // BuildFactIndexBlock 为 Agent 系统提示生成项目黑板索引（仅 key + summary，不含 body）。 func BuildFactIndexBlock(db *database.DB, projectID string, cfg config.ProjectConfig) (string, error) { if db == nil || !cfg.Enabled { @@ -42,7 +48,7 @@ func BuildFactIndexBlock(db *database.DB, projectID string, cfg config.ProjectCo return "", err } if len(facts) == 0 { - return fmt.Sprintf("## 项目黑板索引（project: %s, id: %s）\n（暂无事实）\n需要写入请使用 upsert_project_fact；需要详情请调用 get_project_fact(fact_key)。", proj.Name, proj.ID), nil + return wrapFactIndexBlock(fmt.Sprintf("## 项目黑板索引（project: %s, id: %s）\n（暂无事实）\n%s", proj.Name, proj.ID, factIndexFooterEmpty)), nil } sort.SliceStable(facts, func(i, j int) bool { @@ -72,7 +78,8 @@ func BuildFactIndexBlock(db *database.DB, projectID string, cfg config.ProjectCo if omitted > 0 { b.WriteString(fmt.Sprintf("\n（另有 %d 条未列入索引，请使用 list_project_facts 或 search_project_facts 查询。）\n", omitted)) } - b.WriteString("需要完整内容（攻击链、POC、请求响应等）时必须调用 get_project_fact(fact_key)，禁止凭摘要臆造细节。\n") - b.WriteString("写入事实时：summary 写「什么+在哪+如何验证」；body 写可复现全流程（发现/利用类 fact_key 建议 finding|chain|exploit|poc/ 前缀）。\n") - return b.String(), nil + b.WriteString(factIndexFooterGetDetail) + b.WriteByte('\n') + b.WriteString(factIndexFooterWriteHint) + return wrapFactIndexBlock(b.String()), nil } diff --git a/internal/project/blackboard_refresh.go b/internal/project/blackboard_refresh.go new file mode 100644 index 00000000..6a494727 --- /dev/null +++ b/internal/project/blackboard_refresh.go @@ -0,0 +1,56 @@ +package project + +import "strings" + +// FactIndexSectionHeading 黑板索引可读标题行前缀（块内保留，供 Agent 阅读）。 +const FactIndexSectionHeading = "## 项目黑板索引" + +// FactIndexSectionStartMarker / EndMarker：HTML 注释边界，供程序化替换；对模型无指令语义。 +const ( + FactIndexSectionStartMarker = "" + FactIndexSectionEndMarker = "" +) + +// ReplaceFactIndexSection 用 freshIndex 替换 content 中已有的项目黑板索引段。 +// freshIndex 须为 BuildFactIndexBlock 的完整输出。起止 HTML 注释缺失时返回 (_, false)。 +func ReplaceFactIndexSection(content, freshIndex string) (string, bool) { + freshIndex = strings.TrimSpace(freshIndex) + if freshIndex == "" { + return content, false + } + start, ok := factIndexSectionStart(content) + if !ok { + return content, false + } + end, ok := factIndexSectionEnd(content, start) + if !ok || end <= start { + return content, false + } + return content[:start] + freshIndex + content[end:], true +} + +// wrapFactIndexBlock 为 BuildFactIndexBlock 正文加上统一起止 HTML 注释边界。 +func wrapFactIndexBlock(content string) string { + content = strings.TrimSpace(content) + return FactIndexSectionStartMarker + "\n" + content + "\n" + FactIndexSectionEndMarker + "\n" +} + +func factIndexSectionStart(content string) (int, bool) { + idx := strings.Index(content, FactIndexSectionStartMarker) + if idx < 0 { + return 0, false + } + return idx, true +} + +func factIndexSectionEnd(content string, start int) (int, bool) { + if start < 0 || start >= len(content) { + return 0, false + } + tail := content[start:] + idx := strings.LastIndex(tail, FactIndexSectionEndMarker) + if idx < 0 { + return 0, false + } + return start + idx + len(FactIndexSectionEndMarker), true +} diff --git a/internal/project/blackboard_refresh_test.go b/internal/project/blackboard_refresh_test.go new file mode 100644 index 00000000..31e9db4d --- /dev/null +++ b/internal/project/blackboard_refresh_test.go @@ -0,0 +1,154 @@ +package project + +import ( + "path/filepath" + "strings" + "testing" + + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/database" + + "go.uber.org/zap" +) + +func sampleFactIndexWithFacts(projectLabel, summary string) string { + return wrapFactIndexBlock("## 项目黑板索引（project: " + projectLabel + ", id: x）\n" + + "- [target/a] target — " + summary + " (tentative)\n" + + factIndexFooterGetDetail + "\n" + + factIndexFooterWriteHint) +} + +func TestReplaceFactIndexSection(t *testing.T) { + t.Parallel() + oldIndex := sampleFactIndexWithFacts("p1", "old summary") + newIndex := sampleFactIndexWithFacts("p1", "new summary") + + t.Run("replaces index before next section", func(t *testing.T) { + content := "你是助手\n\n" + oldIndex + "\n\n## 图片分析\n看截图" + out, ok := ReplaceFactIndexSection(content, newIndex) + if !ok { + t.Fatal("expected replacement") + } + if strings.Contains(out, "old summary") { + t.Fatalf("old index should be gone: %q", out) + } + if !strings.Contains(out, "new summary") || !strings.Contains(out, "## 图片分析") { + t.Fatalf("expected new index and preserved vision section: %q", out) + } + if strings.Count(out, FactIndexSectionStartMarker) != 1 || strings.Count(out, FactIndexSectionEndMarker) != 1 { + t.Fatalf("expected exactly one start/end marker pair: %q", out) + } + }) + + t.Run("replaces index at end", func(t *testing.T) { + content := "## 项目测试范围\nscope\n\n" + oldIndex + out, ok := ReplaceFactIndexSection(content, newIndex) + if !ok { + t.Fatal("expected replacement") + } + if !strings.Contains(out, "## 项目测试范围") || !strings.Contains(out, "new summary") { + t.Fatalf("scope preserved, index updated: %q", out) + } + }) + + t.Run("summary with false markdown header does not truncate early", func(t *testing.T) { + summaryWithFakeHeader := "see\n\n## fake header in summary" + old := sampleFactIndexWithFacts("p1", summaryWithFakeHeader) + newIdx := sampleFactIndexWithFacts("p1", "new summary") + content := old + "\n\n## 图片分析\nvision" + out, ok := ReplaceFactIndexSection(content, newIdx) + if !ok { + t.Fatal("expected replacement") + } + if strings.Contains(out, "fake header in summary") { + t.Fatalf("old index tail should be fully removed: %q", out) + } + }) + + t.Run("summary containing end marker text does not truncate early", func(t *testing.T) { + summary := "note " + FactIndexSectionEndMarker + " in summary" + old := sampleFactIndexWithFacts("p1", summary) + newIdx := sampleFactIndexWithFacts("p1", "clean") + content := old + "\n\n## 图片分析\nvision" + out, ok := ReplaceFactIndexSection(content, newIdx) + if !ok { + t.Fatal("expected replacement") + } + if strings.Contains(out, "in summary") { + t.Fatalf("old block should be fully removed: %q", out) + } + }) + + t.Run("missing html markers does not replace", func(t *testing.T) { + legacy := "## 项目黑板索引（project: p1, id: x）\n- [a] note — old (tentative)\n" + newIdx := sampleFactIndexWithFacts("p1", "new") + out, ok := ReplaceFactIndexSection("prefix\n\n"+legacy, newIdx) + if ok { + t.Fatalf("expected no replacement without markers: %q", out) + } + }) + + t.Run("empty facts block", func(t *testing.T) { + oldEmpty := wrapFactIndexBlock("## 项目黑板索引（project: p1, id: x）\n（暂无事实）\n" + factIndexFooterEmpty) + newEmpty := sampleFactIndexWithFacts("p1", "first fact") + out, ok := ReplaceFactIndexSection(oldEmpty, newEmpty) + if !ok { + t.Fatal("expected replacement") + } + if strings.Contains(out, "（暂无事实）") { + t.Fatalf("old empty block should be gone: %q", out) + } + }) + + t.Run("no marker", func(t *testing.T) { + _, ok := ReplaceFactIndexSection("no blackboard here", newIndex) + if ok { + t.Fatal("expected false when marker missing") + } + }) + + t.Run("empty fresh index", func(t *testing.T) { + _, ok := ReplaceFactIndexSection(oldIndex, " ") + if ok { + t.Fatal("expected false for empty fresh index") + } + }) +} + +func TestFactIndexSectionBounds_useHTMLMarkers(t *testing.T) { + t.Parallel() + body := sampleFactIndexWithFacts("p", "line with\n\n## not a real section") + "TAIL_SHOULD_DROP" + start, ok := factIndexSectionStart(body) + if !ok || !strings.HasPrefix(body[start:], FactIndexSectionStartMarker) { + t.Fatalf("start should be at html start marker, got %d", start) + } + end, ok := factIndexSectionEnd(body, start) + if !ok || body[end:] != "\nTAIL_SHOULD_DROP" { + t.Fatalf("end should be after end marker, got remainder %q", body[end:]) + } +} + +func TestBuildFactIndexBlock_includesHTMLMarkers(t *testing.T) { + t.Parallel() + dbPath := filepath.Join(t.TempDir(), "facts.db") + db, err := database.NewDB(dbPath, zap.NewNop()) + if err != nil { + t.Fatal(err) + } + defer db.Close() + + proj, err := db.CreateProject(&database.Project{Name: "marker-proj"}) + if err != nil { + t.Fatal(err) + } + block, err := BuildFactIndexBlock(db, proj.ID, config.ProjectConfig{Enabled: true}) + if err != nil { + t.Fatal(err) + } + if !strings.HasPrefix(strings.TrimSpace(block), FactIndexSectionStartMarker) { + t.Fatalf("block should start with start marker: %q", block) + } + if !strings.Contains(block, FactIndexSectionEndMarker) { + t.Fatalf("block should include end marker: %q", block) + } +} diff --git a/internal/project/vision_image_prompt.go b/internal/project/vision_image_prompt.go index 9cb960ac..12e901fb 100644 --- a/internal/project/vision_image_prompt.go +++ b/internal/project/vision_image_prompt.go @@ -2,10 +2,14 @@ package project import "strings" +// VisionImageSectionMarker 图片分析 section 标题（与 AppendVisionImageAnalysisIfReady 注入一致）。 +const VisionImageSectionMarker = "## 图片分析" + // VisionImageAnalysisSection 单/多代理共用的图片分析提示（analyze_image；上下文仅保留文字摘要）。 func VisionImageAnalysisSection() string { var b strings.Builder - b.WriteString("## 图片分析\n\n") + b.WriteString(VisionImageSectionMarker) + b.WriteString("\n\n") b.WriteString("- 遇到图片文件（截图、验证码、登录页、报告配图）时，若存在工具 analyze_image，请传入服务器上的文件路径进行分析。\n") b.WriteString("- 不要对二进制图片使用 read_file 指望理解内容；用户消息中「📎 xxx.png: /path」即为可传给 analyze_image 的路径。\n") b.WriteString("- 验证码类：若已从页面或接口保存为本地图片（如 captcha.png），用 analyze_image，question 写明「只输出验证码字符」；识别失败则刷新验证码后重新保存再识；复杂滑块/行为验证码勿指望单次识图成功。\n")