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
52 changes: 45 additions & 7 deletions internal/memo/service.go
Original file line number Diff line number Diff line change
Expand Up @@ -265,7 +265,9 @@ func (s *Service) saveEntryLocked(ctx context.Context, entry Entry) error {
for _, removed := range removedEntries {
s.removeAutoExtractTopicLocked(scope, removed.TopicFile)
}
s.trackAutoExtractEntryLocked(scope, entry)
if indexContainsEntryID(working, entry.ID) {
s.trackAutoExtractEntryLocked(scope, entry)
}
}

s.invalidateCache()
Expand Down Expand Up @@ -406,9 +408,7 @@ func (s *Service) listLocked(ctx context.Context, scope Scope) ([]Entry, error)
}
results = append(results, index.Entries...)
}
cloned := make([]Entry, len(results))
copy(cloned, results)
return cloned, nil
return append([]Entry(nil), results...), nil
}

// invalidateCache 触发上下文源缓存失效回调。
Expand Down Expand Up @@ -488,13 +488,51 @@ func trimIndexEntries(index *Index, maxEntries int, maxIndexBytes int) []Entry {
removed = append(removed, index.Entries[0])
index.Entries = index.Entries[1:]
}
for maxIndexBytes > 0 && len(index.Entries) > 0 && len(RenderIndex(index)) > maxIndexBytes {
removed = append(removed, index.Entries[0])
index.Entries = index.Entries[1:]
if maxIndexBytes > 0 && len(index.Entries) > 0 {
removed = append(removed, trimIndexEntriesByBytes(index, maxIndexBytes)...)
}
return removed
}

// trimIndexEntriesByBytes 在索引超过字节阈值时,通过二分定位最小移除数量并返回被删除条目。
func trimIndexEntriesByBytes(index *Index, maxIndexBytes int) []Entry {
if index == nil || len(index.Entries) == 0 || maxIndexBytes <= 0 {
return nil
}
if len(RenderIndex(index)) <= maxIndexBytes {
return nil
}

entries := index.Entries
lo, hi := 0, len(entries)
for lo < hi {
mid := lo + (hi-lo)/2
candidate := &Index{Entries: entries[mid:], UpdatedAt: index.UpdatedAt}
if len(RenderIndex(candidate)) > maxIndexBytes {
lo = mid + 1
continue
}
hi = mid
}

removed := append([]Entry(nil), entries[:lo]...)
index.Entries = entries[lo:]
return removed
}

// indexContainsEntryID 判断索引中是否仍保留目标 ID,用于避免为已裁剪条目建立去重索引。
func indexContainsEntryID(index *Index, entryID string) bool {
if index == nil || strings.TrimSpace(entryID) == "" {
return false
}
for _, item := range index.Entries {
if item.ID == entryID {
return true
}
}
return false
}

// scopesForQuery 将查询范围展开为实际存储分层列表。
func scopesForQuery(scope Scope) []Scope {
switch NormalizeScope(scope) {
Expand Down
77 changes: 65 additions & 12 deletions internal/memo/service_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -188,12 +188,9 @@ func TestServiceRemoveRejectsInvalidScope(t *testing.T) {
}

func TestServiceMaxEntriesTrim(t *testing.T) {
svc := NewService(newMemoryTestStore(), config.MemoConfig{
MaxEntries: 2,
MaxIndexBytes: 16 * 1024,
ExtractTimeoutSec: 15,
ExtractRecentMessages: 10,
}, nil)
cfg := testMemoConfig()
cfg.MaxEntries = 2
svc := NewService(newMemoryTestStore(), cfg, nil)

for _, title := range []string{"A", "B", "C"} {
if err := svc.Add(context.Background(), Entry{
Expand All @@ -216,12 +213,10 @@ func TestServiceMaxEntriesTrim(t *testing.T) {
}

func TestServiceMaxIndexBytesTrim(t *testing.T) {
svc := NewService(newMemoryTestStore(), config.MemoConfig{
MaxEntries: 10,
MaxIndexBytes: 40,
ExtractTimeoutSec: 15,
ExtractRecentMessages: 10,
}, nil)
cfg := testMemoConfig()
cfg.MaxEntries = 10
cfg.MaxIndexBytes = 40
svc := NewService(newMemoryTestStore(), cfg, nil)

for _, title := range []string{"one", "two", "three"} {
if err := svc.Add(context.Background(), Entry{
Expand Down Expand Up @@ -303,6 +298,43 @@ func TestServiceAutoExtractDedupAcrossScopes(t *testing.T) {
}
}

func TestServiceAutoExtractTrimmedEntryDoesNotPolluteDedupIndex(t *testing.T) {
svc := NewService(newMemoryTestStore(), config.MemoConfig{
MaxEntries: 10,
MaxIndexBytes: 1,
ExtractTimeoutSec: 15,
ExtractRecentMessages: 10,
}, nil)
entry := Entry{
Type: TypeUser,
Title: "reply in chinese",
Content: "reply in chinese",
Source: SourceAutoExtract,
}

added, err := svc.addAutoExtractIfAbsent(context.Background(), entry)
if err != nil || !added {
t.Fatalf("first addAutoExtractIfAbsent() = (%v, %v), want (true, nil)", added, err)
}
added, err = svc.addAutoExtractIfAbsent(context.Background(), entry)
if err != nil || !added {
t.Fatalf("second addAutoExtractIfAbsent() = (%v, %v), want (true, nil)", added, err)
}

entries, err := svc.List(context.Background(), ScopeUser)
if err != nil {
t.Fatalf("List() error = %v", err)
}
if len(entries) != 0 {
t.Fatalf("len(entries) = %d, want 0 after byte trim", len(entries))
}

key := autoExtractDedupKey(entry)
if refs := svc.autoExtractKeyRefs[key]; refs != 0 {
t.Fatalf("autoExtractKeyRefs[%q] = %d, want 0", key, refs)
}
}

func TestServiceEnsureAutoExtractIndexLoadsExistingEntries(t *testing.T) {
store := newMemoryTestStore()
store.indexes[ScopeUser] = &Index{Entries: []Entry{{Type: TypeUser, Title: "A", TopicFile: "a.md"}}}
Expand Down Expand Up @@ -342,3 +374,24 @@ func TestMatchesKeywordIncludesContent(t *testing.T) {
t.Fatal("unexpected match for missing keyword")
}
}

func TestTrimIndexEntriesByBytesRemovesMinimalPrefix(t *testing.T) {
index := &Index{
Entries: []Entry{
{Type: TypeUser, Title: "one", TopicFile: "one.md"},
{Type: TypeUser, Title: "two", TopicFile: "two.md"},
{Type: TypeUser, Title: "three", TopicFile: "three.md"},
},
}

target := &Index{Entries: append([]Entry(nil), index.Entries[1:]...)}
maxIndexBytes := len(RenderIndex(target))
removed := trimIndexEntries(index, 10, maxIndexBytes)

if len(removed) != 1 || removed[0].Title != "one" {
t.Fatalf("removed = %#v, want only first entry", removed)
}
if len(index.Entries) != 2 || index.Entries[0].Title != "two" || index.Entries[1].Title != "three" {
t.Fatalf("remaining entries = %#v", index.Entries)
}
}
6 changes: 3 additions & 3 deletions internal/tools/memo/recall.go
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,7 @@ func (t *RecallTool) Schema() map[string]any {
"properties": map[string]any{
"keyword": map[string]any{
"type": "string",
"description": "Search keyword to find matching memory entries (searches title, type, and keywords).",
"description": "Search keyword to find matching memory entries (searches title, type, content, and keywords).",
},
"scope": map[string]any{
"type": "string",
Expand All @@ -69,8 +69,8 @@ func (t *RecallTool) MicroCompactPolicy() tools.MicroCompactPolicy {
func (t *RecallTool) Execute(ctx context.Context, call tools.ToolCallInput) (tools.ToolResult, error) {
var args recallInput
if err := json.Unmarshal(call.Arguments, &args); err != nil {
err = fmt.Errorf("%s: %w", recallToolName, err)
return tools.NewErrorResult(recallToolName, "invalid arguments", err.Error(), nil), err
wrappedErr := fmt.Errorf("%s: %w", recallToolName, err)
return tools.NewErrorResult(recallToolName, "invalid arguments", wrappedErr.Error(), nil), wrappedErr
}

args.Keyword = strings.TrimSpace(args.Keyword)
Expand Down
Loading