diff --git a/internal/context/microcompact.go b/internal/context/microcompact.go index 1e50a7fc..37993eb2 100644 --- a/internal/context/microcompact.go +++ b/internal/context/microcompact.go @@ -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 } @@ -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 { diff --git a/internal/context/microcompact_summarizer_test.go b/internal/context/microcompact_summarizer_test.go index d2f26553..12ffc2e9 100644 --- a/internal/context/microcompact_summarizer_test.go +++ b/internal/context/microcompact_summarizer_test.go @@ -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") + } + }) +} diff --git a/internal/tools/micro_compact_summarizer_test.go b/internal/tools/micro_compact_summarizer_test.go index 90ee6d7b..d01901e6 100644 --- a/internal/tools/micro_compact_summarizer_test.go +++ b/internal/tools/micro_compact_summarizer_test.go @@ -2,6 +2,7 @@ package tools import ( "strings" + "sync" "testing" "unicode/utf8" ) @@ -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() diff --git a/internal/tools/micro_compact_summarizers_builtin.go b/internal/tools/micro_compact_summarizers_builtin.go index dfdf271b..d18481ae 100644 --- a/internal/tools/micro_compact_summarizers_builtin.go +++ b/internal/tools/micro_compact_summarizers_builtin.go @@ -22,7 +22,7 @@ var builtinSummarizers = []builtinSummarizerRegistration{ } // RegisterBuiltinSummarizers 将所有内置工具的内容摘要器注册到 Registry。 -// 应在所有工具注册完成后调用一次。 +// 建议在启动装配阶段调用；可重复调用并覆盖同名摘要器。 func RegisterBuiltinSummarizers(registry *Registry) { if registry == nil { return diff --git a/internal/tools/registry.go b/internal/tools/registry.go index dd7e4331..45a1e3fd 100644 --- a/internal/tools/registry.go +++ b/internal/tools/registry.go @@ -5,6 +5,7 @@ import ( "errors" "sort" "strings" + "sync" providertypes "neo-code/internal/provider/types" "neo-code/internal/security" @@ -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 @@ -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 @@ -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))]