mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-10-20 22:41:23 +02:00
feat: persist compaction checkpoints in the project blackboard
This commit is contained in:
1 parent
deda642bd6
commit
616726d452
14 files changed
+586
-27
No files matched your search
@@ -584,6 +584,8 @@ roles_dir: roles # 角色配置文件目录(相对于配置文件所在目录
|
||||
# ============================================
|
||||
project:
|
||||
enabled: true
|
||||
# 压缩时自动保存续接摘要为 tentative note(每会话/代理更新一条,不替代已确认事实)
|
||||
context_checkpoints: true
|
||||
# default_project_id: "" # 可选:机器人/批量任务创建对话时的默认项目 ID
|
||||
fact_index_max_runes: 65000
|
||||
# 事实关系速览段预算(从索引总预算中预留)
|
||||
|
||||
@@ -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-<hash>.txt`,相同内容的重试复用快照,不重复建文件。`transcript.txt` 仍保留为最新记录的兼容入口,续接摘要引用的是具体快照。快照沿用现有会话记录格式,保留用户、助手与工具记录,省略不变的系统提示与工具/技能索引。
|
||||
|
||||
快照沿用 `conversation_artifacts/<conversationId>/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/)
|
||||
|
||||
回归覆盖多次压缩保留早期证据、重复快照幂等、原子文件权限、主/子代理记录隔离、已确认事实保持不变、配置关闭、项目不匹配、保存失败不宣称成功、多段文本摘要、普通与原生中间件自动落库/刷新,以及索引硬预算和正文不注入。
|
||||
@@ -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)。
|
||||
|
||||
## 机器人
|
||||
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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 = ?`
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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])
|
||||
}
|
||||
@@ -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: "<analysis>internal</analysis><summary>keep important finding</summary>"}}}
|
||||
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 = "<analysis>do not persist scratchpad</analysis><summary>Important endpoint /api/check; next validate response</summary>"
|
||||
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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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 导致上下文丢失。`
|
||||
}
|
||||
|
||||
Reference in new issue
Block a user