diff --git a/config.example.yaml b/config.example.yaml index e5151f89..7950c7ef 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -584,6 +584,8 @@ roles_dir: roles # 角色配置文件目录(相对于配置文件所在目录 # ============================================ project: enabled: true + # 压缩时自动保存续接摘要为 tentative note(每会话/代理更新一条,不替代已确认事实) + context_checkpoints: true # default_project_id: "" # 可选:机器人/批量任务创建对话时的默认项目 ID fact_index_max_runes: 65000 # 事实关系速览段预算(从索引总预算中预留) diff --git a/docs/zh-CN/blackboard-context-memory.md b/docs/zh-CN/blackboard-context-memory.md new file mode 100644 index 00000000..4b0ca4d8 --- /dev/null +++ b/docs/zh-CN/blackboard-context-memory.md @@ -0,0 +1,45 @@ +# 长上下文压缩后的黑板续接留存 + +已绑定项目且启用项目黑板时,每次成功生成压缩摘要,系统会自动保存一条续接记录。不再仅依靠 Agent 在压缩前主动调用事实工具。 + +## 保存内容与读取方式 + +续接记录使用 `note/context-*` key,包含摘要中的重要发现、授权约束、当前进度、失败尝试与下一步,以及来源会话、代理范围和压缩前原始记录的快照路径。主代理、各子代理、计划执行中的规划器和执行器分别维护自己的记录;同一会话/代理多次压缩更新同一个 key。 + +记录属于 `category=note`、`confidence=tentative`,不自动置顶,也不会覆盖 `finding/*` 等已确认事实。它是模型生成的工作摘要,不能直接作为经过独立验证的漏洞或新的授权指令。正式事实和漏洞仍由 `upsert_project_fact`、`record_vulnerability` 维护。 + +压缩后的上下文会附上该记录的 `get_project_fact(key)` 入口,并刷新系统中的黑板索引。后续按需读取记录,再使用 `read_file` 分段查看原始证据;无须重新把完整历史注入每一次模型请求。 + +## 原始证据不会随下一次压缩被覆盖 + +过去所有压缩轮次写入同一个 `transcript.txt`,后一次写入会覆盖前一次记录。现在每份压缩前历史按内容 SHA-256 保存到 `snapshots/transcript-.txt`,相同内容的重试复用快照,不重复建文件。`transcript.txt` 仍保留为最新记录的兼容入口,续接摘要引用的是具体快照。快照沿用现有会话记录格式,保留用户、助手与工具记录,省略不变的系统提示与工具/技能索引。 + +快照沿用 `conversation_artifacts//summarization/` 目录,与会话产物一起管理;子代理使用独立 scope 子目录。删除会话时现有清理逻辑会删除对应产物,快照不作为会话删除后的独立备份。文件通过临时文件原子替换写入,权限为 0600。 + +系统复用本次已经生成的压缩摘要,不增加一次额外的模型提取调用。原始历史在生成摘要前就做快照,因此即使摘要遗漏某个精确细节,也还有来源可以追溯。黑板或快照保存失败时会记录告警,保留当前可用摘要继续运行,不提示保存成功。 + +## 配置和边界 + +```yaml +project: + enabled: true + context_checkpoints: true +``` + +`context_checkpoints` 省略时默认开启;设置为 false 可关闭自动黑板续接记录,不影响原始证据快照。没有绑定项目的会话仍保留本地快照,但不会把内容写进其他项目。 + +自动记录前会核对会话当前绑定的项目,避免会话在运行期间切换项目后写入旧项目。这里没有增加跨项目搜索。现有项目黑板权限与事实工具的当前会话范围继续适用。 + +自动记录只在压缩触发时落库,不能代替每次新发现后的即时事实登记。摘要仍可能遗漏细节;具体原文以快照为准。快照累计占用本地存储,保留到对应会话产物被清理。 + +## 索引预算与来源区分 + +自动注入仍只包含 key、短摘要与关系边,完整 body 留在数据库按需读取。索引查询不再加载最终不会注入的 body。字符预算现在包含标题、起止标记、查询提示、遗漏说明和关系速览;特别小的预算会使用简短查询提示或不注入索引,不截断 HTML 边界。 + +遗留超长 summary 在展示索引时缩短为单行,库中的完整 summary/body 不改动。关系速览仅使用实际列入索引的事实,避免预算外事实继续通过路径视图占用上下文。置顶与事实类型的原有优先级保持不变。 + +## 实现依据与验证 + +Eino 的压缩流程支持在 `Finalize` 中处理摘要并重写上下文,`Callback` 在后处理之后触发;本项目在既有自定义 Finalize 中落库、附上快照引用,然后更新索引。普通消息与原生 AgenticMessage 路径均使用同一留存函数。[Eino Summarization 文档](https://www.cloudwego.io/docs/eino/core_modules/eino_adk/eino_adk_chatmodelagentmiddleware/middleware_summarization/) + +回归覆盖多次压缩保留早期证据、重复快照幂等、原子文件权限、主/子代理记录隔离、已确认事实保持不变、配置关闭、项目不匹配、保存失败不宣称成功、多段文本摘要、普通与原生中间件自动落库/刷新,以及索引硬预算和正文不注入。 diff --git a/docs/zh-CN/configuration.md b/docs/zh-CN/configuration.md index 8b9aed3c..edca05d7 100644 --- a/docs/zh-CN/configuration.md +++ b/docs/zh-CN/configuration.md @@ -256,12 +256,14 @@ c2: enabled: true project: enabled: true + context_checkpoints: true fact_index_max_runes: 65000 ``` - `c2.enabled`:关闭后不启动 C2 监听器,也不注册 C2 MCP 工具。 - WebShell 连接配置存 SQLite,没有单独的主配置开关。 - `project` 控制跨对话事实黑板注入预算。 +- `project.context_checkpoints` 默认开启:已绑定项目的会话成功压缩时,将续接摘要自动保存为 `note/context-*`(tentative),按会话及主/子代理分别更新,不覆盖已确认事实。完整原始记录保存为不可变快照,后续可通过 `get_project_fact` 与 `read_file` 找回。关闭此项仅停止自动黑板记录,仍保留原始上下文快照。详见 [长上下文续接留存](blackboard-context-memory.md)。 ## 机器人 diff --git a/internal/config/config.go b/internal/config/config.go index dc5c9165..2249919b 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -78,6 +78,7 @@ const ( // ProjectConfig 项目黑板(跨对话共享事实)配置。 type ProjectConfig struct { + ContextCheckpoints *bool `yaml:"context_checkpoints,omitempty" json:"context_checkpoints,omitempty"` Enabled bool `yaml:"enabled" json:"enabled"` DefaultProjectID string `yaml:"default_project_id,omitempty" json:"default_project_id,omitempty"` // 机器人/批量等无显式项目时绑定的默认项目 FactIndexMaxRunes int `yaml:"fact_index_max_runes,omitempty" json:"fact_index_max_runes,omitempty"` @@ -86,6 +87,11 @@ type ProjectConfig struct { DefaultInjectDeprecated bool `yaml:"default_inject_deprecated,omitempty" json:"default_inject_deprecated,omitempty"` } +// ContextCheckpointsEnabled persists generated continuation notes on compaction. +func (c ProjectConfig) ContextCheckpointsEnabled() bool { + return c.ContextCheckpoints == nil || *c.ContextCheckpoints +} + // FactIndexMaxRunesEffective 自动注入黑板索引的最大 rune 数。 func (c ProjectConfig) FactIndexMaxRunesEffective() int { if c.FactIndexMaxRunes <= 0 { diff --git a/internal/database/project.go b/internal/database/project.go index c2201267..68620931 100644 --- a/internal/database/project.go +++ b/internal/database/project.go @@ -337,7 +337,7 @@ func (db *DB) SetConversationProjectID(conversationID, projectID string) error { // ListProjectFactsForIndex 列出用于黑板索引注入的事实(不含 deprecated,除非 includeDeprecated)。 func (db *DB) ListProjectFactsForIndex(projectID string, includeDeprecated bool) ([]*ProjectFact, error) { - query := `SELECT id, project_id, fact_key, category, summary, COALESCE(body,''), confidence, + query := `SELECT id, project_id, fact_key, category, summary, '', confidence, COALESCE(source_conversation_id,''), COALESCE(source_message_id,''), pinned, COALESCE(related_vulnerability_id,''), created_at, updated_at FROM project_facts WHERE project_id = ?` diff --git a/internal/multiagent/eino_agentic_summarize.go b/internal/multiagent/eino_agentic_summarize.go index a2423370..0526beb1 100644 --- a/internal/multiagent/eino_agentic_summarize.go +++ b/internal/multiagent/eino_agentic_summarize.go @@ -28,6 +28,7 @@ func newEinoAgenticSummarizationMiddleware( db *database.DB, projectID string, logger *zap.Logger, + memoryScopes ...string, ) (adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage], error) { if summaryModel == nil || appCfg == nil { return nil, fmt.Errorf("multiagent: agentic summarization 需要 model 与配置") @@ -90,12 +91,16 @@ func newEinoAgenticSummarizationMiddleware( summaryInputMax = 4096 } + memoryScope := summarizationMemoryScope(memoryScopes) transcriptPath := "" if conv := strings.TrimSpace(conversationID); conv != "" { baseRoot := filepath.Join(os.TempDir(), "cyberstrike-summarization") if dbPath := strings.TrimSpace(appCfg.Database.Path); dbPath != "" { baseRoot = filepath.Join(filepath.Dir(dbPath), "conversation_artifacts", sanitizeEinoPathSegment(conv), "summarization") } + if memoryScope != "main" { + baseRoot = filepath.Join(baseRoot, "scopes", summarizationScopeID(memoryScope)) + } base := baseRoot if abs, err := filepath.Abs(base); err == nil { base = abs @@ -202,8 +207,7 @@ func newEinoAgenticSummarizationMiddleware( if len(defaultFinalized) == 0 { return nil, fmt.Errorf("agentic summarization default finalize returned no messages") } - summaryMsg := appendTranscriptPathToSummarizationMessage(defaultFinalized[len(defaultFinalized)-1], transcriptPath) - summaryMsg = stripAnalysisFromSummarizationMessage(summaryMsg) + summaryMsg := finalizeSummarizationMemory(defaultFinalized[len(defaultFinalized)-1], classicOriginal, transcriptPath, db, projectID, conversationID, memoryScope, appCfg.Project, logger) userLedger := buildOriginalUserIntentLedgerMessage(classicOriginal, userLedgerMaxRunes, userLedgerEntryMaxRunes) out, ferr := summarizeFinalizeWithRecentAssistantToolTrail(ctx, compactionMessages, summaryMsg, classicTokenCounter, recentTrailMax) if ferr != nil { diff --git a/internal/multiagent/eino_context_checkpoint.go b/internal/multiagent/eino_context_checkpoint.go new file mode 100644 index 00000000..d0c85e00 --- /dev/null +++ b/internal/multiagent/eino_context_checkpoint.go @@ -0,0 +1,181 @@ +package multiagent + +import ( + "crypto/sha256" + "fmt" + "os" + "path/filepath" + "strings" + + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/database" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" + "go.uber.org/zap" +) + +// Content-addressed snapshots keep earlier evidence reachable after the mutable +// transcript.txt alias advances to a later compaction round. +func writeSummarizationTranscriptSnapshot(path string, msgs []adk.Message) (string, error) { + path = strings.TrimSpace(path) + if path == "" { + return "", nil + } + body := []byte(formatSummarizationTranscript(msgs)) + dir := filepath.Join(filepath.Dir(path), "snapshots") + if err := os.MkdirAll(dir, 0o700); err != nil { + return "", err + } + digest := sha256.Sum256(body) + snapshot := filepath.Join(dir, fmt.Sprintf("transcript-%x.txt", digest)) + if _, err := os.Stat(snapshot); os.IsNotExist(err) { + if err := writeContextArtifactAtomically(snapshot, body); err != nil { + return "", err + } + } else if err != nil { + return "", err + } + if err := writeContextArtifactAtomically(path, body); err != nil { + return "", err + } + return snapshot, nil +} + +func writeContextArtifactAtomically(path string, body []byte) error { + file, err := os.CreateTemp(filepath.Dir(path), ".context-*") + if err != nil { + return err + } + temp := file.Name() + defer os.Remove(temp) + if _, err = file.Write(body); err != nil { + file.Close() + return err + } + if err = file.Close(); err != nil { + return err + } + return os.Rename(temp, path) +} + +// Persist generated working memory separately from verified project findings. +// One key per conversation/agent scope is updated; retries do not grow the fact index. +func persistSummarizationCheckpoint(db *database.DB, projectID, conversationID, memoryScope, summary, snapshot string, cfg config.ProjectConfig) (string, error) { + if db == nil || !cfg.Enabled || !cfg.ContextCheckpointsEnabled() || strings.TrimSpace(projectID) == "" || strings.TrimSpace(conversationID) == "" || strings.TrimSpace(summary) == "" { + return "", nil + } + // Bind storage to the current conversation rather than trusting a stale run's project ID. + bound, err := db.GetConversationProjectID(conversationID) + if err != nil { + return "", err + } + if bound != projectID { + return "", fmt.Errorf("checkpoint project does not match conversation") + } + hash := sha256.Sum256([]byte(conversationID + "\x00" + memoryScope)) + key := fmt.Sprintf("note/context-%x", hash[:16]) + body := "## 会话续接记录(自动摘要,尚未经独立验证)\n来源对话: " + conversationID + "\n代理范围: " + memoryScope + "\n\n" + summary + if snapshot != "" { + body += "\n\n压缩前原始证据快照(按需 read_file): " + snapshot + } + _, err = db.UpsertProjectFact(&database.ProjectFact{ + ProjectID: projectID, FactKey: key, Category: "note", Confidence: "tentative", + SourceConversationID: conversationID, + Summary: checkpointSummaryLabel(memoryScope, summary, cfg.FactSummaryMaxRunesEffective()), + Body: body, + }) + if err != nil { + return "", err + } + return key, nil +} + +func checkpointSummaryLabel(scope, summary string, cap int) string { + // A generated note must not turn a permissive fact-summary setting into + // automatic reinjection of the entire compaction output. + if cap > 200 { + cap = 200 + } + var lines []string + for _, line := range strings.Split(summary, "\n") { + line = strings.TrimSpace(line) + if line == "" || strings.HasPrefix(line, "#") || strings.HasPrefix(line, "<") { + continue + } + lines = append(lines, strings.TrimLeft(line, "- *")) + if len(lines) == 4 { + break + } + } + label := "自动续接(" + strings.Join(strings.Fields(scope), " ") + "):" + strings.Join(lines, ";") + runes := []rune(label) + if len(runes) > cap { + return string(runes[:cap-1]) + "…" + } + return label +} + +func finalizeSummarizationMemory(msg adk.Message, originals []adk.Message, transcriptPath string, db *database.DB, projectID, conversationID, memoryScope string, cfg config.ProjectConfig, logger *zap.Logger) adk.Message { + msg = stripAnalysisFromSummarizationMessage(msg) + snapshot, err := writeSummarizationTranscriptSnapshot(transcriptPath, originals) + if err != nil && logger != nil { + logger.Warn("保存上下文证据快照失败", zap.Error(err)) + } + // Only advertise an artifact that was successfully written. + if err == nil { + msg = appendTranscriptPathToSummarizationMessage(msg, snapshot) + } + text := "" + if msg != nil { + text = msg.Content + if text == "" { + for _, part := range msg.UserInputMultiContent { + if part.Type == schema.ChatMessagePartTypeText { + text = part.Text + break + } + } + } + } + if inner, ok := extractSummarizationSummaryBody(text); ok { + text = inner + } + key, err := persistSummarizationCheckpoint(db, projectID, conversationID, memoryScope, text, snapshot, cfg) + if err != nil { + if logger != nil { + logger.Warn("保存黑板续接记录失败", zap.String("conversationId", conversationID), zap.Error(err)) + } + return msg + } + if key == "" || msg == nil { + return msg + } + section := "本轮续接摘要已保存到项目黑板 get_project_fact(" + key + ");这是未独立验证的工作记录,使用前核对证据与当前任务范围。" + clone := *msg + if clone.Content != "" { + clone.Content = appendSummarizationSection(clone.Content, section) + } + if len(clone.UserInputMultiContent) > 0 { + clone.UserInputMultiContent = append([]schema.MessageInputPart(nil), clone.UserInputMultiContent...) + for i, part := range clone.UserInputMultiContent { + if part.Type == schema.ChatMessagePartTypeText { + clone.UserInputMultiContent[i].Text = appendSummarizationSection(part.Text, section) + break + } + } + } + return &clone +} + +func summarizationMemoryScope(scopes []string) string { + if len(scopes) > 0 && strings.TrimSpace(scopes[0]) != "" { + return strings.TrimSpace(scopes[0]) + } + return "main" +} + +func summarizationScopeID(scope string) string { + hash := sha256.Sum256([]byte(scope)) + return fmt.Sprintf("%x", hash[:16]) +} diff --git a/internal/multiagent/eino_context_checkpoint_test.go b/internal/multiagent/eino_context_checkpoint_test.go new file mode 100644 index 00000000..ebc8711f --- /dev/null +++ b/internal/multiagent/eino_context_checkpoint_test.go @@ -0,0 +1,199 @@ +package multiagent + +import ( + "context" + "fmt" + "os" + "path/filepath" + "strings" + "testing" + + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/project" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" + "go.uber.org/zap" +) + +func TestTranscriptSnapshotsPreserveEarlierEvidence(t *testing.T) { + path := filepath.Join(t.TempDir(), "transcript.txt") + first := []adk.Message{schema.UserMessage("first target"), schema.AssistantMessage("first evidence", nil)} + before, err := writeSummarizationTranscriptSnapshot(path, first) + if err != nil { + t.Fatal(err) + } + again, err := writeSummarizationTranscriptSnapshot(path, first) + if err != nil || before != again { + t.Fatalf("snapshot is not idempotent: %s %s %v", before, again, err) + } + after, err := writeSummarizationTranscriptSnapshot(path, []adk.Message{schema.UserMessage("second target")}) + if err != nil || before == after { + t.Fatalf("new history reused old snapshot: %v", err) + } + raw, err := os.ReadFile(before) + if err != nil || !strings.Contains(string(raw), "first evidence") { + t.Fatalf("earlier evidence lost: %s %v", raw, err) + } + alias, _ := os.ReadFile(path) + if !strings.Contains(string(alias), "second target") { + t.Fatal("latest alias not updated") + } + stat, _ := os.Stat(before) + if stat.Mode().Perm() != 0o600 { + t.Fatalf("snapshot permissions: %v", stat.Mode()) + } + files, _ := filepath.Glob(filepath.Join(filepath.Dir(before), "transcript-*.txt")) + if len(files) != 2 { + t.Fatalf("retries duplicated snapshots: %v", files) + } +} + +func checkpointTestProject(t *testing.T) (*database.DB, string, string) { + t.Helper() + db, err := database.NewDB(filepath.Join(t.TempDir(), "test.db"), zap.NewNop()) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { db.Close() }) + proj, err := db.CreateProject(&database.Project{Name: "memory"}) + if err != nil { + t.Fatal(err) + } + conv, err := db.CreateConversation("test", database.ConversationCreateMeta{ProjectID: proj.ID}) + if err != nil { + t.Fatal(err) + } + return db, proj.ID, conv.ID +} + +func TestContextCheckpointsSeparateAgentsAndPreserveConfirmedFacts(t *testing.T) { + db, pid, cid := checkpointTestProject(t) + cfg := config.ProjectConfig{Enabled: true} + db.UpsertProjectFact(&database.ProjectFact{ProjectID: pid, FactKey: "finding/original", Category: "finding", Confidence: "confirmed", Body: "original evidence", Summary: "confirmed result", Pinned: true}) + main, err := persistSummarizationCheckpoint(db, pid, cid, "main", "first summary", "/old/snapshot", cfg) + if err != nil { + t.Fatal(err) + } + updated, err := persistSummarizationCheckpoint(db, pid, cid, "main", "new summary", "/new/snapshot", cfg) + if err != nil || updated != main { + t.Fatalf("unstable key: %s %s %v", main, updated, err) + } + sub, err := persistSummarizationCheckpoint(db, pid, cid, "sub/scanner", "sub findings", "/sub/snapshot", cfg) + if err != nil || sub == main { + t.Fatalf("subagent overwrites coordinator: %v", err) + } + fact, _ := db.GetProjectFactByKey(pid, main) + if fact.Confidence != "tentative" || fact.Pinned || fact.SourceConversationID != cid || !strings.Contains(fact.Body, "new summary") || !strings.Contains(fact.Body, "/new/snapshot") { + t.Fatalf("invalid checkpoint: %#v", fact) + } + existing, _ := db.GetProjectFactByKey(pid, "finding/original") + if existing.Confidence != "confirmed" || existing.Body != "original evidence" || !existing.Pinned { + t.Fatalf("confirmed fact altered: %#v", existing) + } + facts, err := db.ListProjectFacts(pid, database.ProjectFactListFilter{}, 100, 0) + if err != nil || len(facts) != 3 { + t.Fatalf("checkpoint updates grow fact list: %d %v", len(facts), err) + } +} + +func TestContextCheckpointDisabledAndProjectMismatch(t *testing.T) { + db, pid, cid := checkpointTestProject(t) + off := false + for _, cfg := range []config.ProjectConfig{{Enabled: false}, {Enabled: true, ContextCheckpoints: &off}} { + key, err := persistSummarizationCheckpoint(db, pid, cid, "main", "summary", "", cfg) + if err != nil || key != "" { + t.Fatalf("disabled checkpoint wrote data: %s %v", key, err) + } + } + if key, err := persistSummarizationCheckpoint(db, "other-project", cid, "main", "summary", "", config.ProjectConfig{Enabled: true}); err == nil || key != "" { + t.Fatalf("wrong-project write: %s %v", key, err) + } +} + +func TestFinalizeMemorySupportsMultipartSummaryAndReportsOnlySuccessfulWrites(t *testing.T) { + db, pid, cid := checkpointTestProject(t) + original := []adk.Message{schema.UserMessage("exact evidence from history")} + summary := &schema.Message{Role: schema.User, UserInputMultiContent: []schema.MessageInputPart{{Type: schema.ChatMessagePartTypeText, Text: "internalkeep important finding"}}} + out := finalizeSummarizationMemory(summary, original, filepath.Join(t.TempDir(), "transcript.txt"), db, pid, cid, "main", config.ProjectConfig{Enabled: true}, nil) + text := out.UserInputMultiContent[0].Text + if strings.Contains(text, "internal") || !strings.Contains(text, "snapshots/transcript-") || !strings.Contains(text, "get_project_fact(note/context-") { + t.Fatalf("continuation missing durable references: %s", text) + } + if !strings.Contains(summary.UserInputMultiContent[0].Text, "internal") { + t.Fatal("input summary was mutated") + } + bad := finalizeSummarizationMemory(schema.UserMessage("still usable summary"), original, "", db, "wrong-project", cid, "main", config.ProjectConfig{Enabled: true}, nil) + if strings.Contains(bad.Content, "已保存") || !strings.Contains(bad.Content, "still usable summary") { + t.Fatalf("failed write claimed success or lost summary: %s", bad.Content) + } +} + +func TestSummarizationMiddlewaresAutomaticallySaveCheckpoint(t *testing.T) { + for _, native := range []bool{false, true} { + t.Run(fmt.Sprintf("native=%v", native), func(t *testing.T) { + db, pid, cid := checkpointTestProject(t) + cfg := &config.Config{Project: config.ProjectConfig{Enabled: true}} + cfg.Database.Path = filepath.Join(t.TempDir(), "runtime.db") + cfg.OpenAI.Model = "gpt-4o" + cfg.OpenAI.MaxTotalTokens = 5000 + emit := false + mwCfg := &config.MultiAgentEinoMiddlewareConfig{SummarizationEmitInternalEvents: &emit, SummarizationOutputReserveTokens: 1024} + index, err := project.BuildFactIndexBlock(db, pid, cfg.Project) + if err != nil { + t.Fatal(err) + } + originals := []adk.Message{ + schema.SystemMessage("system root\n" + index), + schema.UserMessage("precise original evidence\n" + strings.Repeat("old scan output ", 12000)), + schema.AssistantMessage("found endpoint", nil), + schema.UserMessage("continue from confirmed endpoint"), + } + const summary = "do not persist scratchpadImportant endpoint /api/check; next validate response" + var after []adk.Message + ctx := context.Background() + if native { + m := &capturingAgenticChatModel{output: agenticAssistantTextMessage(summary)} + mw, err := newEinoAgenticSummarizationMiddleware(ctx, m, cfg, mwCfg, cid, db, pid, nil) + if err != nil { + t.Fatal(err) + } + _, state, err := mw.BeforeModelRewriteState(ctx, &adk.TypedChatModelAgentState[*schema.AgenticMessage]{Messages: EinoMessagesToAgentic(originals)}, nil) + if err != nil { + t.Fatal(err) + } + after = AgenticMessagesToEino(state.Messages) + } else { + m := &capturingClassicChatModel{output: &schema.Message{Role: schema.Assistant, Content: summary, ResponseMeta: &schema.ResponseMeta{FinishReason: "stop"}}} + mw, err := newEinoSummarizationMiddleware(ctx, m, cfg, mwCfg, cid, db, pid, nil) + if err != nil { + t.Fatal(err) + } + _, state, err := mw.BeforeModelRewriteState(ctx, &adk.ChatModelAgentState{Messages: originals}, nil) + if err != nil { + t.Fatal(err) + } + after = state.Messages + } + facts, err := db.ListProjectFacts(pid, database.ProjectFactListFilter{}, 100, 0) + if err != nil || len(facts) != 1 { + t.Fatalf("automatic checkpoint missing: %v %#v", err, facts) + } + if !strings.Contains(facts[0].Body, "/api/check") || strings.Contains(facts[0].Body, "scratchpad") { + t.Fatalf("wrong working memory: %s", facts[0].Body) + } + joined := joinClassicMessageContent(after) + if !strings.Contains(joined, facts[0].FactKey) || strings.Contains(joined, "(暂无事实)") { + t.Fatalf("live index did not refresh after save: %s", joined) + } + }) + } +} + +func TestCheckpointIndexLabelStaysShortWithLargeSummarySetting(t *testing.T) { + label := checkpointSummaryLabel("main", "## findings\n- Important endpoint /api/check\n- "+strings.Repeat("raw output ", 1000), 24000) + if len([]rune(label)) > 200 || !strings.Contains(label, "/api/check") { + t.Fatalf("checkpoint reinjects long history or loses key hint: %s", label) + } +} diff --git a/internal/multiagent/eino_orchestration.go b/internal/multiagent/eino_orchestration.go index 051984e2..26b6f61a 100644 --- a/internal/multiagent/eino_orchestration.go +++ b/internal/multiagent/eino_orchestration.go @@ -118,7 +118,7 @@ func buildPlanExecuteAgenticExecutorHandlers(ctx context.Context, a *PlanExecute execHandlers = append(execHandlers, a.AgenticSkillMiddleware) } if a.AppCfg != nil { - sumMw, sumErr := newEinoAgenticSummarizationMiddleware(ctx, a.AgenticExecModel, a.AppCfg, a.MwCfg, a.ConversationID, a.DB, a.ProjectID, a.Logger) + sumMw, sumErr := newEinoAgenticSummarizationMiddleware(ctx, a.AgenticExecModel, a.AppCfg, a.MwCfg, a.ConversationID, a.DB, a.ProjectID, a.Logger, "plan_execute/executor") if sumErr != nil { return nil, fmt.Errorf("plan_execute agentic executor summarization: %w", sumErr) } diff --git a/internal/multiagent/eino_summarize.go b/internal/multiagent/eino_summarize.go index 63eede98..180a2347 100644 --- a/internal/multiagent/eino_summarize.go +++ b/internal/multiagent/eino_summarize.go @@ -88,6 +88,7 @@ func newEinoSummarizationMiddleware( db *database.DB, projectID string, logger *zap.Logger, + memoryScopes ...string, ) (adk.ChatModelAgentMiddleware, error) { if summaryModel == nil || appCfg == nil { return nil, fmt.Errorf("multiagent: summarization 需要 model 与配置") @@ -144,6 +145,7 @@ func newEinoSummarizationMiddleware( if summaryInputMax < 4096 { summaryInputMax = 4096 } + memoryScope := summarizationMemoryScope(memoryScopes) transcriptPath := "" if conv := strings.TrimSpace(conversationID); conv != "" { baseRoot := filepath.Join(os.TempDir(), "cyberstrike-summarization") @@ -151,6 +153,9 @@ func newEinoSummarizationMiddleware( // Persist with the same lifecycle as local conversation storage. baseRoot = filepath.Join(filepath.Dir(dbPath), "conversation_artifacts", sanitizeEinoPathSegment(conv), "summarization") } + if memoryScope != "main" { + baseRoot = filepath.Join(baseRoot, "scopes", summarizationScopeID(memoryScope)) + } base := baseRoot if abs, err := filepath.Abs(base); err == nil { base = abs @@ -247,8 +252,7 @@ func newEinoSummarizationMiddleware( if len(defaultFinalized) == 0 { return nil, fmt.Errorf("summarization default finalize returned no messages") } - summary = appendTranscriptPathToSummarizationMessage(defaultFinalized[len(defaultFinalized)-1], transcriptPath) - summary = stripAnalysisFromSummarizationMessage(summary) + summary = finalizeSummarizationMemory(defaultFinalized[len(defaultFinalized)-1], originalMessages, transcriptPath, db, projectID, conversationID, memoryScope, appCfg.Project, logger) userLedger := buildOriginalUserIntentLedgerMessage(originalMessages, userLedgerMaxRunes, userLedgerEntryMaxRunes) out, ferr := summarizeFinalizeWithRecentAssistantToolTrail(ctx, compactionMessages, summary, tokenCounter, recentTrailMax) if ferr != nil { @@ -679,14 +683,8 @@ func writeSummarizationTranscript(path string, msgs []adk.Message) error { if path == "" { return nil } - body := formatSummarizationTranscript(msgs) - if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { - return fmt.Errorf("mkdir transcript dir: %w", err) - } - if err := os.WriteFile(path, []byte(body), 0o600); err != nil { - return fmt.Errorf("write transcript: %w", err) - } - return nil + _, err := writeSummarizationTranscriptSnapshot(path, msgs) + return err } func einoSummarizationTokenCounter(openAIModel string) summarization.TokenCounterFunc { diff --git a/internal/multiagent/runner.go b/internal/multiagent/runner.go index a4bbb568..c4ed82c2 100644 --- a/internal/multiagent/runner.go +++ b/internal/multiagent/runner.go @@ -222,7 +222,7 @@ func RunDeepAgent( subMax := resolveMaxIterations(appCfg, sub.MaxIterations) - subSumMw, err := newEinoAgenticSummarizationMiddleware(ctx, subModel, appCfg, &ma.EinoMiddleware, conversationID, db, projectID, logger) + subSumMw, err := newEinoAgenticSummarizationMiddleware(ctx, subModel, appCfg, &ma.EinoMiddleware, conversationID, db, projectID, logger, "sub/"+id) if err != nil { return nil, fmt.Errorf("子代理 %q agentic summarization 中间件: %w", id, err) } @@ -488,7 +488,7 @@ func RunDeepAgent( if perr != nil { return nil, fmt.Errorf("plan_execute 执行器 AgenticModel: %w", perr) } - planRewriteSumMw, perr := newEinoSummarizationMiddleware(ctx, execModel, appCfg, &ma.EinoMiddleware, conversationID, db, projectID, logger) + planRewriteSumMw, perr := newEinoSummarizationMiddleware(ctx, execModel, appCfg, &ma.EinoMiddleware, conversationID, db, projectID, logger, "plan_execute/planning") if perr != nil { return nil, fmt.Errorf("plan_execute planner/replanner summarization: %w", perr) } diff --git a/internal/project/blackboard.go b/internal/project/blackboard.go index d1e2aec9..f93a0a98 100644 --- a/internal/project/blackboard.go +++ b/internal/project/blackboard.go @@ -49,31 +49,44 @@ func BuildFactIndexBlock(db *database.DB, projectID string, cfg config.ProjectCo allEdges, _ := db.ListProjectFactEdgesByProject(projectID) _, incomingByTarget := indexEdgeGroupMaps(allEdges) + maxRunes := cfg.FactIndexMaxRunesEffective() + header := fmt.Sprintf("## 项目黑板索引(project: %s, id: %s)\n", indexSingleLine(proj.Name), proj.ID) if len(facts) == 0 { - return wrapFactIndexBlock(fmt.Sprintf("## 项目黑板索引(project: %s, id: %s)\n(暂无事实)\n%s", proj.Name, proj.ID, factIndexFooterEmpty)), nil + block := wrapFactIndexBlock(header + "(暂无事实)\n" + factIndexFooterEmpty) + if len([]rune(block)) <= maxRunes { + return block, nil + } + return minimalFactIndexBlock(maxRunes), nil } sortFactsForIndex(facts) - maxRunes := cfg.FactIndexMaxRunesEffective() pathMaxRunes := cfg.FactIndexPathMaxRunesEffective() footer := factIndexFooterGetDetail + "\n" + factIndexFooterWriteHint - footerRunes := len([]rune(footer)) - factsBudget := maxRunes - pathMaxRunes - footerRunes + // Include markers, separators and the worst-case omission notice in the cap. + omissionReserve := len([]rune(factIndexOmissionNotice(len(facts)))) + framingRunes := len([]rune(wrapFactIndexBlock(header+footer))) + omissionReserve + available := maxRunes - framingRunes + if available <= 0 { + return minimalFactIndexBlock(maxRunes), nil + } + if pathMaxRunes > available { + pathMaxRunes = available + } + factsBudget := available - pathMaxRunes if factsBudget < 800 { - factsBudget = maxRunes - footerRunes + factsBudget = available pathMaxRunes = 0 } indexedKeys := make(map[string]struct{}, len(facts)) var b strings.Builder - b.WriteString(fmt.Sprintf("## 项目黑板索引(project: %s, id: %s)\n", proj.Name, proj.ID)) - used := len([]rune(b.String())) + b.WriteString(header) + used := 0 omitted := 0 for _, f := range facts { - indexedKeys[f.FactKey] = struct{}{} - line := fmt.Sprintf("- [%s] %s — %s (%s)", f.FactKey, f.Category, strings.TrimSpace(f.Summary), f.Confidence) + line := fmt.Sprintf("- [%s] %s — %s (%s)", f.FactKey, f.Category, truncateIndexSummary(indexSingleLine(f.Summary), cfg.FactSummaryMaxRunesEffective()), f.Confidence) line += FormatFactIndexLinksHint(f.FactKey, incomingByTarget[f.FactKey]) line += "\n" lineRunes := len([]rune(line)) @@ -81,15 +94,16 @@ func BuildFactIndexBlock(db *database.DB, projectID string, cfg config.ProjectCo omitted++ continue } + indexedKeys[f.FactKey] = struct{}{} b.WriteString(line) used += lineRunes } if omitted > 0 { - b.WriteString(fmt.Sprintf("\n(另有 %d 条未列入索引,请使用 list_project_facts 或 search_project_facts 查询。)\n", omitted)) + b.WriteString(factIndexOmissionNotice(omitted)) } - if pathSection := BuildFactPathOverviewSection(allEdges, indexedKeys, pathMaxRunes); pathSection != "" { + if pathSection := BuildFactPathOverviewSection(allEdges, indexedKeys, pathMaxRunes-1); len(indexedKeys) > 0 && pathSection != "" { b.WriteString("\n") b.WriteString(pathSection) } @@ -97,3 +111,25 @@ func BuildFactIndexBlock(db *database.DB, projectID string, cfg config.ProjectCo b.WriteString(footer) return wrapFactIndexBlock(b.String()), nil } + +func indexSingleLine(s string) string { return strings.Join(strings.Fields(s), " ") } + +func truncateIndexSummary(s string, maxRunes int) string { + runes := []rune(s) + if maxRunes > 0 && len(runes) > maxRunes { + return string(runes[:maxRunes-1]) + "…" + } + return s +} + +func factIndexOmissionNotice(count int) string { + return fmt.Sprintf("\n(另有 %d 条未列入索引,请使用 list_project_facts 或 search_project_facts 查询。)\n", count) +} + +func minimalFactIndexBlock(maxRunes int) string { + block := wrapFactIndexBlock("## 项目黑板索引\n索引预算不足;使用 list_project_facts 或 search_project_facts 按需查询。") + if len([]rune(block)) > maxRunes { + return "" + } + return block +} diff --git a/internal/project/blackboard_budget_test.go b/internal/project/blackboard_budget_test.go new file mode 100644 index 00000000..d3f67e6b --- /dev/null +++ b/internal/project/blackboard_budget_test.go @@ -0,0 +1,85 @@ +package project + +import ( + "path/filepath" + "strings" + "testing" + + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/database" + "go.uber.org/zap" +) + +func TestFactIndexCapIncludesFramingAndDoesNotInjectBodies(t *testing.T) { + db, err := database.NewDB(filepath.Join(t.TempDir(), "facts.db"), zap.NewNop()) + if err != nil { + t.Fatal(err) + } + defer db.Close() + proj, err := db.CreateProject(&database.Project{Name: strings.Repeat("very long project ", 100)}) + if err != nil { + t.Fatal(err) + } + for i, key := range []string{"target/core", "finding/key", "note/progress"} { + _, err = db.UpsertProjectFact(&database.ProjectFact{ProjectID: proj.ID, FactKey: key, Summary: strings.Repeat("old summary\n", 1000), Body: "SECRET_FULL_BODY", Category: strings.Split(key, "/")[0], Pinned: i == 0, Confidence: "confirmed"}) + if err != nil { + t.Fatal(err) + } + } + for _, cap := range []int{30, 100, 300, 700, 3500, 6500} { + block, err := BuildFactIndexBlock(db, proj.ID, config.ProjectConfig{Enabled: true, FactIndexMaxRunes: cap, FactIndexPathMaxRunes: 1000, FactSummaryMaxRunes: 120}) + if err != nil { + t.Fatal(err) + } + if n := len([]rune(block)); n > cap { + t.Fatalf("index has %d runes for cap %d", n, cap) + } + if strings.Contains(block, "SECRET_FULL_BODY") { + t.Fatal("full body leaked into index") + } + if block != "" && (!strings.HasPrefix(block, FactIndexSectionStartMarker) || !strings.Contains(block, FactIndexSectionEndMarker)) { + t.Fatalf("broken markers: %s", block) + } + } + projected, err := db.ListProjectFactsForIndex(proj.ID, false) + if err != nil { + t.Fatal(err) + } + if projected[0].Body != "" { + t.Fatal("index query loads unused full bodies") + } + full, err := db.GetProjectFactByKey(proj.ID, "target/core") + if err != nil || full.Body != "SECRET_FULL_BODY" { + t.Fatalf("stored evidence altered: %v", err) + } +} + +func TestFactIndexPathOverviewExcludesOmittedFact(t *testing.T) { + db, err := database.NewDB(filepath.Join(t.TempDir(), "facts.db"), zap.NewNop()) + if err != nil { + t.Fatal(err) + } + defer db.Close() + proj, err := db.CreateProject(&database.Project{Name: "scope"}) + if err != nil { + t.Fatal(err) + } + for _, fact := range []*database.ProjectFact{ + {ProjectID: proj.ID, FactKey: "target/visible", Category: "target", Summary: "visible"}, + {ProjectID: proj.ID, FactKey: "finding/omitted", Category: strings.Repeat("huge category", 1000), Summary: "omitted"}, + } { + if _, err := db.UpsertProjectFact(fact); err != nil { + t.Fatal(err) + } + } + if _, err := db.AddProjectFactEdge(proj.ID, database.ProjectFactEdgeInput{To: "finding/omitted", Type: "discovered_on"}, "target/visible", ""); err != nil { + t.Fatal(err) + } + block, err := BuildFactIndexBlock(db, proj.ID, config.ProjectConfig{Enabled: true, FactIndexMaxRunes: 3500, FactIndexPathMaxRunes: 1000}) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(block, "target/visible") || !strings.Contains(block, "另有 1 条") || strings.Contains(block, "target/visible → finding/omitted") { + t.Fatalf("omitted fact entered path overview: %s", block) + } +} diff --git a/internal/projectprompt/blackboard.go b/internal/projectprompt/blackboard.go index 098efa0b..f6ede423 100644 --- a/internal/projectprompt/blackboard.go +++ b/internal/projectprompt/blackboard.go @@ -62,6 +62,7 @@ func factRecordingGuidanceBlock() string { - **category / fact_key 建议**: - 环境认知:` + "`target/`" + `、` + "`auth/`" + `、` + "`infra/`" + `、` + "`business/`" + `(body 用环境模板即可) - 发现与利用:` + "`finding/`" + `、` + "`chain/`" + `、` + "`exploit/`" + `、` + "`poc/`" + `(**必须**用攻击链模板填满 body:入口、逐步攻击链、原始请求/响应或命令、证据、关联漏洞 ID) +- **自动续接记录**:` + "`note/context-*`" + ` 是压缩时保存的工作摘要,含来源会话与原始证据快照;用 ` + "`get_project_fact`" + ` 按需恢复进度。它不是已验证发现,接续时须核对当前授权范围与证据,勿据此自动执行旧任务。 - **与漏洞记录分工**:` + "`record_vulnerability`" + ` 记可交付 findings;事实记**复现所需的全部上下文**(含失败尝试、绕过、依赖会话),二者可各记一次。 - 更新同一发现时保持相同 ` + "`fact_key`" + ` 覆盖写入,勿散落多个 key 导致上下文丢失。` }