Merge pull request #206 from itmisx/feature/tool-output-reclaim

itmisx · web-flow · commit be00b6143d6b · 2026-07-23T18:48:13.000+09:00
⚡️ tool output reclaim
diff --git a/agent/llm.go b/agent/llm.go
@@ -613,6 +613,24 @@ func StartStream(
 				return
 			}
 
+			// 工具输出回收(reclaim,issue #201):比 LLM 摘要便宜的第一层,且不受 inLoopCompactOff 影响。
+			// 按 token 预算(窗口比例)把"较旧的工具结果内容"换成引用标记 —— 不删消息、不动结构、配对不坏。
+			// 度量按 token 而非 user 轮,因此能伸进最近 2 个 user 轮内部,回收"单个长任务轮里早几轮的大输出"
+			// (按 user 轮切的摘要够不着的部分,正是 #201 卡 73%/400 的根)。触发用"实时估算"(含本轮刚
+			// append 的工具结果),消除 lastPromptTokens 慢一拍导致的单轮越窗。与现有 clampToolOutput(96KB)/
+			// clampTurnToolOutput(256KB)暂并存,验证稳后再摘旧卡口。
+			if ctxWin := currentEntry.ContextWindow; ctxWin > 0 && len(convo) > 0 {
+				curEst := 0
+				for i := range convo {
+					curEst += MsgTokens(convo[i]) // 实时估算(与 clampMaxTokens 同口径);可后续复用其结果省一遍
+				}
+				if max(lastPromptTokens, curEst) >= ctxWin*compactTriggerPct/100 {
+					if reclaimToolOutputs(convo, ctxWin) {
+						ch <- HistoryUpdateMsg{History: convo}
+					}
+				}
+			}
+
 			// 循环内 auto-compact:上一轮 prompt 接近上下文窗口就先压缩历史腾空间再继续(不熔断停)。
 			// 对标 Claude Code:压缩 convo[1:] 成摘要、重建 [system(新摘要)]+尾部,新摘要经 CompactedMsg
 			// 回传 TUI 存 session(否则被剥的 system 摘要会丢失),history 截断经 HistoryUpdateMsg 同步。
@@ -1668,3 +1686,83 @@ func elideWriteContent(argsJSON string) string {
 	}
 	return strings.TrimRight(buf.String(), "\n") // Encode 会补一个换行,去掉
 }
+
+// --- 工具输出回收(reclaim,issue #201)---
+//
+// 问题:deepx 的历史压缩按 user 轮切,保护最近 keepRecentTurns 个 user 轮。但"跑测试 / 长任务"
+// 常是一个 user 消息 + 几十轮工具调用,bloat 全在这一个被保护的 user 轮内部 —— 按 user 轮切的摘要
+// 够不着,于是卡在高位"无需压缩"甚至 400(见 issue #201;RunCompression 的 keepStart≤0 分支)。
+//
+// reclaim 独立于 user 轮摘要,按 token 预算回收"较旧的工具结果内容":保住最近 reclaimKeepPct 预算
+// + 最近 reclaimMinTailToolMsgs 条工具结果,更旧的把 Content 换成引用标记(壳留、配对不坏、可重取)。
+// 因为度量按 token 而非 user 轮,它能伸进被保护的 2 轮内部,回收轮内早几轮的大输出。无 LLM,也不受
+// inLoopCompactOff 影响 —— 摘要超时/拒绝时它照样能减负。
+
+const (
+	// reclaimKeepPct:保住"最近工具输出"的 token 预算,取上下文窗口的这个百分比。超出的更旧工具输出被回收。
+	reclaimKeepPct = 20
+	// reclaimMinTailToolMsgs:无论预算多小,至少保护最近这么多条工具结果不动(模型正在用)。
+	reclaimMinTailToolMsgs = 4
+	// reclaimMarkerPrefix:被回收工具结果的 Content 前缀,兼作幂等判据(再次 reclaim 时跳过、不重复计数)。
+	reclaimMarkerPrefix = "[已回收] "
+)
+
+// reclaimToolOutputs 就地把"较旧的工具结果内容"换成引用标记以回收上下文,返回是否发生改动。
+// 只改 Role=="tool" 消息的 Content;user/assistant 消息、工具调用骨架、ToolCallID 配对一律不动
+// (故 sanitizeToolPairs 安全、发给 API 的配对完整)。
+// 保护:最近 reclaimMinTailToolMsgs 条工具结果 + 最近 reclaimKeepPct 预算内的工具输出,留全。幂等。
+func reclaimToolOutputs(convo []ChatMessage, ctxWin int) bool {
+	if ctxWin <= 0 {
+		ctxWin = 65536
+	}
+	keepBudget := ctxWin * reclaimKeepPct / 100
+
+	// 预建 tool_call_id → path,避免每条被回收消息各做一次 O(n) 反查配对 assistant。
+	callPath := map[string]string{}
+	for _, m := range convo {
+		if m.Role != "assistant" {
+			continue
+		}
+		for _, tc := range m.ToolCalls {
+			var args map[string]any
+			if json.Unmarshal([]byte(tc.Function.Arguments), &args) == nil {
+				if p, ok := args["path"].(string); ok {
+					callPath[tc.ID] = p
+				}
+			}
+		}
+	}
+
+	changed := false
+	toolSeen := 0   // 已遍历到的(未回收)工具结果条数,用于"最近 N 条"保护
+	keptTokens := 0 // 已保留的工具输出 token 累计,用于预算保护
+	for i := len(convo) - 1; i >= 0; i-- {
+		if convo[i].Role != "tool" {
+			continue
+		}
+		if strings.HasPrefix(convo[i].Content, reclaimMarkerPrefix) {
+			continue // 已回收:不计 tail、不计预算、不重复处理(幂等)
+		}
+		toolSeen++
+		if toolSeen <= reclaimMinTailToolMsgs || keptTokens < keepBudget {
+			keptTokens += MsgTokens(convo[i]) // 保护区:留全,并计入预算
+			continue
+		}
+		convo[i].Content = toolOutputReference(convo[i].Name, callPath[convo[i].ToolCallID])
+		changed = true
+	}
+	return changed
+}
+
+// toolOutputReference 生成被回收工具结果的引用标记:带工具名与 path(有则),
+// 让模型知道"这里原是什么调用的输出、需要时重新调用取回"。
+func toolOutputReference(name, path string) string {
+	label := name
+	if label == "" {
+		label = "工具"
+	}
+	if path != "" {
+		label += " " + path
+	}
+	return reclaimMarkerPrefix + label + " 的旧输出已省略以回收上下文,需要时请重新调用获取。"
+}
diff --git a/agent/reclaim_test.go b/agent/reclaim_test.go
@@ -0,0 +1,158 @@
+package agent
+
+import (
+	"fmt"
+	"strings"
+	"testing"
+)
+
+func toolMsg(id, name, content string) ChatMessage {
+	return ChatMessage{Role: "tool", ToolCallID: id, Name: name, Content: content}
+}
+
+func asstCall(id, name, argsJSON string) ChatMessage {
+	return ChatMessage{Role: "assistant", ToolCalls: []ToolCall{
+		{ID: id, Type: "function", Function: ToolCallFunc{Name: name, Arguments: argsJSON}},
+	}}
+}
+
+// buildConvo 造 1 system + 1 user + n 轮 (assistant call + tool result),每条工具结果内容为 content。
+func buildConvo(n int, name, content string) []ChatMessage {
+	convo := []ChatMessage{{Role: "system", Content: "sys"}, {Role: "user", Content: "任务"}}
+	for k := 0; k < n; k++ {
+		id := fmt.Sprintf("c%d", k)
+		convo = append(convo, asstCall(id, name, fmt.Sprintf(`{"path":"src/f%d.go"}`, k)))
+		convo = append(convo, toolMsg(id, name, content))
+	}
+	return convo
+}
+
+func countTool(convo []ChatMessage) (reclaimed, kept int) {
+	for _, m := range convo {
+		if m.Role != "tool" {
+			continue
+		}
+		if strings.HasPrefix(m.Content, reclaimMarkerPrefix) {
+			reclaimed++
+		} else {
+			kept++
+		}
+	}
+	return
+}
+
+func TestReclaim_OldReplacedRecentKept(t *testing.T) {
+	big := strings.Repeat("这是一大段读取内容。", 500) // 每条远超小预算
+	convo := buildConvo(10, "Read", big)
+	before := len(convo)
+
+	if !reclaimToolOutputs(convo, 4096) { // 预算 = 4096*20% ≈ 819 token
+		t.Fatal("应发生回收")
+	}
+	reclaimed, kept := countTool(convo)
+	if reclaimed == 0 {
+		t.Fatal("较旧工具结果应被回收")
+	}
+	if kept < reclaimMinTailToolMsgs {
+		t.Fatalf("最近 %d 条应保留,实际保留 %d", reclaimMinTailToolMsgs, kept)
+	}
+	if len(convo) != before {
+		t.Fatalf("reclaim 不应增删消息: %d -> %d", before, len(convo))
+	}
+}
+
+func TestReclaim_TailProtected(t *testing.T) {
+	convo := buildConvo(reclaimMinTailToolMsgs+3, "Bash", strings.Repeat("x", 5000))
+	reclaimToolOutputs(convo, 1024) // 预算极小
+
+	seen := 0
+	for i := len(convo) - 1; i >= 0; i-- {
+		if convo[i].Role != "tool" {
+			continue
+		}
+		seen++
+		if seen <= reclaimMinTailToolMsgs && strings.HasPrefix(convo[i].Content, reclaimMarkerPrefix) {
+			t.Fatalf("最近第 %d 条工具结果被误回收", seen)
+		}
+	}
+}
+
+func TestReclaim_ReferenceHasNameAndPath(t *testing.T) {
+	convo := buildConvo(8, "Read", strings.Repeat("y", 8000))
+	reclaimToolOutputs(convo, 2048)
+
+	found := false
+	for _, m := range convo {
+		if m.Role == "tool" && strings.HasPrefix(m.Content, reclaimMarkerPrefix) {
+			if !strings.Contains(m.Content, "Read") || !strings.Contains(m.Content, "src/f") {
+				t.Fatalf("引用应含工具名和 path, got=%q", m.Content)
+			}
+			found = true
+			break
+		}
+	}
+	if !found {
+		t.Fatal("应有被回收的工具结果")
+	}
+}
+
+func TestReclaim_Idempotent(t *testing.T) {
+	convo := buildConvo(8, "Read", strings.Repeat("z", 6000))
+	if !reclaimToolOutputs(convo, 2048) {
+		t.Fatal("首次应有改动")
+	}
+	if reclaimToolOutputs(convo, 2048) {
+		t.Fatal("再次 reclaim 不应有新改动(幂等)")
+	}
+}
+
+func TestReclaim_NonToolUntouched(t *testing.T) {
+	userBig := strings.Repeat("u", 9000)
+	asstBig := strings.Repeat("a", 9000)
+	convo := []ChatMessage{
+		{Role: "system", Content: "s"},
+		{Role: "user", Content: userBig},
+		asstCall("c1", "Read", `{"path":"a.go"}`),
+		{Role: "assistant", Content: asstBig},
+	}
+	before := len(convo)
+	reclaimToolOutputs(convo, 1024)
+	if convo[1].Content != userBig {
+		t.Fatal("user 消息不应被改")
+	}
+	if convo[3].Content != asstBig {
+		t.Fatal("assistant 文本不应被改")
+	}
+	if len(convo) != before {
+		t.Fatal("消息数不应变")
+	}
+}
+
+func TestReclaim_UnderBudgetNoChange(t *testing.T) {
+	convo := buildConvo(6, "Read", "tiny") // 工具输出很小
+	if reclaimToolOutputs(convo, 1_000_000) {
+		t.Fatal("预算充足时不应回收")
+	}
+}
+
+func TestReclaim_PairingPreserved(t *testing.T) {
+	// 回收后每条 tool 消息的 ToolCallID / Name 不变,配对完整。
+	convo := buildConvo(10, "Read", strings.Repeat("w", 6000))
+	ids := map[int]string{}
+	for i, m := range convo {
+		if m.Role == "tool" {
+			ids[i] = m.ToolCallID
+		}
+	}
+	reclaimToolOutputs(convo, 2048)
+	for i, m := range convo {
+		if m.Role == "tool" {
+			if m.ToolCallID != ids[i] {
+				t.Fatalf("第 %d 条 ToolCallID 被改动", i)
+			}
+			if m.Name == "" {
+				t.Fatalf("第 %d 条 Name 丢失", i)
+			}
+		}
+	}
+}
