feat: simplify snapshot retrieval and add readable timestamps

This commit is contained in:
Sycun committed 2026-10-12 00:33:17 +08:00
1 parent e073709e8f
commit 1d2e900e88
9 files changed
+338 -441

No files matched your search

+11 -8
View File
@@ -10,13 +10,15 @@
压缩后的上下文会附上该记录的 `get_project_fact(key)` 入口,并刷新系统中的黑板索引。后续按需读取记录,再使用 `read_file` 分段查看原始证据;无须重新把完整历史注入每一次模型请求。
新上下文和黑板记录都包含程序生成的证据索引:按实际写入文件的行号标出 Cookie / Set-Cookie、认证 / Token、请求 / 参数、错误 / 失败的关键词命中位置。索引只包含固定类别与位置,不复制凭据值;它是检索线索,不代表相关凭据仍有效或结论已验证。需要精确值而摘要缺失时,Agent 应先检查索引,使用 `grep` 定位,再用 `read_file` 分段核对附近记录的域名、请求和时间,不猜测原值。
新上下文和黑板记录保存同一份续接摘要,并融合原始快照的绝对路径、历史快照目录和回读提示。摘要保留关键证据的来源目标、用途与发现经过;当前任务缺少 Cookie、认证信息、精确参数或错误原文时,Agent 可直接搜索对应快照,再用 `read_file` 分段核对附近原文,不猜测原值。
每类在上下文中最多列出最近 3 个命中行号,完整索引文件每类最多列出最近 32 个;超出时提示总命中数与搜索入口。关键词索引不是完整语义索引,未命中也可以搜索快照原文。
原有摘要章节及原始用户输入账本保持不变,只在同一次摘要中追加「原始证据检索线索」,记录工具结果或关键进展的来源与关键词,不新增对用户输入的归纳,也不由模型编造文件名或行号。
不再生成独立 `.index.md` 文件,不使用正则或关键词分类扫描快照,不增加模型调用。摘要本身承担证据检索入口;它仍可能遗漏细节,原文以快照为准。
## 原始证据不会随下一次压缩被覆盖
过去所有压缩轮次写入同一个 `transcript.txt`,后一次写入会覆盖前一次记录。现在每份压缩前历史按内容 SHA-256 保存到 `snapshots/transcript-<hash>.txt`,相同内容的重试复用快照,不重复建文件。`transcript.txt` 仍保留为最新记录的兼容入口,续接摘要引用的是具体快照。快照沿用现有会话记录格式,保留用户、助手与工具记录,省略不变的系统提示与工具/技能索引。
过去所有压缩轮次写入同一个 `transcript.txt`,后一次写入会覆盖前一次记录。现在每份压缩前历史保存到 `snapshots/transcript-<MD5前5位>-<时间戳>.txt`,例如 `transcript-a1b2c-20261012-010203.000000004+0800.txt`。时间戳带纳秒与时区,记录首次保存时间;相同内容的重试通过完整文件内容核对复用快照,不改变原快照的时间或路径。短 MD5 仅用于识别,不用作唯一性判断;前缀或时间戳碰撞时不覆盖旧文件。旧 SHA-256 文件名保持可读且可复用,不改名以免旧摘要的引用失效。`transcript.txt` 仍保留为最新记录的兼容入口,续接摘要引用的是具体快照。快照沿用现有会话记录格式,保留用户、助手与工具记录,省略不变的系统提示与工具/技能索引。
快照沿用 `conversation_artifacts/<conversationId>/summarization/` 目录,与会话产物一起管理;子代理使用独立 scope 子目录。删除会话时现有清理逻辑会删除对应产物,快照不作为会话删除后的独立备份。文件通过临时文件原子替换写入,权限为 0600。
@@ -26,11 +28,12 @@
data/conversation_artifacts/<conversationId>/summarization/
├── transcript.txt
└── snapshots/
├── transcript-<hash>.txt
└── transcript-<hash>.txt.index.md
└── transcript-<MD5前5位>-<时间戳>.txt
```
子代理、规划器和执行器使用 `summarization/scopes/<scopeHash>/snapshots/`。每份快照有独立索引文件,相同内容重试复用同一文件;历史索引不随最新记录更新而消失。摘要提供历史快照目录,Agent 可用 `glob` 查找 `transcript-*.txt.index.md`,按需查找旧证据,无须维护多个代理并发覆盖的共享目录清单。
可按文件名中的时间排序判断快照先后。快照可能在生成摘要前保存,包括失败尝试;相同历史的重试复用文件,因此文件序号不严格等同于成功压缩次数。
子代理、规划器和执行器使用 `summarization/scopes/<scopeHash>/snapshots/`。摘要提供历史快照目录,Agent 可用 `glob` 查找 `transcript-*.txt`,按需回读更早证据。已有旧索引文件不参与新流程,也不会在升级时被自动删除。
系统复用本次已经生成的压缩摘要,不增加一次额外的模型提取调用。原始历史在生成摘要前就做快照,因此即使摘要遗漏某个精确细节,也还有来源可以追溯。黑板或快照保存失败时会记录告警,保留当前可用摘要继续运行,不提示保存成功。
@@ -52,9 +55,9 @@ project:
自动注入仍只包含 key、短摘要与关系边,完整 body 留在数据库按需读取。索引查询不再加载最终不会注入的 body。字符预算现在包含标题、起止标记、查询提示、遗漏说明和关系速览;特别小的预算会使用简短查询提示或不注入索引,不截断 HTML 边界。
黑板不需要再次调用模型压缩:同一会话/代理的自动续接正文更新为本次摘要,短标签最多 200 字符并保留可检索的证据类别,不自动注入完整正文或快照索引。每次生成新证据索引时跳过历史中已有的程序索引块,防止把检索提示本身重复当作证据。历史原文及其索引留在文件中,已确认事实不因续接记录更新而被改写。
黑板不需要再次调用模型压缩:同一会话/代理的自动续接正文更新为本次摘要,短标签最多 200 字符,不自动注入完整正文或原始快照。读取入口只附在续接摘要与黑板正文中,不复制完整历史,也不维护额外的证据类别索引。历史原文保留在文件中,已确认事实不因续接记录更新而被改写。
跳过和替换仅处理同一消息内带完整起止标记的索引块;模型复述时遗漏结束标记,不会导致后续消息中的原始证据被跳过或摘要正文被误删。快照写入失败时,不注入框架默认的 `transcript.txt` 路径提示,避免 Agent 读取不存在或属于上一轮的文件。
快照写入失败时,不注入框架默认的 `transcript.txt` 路径提示,避免 Agent 读取不存在或属于上一轮的文件。普通文本、带 `<summary>` 标签的文本与多段消息均保留同一读取入口,避免正文重复附加。
遗留超长 summary 在展示索引时缩短为单行,库中的完整 summary/body 不改动。关系速览仅使用实际列入索引的事实,避免预算外事实继续通过路径视图占用上下文。置顶与事实类型的原有优先级保持不变。
+73 -55
View File
@@ -1,11 +1,15 @@
package multiagent
import (
"bytes"
"crypto/md5"
"crypto/sha256"
"fmt"
"os"
"path/filepath"
"strings"
"sync"
"time"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/database"
@@ -15,41 +19,27 @@ import (
"go.uber.org/zap"
)
// Content-addressed snapshots keep earlier evidence reachable after the mutable
// transcript.txt alias advances to a later compaction round.
var transcriptSnapshotWriteMu sync.Mutex
const transcriptSnapshotTimeLayout = "20060102-150405.000000000-0700"
// Time-stamped snapshots keep earlier evidence reachable after the mutable
// alias advances. Short MD5 prefixes are labels, never evidence of equality.
func writeSummarizationTranscriptSnapshot(path string, msgs []adk.Message) (string, error) {
path = strings.TrimSpace(path)
if path == "" {
return "", nil
}
body := []byte(formatSummarizationTranscript(msgs))
transcriptSnapshotWriteMu.Lock()
defer transcriptSnapshotWriteMu.Unlock()
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 info, err := os.Stat(snapshot); os.IsNotExist(err) {
if err := writeContextArtifactAtomically(snapshot, body); err != nil {
return "", err
}
} else if err != nil {
snapshot, err := findOrCreateTranscriptSnapshot(dir, body, time.Now())
if err != nil {
return "", err
} else if !info.Mode().IsRegular() {
return "", fmt.Errorf("snapshot is not a regular file: %s", snapshot)
}
// One immutable companion per snapshot avoids a shared mutable catalog and
// keeps discovery scoped to this conversation/agent, including concurrent runs.
indexPath := snapshot + ".index.md"
if info, err := os.Stat(indexPath); os.IsNotExist(err) {
index := buildTranscriptEvidenceIndex(snapshot, string(body), 32)
if err := writeContextArtifactAtomically(indexPath, []byte(index)); err != nil {
return "", err
}
} else if err != nil {
return "", err
} else if !info.Mode().IsRegular() {
return "", fmt.Errorf("evidence index is not a regular file: %s", indexPath)
}
if err := writeContextArtifactAtomically(path, body); err != nil {
return "", err
@@ -57,6 +47,59 @@ func writeSummarizationTranscriptSnapshot(path string, msgs []adk.Message) (stri
return snapshot, nil
}
func findOrCreateTranscriptSnapshot(dir string, body []byte, timestamp time.Time) (string, error) {
prefix := fmt.Sprintf("%x", md5.Sum(body))[:5]
entries, err := os.ReadDir(dir)
if err != nil {
return "", err
}
var candidates []string
for _, entry := range entries {
if strings.HasPrefix(entry.Name(), "transcript-"+prefix+"-") && strings.HasSuffix(entry.Name(), ".txt") {
candidates = append(candidates, filepath.Join(dir, entry.Name()))
}
}
// Existing SHA-256 names remain valid after upgrading; never rename evidence
// referenced by older summaries or create a duplicate on an identical retry.
legacy := filepath.Join(dir, fmt.Sprintf("transcript-%x.txt", sha256.Sum256(body)))
if _, err := os.Stat(legacy); err == nil {
candidates = append(candidates, legacy)
} else if !os.IsNotExist(err) {
return "", err
}
for _, candidate := range candidates {
info, err := os.Lstat(candidate)
if err != nil {
return "", err
}
if !info.Mode().IsRegular() {
return "", fmt.Errorf("snapshot is not a regular file: %s", candidate)
}
stored, err := os.ReadFile(candidate)
if err != nil {
return "", err
}
if bytes.Equal(stored, body) {
return candidate, nil
}
}
for {
name := "transcript-" + prefix + "-" + timestamp.Format(transcriptSnapshotTimeLayout) + ".txt"
path := filepath.Join(dir, name)
if _, err := os.Lstat(path); os.IsNotExist(err) {
if err := writeContextArtifactAtomically(path, body); err != nil {
return "", err
}
return path, nil
} else if err != nil {
return "", err
}
// A colliding timestamp must not replace an existing file, even if the
// short digest prefix happens to match different content.
timestamp = timestamp.Add(time.Nanosecond)
}
}
func writeContextArtifactAtomically(path string, body []byte) error {
file, err := os.CreateTemp(filepath.Dir(path), ".context-*")
if err != nil {
@@ -92,7 +135,7 @@ func persistSummarizationCheckpoint(db *database.DB, projectID, conversationID,
key := fmt.Sprintf("note/context-%x", hash[:16])
body := "## 会话续接记录(自动摘要,尚未经独立验证)\n来源对话: " + conversationID + "\n代理范围: " + memoryScope + "\n\n" + summary
if snapshot != "" {
body += "\n\n压缩前原始证据快照(按需 read_file): " + snapshot
body = appendSummarizationSection(body, summarizationTranscriptRetrievalSection(snapshot))
}
_, err = db.UpsertProjectFact(&database.ProjectFact{
ProjectID: projectID, FactKey: key, Category: "note", Confidence: "tentative",
@@ -112,23 +155,7 @@ func checkpointSummaryLabel(scope, summary string, cap int) string {
if cap > 200 {
cap = 200
}
// Keep retrieval cues in the short blackboard label without injecting paths
// or the full evidence index into every model request.
evidenceHint := ""
if start := strings.Index(summary, evidenceIndexMarker); start >= 0 {
section := summary[start:]
if end := strings.Index(section, evidenceIndexEnd); end >= 0 {
section = section[:end]
}
var topics []string
for _, topic := range evidenceTopics {
if strings.Contains(section, "- "+topic.label+":行 ") {
topics = append(topics, topic.label)
}
}
if len(topics) > 0 {
evidenceHint = ";证据索引:" + strings.Join(topics, "、")
}
if start := strings.Index(summary, summarizationSnapshotHeading); start >= 0 {
summary = summary[:start]
}
var lines []string
@@ -143,11 +170,7 @@ func checkpointSummaryLabel(scope, summary string, cap int) string {
}
}
label := "自动续接(" + strings.Join(strings.Fields(scope), " ") + "):" + strings.Join(lines, ";")
available := cap - len([]rune(evidenceHint))
if available < 1 {
return truncateIndexLabel(label+evidenceHint, cap)
}
return truncateIndexLabel(label, available) + evidenceHint
return truncateIndexLabel(label, cap)
}
func truncateIndexLabel(text string, cap int) string {
@@ -168,13 +191,8 @@ func finalizeSummarizationMemory(msg adk.Message, originals []adk.Message, trans
logger.Warn("保存上下文证据快照失败", zap.Error(err))
}
// Only advertise an artifact that was successfully written.
evidenceIndex := ""
if err == nil {
msg = appendTranscriptPathToSummarizationMessage(msg, snapshot)
if snapshot != "" {
evidenceIndex = buildTranscriptEvidenceIndex(snapshot, formatSummarizationTranscript(originals), 3)
msg = appendEvidenceIndexToMessage(msg, evidenceIndex)
}
}
text := ""
if msg != nil {
@@ -188,12 +206,12 @@ func finalizeSummarizationMemory(msg adk.Message, originals []adk.Message, trans
}
}
}
if snapshot != "" {
text = strings.TrimSuffix(text, "\n\n"+summarizationTranscriptRetrievalSection(snapshot))
}
if inner, ok := extractSummarizationSummaryBody(text); ok {
text = inner
}
if evidenceIndex != "" {
text = appendEvidenceIndex(text, evidenceIndex)
}
key, err := persistSummarizationCheckpoint(db, projectID, conversationID, memoryScope, text, snapshot, cfg)
if err != nil {
if logger != nil {
@@ -187,6 +187,10 @@ func TestSummarizationMiddlewaresAutomaticallySaveCheckpoint(t *testing.T) {
if !strings.Contains(joined, facts[0].FactKey) || strings.Contains(joined, "(暂无事实)") {
t.Fatalf("live index did not refresh after save: %s", joined)
}
snapshots, err := filepath.Glob(filepath.Join(filepath.Dir(cfg.Database.Path), "conversation_artifacts", cid, "summarization", "snapshots", "*.txt"))
if err != nil || len(snapshots) != 1 {
t.Fatalf("one compaction should reuse one original transcript across preflight/finalize/callback: %v %v", snapshots, err)
}
})
}
}
-122
View File
@@ -1,122 +0,0 @@
package multiagent
import (
"fmt"
"path/filepath"
"regexp"
"strings"
"github.com/cloudwego/eino/adk"
"github.com/cloudwego/eino/schema"
)
const evidenceIndexMarker = "## 原始证据索引(程序定位)"
const evidenceIndexEnd = "<!-- evidence-index-end -->"
var evidenceTopics = []struct {
label string
pattern *regexp.Regexp
}{
{"Cookie / Set-Cookie", regexp.MustCompile(`(?i)\b(?:set-cookie|cookie)\b`)},
{"认证 / Token", regexp.MustCompile(`(?i)\b(?:authorization|bearer|access_token|refresh_token|api[_-]?key|session[_-]?id)\b`)},
{"请求 / 参数", regexp.MustCompile(`(?i)\b(?:https?://|curl|request|headers?|params?|payload)\b|请求|参数`)},
{"错误 / 失败", regexp.MustCompile(`(?i)\b(?:error|exception|traceback|failed|failure)\b|错误|失败`)},
}
// Locations are derived from the exact serialized file, never from model output.
// Only fixed category labels and line numbers are exposed, not credential values.
func buildTranscriptEvidenceIndex(snapshot, transcript string, limit int) string {
var out strings.Builder
fmt.Fprintf(&out, "%s\n快照: %s\n历史快照目录: %s\n索引文件: %s.index.md\n", evidenceIndexMarker, snapshot, filepath.Dir(snapshot), snapshot)
lines := strings.Split(transcript, "\n")
skip := make([]bool, len(lines))
for _, span := range completeEvidenceIndexLineSpans(lines) {
for i := span[0]; i <= span[1]; i++ {
skip[i] = true
}
}
for _, topic := range evidenceTopics {
var locations []string
count := 0
for i, line := range lines {
if skip[i] {
continue
}
if topic.pattern.MatchString(line) {
count++
locations = append(locations, fmt.Sprintf("%d", i+1))
if len(locations) > limit {
locations = locations[1:]
}
}
}
if count > 0 {
fmt.Fprintf(&out, "- %s:行 %s", topic.label, strings.Join(locations, ", "))
if count > len(locations) {
fmt.Fprintf(&out, "(最近 %d 行,共 %d 行命中,其余用 grep 定位)", len(locations), count)
}
out.WriteByte('\n')
}
}
out.WriteString("以上为关键词定位,不代表凭据有效或结论已验证。当前任务缺少 Cookie、认证信息、精确参数或错误原文时,先检查索引,用 grep 定位,再用 read_file 分段核对附近记录的域名、请求与时间;不要猜测原值。历史索引用 glob 在该目录查找 transcript-*.txt.index.md,再按需读取;未命中仍可搜索快照原文。不要整批重新载入历史,也不要将旧记录当作新授权。\n")
out.WriteString(evidenceIndexEnd)
return out.String()
}
func appendEvidenceIndex(msgText, section string) string {
// Replace a previously generated section rather than nesting evidence indexes.
lines := strings.SplitAfter(msgText, "\n")
spans := completeEvidenceIndexLineSpans(lines)
var out strings.Builder
next := 0
for _, span := range spans {
for _, line := range lines[next:span[0]] {
out.WriteString(line)
}
next = span[1] + 1
}
for _, line := range lines[next:] {
out.WriteString(line)
}
return appendSummarizationSection(out.String(), section)
}
// Only complete blocks within a single transcript turn are generated indexes.
// A truncated or nested heading must never swallow subsequent conversation data.
func completeEvidenceIndexLineSpans(lines []string) [][2]int {
var spans [][2]int
start := -1
for i, line := range lines {
line = strings.TrimSpace(line)
if strings.HasPrefix(line, "--- [") && strings.HasSuffix(line, "] ---") {
start = -1
}
if line == evidenceIndexMarker {
start = i
} else if line == evidenceIndexEnd && start >= 0 {
spans = append(spans, [2]int{start, i})
start = -1
}
}
return spans
}
func appendEvidenceIndexToMessage(msg adk.Message, section string) adk.Message {
if msg == nil {
return nil
}
clone := *msg
if clone.Content != "" {
clone.Content = appendEvidenceIndex(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 && part.Text != "" {
clone.UserInputMultiContent[i].Text = appendEvidenceIndex(part.Text, section)
break
}
}
}
return &clone
}
@@ -1,252 +0,0 @@
package multiagent
import (
"context"
"crypto/sha256"
"fmt"
"os"
"path/filepath"
"strconv"
"strings"
"sync"
"testing"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/database"
"github.com/cloudwego/eino/adk"
"github.com/cloudwego/eino/schema"
)
func TestEvidenceIndexLocatesSerializedEvidenceWithoutRepeatingSecrets(t *testing.T) {
path := filepath.Join(t.TempDir(), "transcript.txt")
msgs := []adk.Message{
schema.SystemMessage("static Cookie: ignored-secret"),
schema.UserMessage("https://example.test/login"),
{Role: schema.Tool, Content: "Set-Cookie: sid=test-secret; Path=/\nAuthorization: Bearer test-token\nError: invalid session"},
}
snapshot, err := writeSummarizationTranscriptSnapshot(path, msgs)
if err != nil {
t.Fatal(err)
}
raw, _ := os.ReadFile(snapshot)
index, err := os.ReadFile(snapshot + ".index.md")
if err != nil {
t.Fatal(err)
}
text := string(index)
for _, secret := range []string{"test-secret", "test-token", "ignored-secret"} {
if strings.Contains(text, secret) {
t.Fatalf("index repeats secret %q", secret)
}
}
// Check locations against the actual on-disk transcript, including its header.
for _, topic := range evidenceTopics {
var expected []string
for i, line := range strings.Split(string(raw), "\n") {
if topic.pattern.MatchString(line) {
expected = append(expected, strconv.Itoa(i+1))
}
}
if len(expected) == 0 || !strings.Contains(text, topic.label+":行 "+strings.Join(expected, ", ")) {
t.Fatalf("incorrect %s locations: %s", topic.label, text)
}
}
stat, _ := os.Stat(snapshot + ".index.md")
if stat.Mode().Perm() != 0o600 {
t.Fatalf("index permissions: %v", stat.Mode())
}
if _, err := writeSummarizationTranscriptSnapshot(path, msgs); err != nil {
t.Fatal(err)
}
files, _ := filepath.Glob(filepath.Join(filepath.Dir(snapshot), "*.index.md"))
if len(files) != 1 {
t.Fatalf("retry duplicated indexes: %v", files)
}
}
func TestFailedTranscriptWriteDoesNotLeaveFrameworkAliasHint(t *testing.T) {
for _, native := range []bool{false, true} {
t.Run(fmt.Sprintf("native=%v", native), func(t *testing.T) {
cfg := &config.Config{}
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}
ctx := context.Background()
base := filepath.Join(filepath.Dir(cfg.Database.Path), "conversation_artifacts", "failure", "summarization")
if err := os.MkdirAll(base, 0o700); err != nil {
t.Fatal(err)
}
// Make snapshot creation fail while the middleware can still create its base directory.
if err := os.WriteFile(filepath.Join(base, "snapshots"), []byte("blocked"), 0o600); err != nil {
t.Fatal(err)
}
original := []adk.Message{schema.SystemMessage("system"), schema.UserMessage(strings.Repeat("history ", 20000)), schema.UserMessage("continue")}
const summary = "<summary>usable continuation</summary>"
var after []adk.Message
if native {
mw, err := newEinoAgenticSummarizationMiddleware(ctx, &capturingAgenticChatModel{output: agenticAssistantTextMessage(summary)}, cfg, mwCfg, "failure", nil, "", nil)
if err != nil {
t.Fatal(err)
}
_, state, err := mw.BeforeModelRewriteState(ctx, &adk.TypedChatModelAgentState[*schema.AgenticMessage]{Messages: EinoMessagesToAgentic(original)}, nil)
if err != nil {
t.Fatal(err)
}
after = AgenticMessagesToEino(state.Messages)
} else {
mw, err := newEinoSummarizationMiddleware(ctx, &capturingClassicChatModel{output: &schema.Message{Role: schema.Assistant, Content: summary, ResponseMeta: &schema.ResponseMeta{FinishReason: "stop"}}}, cfg, mwCfg, "failure", nil, "", nil)
if err != nil {
t.Fatal(err)
}
_, state, err := mw.BeforeModelRewriteState(ctx, &adk.ChatModelAgentState{Messages: original}, nil)
if err != nil {
t.Fatal(err)
}
after = state.Messages
}
text := joinClassicMessageContent(after)
if !strings.Contains(text, "usable continuation") || strings.Contains(text, "transcript.txt") || strings.Contains(text, evidenceIndexMarker) {
t.Fatalf("failed snapshot has a misleading retrieval hint: %s", text)
}
})
}
}
func TestConcurrentSnapshotWritesRemainCompleteAndDeduplicated(t *testing.T) {
path := filepath.Join(t.TempDir(), "transcript.txt")
var wg sync.WaitGroup
errors := make(chan error, 20)
for i := 0; i < 20; i++ {
wg.Add(1)
go func(i int) {
defer wg.Done()
msgs := []adk.Message{schema.UserMessage(fmt.Sprintf("Cookie: fixture-%d", i%2))}
snapshot, err := writeSummarizationTranscriptSnapshot(path, msgs)
if err == nil {
var body []byte
body, err = os.ReadFile(snapshot)
if err == nil && filepath.Base(snapshot) != fmt.Sprintf("transcript-%x.txt", sha256.Sum256(body)) {
err = fmt.Errorf("partial or mismatched snapshot: %s", snapshot)
}
}
errors <- err
}(i)
}
wg.Wait()
close(errors)
for err := range errors {
if err != nil {
t.Fatal(err)
}
}
files, _ := filepath.Glob(filepath.Join(filepath.Dir(path), "snapshots", "*.txt"))
indexes, _ := filepath.Glob(filepath.Join(filepath.Dir(path), "snapshots", "*.index.md"))
if len(files) != 2 || len(indexes) != 2 {
t.Fatalf("concurrent retries not deduplicated: %v %v", files, indexes)
}
}
func TestEvidenceIndexIsBoundedAndSkipsEarlierGeneratedIndexes(t *testing.T) {
old := buildTranscriptEvidenceIndex("/old/transcript.txt", "Cookie: earlier", 3)
transcript := old + "\n" + strings.Repeat("Cookie: current\n", 100)
index := buildTranscriptEvidenceIndex("/new/transcript.txt", transcript, 3)
if !strings.Contains(index, "共 100 行命中") || strings.Contains(index, "共 101 行命中") {
t.Fatalf("indexed its own earlier hints: %s", index)
}
if strings.Contains(index, "认证 / Token:行") || strings.Contains(index, "错误 / 失败:行") {
t.Fatalf("generated hints counted as evidence: %s", index)
}
lineCount := strings.Count(old, "\n") + 1
want := "Cookie / Set-Cookie:行 " + strconv.Itoa(lineCount+98) + ", " + strconv.Itoa(lineCount+99) + ", " + strconv.Itoa(lineCount+100)
if !strings.Contains(index, want) {
t.Fatalf("did not retain latest locations %s: %s", want, index)
}
if len(index) > 3000 {
t.Fatalf("unbounded inline index: %d", len(index))
}
updated := appendEvidenceIndex("summary\n"+old+"\n"+old, index)
if strings.Count(updated, evidenceIndexMarker) != 1 || strings.Contains(updated, "/old/") {
t.Fatalf("indexes accumulate on replacement: %s", updated)
}
}
func TestEvidenceIndexFailureDoesNotAdvertiseUnavailableArtifacts(t *testing.T) {
path := filepath.Join(t.TempDir(), "transcript.txt")
originals := []adk.Message{schema.UserMessage("Cookie: sid=fixture")}
snapshot, err := writeSummarizationTranscriptSnapshot(path, originals)
if err != nil {
t.Fatal(err)
}
if err := os.Remove(snapshot + ".index.md"); err != nil {
t.Fatal(err)
}
if err := os.Mkdir(snapshot+".index.md", 0o700); err != nil {
t.Fatal(err)
}
out := finalizeSummarizationMemory(schema.UserMessage("continue"), originals, path, nil, "", "", "main", config.ProjectConfig{}, nil)
if strings.Contains(out.Content, evidenceIndexMarker) || strings.Contains(out.Content, snapshot) || out.Content != "continue" {
t.Fatalf("failed artifact advertised or summary lost: %s", out.Content)
}
}
func TestEvidenceIndexSurvivesSummaryExtractionIntoBlackboard(t *testing.T) {
db, pid, cid := checkpointTestProject(t)
for _, multipart := range []bool{false, true} {
msg := schema.UserMessage("<summary>login completed; continue</summary>")
if multipart {
msg.UserInputMultiContent = []schema.MessageInputPart{{Type: schema.ChatMessagePartTypeText, Text: msg.Content}}
msg.Content = ""
}
out := finalizeSummarizationMemory(msg, []adk.Message{schema.UserMessage("Cookie: sid=fixture-value")}, filepath.Join(t.TempDir(), "transcript.txt"), db, pid, cid, "main", config.ProjectConfig{Enabled: true}, nil)
text := out.Content
if multipart {
text = out.UserInputMultiContent[0].Text
}
facts, err := db.ListProjectFacts(pid, database.ProjectFactListFilter{}, 10, 0)
if err != nil || len(facts) != 1 {
t.Fatalf("checkpoint count: %v %v", facts, err)
}
for _, body := range []string{text, facts[0].Body} {
if !strings.Contains(body, "Cookie / Set-Cookie:行") || !strings.Contains(body, ".index.md") || strings.Contains(body, "fixture-value") {
t.Fatalf("missing index or repeated secret: %s", body)
}
if strings.Count(body, evidenceIndexMarker) != 1 {
t.Fatalf("duplicate generated index: %s", body)
}
}
if !strings.Contains(facts[0].Summary, "证据索引:Cookie / Set-Cookie") || strings.Contains(facts[0].Summary, ".index.md") {
t.Fatalf("blackboard label lost retrieval cue or includes paths: %s", facts[0].Summary)
}
}
}
func TestEvidenceHintFitsBlackboardLabelBudget(t *testing.T) {
index := buildTranscriptEvidenceIndex("/snapshot.txt", "Cookie: fixture\nAuthorization: fixture\nhttps://example.test\nError: fixture", 3)
label := checkpointSummaryLabel("main", strings.Repeat("long summary ", 1000)+"\n"+index, 24000)
if len([]rune(label)) > 200 || !strings.Contains(label, "Cookie / Set-Cookie") || !strings.Contains(label, "错误 / 失败") {
t.Fatalf("unbounded label or lost evidence categories: %s", label)
}
}
func TestTruncatedEvidenceIndexDoesNotHideLaterConversationEvidence(t *testing.T) {
transcript := formatSummarizationTranscript([]adk.Message{
schema.UserMessage("old summary\n" + evidenceIndexMarker + "\nindex interrupted"),
{Role: schema.Tool, Content: "Cookie: sid=new-fixture\nError: later request failed"},
})
index := buildTranscriptEvidenceIndex("/new/snapshot.txt", transcript, 3)
if !strings.Contains(index, "Cookie / Set-Cookie:行") || !strings.Contains(index, "错误 / 失败:行") {
t.Fatalf("truncated old index suppressed later evidence: %s", index)
}
}
func TestIncompleteEvidenceIndexDoesNotConsumeLaterSummaryWhenReplaced(t *testing.T) {
complete := buildTranscriptEvidenceIndex("/old/snapshot.txt", "Cookie: fixture", 3)
input := "summary\n" + evidenceIndexMarker + "\nincomplete\nkeep this progress\n" + complete
updated := appendEvidenceIndex(input, "new index")
if !strings.Contains(updated, "keep this progress") || strings.Contains(updated, "/old/snapshot.txt") {
t.Fatalf("replacing a completed index consumed unrelated summary: %s", updated)
}
}
@@ -0,0 +1,230 @@
package multiagent
import (
"bytes"
"context"
"crypto/md5"
"crypto/sha256"
"fmt"
"os"
"path/filepath"
"strings"
"sync"
"testing"
"time"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/database"
"github.com/cloudwego/eino/adk"
"github.com/cloudwego/eino/schema"
)
func TestFailedTranscriptWriteDoesNotLeaveFrameworkAliasHint(t *testing.T) {
for _, native := range []bool{false, true} {
t.Run(fmt.Sprintf("native=%v", native), func(t *testing.T) {
cfg := &config.Config{}
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}
ctx := context.Background()
base := filepath.Join(filepath.Dir(cfg.Database.Path), "conversation_artifacts", "failure", "summarization")
if err := os.MkdirAll(base, 0o700); err != nil {
t.Fatal(err)
}
// Make snapshot creation fail while the middleware can still create its base directory.
if err := os.WriteFile(filepath.Join(base, "snapshots"), []byte("blocked"), 0o600); err != nil {
t.Fatal(err)
}
original := []adk.Message{schema.SystemMessage("system"), schema.UserMessage(strings.Repeat("history ", 20000)), schema.UserMessage("continue")}
const summary = "<summary>usable continuation</summary>"
var after []adk.Message
if native {
mw, err := newEinoAgenticSummarizationMiddleware(ctx, &capturingAgenticChatModel{output: agenticAssistantTextMessage(summary)}, cfg, mwCfg, "failure", nil, "", nil)
if err != nil {
t.Fatal(err)
}
_, state, err := mw.BeforeModelRewriteState(ctx, &adk.TypedChatModelAgentState[*schema.AgenticMessage]{Messages: EinoMessagesToAgentic(original)}, nil)
if err != nil {
t.Fatal(err)
}
after = AgenticMessagesToEino(state.Messages)
} else {
mw, err := newEinoSummarizationMiddleware(ctx, &capturingClassicChatModel{output: &schema.Message{Role: schema.Assistant, Content: summary, ResponseMeta: &schema.ResponseMeta{FinishReason: "stop"}}}, cfg, mwCfg, "failure", nil, "", nil)
if err != nil {
t.Fatal(err)
}
_, state, err := mw.BeforeModelRewriteState(ctx, &adk.ChatModelAgentState{Messages: original}, nil)
if err != nil {
t.Fatal(err)
}
after = state.Messages
}
text := joinClassicMessageContent(after)
if !strings.Contains(text, "usable continuation") || strings.Contains(text, "transcript.txt") || strings.Contains(text, summarizationSnapshotHeading) {
t.Fatalf("failed snapshot has a misleading retrieval hint: %s", text)
}
})
}
}
func TestConcurrentSnapshotWritesRemainCompleteAndDeduplicated(t *testing.T) {
path := filepath.Join(t.TempDir(), "transcript.txt")
var wg sync.WaitGroup
errors := make(chan error, 20)
for i := 0; i < 20; i++ {
wg.Add(1)
go func(i int) {
defer wg.Done()
msgs := []adk.Message{schema.UserMessage(fmt.Sprintf("Cookie: fixture-%d", i%2))}
snapshot, err := writeSummarizationTranscriptSnapshot(path, msgs)
if err == nil {
var body []byte
body, err = os.ReadFile(snapshot)
if err == nil && !bytes.Equal(body, []byte(formatSummarizationTranscript(msgs))) {
err = fmt.Errorf("partial or mismatched snapshot: %s", snapshot)
}
}
errors <- err
}(i)
}
wg.Wait()
close(errors)
for err := range errors {
if err != nil {
t.Fatal(err)
}
}
files, _ := filepath.Glob(filepath.Join(filepath.Dir(path), "snapshots", "*.txt"))
indexes, _ := filepath.Glob(filepath.Join(filepath.Dir(path), "snapshots", "*.index.md"))
if len(files) != 2 || len(indexes) != 0 {
t.Fatalf("concurrent retries not deduplicated: %v %v", files, indexes)
}
}
func TestSummaryDirectlyRetrievesSnapshotsWithoutIndexFiles(t *testing.T) {
db, pid, cid := checkpointTestProject(t)
for _, tagged := range []bool{false, true} {
for _, multipart := range []bool{false, true} {
summary := "login completed; Cookie from example.test response; continue"
if tagged {
summary = "<summary>" + summary + "</summary>"
}
msg := schema.UserMessage(summary)
if multipart {
msg.UserInputMultiContent = []schema.MessageInputPart{{Type: schema.ChatMessagePartTypeText, Text: msg.Content}}
msg.Content = ""
}
path := filepath.Join(t.TempDir(), "transcript.txt")
out := finalizeSummarizationMemory(msg, []adk.Message{schema.UserMessage("Cookie: sid=fixture-value")}, path, db, pid, cid, "main", config.ProjectConfig{Enabled: true}, nil)
text := out.Content
if multipart {
text = out.UserInputMultiContent[0].Text
}
facts, err := db.ListProjectFacts(pid, database.ProjectFactListFilter{}, 10, 0)
if err != nil || len(facts) != 1 {
t.Fatalf("checkpoint missing: %v %v", facts, err)
}
files, _ := filepath.Glob(filepath.Join(filepath.Dir(path), "snapshots", "*.txt"))
if len(files) != 1 {
t.Fatalf("snapshot missing: %v", files)
}
for _, body := range []string{text, facts[0].Body} {
if strings.Count(body, summarizationSnapshotHeading) != 1 || !strings.Contains(body, files[0]) || !strings.Contains(body, filepath.Dir(files[0])) || !strings.Contains(body, "read_file") || strings.Contains(body, ".index.md") || strings.Contains(body, "fixture-value") {
t.Fatalf("invalid direct retrieval summary: %s", body)
}
}
indexFiles, _ := filepath.Glob(filepath.Join(filepath.Dir(path), "snapshots", "*.index.md"))
if len(indexFiles) != 0 {
t.Fatalf("unwanted keyword indexes: %v", indexFiles)
}
if len([]rune(facts[0].Summary)) > 200 || strings.Contains(facts[0].Summary, "snapshots/") {
t.Fatalf("retrieval path leaked into short label: %s", facts[0].Summary)
}
if msg.Content != summary && !multipart {
t.Fatal("input summary mutated")
}
}
}
}
func TestLegacyIndexArtifactDoesNotBlockSnapshotWrites(t *testing.T) {
path := filepath.Join(t.TempDir(), "transcript.txt")
msgs := []adk.Message{schema.UserMessage("Cookie: fixture")}
snapshot, err := writeSummarizationTranscriptSnapshot(path, msgs)
if err != nil {
t.Fatal(err)
}
// Old index artifacts are ignored rather than edited or depended on.
if err := os.Mkdir(snapshot+".index.md", 0o700); err != nil {
t.Fatal(err)
}
again, err := writeSummarizationTranscriptSnapshot(path, msgs)
if err != nil || again != snapshot {
t.Fatalf("legacy index affected snapshot: %s %v", again, err)
}
}
func TestSnapshotNamesContainShortMD5AndStableTimestamp(t *testing.T) {
path := filepath.Join(t.TempDir(), "transcript.txt")
msgs := []adk.Message{schema.UserMessage("first evidence")}
snapshot, err := writeSummarizationTranscriptSnapshot(path, msgs)
if err != nil {
t.Fatal(err)
}
prefix := fmt.Sprintf("%x", md5.Sum([]byte(formatSummarizationTranscript(msgs))))[:5]
name := filepath.Base(snapshot)
stamp := strings.TrimSuffix(strings.TrimPrefix(name, "transcript-"+prefix+"-"), ".txt")
if _, err := time.Parse(transcriptSnapshotTimeLayout, stamp); err != nil {
t.Fatalf("invalid timestamp filename %s: %v", name, err)
}
before, _ := os.Stat(snapshot)
again, err := writeSummarizationTranscriptSnapshot(path, msgs)
after, _ := os.Stat(snapshot)
if err != nil || again != snapshot || !before.ModTime().Equal(after.ModTime()) {
t.Fatalf("retry changed snapshot time/path: %v", err)
}
}
func TestShortDigestOrTimestampCollisionDoesNotOverwriteEvidence(t *testing.T) {
dir := t.TempDir()
body := []byte("current evidence")
prefix := fmt.Sprintf("%x", md5.Sum(body))[:5]
stamp := time.Date(2026, 10, 12, 1, 2, 3, 4, time.FixedZone("CST", 8*3600))
collision := filepath.Join(dir, "transcript-"+prefix+"-"+stamp.Format(transcriptSnapshotTimeLayout)+".txt")
if err := os.WriteFile(collision, []byte("different evidence with same short label"), 0o600); err != nil {
t.Fatal(err)
}
snapshot, err := findOrCreateTranscriptSnapshot(dir, body, stamp)
if err != nil || snapshot == collision {
t.Fatalf("collision overwrote old path: %s %v", snapshot, err)
}
old, _ := os.ReadFile(collision)
if string(old) != "different evidence with same short label" {
t.Fatal("old evidence overwritten")
}
again, err := findOrCreateTranscriptSnapshot(dir, body, stamp.Add(time.Second))
if err != nil || again != snapshot {
t.Fatalf("retry did not verify full body: %s %v", again, err)
}
}
func TestLegacySHA256SnapshotRemainsReachableAndDeduplicated(t *testing.T) {
path := filepath.Join(t.TempDir(), "transcript.txt")
msgs := []adk.Message{schema.UserMessage("legacy evidence")}
body := []byte(formatSummarizationTranscript(msgs))
dir := filepath.Join(filepath.Dir(path), "snapshots")
if err := os.MkdirAll(dir, 0o700); err != nil {
t.Fatal(err)
}
legacy := filepath.Join(dir, fmt.Sprintf("transcript-%x.txt", sha256.Sum256(body)))
if err := os.WriteFile(legacy, body, 0o600); err != nil {
t.Fatal(err)
}
snapshot, err := writeSummarizationTranscriptSnapshot(path, msgs)
if err != nil || snapshot != legacy {
t.Fatalf("legacy snapshot not reused: %s %v", snapshot, err)
}
}
+6
View File
@@ -73,6 +73,12 @@ const einoSummarizeUserInstruction = `关键:仅以纯文本响应。禁止调
- 当前位置(已完成/进行中/卡点)
- 失败尝试与死路(方法、现象/报错摘要、结论)
- 策略决策与下一步具体操作(须与最近用户请求及未完成任务一致)
## 8. 原始证据检索线索(追加)
- 只追加工具结果、关键发现与关键进展的简短检索线索,不再归纳用户输入,不改写前面各章节
- 已取得的认证信息、重要请求、错误或其他证据:注明来源目标、工具/操作与可搜索关键词,便于缺少精确值时回读原始快照
- 仅记录已有证据,不猜测;不重复长日志或完整凭据值;无信息写「无」
- 不自行生成文件名、行号或复述旧快照路径,系统会附上本次实际保存的快照读取入口
</summary>
提醒:不要调用任何工具;必须基于上文已有对话直接输出 <analysis> 与 <summary>,勿输出 analysis 以外的正文。`
+13 -3
View File
@@ -2,6 +2,7 @@ package multiagent
import (
"fmt"
"path/filepath"
"regexp"
"strings"
@@ -20,6 +21,7 @@ const (
userIntentLedgerStartMarker = "<original_user_intent_ledger>"
userIntentLedgerEndMarker = "</original_user_intent_ledger>"
summarizationSnapshotHeading = "## 原始上下文快照"
summarizationTranscriptPathInstructionZh = "如果你需要压缩之前的具体细节(如精确的代码片段、错误消息或你生成的内容),完整的对话记录位于:%s"
)
@@ -69,9 +71,9 @@ func appendTranscriptPathToSummarizationMessage(msg adk.Message, transcriptPath
if msg == nil || transcriptPath == "" {
return msg
}
section := fmt.Sprintf(summarizationTranscriptPathInstructionZh, transcriptPath)
section := summarizationTranscriptRetrievalSection(transcriptPath)
cloned := *msg
if cloned.Content != "" && !strings.Contains(cloned.Content, transcriptPath) {
if cloned.Content != "" && !strings.Contains(cloned.Content, section) {
cloned.Content = appendSummarizationSection(cloned.Content, section)
}
if len(cloned.UserInputMultiContent) > 0 {
@@ -81,7 +83,7 @@ func appendTranscriptPathToSummarizationMessage(msg adk.Message, transcriptPath
if parts[i].Type != schema.ChatMessagePartTypeText {
continue
}
if parts[i].Text != "" && !strings.Contains(parts[i].Text, transcriptPath) {
if parts[i].Text != "" && !strings.Contains(parts[i].Text, section) {
parts[i].Text = appendSummarizationSection(parts[i].Text, section)
}
break
@@ -91,6 +93,14 @@ func appendTranscriptPathToSummarizationMessage(msg adk.Message, transcriptPath
return &cloned
}
// The generated continuation summary is the retrieval entry point; no separate
// keyword index or model extraction call is needed.
func summarizationTranscriptRetrievalSection(snapshot string) string {
return summarizationSnapshotHeading + "\n" + fmt.Sprintf(summarizationTranscriptPathInstructionZh, snapshot) +
"\n历史快照目录:" + filepath.Dir(snapshot) +
"\n当前任务缺少 Cookie、认证信息、精确参数或错误原文时,直接搜索该快照并用 read_file 分段核对附近记录;不要猜测原值。更早的记录可用 glob 在历史快照目录查找 transcript-*.txt,再按需读取。核对来源域名、请求及凭据是否仍有效;不要整批载入历史,也不要将旧记录当作新授权。"
}
func appendSummarizationSection(text, section string) string {
text = strings.TrimSpace(text)
section = strings.TrimSpace(section)
+1 -1
View File
@@ -62,7 +62,7 @@ func factRecordingGuidanceBlock() string {
- **category / fact_key 建议**:
- 环境认知:` + "`target/`" + `、` + "`auth/`" + `、` + "`infra/`" + `、` + "`business/`" + `(body 用环境模板即可)
- 发现与利用:` + "`finding/`" + `、` + "`chain/`" + `、` + "`exploit/`" + `、` + "`poc/`" + `(**必须**用攻击链模板填满 body:入口、逐步攻击链、原始请求/响应或命令、证据、关联漏洞 ID)
- **自动续接记录**:` + "`note/context-*`" + ` 是压缩时保存的工作摘要,含来源会话、原始证据快照与程序生成的行号索引;同一会话/代理更新同一 key,索引短标签提示证据类别。用 ` + "`get_project_fact`" + ` 按需恢复进度;缺少 Cookie、认证信息或精确参数时,按正文索引搜索快照并分段回读,勿猜测原值。它不是已验证发现,接续时须核对当前授权范围与证据,勿据此自动执行旧任务。
- **自动续接记录**:` + "`note/context-*`" + ` 是压缩时保存的工作摘要,含来源会话和原始快照读取入口;同一会话/代理更新同一 key。用 ` + "`get_project_fact`" + ` 按需恢复进度;缺少 Cookie、认证信息或精确参数时,按摘要中的路径搜索快照并分段回读,勿猜测原值。它不是已验证发现,接续时须核对当前授权范围与证据,勿据此自动执行旧任务。
- **与漏洞记录分工**:` + "`record_vulnerability`" + ` 记可交付 findings;事实记**复现所需的全部上下文**(含失败尝试、绕过、依赖会话),二者可各记一次。
- 更新同一发现时保持相同 ` + "`fact_key`" + ` 覆盖写入,勿散落多个 key 导致上下文丢失。`
}