Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 17 additions & 4 deletions internal/context/microcompact.go
Original file line number Diff line number Diff line change
Expand Up @@ -50,12 +50,15 @@ func microCompactMessagesWithPolicies(messages []providertypes.Message, policies
if len(compactableIDs) == 0 {
continue
}
compactableContents := compactableToolMessageContents(cloned, span, compactableIDs)
if len(compactableContents) == 0 {
if retainedCompactableSpans < retainedToolSpans {
if hasCompactableToolMessage(cloned, span, compactableIDs) {
retainedCompactableSpans++
}
continue
}
if retainedCompactableSpans < retainedToolSpans {
retainedCompactableSpans++

compactableContents := compactableToolMessageContents(cloned, span, compactableIDs)
if len(compactableContents) == 0 {
continue
}

Expand Down Expand Up @@ -152,6 +155,16 @@ func compactableToolMessageContents(messages []providertypes.Message, span inter
return contents
}

// hasCompactableToolMessage 判断工具块中是否存在至少一条可压缩的工具消息。
func hasCompactableToolMessage(messages []providertypes.Message, span internalcompact.MessageSpan, compactableIDs map[string]struct{}) bool {
for messageIndex := span.Start + 1; messageIndex < span.End; messageIndex++ {
if _, ok := compactableToolMessageContent(messages[messageIndex], compactableIDs); ok {
return true
}
}
return false
}

// compactableToolMessageContent 判断 tool 消息是否可压缩,并返回渲染后的内容文本。
func compactableToolMessageContent(message providertypes.Message, compactableIDs map[string]struct{}) (string, bool) {
if message.Role != providertypes.RoleTool || message.IsError {
Expand Down
30 changes: 30 additions & 0 deletions internal/context/microcompact_summarizer_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -347,3 +347,33 @@ func TestCompactableToolCallIDsEmptyInput(t *testing.T) {
t.Fatalf("expected nil maps for empty input, got ids=%v names=%v", ids, names)
}
}

// TestHasCompactableToolMessage 验证工具块可压缩消息探测逻辑。
func TestHasCompactableToolMessage(t *testing.T) {
t.Parallel()

span := internalcompact.MessageSpan{Start: 0, End: 3}
ids := map[string]struct{}{"call-1": {}}

t.Run("true_when_matching_tool_message_exists", func(t *testing.T) {
messages := []providertypes.Message{
{Role: providertypes.RoleAssistant, ToolCalls: []providertypes.ToolCall{{ID: "call-1", Name: "bash"}}},
{Role: providertypes.RoleTool, ToolCallID: "call-1", Parts: []providertypes.ContentPart{providertypes.NewTextPart("output")}},
{Role: providertypes.RoleUser, Parts: []providertypes.ContentPart{providertypes.NewTextPart("u")}},
}
if !hasCompactableToolMessage(messages, span, ids) {
t.Fatal("expected compactable tool message to be found")
}
})

t.Run("false_when_tool_messages_are_not_compactable", func(t *testing.T) {
messages := []providertypes.Message{
{Role: providertypes.RoleAssistant, ToolCalls: []providertypes.ToolCall{{ID: "call-1", Name: "bash"}}},
{Role: providertypes.RoleTool, ToolCallID: "call-1", IsError: true, Parts: []providertypes.ContentPart{providertypes.NewTextPart("error")}},
{Role: providertypes.RoleTool, ToolCallID: "call-2", Parts: []providertypes.ContentPart{providertypes.NewTextPart("other")}},
}
if hasCompactableToolMessage(messages, span, ids) {
t.Fatal("expected no compactable tool message")
}
})
}
31 changes: 31 additions & 0 deletions internal/tools/micro_compact_summarizer_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package tools

import (
"strings"
"sync"
"testing"
"unicode/utf8"
)
Expand Down Expand Up @@ -354,6 +355,36 @@ func TestRegisterSummarizerNilRegistry(t *testing.T) {
t.Fatal("expected nil summarizer on nil registry")
}
}

func TestRegisterSummarizerConcurrentAccess(t *testing.T) {
t.Parallel()

registry := NewRegistry()
var wg sync.WaitGroup

for i := 0; i < 8; i++ {
wg.Add(1)
go func() {
defer wg.Done()
for j := 0; j < 200; j++ {
if j%3 == 0 {
registry.RegisterSummarizer("concurrent_tool", nil)
continue
}
registry.RegisterSummarizer("concurrent_tool", func(content string, metadata map[string]string, isError bool) string {
return "worker"
})
s := registry.MicroCompactSummarizer("concurrent_tool")
if s != nil {
_ = s("content", nil, false)
}
}
}()
}

wg.Wait()
}

func TestTruncateRunes(t *testing.T) {
t.Parallel()

Expand Down
2 changes: 1 addition & 1 deletion internal/tools/micro_compact_summarizers_builtin.go
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ var builtinSummarizers = []builtinSummarizerRegistration{
}

// RegisterBuiltinSummarizers 将所有内置工具的内容摘要器注册到 Registry。
// 应在所有工具注册完成后调用一次
// 建议在启动装配阶段调用;可重复调用并覆盖同名摘要器
func RegisterBuiltinSummarizers(registry *Registry) {
if registry == nil {
return
Expand Down
11 changes: 10 additions & 1 deletion internal/tools/registry.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import (
"errors"
"sort"
"strings"
"sync"

providertypes "neo-code/internal/provider/types"
"neo-code/internal/security"
Expand All @@ -15,6 +16,7 @@ type Registry struct {
tools map[string]Tool
microCompactPolicies map[string]MicroCompactPolicy
microCompactSummarizers map[string]ContentSummarizer
microCompactSummaryMu sync.RWMutex
mcpRegistry *mcp.Registry
mcpFactory *mcp.AdapterFactory
mcpExposureFilter mcp.ExposureFilter
Expand Down Expand Up @@ -110,6 +112,8 @@ func (r *Registry) RegisterSummarizer(toolName string, summarizer ContentSummari
return
}
name := strings.ToLower(strings.TrimSpace(toolName))
r.microCompactSummaryMu.Lock()
defer r.microCompactSummaryMu.Unlock()
if summarizer == nil {
delete(r.microCompactSummarizers, name)
return
Expand All @@ -119,7 +123,12 @@ func (r *Registry) RegisterSummarizer(toolName string, summarizer ContentSummari

// MicroCompactSummarizer 返回指定工具的内容摘要器;无注册时返回 nil。
func (r *Registry) MicroCompactSummarizer(name string) ContentSummarizer {
if r == nil || r.microCompactSummarizers == nil {
if r == nil {
return nil
}
r.microCompactSummaryMu.RLock()
defer r.microCompactSummaryMu.RUnlock()
if r.microCompactSummarizers == nil {
return nil
}
return r.microCompactSummarizers[strings.ToLower(strings.TrimSpace(name))]
Expand Down
Loading