From 233be28b5ea2aaffea59c96e57d9f90d42ac1eef Mon Sep 17 00:00:00 2001 From: xgopilot Date: Mon, 13 Apr 2026 00:54:18 +0000 Subject: [PATCH 1/6] fix(memo): address review findings and simplify extractor/runtime integration - use rune-safe title truncation for auto extracted memo entries - guard memo remember/recall tools against nil service and cap recall output - limit runtime memo extraction input and add bounded async execution - add regression tests for utf-8 truncation, nil guards, output limit, and runtime hook Generated with [codeagent](https://github.com/qbox/codeagent) Co-authored-by: Yumiue <188874804+Yumiue@users.noreply.github.com> --- internal/memo/extractor.go | 41 +++++++++++++--- internal/memo/extractor_test.go | 12 +++-- internal/runtime/runtime.go | 73 +++++++++++++++++++++++----- internal/runtime/runtime_test.go | 71 +++++++++++++++++++++++++++ internal/tools/memo/common.go | 13 +++++ internal/tools/memo/recall.go | 11 +++-- internal/tools/memo/recall_test.go | 44 +++++++++++++++++ internal/tools/memo/remember.go | 3 ++ internal/tools/memo/remember_test.go | 19 ++++++++ 9 files changed, 260 insertions(+), 27 deletions(-) create mode 100644 internal/tools/memo/common.go diff --git a/internal/memo/extractor.go b/internal/memo/extractor.go index 900c2a05..8f3490fc 100644 --- a/internal/memo/extractor.go +++ b/internal/memo/extractor.go @@ -9,13 +9,15 @@ import ( // signalPhrases 包含规则提取器识别的显式记忆信号词。 var signalPhrases = []string{ - "记住", "记下来", "以后都这样", "以后都这样", + "记住", "记下来", "以后都这样", "我偏好", "我喜欢", "我习惯", "我希望", "别再", "不要再", "不要使用", "避免", "always", "never", "prefer", "avoid", "remember", "make sure", "from now on", } +var lowerSignalPhrases = normalizeSignalPhrases(signalPhrases) + // RuleExtractor 基于规则的轻量记忆提取器，检测用户消息中的显式信号词。 // 无外部依赖，适合作为默认提取器。 type RuleExtractor struct{} @@ -48,11 +50,9 @@ func (r *RuleExtractor) Extract(ctx context.Context, messages []providertypes.Me return nil, nil } - // 截断过长内容作为标题 + // 截断过长内容作为标题（按 rune 截断，避免破坏 UTF-8）。 title := NormalizeTitle(lastUserMsg) - if len(title) > 150 { - title = title[:147] + "..." - } + title = truncateWithEllipsis(title, 150) return []Entry{ { @@ -67,14 +67,41 @@ func (r *RuleExtractor) Extract(ctx context.Context, messages []providertypes.Me // containsSignal 检查文本是否包含任意信号词。 func containsSignal(text string) bool { lower := strings.ToLower(text) - for _, phrase := range signalPhrases { - if strings.Contains(lower, strings.ToLower(phrase)) { + for _, phrase := range lowerSignalPhrases { + if strings.Contains(lower, phrase) { return true } } return false } +// normalizeSignalPhrases 将信号词标准化为小写，避免运行时重复转换。 +func normalizeSignalPhrases(phrases []string) []string { + result := make([]string, 0, len(phrases)) + for _, phrase := range phrases { + text := strings.TrimSpace(strings.ToLower(phrase)) + if text != "" { + result = append(result, text) + } + } + return result +} + +// truncateWithEllipsis 按 rune 截断字符串，并在超长时追加省略号。 +func truncateWithEllipsis(text string, maxRunes int) string { + if maxRunes <= 0 { + return "" + } + runes := []rune(text) + if len(runes) <= maxRunes { + return text + } + if maxRunes <= 3 { + return string(runes[:maxRunes]) + } + return string(runes[:maxRunes-3]) + "..." +} + // ExtractAndStore 从消息中提取记忆并保存到 Service。 // 提取失败静默处理，不影响主循环。 func ExtractAndStore(ctx context.Context, extractor Extractor, svc *Service, messages []providertypes.Message) { diff --git a/internal/memo/extractor_test.go b/internal/memo/extractor_test.go index 6f7ac809..74f1194a 100644 --- a/internal/memo/extractor_test.go +++ b/internal/memo/extractor_test.go @@ -4,6 +4,7 @@ import ( "context" "strings" "testing" + "unicode/utf8" "neo-code/internal/config" providertypes "neo-code/internal/provider/types" @@ -123,8 +124,8 @@ func TestRuleExtractorExtractEnglishSignals(t *testing.T) { func TestRuleExtractorExtractLongContent(t *testing.T) { extractor := NewRuleExtractor() - // 超过 150 字符 - longContent := "记住" + strings.Repeat("a", 200) + // 超过 150 rune，验证按 rune 截断且保持 UTF-8 合法。 + longContent := "记住" + strings.Repeat("中", 200) messages := []providertypes.Message{ {Role: providertypes.RoleUser, Content: longContent}, } @@ -136,12 +137,15 @@ func TestRuleExtractorExtractLongContent(t *testing.T) { if len(entries) != 1 { t.Fatalf("entries = %d, want 1", len(entries)) } - if len(entries[0].Title) > 150 { - t.Errorf("Title length = %d, should be <= 150", len(entries[0].Title)) + if got := len([]rune(entries[0].Title)); got > 150 { + t.Errorf("Title rune length = %d, should be <= 150", got) } if !strings.HasSuffix(entries[0].Title, "...") { t.Error("Title should be truncated with ...") } + if !utf8.ValidString(entries[0].Title) { + t.Errorf("Title should remain valid UTF-8: %q", entries[0].Title) + } } func TestRuleExtractorExtractOnlyLastUserMessage(t *testing.T) { diff --git a/internal/runtime/runtime.go b/internal/runtime/runtime.go index 6bd7230d..9707023e 100644 --- a/internal/runtime/runtime.go +++ b/internal/runtime/runtime.go @@ -29,6 +29,10 @@ const ( providerRetryBaseWait = 1 * time.Second // providerRetryMaxWait 是 runtime 层重试的最大等待时间。 providerRetryMaxWait = 5 * time.Second + // memoExtractTimeout 控制单次异步记忆提取的最长执行时间。 + memoExtractTimeout = 3 * time.Second + // memoExtractQueueSize 控制并发记忆提取任务上限，避免高频请求放大后台负载。 + memoExtractQueueSize = 4 ) // streamAccumulator 在流式事件处理过程中累积本轮对话需要持久化的助手消息状态， @@ -151,6 +155,7 @@ type Service struct { contextBuilder agentcontext.Builder // 上下文构建器，负责组装 system prompt 与本轮发给模型的消息上下文。 compactRunner contextcompact.Runner memoExtractor MemoExtractor // 可选的记忆提取钩子，完成后异步提取记忆。 + memoExtractSlots chan struct{} // 异步记忆提取并发限流槽位。 events chan RuntimeEvent operationMu sync.Mutex // 运行级互斥：串行化 Run 与 Compact，避免并发写同一会话。 runMu sync.Mutex // 仅保护 activeRun* 字段的并发读写。 @@ -179,12 +184,13 @@ func NewWithFactory( } return &Service{ - configManager: configManager, - sessionStore: sessionStore, - toolManager: toolManager, - providerFactory: providerFactory, - contextBuilder: contextBuilder, - events: make(chan RuntimeEvent, 128), + configManager: configManager, + sessionStore: sessionStore, + toolManager: toolManager, + providerFactory: providerFactory, + contextBuilder: contextBuilder, + memoExtractSlots: make(chan struct{}, memoExtractQueueSize), + events: make(chan RuntimeEvent, 128), } } @@ -344,12 +350,7 @@ func (s *Service) Run(ctx context.Context, input UserInput) error { } if len(assistant.ToolCalls) == 0 { s.emit(ctx, EventAgentDone, input.RunID, session.ID, assistant) - // 异步提取记忆：不影响主循环，失败静默处理。 - if s.memoExtractor != nil { - msgs := make([]providertypes.Message, len(session.Messages)) - copy(msgs, session.Messages) - go s.memoExtractor.ExtractAndStore(context.Background(), msgs) - } + s.scheduleMemoExtract(ctx, userMessage) return nil } @@ -875,3 +876,51 @@ func resolveWorkdirForSession(defaultWorkdir string, currentWorkdir string, requ func normalizeExistingWorkdir(workdir string) (string, error) { return agentsession.ResolveExistingDir(workdir) } + +// scheduleMemoExtract 在回答完成后异步触发记忆提取，不阻塞主流程。 +// 调用方仅传入本轮用户消息，并通过超时与限流控制后台资源占用。 +func (s *Service) scheduleMemoExtract(parent context.Context, userMessage providertypes.Message) { + if s == nil || s.memoExtractor == nil { + return + } + if strings.TrimSpace(userMessage.Role) == "" { + userMessage.Role = providertypes.RoleUser + } + if userMessage.Role != providertypes.RoleUser { + return + } + if !s.tryAcquireMemoExtractSlot() { + return + } + + extractCtx, cancel := context.WithTimeout(context.WithoutCancel(parent), memoExtractTimeout) + go func(message providertypes.Message) { + defer cancel() + defer s.releaseMemoExtractSlot() + s.memoExtractor.ExtractAndStore(extractCtx, []providertypes.Message{message}) + }(userMessage) +} + +// tryAcquireMemoExtractSlot 以非阻塞方式申请记忆提取槽位。 +func (s *Service) tryAcquireMemoExtractSlot() bool { + if s == nil || s.memoExtractSlots == nil { + return false + } + select { + case s.memoExtractSlots <- struct{}{}: + return true + default: + return false + } +} + +// releaseMemoExtractSlot 释放记忆提取槽位。 +func (s *Service) releaseMemoExtractSlot() { + if s == nil || s.memoExtractSlots == nil { + return + } + select { + case <-s.memoExtractSlots: + default: + } +} diff --git a/internal/runtime/runtime_test.go b/internal/runtime/runtime_test.go index e92f7499..c35ad5ea 100644 --- a/internal/runtime/runtime_test.go +++ b/internal/runtime/runtime_test.go @@ -310,6 +310,23 @@ func (m *stubToolManager) RememberSessionDecision(sessionID string, action secur return m.rememberErr } +type stubMemoExtractor struct { + calls int + messages []providertypes.Message + done chan struct{} +} + +func (s *stubMemoExtractor) ExtractAndStore(ctx context.Context, messages []providertypes.Message) { + s.calls++ + s.messages = append([]providertypes.Message(nil), messages...) + if s.done != nil { + select { + case s.done <- struct{}{}: + default: + } + } +} + func TestServiceRun(t *testing.T) { tests := []struct { name string @@ -1385,6 +1402,60 @@ func TestServiceNewWithFactoryDefaultsToolManager(t *testing.T) { assertEventSequence(t, events, []EventType{EventUserMessage, EventAgentDone}) } +func TestServiceRunMemoExtractorOnlyReceivesLatestUserMessage(t *testing.T) { + manager := newRuntimeConfigManager(t) + store := newMemoryStore() + now := time.Now() + store.sessions["session-memo"] = agentsession.Session{ + ID: "session-memo", + Title: "history", + CreatedAt: now.Add(-time.Minute), + UpdatedAt: now.Add(-time.Minute), + Messages: []providertypes.Message{ + {Role: providertypes.RoleUser, Content: "旧偏好"}, + {Role: providertypes.RoleAssistant, Content: "收到"}, + }, + } + + service := NewWithFactory(manager, &stubToolManager{}, store, &scriptedProviderFactory{ + provider: &scriptedProvider{ + streams: [][]providertypes.StreamEvent{ + {providertypes.NewTextDeltaStreamEvent("done")}, + }, + }, + }, nil) + + extractor := &stubMemoExtractor{done: make(chan struct{}, 1)} + service.SetMemoExtractor(extractor) + if err := service.Run(context.Background(), UserInput{ + SessionID: "session-memo", + RunID: "run-memo-extract", + Content: "记住以后都用 tab", + }); err != nil { + t.Fatalf("Run() error = %v", err) + } + + select { + case <-extractor.done: + case <-time.After(2 * time.Second): + t.Fatal("memo extractor was not called") + } + + if extractor.calls != 1 { + t.Fatalf("memo extractor calls = %d, want 1", extractor.calls) + } + if len(extractor.messages) != 1 { + t.Fatalf("memo extractor message count = %d, want 1", len(extractor.messages)) + } + got := extractor.messages[0] + if got.Role != providertypes.RoleUser { + t.Fatalf("memo extractor role = %q, want %q", got.Role, providertypes.RoleUser) + } + if got.Content != "记住以后都用 tab" { + t.Fatalf("memo extractor content = %q", got.Content) + } +} + func TestServiceRunErrorPaths(t *testing.T) { tests := []struct { name string diff --git a/internal/tools/memo/common.go b/internal/tools/memo/common.go new file mode 100644 index 00000000..5534b3d5 --- /dev/null +++ b/internal/tools/memo/common.go @@ -0,0 +1,13 @@ +package memo + +import ( + "fmt" + + "neo-code/internal/tools" +) + +// nilServiceError 构造 memo 工具缺少 service 依赖时的统一错误结果。 +func nilServiceError(toolName string) (tools.ToolResult, error) { + err := fmt.Errorf("%s: service is nil", toolName) + return tools.NewErrorResult(toolName, tools.NormalizeErrorReason(toolName, err), "", nil), err +} diff --git a/internal/tools/memo/recall.go b/internal/tools/memo/recall.go index 52142ade..e921787b 100644 --- a/internal/tools/memo/recall.go +++ b/internal/tools/memo/recall.go @@ -72,6 +72,9 @@ func (t *RecallTool) Execute(ctx context.Context, call tools.ToolCallInput) (too err := fmt.Errorf("%s: keyword is required", recallToolName) return tools.NewErrorResult(recallToolName, tools.NormalizeErrorReason(recallToolName, err), "", nil), err } + if t.svc == nil { + return nilServiceError(recallToolName) + } results, err := t.svc.Recall(ctx, args.Keyword) if err != nil { @@ -79,10 +82,10 @@ func (t *RecallTool) Execute(ctx context.Context, call tools.ToolCallInput) (too } if len(results) == 0 { - return tools.ToolResult{ + return tools.ApplyOutputLimit(tools.ToolResult{ Name: recallToolName, Content: fmt.Sprintf("No memories found matching %q.", args.Keyword), - }, nil + }, tools.DefaultOutputLimitBytes), nil } // 按 key 排序保证输出稳定性 @@ -98,8 +101,8 @@ func (t *RecallTool) Execute(ctx context.Context, call tools.ToolCallInput) (too fmt.Fprintf(&builder, "--- %s ---\n%s\n\n", k, results[k]) } - return tools.ToolResult{ + return tools.ApplyOutputLimit(tools.ToolResult{ Name: recallToolName, Content: builder.String(), - }, nil + }, tools.DefaultOutputLimitBytes), nil } diff --git a/internal/tools/memo/recall_test.go b/internal/tools/memo/recall_test.go index 22487263..239a0fd4 100644 --- a/internal/tools/memo/recall_test.go +++ b/internal/tools/memo/recall_test.go @@ -91,6 +91,21 @@ func TestRecallToolExecuteInvalidJSON(t *testing.T) { } } +func TestRecallToolExecuteNilService(t *testing.T) { + tool := NewRecallTool(nil) + args, _ := json.Marshal(recallInput{Keyword: "tab"}) + result, err := tool.Execute(context.Background(), tools.ToolCallInput{Arguments: args}) + if err == nil { + t.Fatal("expected error for nil service") + } + if !result.IsError { + t.Fatal("expected error result") + } + if !strings.Contains(result.Content, "service is nil") { + t.Fatalf("unexpected error content: %q", result.Content) + } +} + func TestRecallToolExecuteEmptyKeyword(t *testing.T) { svc := newTestService(t) tool := NewRecallTool(svc) @@ -145,3 +160,32 @@ func TestRecallToolDescription(t *testing.T) { t.Errorf("Description should mention 'memory': %q", desc) } } + +func TestRecallToolExecuteAppliesOutputLimit(t *testing.T) { + svc := newTestService(t) + if err := svc.Add(context.Background(), memo.Entry{ + Type: memo.TypeReference, + Title: "超长记忆", + Content: strings.Repeat("x", tools.DefaultOutputLimitBytes+1024), + Source: memo.SourceUserManual, + }); err != nil { + t.Fatalf("seed memo entry: %v", err) + } + + tool := NewRecallTool(svc) + args, _ := json.Marshal(recallInput{Keyword: "超长"}) + result, err := tool.Execute(context.Background(), tools.ToolCallInput{Arguments: args}) + if err != nil { + t.Fatalf("Execute error: %v", err) + } + if result.IsError { + t.Fatalf("expected success result, got error: %s", result.Content) + } + if !strings.Contains(result.Content, "...[truncated]") { + t.Fatalf("expected truncated suffix, got content length %d", len(result.Content)) + } + truncated, ok := result.Metadata["truncated"].(bool) + if !ok || !truncated { + t.Fatalf("expected metadata truncated=true, got %+v", result.Metadata) + } +} diff --git a/internal/tools/memo/remember.go b/internal/tools/memo/remember.go index 6edd0b3d..d6cac2db 100644 --- a/internal/tools/memo/remember.go +++ b/internal/tools/memo/remember.go @@ -96,6 +96,9 @@ func (t *RememberTool) Execute(ctx context.Context, call tools.ToolCallInput) (t err := fmt.Errorf("%s: invalid type %q, must be one of user/feedback/project/reference", rememberToolName, args.Type) return tools.NewErrorResult(rememberToolName, tools.NormalizeErrorReason(rememberToolName, err), "", nil), err } + if t.svc == nil { + return nilServiceError(rememberToolName) + } title := memo.NormalizeTitle(args.Title) entry := memo.Entry{ diff --git a/internal/tools/memo/remember_test.go b/internal/tools/memo/remember_test.go index f3b881f6..f2ac7fab 100644 --- a/internal/tools/memo/remember_test.go +++ b/internal/tools/memo/remember_test.go @@ -118,6 +118,25 @@ func TestRememberToolExecuteInvalidJSON(t *testing.T) { } } +func TestRememberToolExecuteNilService(t *testing.T) { + tool := NewRememberTool(nil) + args, _ := json.Marshal(rememberInput{ + Type: "user", + Title: "偏好中文注释", + Content: "用户偏好使用中文注释和 tab 缩进", + }) + result, err := tool.Execute(context.Background(), tools.ToolCallInput{Arguments: args}) + if err == nil { + t.Fatal("expected error for nil service") + } + if !result.IsError { + t.Fatal("expected error result") + } + if !strings.Contains(result.Content, "service is nil") { + t.Fatalf("unexpected error content: %q", result.Content) + } +} + func TestRememberToolExecuteMissingFields(t *testing.T) { svc := newTestService(t) tool := NewRememberTool(svc) From b1d4a3634281abc1fb5b7e9260de5bf5fef12404 Mon Sep 17 00:00:00 2001 From: Yumiue <229866007@qq.com> Date: Mon, 13 Apr 2026 15:45:23 +0800 Subject: [PATCH 2/6] =?UTF-8?q?feat(memo):=20=E6=8E=A5=E5=85=A5=20LLM=20?= =?UTF-8?q?=E8=87=AA=E5=8A=A8=E6=8F=90=E5=8F=96=E4=B8=8E=E5=90=8E=E5=8F=B0?= =?UTF-8?q?=E8=B0=83=E5=BA=A6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # Conflicts: # internal/runtime/runtime.go # internal/runtime/runtime_test.go --- internal/app/bootstrap.go | 99 ++++++- internal/app/provider_text_generator_test.go | 156 +++++++++++ internal/memo/auto_extractor.go | 273 ++++++++++++++++++ internal/memo/auto_extractor_test.go | 280 +++++++++++++++++++ internal/memo/llm_extractor.go | 240 ++++++++++++++++ internal/memo/llm_extractor_test.go | 233 +++++++++++++++ internal/memo/service.go | 108 +++++-- internal/memo/types.go | 68 ++--- internal/runtime/runtime.go | 113 ++++---- internal/runtime/runtime_test.go | 192 ++++++++----- 10 files changed, 1561 insertions(+), 201 deletions(-) create mode 100644 internal/app/provider_text_generator_test.go create mode 100644 internal/memo/auto_extractor.go create mode 100644 internal/memo/auto_extractor_test.go create mode 100644 internal/memo/llm_extractor.go create mode 100644 internal/memo/llm_extractor_test.go diff --git a/internal/app/bootstrap.go b/internal/app/bootstrap.go index 7f43e051..349e1b76 100644 --- a/internal/app/bootstrap.go +++ b/internal/app/bootstrap.go @@ -2,6 +2,7 @@ package app import ( "context" + "fmt" "strings" "time" @@ -118,10 +119,8 @@ func BuildRuntime(ctx context.Context, opts BootstrapOptions) (RuntimeBundle, er // 注入记忆提取钩子：当 AutoExtract 启用且 memoSvc 可用时，ReAct 循环完成后异步提取记忆。 if memoSvc != nil && cfg.Memo.AutoExtract { - runtimeSvc.SetMemoExtractor(&memoExtractorAdapter{ - extractor: memo.NewRuleExtractor(), - svc: memoSvc, - }) + textGen := newProviderTextGenerator(providerRegistry, manager) + runtimeSvc.SetMemoExtractor(memo.NewAutoExtractor(memo.NewLLMExtractor(textGen), memoSvc)) } return RuntimeBundle{ @@ -223,13 +222,91 @@ func buildToolManager(registry *tools.Registry) (tools.Manager, error) { return tools.NewManager(registry, engine, security.NewWorkspaceSandbox()) } -// memoExtractorAdapter 适配 memo.RuleExtractor + memo.Service 到 runtime.MemoExtractor 接口。 -type memoExtractorAdapter struct { - extractor *memo.RuleExtractor - svc *memo.Service +// providerTextGenerator 复用当前 provider 配置，向 memo 提供纯文本生成能力。 +type providerTextGenerator struct { + factory agentruntime.ProviderFactory + configManager *config.Manager } -// ExtractAndStore 实现 runtime.MemoExtractor 接口，从消息中提取记忆并保存。 -func (a *memoExtractorAdapter) ExtractAndStore(ctx context.Context, messages []providertypes.Message) { - memo.ExtractAndStore(ctx, a.extractor, a.svc, messages) +// newProviderTextGenerator 创建基于当前 provider 选择的文本生成适配器。 +func newProviderTextGenerator(factory agentruntime.ProviderFactory, cm *config.Manager) *providerTextGenerator { + return &providerTextGenerator{ + factory: factory, + configManager: cm, + } +} + +// Generate 使用当前选中的 provider/model 发起无工具的独立生成请求。 +func (g *providerTextGenerator) Generate( + ctx context.Context, + prompt string, + messages []providertypes.Message, +) (string, error) { + cfg := g.configManager.Get() + resolved, err := config.ResolveSelectedProvider(cfg) + if err != nil { + return "", err + } + + modelProvider, err := g.factory.Build(ctx, resolved.ToRuntimeConfig()) + if err != nil { + return "", err + } + + events := make(chan providertypes.StreamEvent, 32) + done := make(chan error, 1) + var builder strings.Builder + + go func() { + messageDone := false + var streamErr error + for event := range events { + switch event.Type { + case providertypes.StreamEventTextDelta: + payload, err := event.TextDeltaValue() + if err != nil { + if streamErr == nil { + streamErr = err + } + continue + } + builder.WriteString(payload.Text) + case providertypes.StreamEventMessageDone: + if _, err := event.MessageDoneValue(); err != nil { + if streamErr == nil { + streamErr = err + } + continue + } + messageDone = true + default: + if streamErr == nil { + streamErr = fmt.Errorf("memo: unexpected provider stream event %q", event.Type) + } + } + } + if streamErr == nil && !messageDone { + streamErr = fmt.Errorf("memo: provider stream ended without message_done event") + } + done <- streamErr + }() + + err = modelProvider.Generate(ctx, providertypes.GenerateRequest{ + Model: cfg.CurrentModel, + SystemPrompt: prompt, + Messages: append([]providertypes.Message(nil), messages...), + }, events) + close(events) + + streamErr := <-done + if streamErr != nil { + if err != nil { + return "", fmt.Errorf("memo: provider generate failed: %v: %w", streamErr, err) + } + return "", streamErr + } + if err != nil { + return "", err + } + return builder.String(), nil } diff --git a/internal/app/provider_text_generator_test.go b/internal/app/provider_text_generator_test.go new file mode 100644 index 00000000..8de740fb --- /dev/null +++ b/internal/app/provider_text_generator_test.go @@ -0,0 +1,156 @@ +package app + +import ( + "context" + "errors" + "strings" + "testing" + + "neo-code/internal/config" + "neo-code/internal/provider" + providertypes "neo-code/internal/provider/types" +) + +type stubTextGenProvider struct { + requests []providertypes.GenerateRequest + generate func(ctx context.Context, req providertypes.GenerateRequest, events chan<- providertypes.StreamEvent) error +} + +func (s *stubTextGenProvider) Generate( + ctx context.Context, + req providertypes.GenerateRequest, + events chan<- providertypes.StreamEvent, +) error { + s.requests = append(s.requests, req) + if s.generate != nil { + return s.generate(ctx, req, events) + } + return nil +} + +type stubTextGenFactory struct { + provider provider.Provider + err error + calls int + configs []provider.RuntimeConfig +} + +func (s *stubTextGenFactory) Build(ctx context.Context, cfg provider.RuntimeConfig) (provider.Provider, error) { + s.calls++ + s.configs = append(s.configs, cfg) + if s.err != nil { + return nil, s.err + } + return s.provider, nil +} + +// TestProviderTextGeneratorGenerateSuccess 验证文本生成器可以聚合流式文本并且不传 tools。 +func TestProviderTextGeneratorGenerateSuccess(t *testing.T) { + manager := newLoadedConfigManagerForTextGenerator(t) + providerStub := &stubTextGenProvider{ + generate: func(ctx context.Context, req providertypes.GenerateRequest, events chan<- providertypes.StreamEvent) error { + events <- providertypes.NewTextDeltaStreamEvent("hello ") + events <- providertypes.NewTextDeltaStreamEvent("world") + events <- providertypes.NewMessageDoneStreamEvent("stop", nil) + return nil + }, + } + factory := &stubTextGenFactory{provider: providerStub} + generator := newProviderTextGenerator(factory, manager) + + text, err := generator.Generate(context.Background(), "memo system prompt", []providertypes.Message{ + {Role: providertypes.RoleUser, Content: "记住这个。"}, + }) + if err != nil { + t.Fatalf("Generate() error = %v", err) + } + if text != "hello world" { + t.Fatalf("text = %q, want %q", text, "hello world") + } + if factory.calls != 1 { + t.Fatalf("Build() calls = %d, want 1", factory.calls) + } + if len(providerStub.requests) != 1 { + t.Fatalf("provider requests = %d, want 1", len(providerStub.requests)) + } + request := providerStub.requests[0] + if request.SystemPrompt != "memo system prompt" { + t.Fatalf("SystemPrompt = %q", request.SystemPrompt) + } + if len(request.Tools) != 0 { + t.Fatalf("Tools = %#v, want empty", request.Tools) + } + if request.Messages[0].Content != "记住这个。" { + t.Fatalf("Messages = %#v", request.Messages) + } +} + +// TestProviderTextGeneratorGenerateBuildFailure 验证 provider 构建失败会原样返回。 +func TestProviderTextGeneratorGenerateBuildFailure(t *testing.T) { + manager := newLoadedConfigManagerForTextGenerator(t) + factory := &stubTextGenFactory{err: errors.New("build failed")} + generator := newProviderTextGenerator(factory, manager) + + _, err := generator.Generate(context.Background(), "prompt", nil) + if err == nil || !strings.Contains(err.Error(), "build failed") { + t.Fatalf("Generate() error = %v", err) + } +} + +// TestProviderTextGeneratorGenerateFailure 验证 provider.Generate 失败会向上返回。 +func TestProviderTextGeneratorGenerateFailure(t *testing.T) { + manager := newLoadedConfigManagerForTextGenerator(t) + providerStub := &stubTextGenProvider{ + generate: func(ctx context.Context, req providertypes.GenerateRequest, events chan<- providertypes.StreamEvent) error { + return errors.New("generate failed") + }, + } + generator := newProviderTextGenerator(&stubTextGenFactory{provider: providerStub}, manager) + + _, err := generator.Generate(context.Background(), "prompt", nil) + if err == nil || !strings.Contains(err.Error(), "generate failed") { + t.Fatalf("Generate() error = %v", err) + } +} + +// TestProviderTextGeneratorGenerateRequiresMessageDone 验证缺失 message_done 会视为错误。 +func TestProviderTextGeneratorGenerateRequiresMessageDone(t *testing.T) { + manager := newLoadedConfigManagerForTextGenerator(t) + providerStub := &stubTextGenProvider{ + generate: func(ctx context.Context, req providertypes.GenerateRequest, events chan<- providertypes.StreamEvent) error { + events <- providertypes.NewTextDeltaStreamEvent("partial") + return nil + }, + } + generator := newProviderTextGenerator(&stubTextGenFactory{provider: providerStub}, manager) + + _, err := generator.Generate(context.Background(), "prompt", nil) + if err == nil || !strings.Contains(err.Error(), "message_done") { + t.Fatalf("Generate() error = %v", err) + } +} + +// newLoadedConfigManagerForTextGenerator 创建带默认 provider 选择的配置管理器。 +func newLoadedConfigManagerForTextGenerator(t *testing.T) *config.Manager { + t.Helper() + home := t.TempDir() + t.Setenv("HOME", home) + t.Setenv("USERPROFILE", home) + t.Setenv(config.OpenAIDefaultAPIKeyEnv, "test-key") + + defaults := config.StaticDefaults() + defaults.Workdir = t.TempDir() + loader := config.NewLoader("", defaults) + manager := config.NewManager(loader) + if _, err := manager.Load(context.Background()); err != nil { + t.Fatalf("Load() error = %v", err) + } + if err := manager.Update(context.Background(), func(cfg *config.Config) error { + cfg.SelectedProvider = config.OpenAIName + cfg.CurrentModel = config.OpenAIDefaultModel + return nil + }); err != nil { + t.Fatalf("Update() error = %v", err) + } + return manager +} diff --git a/internal/memo/auto_extractor.go b/internal/memo/auto_extractor.go new file mode 100644 index 00000000..69d5af7f --- /dev/null +++ b/internal/memo/auto_extractor.go @@ -0,0 +1,273 @@ +package memo + +import ( + "context" + "log" + "strings" + "sync" + "time" + + providertypes "neo-code/internal/provider/types" +) + +const autoExtractDebounce = 2 * time.Second + +// AutoExtractor 负责按会话在后台调度自动提取，并处理防抖、互斥和尾随执行。 +type AutoExtractor struct { + extractor Extractor + svc *Service + debounce time.Duration + logf func(format string, args ...any) + + mu sync.Mutex + workers map[string]*autoExtractWorker +} + +type autoExtractWorker struct { + updates chan autoExtractRequest +} + +type autoExtractRequest struct { + messages []providertypes.Message + dueAt time.Time + skip bool +} + +// NewAutoExtractor 创建后台自动提取调度器。 +func NewAutoExtractor(extractor Extractor, svc *Service) *AutoExtractor { + return &AutoExtractor{ + extractor: extractor, + svc: svc, + debounce: autoExtractDebounce, + logf: log.Printf, + workers: make(map[string]*autoExtractWorker), + } +} + +// Schedule 按会话维度安排一次后台自动提取，skip=true 时清空当前待执行请求。 +func (a *AutoExtractor) Schedule(sessionID string, messages []providertypes.Message, skip bool) { + if a == nil || a.extractor == nil || a.svc == nil { + return + } + + sessionID = strings.TrimSpace(sessionID) + if sessionID == "" { + return + } + + req := autoExtractRequest{ + messages: cloneProviderMessages(messages), + dueAt: time.Now().Add(a.debounce), + skip: skip, + } + + worker := a.ensureWorker(sessionID) + worker.updates <- req +} + +// ensureWorker 获取或创建指定会话的后台 worker。 +func (a *AutoExtractor) ensureWorker(sessionID string) *autoExtractWorker { + a.mu.Lock() + defer a.mu.Unlock() + + if worker, ok := a.workers[sessionID]; ok { + return worker + } + + worker := &autoExtractWorker{ + updates: make(chan autoExtractRequest, 32), + } + a.workers[sessionID] = worker + go a.runWorker(worker) + return worker +} + +// runWorker 串行处理单个会话的调度状态，避免并发提取并支持 trailing extraction。 +func (a *AutoExtractor) runWorker(worker *autoExtractWorker) { + var ( + pending *autoExtractRequest + timer *time.Timer + timerCh <-chan time.Time + running bool + doneCh = make(chan struct{}, 1) + ) + + for { + select { + case req := <-worker.updates: + if req.skip { + pending = nil + stopTimer(timer) + timer = nil + timerCh = nil + continue + } + + reqCopy := req + pending = &reqCopy + if running { + continue + } + + timer = resetTimer(timer, time.Until(req.dueAt)) + timerCh = timer.C + + case <-timerCh: + timer = nil + timerCh = nil + + for { + select { + case req := <-worker.updates: + if req.skip { + pending = nil + continue + } + reqCopy := req + pending = &reqCopy + default: + goto launch + } + } + + launch: + if running || pending == nil { + continue + } + if wait := time.Until(pending.dueAt); wait > 0 { + timer = resetTimer(timer, wait) + timerCh = timer.C + continue + } + + current := *pending + pending = nil + running = true + + go func(req autoExtractRequest) { + a.extractAndStore(req.messages) + doneCh <- struct{}{} + }(current) + + case <-doneCh: + running = false + if pending == nil { + continue + } + + timer = resetTimer(timer, time.Until(pending.dueAt)) + timerCh = timer.C + } + } +} + +// extractAndStore 执行提取，并在写入前做本地批次去重和持久化级别的原子去重。 +func (a *AutoExtractor) extractAndStore(messages []providertypes.Message) { + ctx := context.Background() + entries, err := a.extractor.Extract(ctx, messages) + if err != nil { + a.logError("memo: auto extract failed: %v", err) + return + } + if len(entries) == 0 { + return + } + + seen := make(map[string]struct{}, len(entries)) + for _, entry := range entries { + entry.Source = SourceAutoExtract + key := autoExtractDedupKey(entry) + if key == "" { + continue + } + if _, exists := seen[key]; exists { + continue + } + + added, err := a.svc.addAutoExtractIfAbsent(ctx, entry) + if err != nil { + a.logError("memo: auto extract add failed: %v", err) + continue + } + + seen[key] = struct{}{} + if !added { + continue + } + } +} + +// autoExtractDedupKey 生成自动提取条目的精确去重键。 +func autoExtractDedupKey(entry Entry) string { + title := NormalizeTitle(entry.Title) + content := strings.TrimSpace(entry.Content) + if !IsValidType(entry.Type) || title == "" || content == "" { + return "" + } + return strings.Join([]string{string(entry.Type), title, content}, "\x1f") +} + +// parseTopicSourceAndContent 从 topic 文件中解析 source frontmatter 和正文内容。 +func parseTopicSourceAndContent(topic string) (string, string) { + parts := strings.Split(topic, "---") + if len(parts) < 3 { + return "", strings.TrimSpace(topic) + } + + frontmatter := parts[1] + body := strings.TrimSpace(strings.Join(parts[2:], "---")) + for _, line := range strings.Split(frontmatter, "\n") { + line = strings.TrimSpace(line) + if !strings.HasPrefix(line, "source:") { + continue + } + return strings.TrimSpace(strings.TrimPrefix(line, "source:")), body + } + return "", body +} + +// cloneProviderMessages 深拷贝消息切片，避免后台任务读取到后续修改。 +func cloneProviderMessages(messages []providertypes.Message) []providertypes.Message { + if len(messages) == 0 { + return nil + } + + cloned := make([]providertypes.Message, 0, len(messages)) + for _, message := range messages { + cloned = append(cloned, cloneProviderMessage(message)) + } + return cloned +} + +// resetTimer 以安全方式重置定时器，避免旧事件污染新的防抖窗口。 +func resetTimer(timer *time.Timer, wait time.Duration) *time.Timer { + if wait < 0 { + wait = 0 + } + if timer == nil { + return time.NewTimer(wait) + } + stopTimer(timer) + timer.Reset(wait) + return timer +} + +// stopTimer 停止定时器并在必要时清空通道。 +func stopTimer(timer *time.Timer) { + if timer == nil { + return + } + if !timer.Stop() { + select { + case <-timer.C: + default: + } + } +} + +// logError 统一收敛后台提取日志，避免把错误暴露给主链路。 +func (a *AutoExtractor) logError(format string, args ...any) { + if a != nil && a.logf != nil { + a.logf(format, args...) + } +} diff --git a/internal/memo/auto_extractor_test.go b/internal/memo/auto_extractor_test.go new file mode 100644 index 00000000..fc4c6d2f --- /dev/null +++ b/internal/memo/auto_extractor_test.go @@ -0,0 +1,280 @@ +package memo + +import ( + "context" + "errors" + "strings" + "sync" + "testing" + "time" + + "neo-code/internal/config" + providertypes "neo-code/internal/provider/types" +) + +type stubMemoExtractor struct { + mu sync.Mutex + callCount int + calls [][]providertypes.Message + extractFn func(ctx context.Context, messages []providertypes.Message) ([]Entry, error) +} + +func (s *stubMemoExtractor) Extract(ctx context.Context, messages []providertypes.Message) ([]Entry, error) { + s.mu.Lock() + s.callCount++ + s.calls = append(s.calls, cloneProviderMessages(messages)) + extractFn := s.extractFn + s.mu.Unlock() + + if extractFn != nil { + return extractFn(ctx, messages) + } + return nil, nil +} + +func (s *stubMemoExtractor) Calls() int { + s.mu.Lock() + defer s.mu.Unlock() + return s.callCount +} + +func newAutoExtractorTestService(t *testing.T) *Service { + t.Helper() + store := NewFileStore(t.TempDir(), t.TempDir()) + return NewService(store, nil, config.MemoConfig{MaxIndexLines: 200}, nil) +} + +func TestAutoExtractorDebounceMergesRequests(t *testing.T) { + svc := newAutoExtractorTestService(t) + extractor := &stubMemoExtractor{ + extractFn: func(ctx context.Context, messages []providertypes.Message) ([]Entry, error) { + last := messages[len(messages)-1].Content + return []Entry{{Type: TypeProject, Title: last, Content: last, Source: SourceAutoExtract}}, nil + }, + } + auto := NewAutoExtractor(extractor, svc) + auto.debounce = 20 * time.Millisecond + auto.logf = func(string, ...any) {} + + auto.Schedule("session-1", []providertypes.Message{{Role: providertypes.RoleUser, Content: "first"}}, false) + auto.Schedule("session-1", []providertypes.Message{{Role: providertypes.RoleUser, Content: "second"}}, false) + + waitFor(t, time.Second, func() bool { return extractor.Calls() == 1 }) + time.Sleep(60 * time.Millisecond) + + if extractor.Calls() != 1 { + t.Fatalf("extractor calls = %d, want 1", extractor.Calls()) + } + + recall, err := svc.Recall(context.Background(), "second") + if err != nil { + t.Fatalf("Recall() error = %v", err) + } + if len(recall) != 1 { + t.Fatalf("recall = %#v", recall) + } + for _, content := range recall { + if !strings.Contains(content, "second") { + t.Fatalf("recall content = %q", content) + } + } +} + +func TestAutoExtractorSkipCancelsPendingRequest(t *testing.T) { + svc := newAutoExtractorTestService(t) + extractor := &stubMemoExtractor{} + auto := NewAutoExtractor(extractor, svc) + auto.debounce = 40 * time.Millisecond + auto.logf = func(string, ...any) {} + + auto.Schedule("session-1", []providertypes.Message{{Role: providertypes.RoleUser, Content: "remember me"}}, false) + time.Sleep(20 * time.Millisecond) + auto.Schedule("session-1", []providertypes.Message{{Role: providertypes.RoleUser, Content: "remember me"}}, true) + time.Sleep(80 * time.Millisecond) + + if extractor.Calls() != 0 { + t.Fatalf("extractor calls = %d, want 0", extractor.Calls()) + } +} + +func TestAutoExtractorTrailingRun(t *testing.T) { + svc := newAutoExtractorTestService(t) + firstStarted := make(chan struct{}, 1) + secondStarted := make(chan struct{}, 1) + releaseFirst := make(chan struct{}) + + extractor := &stubMemoExtractor{ + extractFn: func(ctx context.Context, messages []providertypes.Message) ([]Entry, error) { + switch messages[len(messages)-1].Content { + case "first": + firstStarted <- struct{}{} + <-releaseFirst + case "second": + secondStarted <- struct{}{} + } + last := messages[len(messages)-1].Content + return []Entry{{Type: TypeProject, Title: last, Content: last, Source: SourceAutoExtract}}, nil + }, + } + auto := NewAutoExtractor(extractor, svc) + auto.debounce = 15 * time.Millisecond + auto.logf = func(string, ...any) {} + + auto.Schedule("session-1", []providertypes.Message{{Role: providertypes.RoleUser, Content: "first"}}, false) + + select { + case <-firstStarted: + case <-time.After(time.Second): + t.Fatal("first extraction did not start") + } + + auto.Schedule("session-1", []providertypes.Message{{Role: providertypes.RoleUser, Content: "second"}}, false) + time.Sleep(40 * time.Millisecond) + close(releaseFirst) + + select { + case <-secondStarted: + case <-time.After(time.Second): + t.Fatal("second trailing extraction did not start") + } + + waitFor(t, time.Second, func() bool { return extractor.Calls() == 2 }) + waitFor(t, time.Second, func() bool { + entries, err := svc.List(context.Background()) + return err == nil && len(entries) == 2 + }) +} + +func TestAutoExtractorSkipDoesNotQueue(t *testing.T) { + svc := newAutoExtractorTestService(t) + extractor := &stubMemoExtractor{} + auto := NewAutoExtractor(extractor, svc) + auto.debounce = 10 * time.Millisecond + auto.logf = func(string, ...any) {} + + auto.Schedule("session-1", []providertypes.Message{{Role: providertypes.RoleUser, Content: "skip"}}, true) + time.Sleep(50 * time.Millisecond) + + if extractor.Calls() != 0 { + t.Fatalf("extractor calls = %d, want 0", extractor.Calls()) + } +} + +func TestAutoExtractorErrorsAreSilent(t *testing.T) { + svc := newAutoExtractorTestService(t) + extractor := &stubMemoExtractor{ + extractFn: func(ctx context.Context, messages []providertypes.Message) ([]Entry, error) { + return nil, errors.New("boom") + }, + } + auto := NewAutoExtractor(extractor, svc) + auto.debounce = 10 * time.Millisecond + auto.logf = func(string, ...any) {} + + auto.Schedule("session-1", []providertypes.Message{{Role: providertypes.RoleUser, Content: "x"}}, false) + waitFor(t, time.Second, func() bool { return extractor.Calls() == 1 }) + + entries, err := svc.List(context.Background()) + if err != nil { + t.Fatalf("List() error = %v", err) + } + if len(entries) != 0 { + t.Fatalf("entries = %#v, want empty", entries) + } +} + +func TestAutoExtractorSuppressesExactDuplicates(t *testing.T) { + svc := newAutoExtractorTestService(t) + if err := svc.Add(context.Background(), Entry{ + Type: TypeUser, + Title: "reply in chinese", + Content: "reply in chinese", + Source: SourceAutoExtract, + }); err != nil { + t.Fatalf("seed Add() error = %v", err) + } + + extractor := &stubMemoExtractor{ + extractFn: func(ctx context.Context, messages []providertypes.Message) ([]Entry, error) { + return []Entry{ + {Type: TypeUser, Title: "reply in chinese", Content: "reply in chinese", Source: SourceAutoExtract}, + {Type: TypeFeedback, Title: "run tests first", Content: "run tests first", Source: SourceAutoExtract}, + {Type: TypeFeedback, Title: "run tests first", Content: "run tests first", Source: SourceAutoExtract}, + }, nil + }, + } + auto := NewAutoExtractor(extractor, svc) + auto.debounce = 10 * time.Millisecond + auto.logf = func(string, ...any) {} + + auto.Schedule("session-1", []providertypes.Message{{Role: providertypes.RoleUser, Content: "dedupe"}}, false) + waitFor(t, time.Second, func() bool { + entries, err := svc.List(context.Background()) + return err == nil && len(entries) == 2 + }) + + entries, err := svc.List(context.Background()) + if err != nil { + t.Fatalf("List() error = %v", err) + } + if len(entries) != 2 { + t.Fatalf("len(entries) = %d, want 2", len(entries)) + } +} + +func TestAutoExtractorSuppressesExactDuplicatesAcrossSessions(t *testing.T) { + svc := newAutoExtractorTestService(t) + started := make(chan struct{}, 2) + release := make(chan struct{}) + + extractor := &stubMemoExtractor{ + extractFn: func(ctx context.Context, messages []providertypes.Message) ([]Entry, error) { + started <- struct{}{} + <-release + return []Entry{ + {Type: TypeProject, Title: "same title", Content: "same content", Source: SourceAutoExtract}, + }, nil + }, + } + auto := NewAutoExtractor(extractor, svc) + auto.debounce = 0 + auto.logf = func(string, ...any) {} + + auto.Schedule("session-1", []providertypes.Message{{Role: providertypes.RoleUser, Content: "one"}}, false) + auto.Schedule("session-2", []providertypes.Message{{Role: providertypes.RoleUser, Content: "two"}}, false) + + for i := 0; i < 2; i++ { + select { + case <-started: + case <-time.After(time.Second): + t.Fatal("concurrent extraction did not start") + } + } + close(release) + + waitFor(t, time.Second, func() bool { return extractor.Calls() == 2 }) + waitFor(t, time.Second, func() bool { + entries, err := svc.List(context.Background()) + return err == nil && len(entries) == 1 + }) + + entries, err := svc.List(context.Background()) + if err != nil { + t.Fatalf("List() error = %v", err) + } + if len(entries) != 1 { + t.Fatalf("len(entries) = %d, want 1", len(entries)) + } +} + +func waitFor(t *testing.T, timeout time.Duration, fn func() bool) { + t.Helper() + deadline := time.Now().Add(timeout) + for time.Now().Before(deadline) { + if fn() { + return + } + time.Sleep(10 * time.Millisecond) + } + t.Fatal("condition not met before timeout") +} diff --git a/internal/memo/llm_extractor.go b/internal/memo/llm_extractor.go new file mode 100644 index 00000000..1aa10542 --- /dev/null +++ b/internal/memo/llm_extractor.go @@ -0,0 +1,240 @@ +package memo + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "strings" + "time" + + providertypes "neo-code/internal/provider/types" +) + +const llmExtractorRecentMessageLimit = 10 + +// LLMExtractor 基于 LLM 分析最近对话，并返回结构化记忆条目。 +type LLMExtractor struct { + generator TextGenerator + now func() time.Time +} + +type extractedEntry struct { + Type string `json:"type"` + Title string `json:"title"` + Content string `json:"content"` + Keywords []string `json:"keywords"` +} + +// NewLLMExtractor 创建基于 TextGenerator 的记忆提取器。 +func NewLLMExtractor(generator TextGenerator) *LLMExtractor { + return &LLMExtractor{ + generator: generator, + now: time.Now, + } +} + +// Extract 从最近对话中提取可跨会话持久化的记忆条目。 +func (e *LLMExtractor) Extract(ctx context.Context, messages []providertypes.Message) ([]Entry, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + if e == nil || e.generator == nil { + return nil, errors.New("memo: text generator is nil") + } + + recent := recentMessagesForExtraction(messages, llmExtractorRecentMessageLimit) + if len(recent) == 0 || !containsUserMessage(recent) { + return nil, nil + } + + response, err := e.generator.Generate(ctx, buildExtractionPrompt(e.now()), recent) + if err != nil { + return nil, err + } + if err := ctx.Err(); err != nil { + return nil, err + } + + jsonText, err := extractJSONArray(response) + if err != nil { + return nil, err + } + + var extracted []extractedEntry + if err := json.Unmarshal([]byte(jsonText), &extracted); err != nil { + return nil, fmt.Errorf("memo: parse extraction response: %w", err) + } + + entries := make([]Entry, 0, len(extracted)) + for _, item := range extracted { + entry, ok := toMemoEntry(item) + if !ok { + continue + } + entries = append(entries, entry) + } + return entries, nil +} + +// buildExtractionPrompt 构造记忆提取专用的 system prompt。 +func buildExtractionPrompt(now time.Time) string { + currentDate := now.In(time.Local).Format("2006-01-02") + return strings.TrimSpace(fmt.Sprintf(` +你是一个记忆提取助手（memory extraction assistant）。 +分析最近对话中值得跨会话持久记住的信息，并返回严格 JSON 数组。 + +当前本地日期：%s +如果对话中出现“今天、明天、下周二”等相对日期，必须先转换为绝对日期再写入 content。 + +只允许以下四种 type： +- user: 用户偏好、习惯、背景、专长 +- feedback: 用户对 Agent 做法的纠正、要求、确认过的工作方式 +- project: 项目事实、项目决策、截止时间、进行中的工作 +- reference: 外部资源、文档、链接、仪表盘、沟通渠道 + +提取规则： +1. 只提取无法从代码仓库直接推导的信息。 +2. 不要提取通用编程知识、代码结构、文件路径、Git 历史。 +3. 每条记忆必须具体、可操作。 +4. 没有值得记住的信息时，返回 []。 +5. 输出必须是 JSON 数组，不要输出任何额外解释。 + +输出格式： +[{"type":"user","title":"...","content":"...","keywords":["..."]}] +`, currentDate)) +} + +// recentMessagesForExtraction 截取最近的非 tool 消息，保留原始顺序。 +func recentMessagesForExtraction(messages []providertypes.Message, limit int) []providertypes.Message { + if len(messages) == 0 || limit <= 0 { + return nil + } + + filtered := make([]providertypes.Message, 0, len(messages)) + for _, message := range messages { + if message.Role == providertypes.RoleTool { + continue + } + filtered = append(filtered, cloneProviderMessage(message)) + } + if len(filtered) <= limit { + return filtered + } + return filtered[len(filtered)-limit:] +} + +// containsUserMessage 检查待提取消息中是否包含用户输入。 +func containsUserMessage(messages []providertypes.Message) bool { + for _, message := range messages { + if message.Role == providertypes.RoleUser && strings.TrimSpace(message.Content) != "" { + return true + } + } + return false +} + +// toMemoEntry 将 LLM 输出条目收敛为合法的 memo.Entry。 +func toMemoEntry(item extractedEntry) (Entry, bool) { + entryType, ok := ParseType(strings.TrimSpace(item.Type)) + if !ok { + return Entry{}, false + } + + title := NormalizeTitle(item.Title) + content := strings.TrimSpace(item.Content) + if title == "" || content == "" { + return Entry{}, false + } + + return Entry{ + Type: entryType, + Title: title, + Content: content, + Keywords: normalizeKeywords(item.Keywords), + Source: SourceAutoExtract, + }, true +} + +// normalizeKeywords 规范化关键词列表，移除空值和重复值。 +func normalizeKeywords(keywords []string) []string { + if len(keywords) == 0 { + return nil + } + + result := make([]string, 0, len(keywords)) + seen := make(map[string]struct{}, len(keywords)) + for _, keyword := range keywords { + trimmed := strings.TrimSpace(keyword) + if trimmed == "" { + continue + } + key := strings.ToLower(trimmed) + if _, exists := seen[key]; exists { + continue + } + seen[key] = struct{}{} + result = append(result, trimmed) + } + if len(result) == 0 { + return nil + } + return result +} + +// extractJSONArray 从模型返回文本中提取最外层 JSON 数组，容忍前后噪声。 +func extractJSONArray(text string) (string, error) { + start := strings.Index(text, "[") + if start < 0 { + return "", errors.New("memo: extraction response does not contain a JSON array") + } + + depth := 0 + inString := false + escaped := false + for index := start; index < len(text); index++ { + ch := text[index] + if inString { + if escaped { + escaped = false + continue + } + switch ch { + case '\\': + escaped = true + case '"': + inString = false + } + continue + } + + switch ch { + case '"': + inString = true + case '[': + depth++ + case ']': + depth-- + if depth == 0 { + return strings.TrimSpace(text[start : index+1]), nil + } + } + } + + return "", errors.New("memo: extraction response contains an incomplete JSON array") +} + +// cloneProviderMessage 深拷贝消息结构，避免后台提取读取到运行时后续修改。 +func cloneProviderMessage(message providertypes.Message) providertypes.Message { + cloned := message + if len(message.ToolCalls) > 0 { + cloned.ToolCalls = append([]providertypes.ToolCall(nil), message.ToolCalls...) + } + if len(message.ToolMetadata) > 0 { + cloned.ToolMetadata = make(map[string]string, len(message.ToolMetadata)) + for key, value := range message.ToolMetadata { + cloned.ToolMetadata[key] = value + } + } + return cloned +} diff --git a/internal/memo/llm_extractor_test.go b/internal/memo/llm_extractor_test.go new file mode 100644 index 00000000..7d278c26 --- /dev/null +++ b/internal/memo/llm_extractor_test.go @@ -0,0 +1,233 @@ +package memo + +import ( + "context" + "errors" + "strings" + "testing" + "time" + + providertypes "neo-code/internal/provider/types" +) + +type stubTextGenerator struct { + response string + err error + calls int + prompt string + messages []providertypes.Message +} + +func (s *stubTextGenerator) Generate( + ctx context.Context, + prompt string, + messages []providertypes.Message, +) (string, error) { + s.calls++ + s.prompt = prompt + s.messages = append([]providertypes.Message(nil), messages...) + if s.err != nil { + return "", s.err + } + return s.response, nil +} + +// TestLLMExtractorExtractValidJSON 验证提取器可以解析合法 JSON 并收敛字段。 +func TestLLMExtractorExtractValidJSON(t *testing.T) { + generator := &stubTextGenerator{ + response: `[{"type":"user","title":" 偏好 Go 代码风格 ","content":"用户偏好使用 Go 惯用写法。","keywords":["go"," style ","go"]}]`, + } + extractor := NewLLMExtractor(generator) + extractor.now = func() time.Time { + return time.Date(2026, 4, 13, 10, 0, 0, 0, time.FixedZone("CST", 8*3600)) + } + + entries, err := extractor.Extract(context.Background(), []providertypes.Message{ + {Role: providertypes.RoleUser, Content: "以后默认按 Go 惯用风格写。"}, + {Role: providertypes.RoleAssistant, Content: "收到。"}, + }) + if err != nil { + t.Fatalf("Extract() error = %v", err) + } + if len(entries) != 1 { + t.Fatalf("len(entries) = %d, want 1", len(entries)) + } + if entries[0].Type != TypeUser { + t.Fatalf("Type = %q, want %q", entries[0].Type, TypeUser) + } + if entries[0].Title != "偏好 Go 代码风格" { + t.Fatalf("Title = %q", entries[0].Title) + } + if entries[0].Content != "用户偏好使用 Go 惯用写法。" { + t.Fatalf("Content = %q", entries[0].Content) + } + if len(entries[0].Keywords) != 2 || entries[0].Keywords[1] != "style" { + t.Fatalf("Keywords = %#v", entries[0].Keywords) + } + if entries[0].Source != SourceAutoExtract { + t.Fatalf("Source = %q, want %q", entries[0].Source, SourceAutoExtract) + } + if generator.calls != 1 { + t.Fatalf("Generate() calls = %d, want 1", generator.calls) + } + if !strings.Contains(generator.prompt, "user: 用户偏好") { + t.Fatalf("prompt should describe user type, got %q", generator.prompt) + } + if !strings.Contains(generator.prompt, "当前本地日期：2026-04-13") { + t.Fatalf("prompt should include absolute local date, got %q", generator.prompt) + } + if !strings.Contains(generator.prompt, "必须先转换为绝对日期") { + t.Fatalf("prompt should require absolute dates, got %q", generator.prompt) + } +} + +// TestLLMExtractorExtractEmptyResult 验证空数组响应会返回零条记忆。 +func TestLLMExtractorExtractEmptyResult(t *testing.T) { + extractor := NewLLMExtractor(&stubTextGenerator{response: `[]`}) + + entries, err := extractor.Extract(context.Background(), []providertypes.Message{ + {Role: providertypes.RoleUser, Content: "这轮没有需要记住的内容。"}, + }) + if err != nil { + t.Fatalf("Extract() error = %v", err) + } + if len(entries) != 0 { + t.Fatalf("len(entries) = %d, want 0", len(entries)) + } +} + +// TestLLMExtractorExtractNoUserMessage 验证没有用户消息时不会调用模型。 +func TestLLMExtractorExtractNoUserMessage(t *testing.T) { + generator := &stubTextGenerator{response: `[]`} + extractor := NewLLMExtractor(generator) + + entries, err := extractor.Extract(context.Background(), []providertypes.Message{ + {Role: providertypes.RoleAssistant, Content: "只有助手消息。"}, + }) + if err != nil { + t.Fatalf("Extract() error = %v", err) + } + if len(entries) != 0 { + t.Fatalf("len(entries) = %d, want 0", len(entries)) + } + if generator.calls != 0 { + t.Fatalf("Generate() calls = %d, want 0", generator.calls) + } +} + +// TestLLMExtractorExtractNoMessages 验证空消息输入直接返回空结果。 +func TestLLMExtractorExtractNoMessages(t *testing.T) { + extractor := NewLLMExtractor(&stubTextGenerator{response: `[]`}) + + entries, err := extractor.Extract(context.Background(), nil) + if err != nil { + t.Fatalf("Extract() error = %v", err) + } + if len(entries) != 0 { + t.Fatalf("len(entries) = %d, want 0", len(entries)) + } +} + +// TestLLMExtractorExtractInvalidJSON 验证无效 JSON 会返回错误。 +func TestLLMExtractorExtractInvalidJSON(t *testing.T) { + extractor := NewLLMExtractor(&stubTextGenerator{response: `[{invalid json}]`}) + + _, err := extractor.Extract(context.Background(), []providertypes.Message{ + {Role: providertypes.RoleUser, Content: "记住这个。"}, + }) + if err == nil { + t.Fatal("expected invalid JSON error") + } +} + +// TestLLMExtractorExtractToleratesWrappedJSON 验证 JSON 前后噪声不会影响解析。 +func TestLLMExtractorExtractToleratesWrappedJSON(t *testing.T) { + extractor := NewLLMExtractor(&stubTextGenerator{ + response: "分析如下：\n[{\"type\":\"feedback\",\"title\":\"以后先跑测试\",\"content\":\"用户要求修改后先跑测试。\"}]\n以上完毕。", + }) + + entries, err := extractor.Extract(context.Background(), []providertypes.Message{ + {Role: providertypes.RoleUser, Content: "以后改完先跑测试。"}, + }) + if err != nil { + t.Fatalf("Extract() error = %v", err) + } + if len(entries) != 1 || entries[0].Type != TypeFeedback { + t.Fatalf("entries = %#v", entries) + } +} + +// TestLLMExtractorExtractFiltersInvalidEntries 验证非法类型和空字段会被过滤。 +func TestLLMExtractorExtractFiltersInvalidEntries(t *testing.T) { + extractor := NewLLMExtractor(&stubTextGenerator{ + response: `[ + {"type":"invalid","title":"bad","content":"bad"}, + {"type":"project","title":" ","content":"missing title"}, + {"type":"reference","title":"文档入口","content":"查看 docs/runtime-provider-event-flow.md"} + ]`, + }) + + entries, err := extractor.Extract(context.Background(), []providertypes.Message{ + {Role: providertypes.RoleUser, Content: "参考文档在 docs/runtime-provider-event-flow.md。"}, + }) + if err != nil { + t.Fatalf("Extract() error = %v", err) + } + if len(entries) != 1 || entries[0].Type != TypeReference { + t.Fatalf("entries = %#v", entries) + } +} + +// TestLLMExtractorExtractCancelledContext 验证已取消上下文会中止提取。 +func TestLLMExtractorExtractCancelledContext(t *testing.T) { + generator := &stubTextGenerator{response: `[]`} + extractor := NewLLMExtractor(generator) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + _, err := extractor.Extract(ctx, []providertypes.Message{ + {Role: providertypes.RoleUser, Content: "记住这个。"}, + }) + if !errors.Is(err, context.Canceled) { + t.Fatalf("Extract() error = %v, want context.Canceled", err) + } + if generator.calls != 0 { + t.Fatalf("Generate() calls = %d, want 0", generator.calls) + } +} + +// TestLLMExtractorExtractUsesRecentNonToolMessages 验证只取最近 10 条非 tool 消息。 +func TestLLMExtractorExtractUsesRecentNonToolMessages(t *testing.T) { + generator := &stubTextGenerator{response: `[]`} + extractor := NewLLMExtractor(generator) + + messages := make([]providertypes.Message, 0, 16) + for index := 0; index < 12; index++ { + messages = append(messages, providertypes.Message{ + Role: providertypes.RoleUser, + Content: "user-" + string(rune('a'+index)), + }) + if index%3 == 0 { + messages = append(messages, providertypes.Message{ + Role: providertypes.RoleTool, + Content: "tool-" + string(rune('a'+index)), + }) + } + } + + _, err := extractor.Extract(context.Background(), messages) + if err != nil { + t.Fatalf("Extract() error = %v", err) + } + if len(generator.messages) != 10 { + t.Fatalf("len(generator.messages) = %d, want 10", len(generator.messages)) + } + for _, message := range generator.messages { + if message.Role == providertypes.RoleTool { + t.Fatalf("unexpected tool message in extraction context: %#v", message) + } + } + if generator.messages[0].Content != "user-c" || generator.messages[9].Content != "user-l" { + t.Fatalf("unexpected recent window: first=%q last=%q", generator.messages[0].Content, generator.messages[9].Content) + } +} diff --git a/internal/memo/service.go b/internal/memo/service.go index 49bb0075..42325bce 100644 --- a/internal/memo/service.go +++ b/internal/memo/service.go @@ -18,11 +18,10 @@ type Service struct { extractor Extractor config config.MemoConfig mu sync.Mutex - sourceInvl func() // 可选的缓存失效回调 + sourceInvl func() } -// NewService 创建 memo Service 实例。 -// extractor 可以为 nil（禁用自动提取时不需要）。 +// NewService 创建 memo Service 实例；extractor 可以为 nil。 func NewService(store Store, extractor Extractor, cfg config.MemoConfig, sourceInvl func()) *Service { return &Service{ store: store, @@ -34,17 +33,43 @@ func NewService(store Store, extractor Extractor, cfg config.MemoConfig, sourceI // Add 添加一条记忆并持久化索引和 topic 文件。 func (s *Service) Add(ctx context.Context, entry Entry) error { - if !IsValidType(entry.Type) { - return fmt.Errorf("memo: invalid type %q", entry.Type) + entry, err := normalizeEntryForPersist(entry) + if err != nil { + return err } - entry.Title = NormalizeTitle(entry.Title) - if entry.Title == "" { - return fmt.Errorf("memo: title is empty") + + s.mu.Lock() + defer s.mu.Unlock() + + return s.saveEntryLocked(ctx, entry) +} + +// addAutoExtractIfAbsent 在同一把锁内完成自动提取条目的查重与写入。 +func (s *Service) addAutoExtractIfAbsent(ctx context.Context, entry Entry) (bool, error) { + entry, err := normalizeEntryForPersist(entry) + if err != nil { + return false, err } s.mu.Lock() defer s.mu.Unlock() + duplicate, err := s.hasExactAutoExtractLocked(ctx, entry) + if err != nil { + return false, err + } + if duplicate { + return false, nil + } + + if err := s.saveEntryLocked(ctx, entry); err != nil { + return false, err + } + return true, nil +} + +// saveEntryLocked 在持有 Service 锁的前提下持久化单条记忆及索引。 +func (s *Service) saveEntryLocked(ctx context.Context, entry Entry) error { now := time.Now() if entry.ID == "" { entry.ID = newEntryID(entry.Type) @@ -64,7 +89,6 @@ func (s *Service) Add(ctx context.Context, entry Entry) error { } working := cloneIndex(index) - // 检查是否已存在相同 ID（更新场景） replaced := false for i, existing := range working.Entries { if existing.ID == entry.ID { @@ -78,11 +102,9 @@ func (s *Service) Add(ctx context.Context, entry Entry) error { } working.UpdatedAt = now - // 截断索引到最大行数 var topicsToDelete []string if s.config.MaxIndexLines > 0 && len(working.Entries) > s.config.MaxIndexLines { excess := len(working.Entries) - s.config.MaxIndexLines - // 记录最旧条目对应的 topic 文件，待索引写入成功后再删除。 for i := 0; i < excess; i++ { topicFile := strings.TrimSpace(working.Entries[i].TopicFile) if topicFile != "" && topicFile != entry.TopicFile { @@ -109,8 +131,42 @@ func (s *Service) Add(ctx context.Context, entry Entry) error { return nil } -// loadIndexLocked 在持有锁的状态下加载索引，供多个 Service 方法复用。 -// 调用方须持有 s.mu 锁。 +// hasExactAutoExtractLocked 检查是否已存在完全相同的自动提取记忆。 +func (s *Service) hasExactAutoExtractLocked(ctx context.Context, target Entry) (bool, error) { + targetKey := autoExtractDedupKey(target) + if targetKey == "" { + return false, nil + } + + index, err := s.loadIndexLocked(ctx) + if err != nil { + return false, err + } + for _, entry := range index.Entries { + if strings.TrimSpace(entry.TopicFile) == "" { + continue + } + + topicContent, err := s.store.LoadTopic(ctx, entry.TopicFile) + if err != nil { + continue + } + + source, content := parseTopicSourceAndContent(topicContent) + if source != SourceAutoExtract { + continue + } + + entry.Source = source + entry.Content = content + if autoExtractDedupKey(entry) == targetKey { + return true, nil + } + } + return false, nil +} + +// loadIndexLocked 在持有锁的状态下加载索引。 func (s *Service) loadIndexLocked(ctx context.Context) (*Index, error) { index, err := s.store.LoadIndex(ctx) if err != nil { @@ -119,8 +175,7 @@ func (s *Service) loadIndexLocked(ctx context.Context) (*Index, error) { return index, nil } -// Remove 按关键词搜索并删除匹配的记忆条目。 -// 返回被删除的条目数量。 +// Remove 按关键词搜索并删除匹配的记忆条目，返回删除数量。 func (s *Service) Remove(ctx context.Context, keyword string) (int, error) { s.mu.Lock() defer s.mu.Unlock() @@ -167,7 +222,7 @@ func (s *Service) Remove(ctx context.Context, keyword string) (int, error) { return removed, nil } -// List 返回所有记忆条目的浅拷贝。 +// List 返回索引中的所有记忆条目浅拷贝。 func (s *Service) List(ctx context.Context) ([]Entry, error) { s.mu.Lock() defer s.mu.Unlock() @@ -181,7 +236,7 @@ func (s *Service) List(ctx context.Context) ([]Entry, error) { return result, nil } -// Search 按关键词搜索记忆条目，返回匹配结果。 +// Search 按关键词搜索记忆条目。 func (s *Service) Search(ctx context.Context, keyword string) ([]Entry, error) { s.mu.Lock() defer s.mu.Unlock() @@ -229,15 +284,14 @@ func (s *Service) Recall(ctx context.Context, keyword string) (map[string]string return results, nil } -// invalidateCache 触发上下文源的缓存失效回调。 +// invalidateCache 触发上下文源缓存失效回调。 func (s *Service) invalidateCache() { if s.sourceInvl != nil { s.sourceInvl() } } -// matchesKeyword 检查条目是否匹配关键词（标题、关键词列表、类型）。 -// 调用方须确保 keyword 已转换为小写。 +// matchesKeyword 检查条目是否匹配关键词。 func matchesKeyword(entry Entry, keyword string) bool { if strings.Contains(strings.ToLower(entry.Title), keyword) { return true @@ -253,6 +307,18 @@ func matchesKeyword(entry Entry, keyword string) bool { return false } +// normalizeEntryForPersist 统一校验并标准化写入前的记忆条目。 +func normalizeEntryForPersist(entry Entry) (Entry, error) { + if !IsValidType(entry.Type) { + return Entry{}, fmt.Errorf("memo: invalid type %q", entry.Type) + } + entry.Title = NormalizeTitle(entry.Title) + if entry.Title == "" { + return Entry{}, fmt.Errorf("memo: title is empty") + } + return entry, nil +} + // newEntryID 生成格式为 __ 的唯一 ID。 func newEntryID(t Type) string { ts := fmt.Sprintf("%x", time.Now().Unix()) @@ -262,7 +328,7 @@ func newEntryID(t Type) string { return fmt.Sprintf("%s_%s_%s", t, ts, randHex) } -// cloneIndex 复制索引结构，避免在持久化失败时污染内存中的原始数据引用。 +// cloneIndex 复制索引结构，避免持久化失败时污染原始数据引用。 func cloneIndex(index *Index) *Index { if index == nil { return &Index{} diff --git a/internal/memo/types.go b/internal/memo/types.go index 9ea4442d..62ab2b9d 100644 --- a/internal/memo/types.go +++ b/internal/memo/types.go @@ -7,82 +7,70 @@ import ( providertypes "neo-code/internal/provider/types" ) -// Type 定义记忆条目的分类。 +// Type 定义记忆条目的闭合分类。 type Type string const ( - // TypeUser 表示用户画像、偏好、专长类记忆。 + // TypeUser 表示用户画像、偏好和背景信息。 TypeUser Type = "user" - // TypeFeedback 表示用户纠正与指导类记忆。 + // TypeFeedback 表示用户对 Agent 行为的纠正与指导。 TypeFeedback Type = "feedback" - // TypeProject 表示项目事实、决策、进行中工作类记忆。 + // TypeProject 表示项目事实、决策和进行中的工作。 TypeProject Type = "project" - // TypeReference 表示外部资源指针类记忆。 + // TypeReference 表示外部资源和信息入口。 TypeReference Type = "reference" ) -// SourceAutoExtract 表示记忆由提取器自动生成。 -const SourceAutoExtract = "extractor_auto" - -// SourceUserManual 表示记忆由用户手动创建。 -const SourceUserManual = "user_manual" - -// SourceToolInitiated 表示记忆由模型通过工具主动创建。 -const SourceToolInitiated = "tool_initiated" +const ( + // SourceAutoExtract 表示记忆由自动提取器生成。 + SourceAutoExtract = "extractor_auto" + // SourceUserManual 表示记忆由用户手动创建。 + SourceUserManual = "user_manual" + // SourceToolInitiated 表示记忆由模型主动调用工具创建。 + SourceToolInitiated = "tool_initiated" +) // Entry 表示单个持久化记忆条目。 type Entry struct { - // ID 唯一标识，格式为 __。 - ID string - // Type 记忆分类。 - Type Type - // Title 索引行显示内容，约 150 字符以内。 - Title string - // Content 详细内容，存入 topic 文件。 - Content string - // Keywords 关键词列表，用于搜索和去重。 - Keywords []string - // Source 记忆来源：extractor_auto / user_manual / tool_initiated。 - Source string - // TopicFile 对应的 topic 文件名（如 user_profile.md）。 + ID string + Type Type + Title string + Content string + Keywords []string + Source string TopicFile string - // CreatedAt 创建时间。 CreatedAt time.Time - // UpdatedAt 最后更新时间。 UpdatedAt time.Time } // Index 表示 MEMO.md 索引文件的内存模型。 type Index struct { - // Entries 索引中的所有记忆条目。 - Entries []Entry - // UpdatedAt 索引最后更新时间。 + Entries []Entry UpdatedAt time.Time } // Store 定义记忆持久化的最小抽象。 type Store interface { - // LoadIndex 加载工作区级别的 MEMO.md 索引。 LoadIndex(ctx context.Context) (*Index, error) - // SaveIndex 持久化索引到 MEMO.md。 SaveIndex(ctx context.Context, index *Index) error - // LoadTopic 加载指定 topic 文件的内容。 LoadTopic(ctx context.Context, filename string) (string, error) - // SaveTopic 持久化 topic 文件内容。 SaveTopic(ctx context.Context, filename string, content string) error - // DeleteTopic 删除指定 topic 文件。 DeleteTopic(ctx context.Context, filename string) error - // ListTopics 列出所有 topic 文件名。 ListTopics(ctx context.Context) ([]string, error) } -// Extractor 定义从对话中提取记忆的最小能力。 +// Extractor 定义从对话消息中提取记忆的最小能力。 type Extractor interface { - // Extract 从对话消息中提取值得记忆的信息。 Extract(ctx context.Context, messages []providertypes.Message) ([]Entry, error) } -// ValidTypes 返回所有合法的记忆类型列表。 +// TextGenerator 定义调用 LLM 生成文本的最小能力，用于记忆提取。 +// 该接口隔离 provider 细节，避免 memo 包直接依赖 provider 基础设施。 +type TextGenerator interface { + Generate(ctx context.Context, prompt string, messages []providertypes.Message) (string, error) +} + +// ValidTypes 返回所有合法的记忆类型。 func ValidTypes() []Type { return []Type{TypeUser, TypeFeedback, TypeProject, TypeReference} } diff --git a/internal/runtime/runtime.go b/internal/runtime/runtime.go index 9707023e..ec67ba23 100644 --- a/internal/runtime/runtime.go +++ b/internal/runtime/runtime.go @@ -29,10 +29,6 @@ const ( providerRetryBaseWait = 1 * time.Second // providerRetryMaxWait 是 runtime 层重试的最大等待时间。 providerRetryMaxWait = 5 * time.Second - // memoExtractTimeout 控制单次异步记忆提取的最长执行时间。 - memoExtractTimeout = 3 * time.Second - // memoExtractQueueSize 控制并发记忆提取任务上限，避免高频请求放大后台负载。 - memoExtractQueueSize = 4 ) // streamAccumulator 在流式事件处理过程中累积本轮对话需要持久化的助手消息状态， @@ -139,7 +135,7 @@ type ProviderFactory interface { // 与 memo.Extractor 解耦，避免 runtime 直接依赖 memo 包的具体类型。 type MemoExtractor interface { // ExtractAndStore 从消息中提取记忆并保存，失败静默处理。 - ExtractAndStore(ctx context.Context, messages []providertypes.Message) + Schedule(sessionID string, messages []providertypes.Message, skip bool) } // SetMemoExtractor 设置可选的记忆提取钩子，完成后由 ReAct 循环调用。 @@ -155,7 +151,6 @@ type Service struct { contextBuilder agentcontext.Builder // 上下文构建器，负责组装 system prompt 与本轮发给模型的消息上下文。 compactRunner contextcompact.Runner memoExtractor MemoExtractor // 可选的记忆提取钩子，完成后异步提取记忆。 - memoExtractSlots chan struct{} // 异步记忆提取并发限流槽位。 events chan RuntimeEvent operationMu sync.Mutex // 运行级互斥：串行化 Run 与 Compact，避免并发写同一会话。 runMu sync.Mutex // 仅保护 activeRun* 字段的并发读写。 @@ -184,13 +179,12 @@ func NewWithFactory( } return &Service{ - configManager: configManager, - sessionStore: sessionStore, - toolManager: toolManager, - providerFactory: providerFactory, - contextBuilder: contextBuilder, - memoExtractSlots: make(chan struct{}, memoExtractQueueSize), - events: make(chan RuntimeEvent, 128), + configManager: configManager, + sessionStore: sessionStore, + toolManager: toolManager, + providerFactory: providerFactory, + contextBuilder: contextBuilder, + events: make(chan RuntimeEvent, 128), } } @@ -230,6 +224,7 @@ func (s *Service) Run(ctx context.Context, input UserInput) error { autoCompacted := false reactiveRetried := false + rememberedThisRun := false for attempt := 0; ; attempt++ { if err := ctx.Err(); err != nil { @@ -350,7 +345,10 @@ func (s *Service) Run(ctx context.Context, input UserInput) error { } if len(assistant.ToolCalls) == 0 { s.emit(ctx, EventAgentDone, input.RunID, session.ID, assistant) - s.scheduleMemoExtract(ctx, userMessage) + // 异步提取记忆：不影响主循环，失败静默处理。 + if s.memoExtractor != nil { + s.memoExtractor.Schedule(session.ID, cloneMessages(session.Messages), rememberedThisRun) + } return nil } @@ -405,6 +403,9 @@ func (s *Service) Run(ctx context.Context, input UserInput) error { } s.emit(ctx, EventToolResult, input.RunID, session.ID, result) + if isSuccessfulRememberToolCall(call.Name, result, execErr) { + rememberedThisRun = true + } if execErr != nil { if err := ctx.Err(); err != nil { return s.handleRunError(ctx, input.RunID, session.ID, err) @@ -779,6 +780,42 @@ func (s *Service) isRunCanceled(err error) bool { return errors.Is(err, context.Canceled) } +// isSuccessfulRememberToolCall 判断工具调用是否已成功完成显式记忆写入。 +func isSuccessfulRememberToolCall(callName string, result tools.ToolResult, execErr error) bool { + if execErr != nil || result.IsError { + return false + } + + toolName := strings.TrimSpace(result.Name) + if toolName == "" { + toolName = strings.TrimSpace(callName) + } + return toolName == tools.ToolNameMemoRemember +} + +// cloneMessages 深拷贝消息切片，避免后台调度读取到后续运行态修改。 +func cloneMessages(messages []providertypes.Message) []providertypes.Message { + if len(messages) == 0 { + return nil + } + + cloned := make([]providertypes.Message, 0, len(messages)) + for _, message := range messages { + next := message + if len(message.ToolCalls) > 0 { + next.ToolCalls = append([]providertypes.ToolCall(nil), message.ToolCalls...) + } + if len(message.ToolMetadata) > 0 { + next.ToolMetadata = make(map[string]string, len(message.ToolMetadata)) + for key, value := range message.ToolMetadata { + next.ToolMetadata[key] = value + } + } + cloned = append(cloned, next) + } + return cloned +} + type permissionEventView struct { toolName string actionType string @@ -876,51 +913,3 @@ func resolveWorkdirForSession(defaultWorkdir string, currentWorkdir string, requ func normalizeExistingWorkdir(workdir string) (string, error) { return agentsession.ResolveExistingDir(workdir) } - -// scheduleMemoExtract 在回答完成后异步触发记忆提取，不阻塞主流程。 -// 调用方仅传入本轮用户消息，并通过超时与限流控制后台资源占用。 -func (s *Service) scheduleMemoExtract(parent context.Context, userMessage providertypes.Message) { - if s == nil || s.memoExtractor == nil { - return - } - if strings.TrimSpace(userMessage.Role) == "" { - userMessage.Role = providertypes.RoleUser - } - if userMessage.Role != providertypes.RoleUser { - return - } - if !s.tryAcquireMemoExtractSlot() { - return - } - - extractCtx, cancel := context.WithTimeout(context.WithoutCancel(parent), memoExtractTimeout) - go func(message providertypes.Message) { - defer cancel() - defer s.releaseMemoExtractSlot() - s.memoExtractor.ExtractAndStore(extractCtx, []providertypes.Message{message}) - }(userMessage) -} - -// tryAcquireMemoExtractSlot 以非阻塞方式申请记忆提取槽位。 -func (s *Service) tryAcquireMemoExtractSlot() bool { - if s == nil || s.memoExtractSlots == nil { - return false - } - select { - case s.memoExtractSlots <- struct{}{}: - return true - default: - return false - } -} - -// releaseMemoExtractSlot 释放记忆提取槽位。 -func (s *Service) releaseMemoExtractSlot() { - if s == nil || s.memoExtractSlots == nil { - return - } - select { - case <-s.memoExtractSlots: - default: - } -} diff --git a/internal/runtime/runtime_test.go b/internal/runtime/runtime_test.go index c35ad5ea..acb8f762 100644 --- a/internal/runtime/runtime_test.go +++ b/internal/runtime/runtime_test.go @@ -310,21 +310,24 @@ func (m *stubToolManager) RememberSessionDecision(sessionID string, action secur return m.rememberErr } -type stubMemoExtractor struct { - calls int - messages []providertypes.Message - done chan struct{} +type stubScheduledMemoExtractor struct { + calls []struct { + sessionID string + messages []providertypes.Message + skip bool + } } -func (s *stubMemoExtractor) ExtractAndStore(ctx context.Context, messages []providertypes.Message) { - s.calls++ - s.messages = append([]providertypes.Message(nil), messages...) - if s.done != nil { - select { - case s.done <- struct{}{}: - default: - } - } +func (s *stubScheduledMemoExtractor) Schedule(sessionID string, messages []providertypes.Message, skip bool) { + s.calls = append(s.calls, struct { + sessionID string + messages []providertypes.Message + skip bool + }{ + sessionID: sessionID, + messages: cloneMessages(messages), + skip: skip, + }) } func TestServiceRun(t *testing.T) { @@ -507,6 +510,115 @@ func TestServiceRun(t *testing.T) { } } +func TestServiceRunSchedulesMemoExtractionAfterFinalReply(t *testing.T) { + t.Parallel() + + manager := newRuntimeConfigManager(t) + store := newMemoryStore() + registry := tools.NewRegistry() + registry.Register(&stubTool{name: "filesystem_read_file", content: "default"}) + + scripted := &scriptedProvider{ + streams: [][]providertypes.StreamEvent{ + { + providertypes.NewTextDeltaStreamEvent("final answer"), + providertypes.NewMessageDoneStreamEvent("stop", nil), + }, + }, + } + + service := NewWithFactory(manager, registry, store, &scriptedProviderFactory{provider: scripted}, &stubContextBuilder{}) + memoExtractor := &stubScheduledMemoExtractor{} + service.SetMemoExtractor(memoExtractor) + + if err := service.Run(context.Background(), UserInput{RunID: "run-memo-schedule", Content: "hello"}); err != nil { + t.Fatalf("Run() error = %v", err) + } + if len(memoExtractor.calls) != 1 { + t.Fatalf("memo schedule calls = %d, want 1", len(memoExtractor.calls)) + } + if memoExtractor.calls[0].skip { + t.Fatalf("skip = true, want false") + } + if len(memoExtractor.calls[0].messages) != 2 { + t.Fatalf("scheduled messages = %#v", memoExtractor.calls[0].messages) + } +} + +func TestServiceRunSkipsAutoMemoExtractionAfterRememberTool(t *testing.T) { + t.Parallel() + + manager := newRuntimeConfigManager(t) + store := newMemoryStore() + registry := tools.NewRegistry() + registry.Register(&stubTool{name: tools.ToolNameMemoRemember, content: "Memory saved"}) + + scripted := &scriptedProvider{ + streams: [][]providertypes.StreamEvent{ + { + providertypes.NewToolCallStartStreamEvent(0, "call-remember", tools.ToolNameMemoRemember), + providertypes.NewToolCallDeltaStreamEvent(0, "call-remember", `{"type":"user","title":"t","content":"c"}`), + providertypes.NewMessageDoneStreamEvent("tool_calls", nil), + }, + { + providertypes.NewTextDeltaStreamEvent("done"), + providertypes.NewMessageDoneStreamEvent("stop", nil), + }, + }, + } + + service := NewWithFactory(manager, registry, store, &scriptedProviderFactory{provider: scripted}, &stubContextBuilder{}) + memoExtractor := &stubScheduledMemoExtractor{} + service.SetMemoExtractor(memoExtractor) + + if err := service.Run(context.Background(), UserInput{RunID: "run-remember-skip", Content: "remember this"}); err != nil { + t.Fatalf("Run() error = %v", err) + } + if len(memoExtractor.calls) != 1 { + t.Fatalf("memo schedule calls = %d, want 1", len(memoExtractor.calls)) + } + if !memoExtractor.calls[0].skip { + t.Fatalf("skip = false, want true") + } +} + +func TestServiceRunSchedulesMemoExtractionOnlyAfterFinalCompletion(t *testing.T) { + t.Parallel() + + manager := newRuntimeConfigManager(t) + store := newMemoryStore() + registry := tools.NewRegistry() + registry.Register(&stubTool{name: "filesystem_edit", content: "tool output"}) + + scripted := &scriptedProvider{ + streams: [][]providertypes.StreamEvent{ + { + providertypes.NewToolCallStartStreamEvent(0, "call-1", "filesystem_edit"), + providertypes.NewToolCallDeltaStreamEvent(0, "call-1", `{"path":"main.go"}`), + providertypes.NewMessageDoneStreamEvent("tool_calls", nil), + }, + { + providertypes.NewTextDeltaStreamEvent("done"), + providertypes.NewMessageDoneStreamEvent("stop", nil), + }, + }, + } + + service := NewWithFactory(manager, registry, store, &scriptedProviderFactory{provider: scripted}, &stubContextBuilder{}) + memoExtractor := &stubScheduledMemoExtractor{} + service.SetMemoExtractor(memoExtractor) + + if err := service.Run(context.Background(), UserInput{RunID: "run-final-only", Content: "edit file"}); err != nil { + t.Fatalf("Run() error = %v", err) + } + if len(memoExtractor.calls) != 1 { + t.Fatalf("memo schedule calls = %d, want 1", len(memoExtractor.calls)) + } + if len(memoExtractor.calls[0].messages) != 4 { + t.Fatalf("scheduled messages = %#v", memoExtractor.calls[0].messages) + } +} + func TestServiceRunMergesLateToolCallMetadata(t *testing.T) { t.Parallel() @@ -1402,60 +1514,6 @@ func TestServiceNewWithFactoryDefaultsToolManager(t *testing.T) { assertEventSequence(t, events, []EventType{EventUserMessage, EventAgentDone}) } -func TestServiceRunMemoExtractorOnlyReceivesLatestUserMessage(t *testing.T) { - manager := newRuntimeConfigManager(t) - store := newMemoryStore() - now := time.Now() - store.sessions["session-memo"] = agentsession.Session{ - ID: "session-memo", - Title: "history", - CreatedAt: now.Add(-time.Minute), - UpdatedAt: now.Add(-time.Minute), - Messages: []providertypes.Message{ - {Role: providertypes.RoleUser, Content: "旧偏好"}, - {Role: providertypes.RoleAssistant, Content: "收到"}, - }, - } - - service := NewWithFactory(manager, &stubToolManager{}, store, &scriptedProviderFactory{ - provider: &scriptedProvider{ - streams: [][]providertypes.StreamEvent{ - {providertypes.NewTextDeltaStreamEvent("done")}, - }, - }, - }, nil) - - extractor := &stubMemoExtractor{done: make(chan struct{}, 1)} - service.SetMemoExtractor(extractor) - if err := service.Run(context.Background(), UserInput{ - SessionID: "session-memo", - RunID: "run-memo-extract", - Content: "记住以后都用 tab", - }); err != nil { - t.Fatalf("Run() error = %v", err) - } - - select { - case <-extractor.done: - case <-time.After(2 * time.Second): - t.Fatal("memo extractor was not called") - } - - if extractor.calls != 1 { - t.Fatalf("memo extractor calls = %d, want 1", extractor.calls) - } - if len(extractor.messages) != 1 { - t.Fatalf("memo extractor message count = %d, want 1", len(extractor.messages)) - } - got := extractor.messages[0] - if got.Role != providertypes.RoleUser { - t.Fatalf("memo extractor role = %q, want %q", got.Role, providertypes.RoleUser) - } - if got.Content != "记住以后都用 tab" { - t.Fatalf("memo extractor content = %q", got.Content) - } -} - func TestServiceRunErrorPaths(t *testing.T) { tests := []struct { name string From 96ba85d90028d0b2810d0b8a1bba4026091c94c7 Mon Sep 17 00:00:00 2001 From: Yumiue <229866007@qq.com> Date: Mon, 13 Apr 2026 16:33:04 +0800 Subject: [PATCH 3/6] fix(memo): harden auto extraction scheduling --- internal/app/bootstrap.go | 82 ++++-- internal/app/provider_text_generator_test.go | 104 +++++++- internal/memo/auto_extractor.go | 254 ++++++++++--------- internal/memo/auto_extractor_test.go | 108 +++++--- internal/memo/service.go | 150 +++++++++-- internal/runtime/runtime.go | 15 +- internal/runtime/runtime_test.go | 52 +++- 7 files changed, 546 insertions(+), 219 deletions(-) diff --git a/internal/app/bootstrap.go b/internal/app/bootstrap.go index 349e1b76..2a509906 100644 --- a/internal/app/bootstrap.go +++ b/internal/app/bootstrap.go @@ -3,6 +3,7 @@ package app import ( "context" "fmt" + "log" "strings" "time" @@ -12,6 +13,7 @@ import ( configstate "neo-code/internal/config/state" agentcontext "neo-code/internal/context" "neo-code/internal/memo" + "neo-code/internal/provider" "neo-code/internal/provider/builtin" providercatalog "neo-code/internal/provider/catalog" providertypes "neo-code/internal/provider/types" @@ -119,8 +121,11 @@ func BuildRuntime(ctx context.Context, opts BootstrapOptions) (RuntimeBundle, er // 注入记忆提取钩子：当 AutoExtract 启用且 memoSvc 可用时，ReAct 循环完成后异步提取记忆。 if memoSvc != nil && cfg.Memo.AutoExtract { - textGen := newProviderTextGenerator(providerRegistry, manager) - runtimeSvc.SetMemoExtractor(memo.NewAutoExtractor(memo.NewLLMExtractor(textGen), memoSvc)) + runtimeSvc.SetMemoExtractor(newMemoExtractorAdapter( + providerRegistry, + manager, + memo.NewAutoExtractor(nil, memoSvc), + )) } return RuntimeBundle{ @@ -222,33 +227,78 @@ func buildToolManager(registry *tools.Registry) (tools.Manager, error) { return tools.NewManager(registry, engine, security.NewWorkspaceSandbox()) } -// providerTextGenerator 复用当前 provider 配置，向 memo 提供纯文本生成能力。 -type providerTextGenerator struct { +type memoExtractorScheduler interface { + ScheduleWithExtractor(sessionID string, messages []providertypes.Message, extractor memo.Extractor) +} + +// memoExtractorAdapter 在调度时绑定 provider 配置快照，避免后台任务读取全局可变配置。 +type memoExtractorAdapter struct { factory agentruntime.ProviderFactory configManager *config.Manager + scheduler memoExtractorScheduler } -// newProviderTextGenerator 创建基于当前 provider 选择的文本生成适配器。 -func newProviderTextGenerator(factory agentruntime.ProviderFactory, cm *config.Manager) *providerTextGenerator { - return &providerTextGenerator{ +// newMemoExtractorAdapter 创建绑定当前 provider 选择的记忆提取调度适配器。 +func newMemoExtractorAdapter( + factory agentruntime.ProviderFactory, + cm *config.Manager, + scheduler memoExtractorScheduler, +) *memoExtractorAdapter { + return &memoExtractorAdapter{ factory: factory, configManager: cm, + scheduler: scheduler, + } +} + +// Schedule 在当前运行结束时绑定 provider/model 快照，再交给后台调度器延后执行。 +func (a *memoExtractorAdapter) Schedule(sessionID string, messages []providertypes.Message) { + if a == nil || a.scheduler == nil { + return + } + + cfg := a.configManager.Get() + resolved, err := config.ResolveSelectedProvider(cfg) + if err != nil { + log.Printf("memo: resolve selected provider failed: %v", err) + return } + + extractor := memo.NewLLMExtractor(newProviderTextGenerator( + a.factory, + resolved.ToRuntimeConfig(), + cfg.CurrentModel, + )) + a.scheduler.ScheduleWithExtractor(sessionID, messages, extractor) +} + +// providerTextGenerator 复用当前 provider 配置快照，向 memo 提供纯文本生成能力。 +type providerTextGenerator struct { + factory agentruntime.ProviderFactory + runtimeCfg provider.RuntimeConfig + model string } -// Generate 使用当前选中的 provider/model 发起无工具的独立生成请求。 +// newProviderTextGenerator 创建绑定固定 provider/model 快照的文本生成适配器。 +func newProviderTextGenerator( + factory agentruntime.ProviderFactory, + runtimeCfg provider.RuntimeConfig, + model string, +) *providerTextGenerator { + return &providerTextGenerator{ + factory: factory, + runtimeCfg: runtimeCfg, + model: model, + } +} + +// Generate 使用预先绑定的 provider/model 快照发起无工具的独立生成请求。 func (g *providerTextGenerator) Generate( ctx context.Context, prompt string, messages []providertypes.Message, ) (string, error) { - cfg := g.configManager.Get() - resolved, err := config.ResolveSelectedProvider(cfg) - if err != nil { - return "", err - } - - modelProvider, err := g.factory.Build(ctx, resolved.ToRuntimeConfig()) + modelProvider, err := g.factory.Build(ctx, g.runtimeCfg) if err != nil { return "", err } @@ -292,7 +342,7 @@ func (g *providerTextGenerator) Generate( }() err = modelProvider.Generate(ctx, providertypes.GenerateRequest{ - Model: cfg.CurrentModel, + Model: g.model, SystemPrompt: prompt, Messages: append([]providertypes.Message(nil), messages...), }, events) diff --git a/internal/app/provider_text_generator_test.go b/internal/app/provider_text_generator_test.go index 8de740fb..b02e82bd 100644 --- a/internal/app/provider_text_generator_test.go +++ b/internal/app/provider_text_generator_test.go @@ -7,6 +7,7 @@ import ( "testing" "neo-code/internal/config" + "neo-code/internal/memo" "neo-code/internal/provider" providertypes "neo-code/internal/provider/types" ) @@ -44,9 +45,28 @@ func (s *stubTextGenFactory) Build(ctx context.Context, cfg provider.RuntimeConf return s.provider, nil } +type stubMemoScheduler struct { + calls int + sessionID string + messages []providertypes.Message + extractor memo.Extractor +} + +func (s *stubMemoScheduler) ScheduleWithExtractor( + sessionID string, + messages []providertypes.Message, + extractor memo.Extractor, +) { + s.calls++ + s.sessionID = sessionID + s.messages = append([]providertypes.Message(nil), messages...) + s.extractor = extractor +} + // TestProviderTextGeneratorGenerateSuccess 验证文本生成器可以聚合流式文本并且不传 tools。 func TestProviderTextGeneratorGenerateSuccess(t *testing.T) { manager := newLoadedConfigManagerForTextGenerator(t) + runtimeCfg, model := resolvedProviderConfigForTextGenerator(t, manager) providerStub := &stubTextGenProvider{ generate: func(ctx context.Context, req providertypes.GenerateRequest, events chan<- providertypes.StreamEvent) error { events <- providertypes.NewTextDeltaStreamEvent("hello ") @@ -56,7 +76,7 @@ func TestProviderTextGeneratorGenerateSuccess(t *testing.T) { }, } factory := &stubTextGenFactory{provider: providerStub} - generator := newProviderTextGenerator(factory, manager) + generator := newProviderTextGenerator(factory, runtimeCfg, model) text, err := generator.Generate(context.Background(), "memo system prompt", []providertypes.Message{ {Role: providertypes.RoleUser, Content: "记住这个。"}, @@ -83,13 +103,17 @@ func TestProviderTextGeneratorGenerateSuccess(t *testing.T) { if request.Messages[0].Content != "记住这个。" { t.Fatalf("Messages = %#v", request.Messages) } + if len(factory.configs) != 1 || factory.configs[0].BaseURL != runtimeCfg.BaseURL { + t.Fatalf("Build() config = %+v, want base_url %q", factory.configs, runtimeCfg.BaseURL) + } } // TestProviderTextGeneratorGenerateBuildFailure 验证 provider 构建失败会原样返回。 func TestProviderTextGeneratorGenerateBuildFailure(t *testing.T) { manager := newLoadedConfigManagerForTextGenerator(t) + runtimeCfg, model := resolvedProviderConfigForTextGenerator(t, manager) factory := &stubTextGenFactory{err: errors.New("build failed")} - generator := newProviderTextGenerator(factory, manager) + generator := newProviderTextGenerator(factory, runtimeCfg, model) _, err := generator.Generate(context.Background(), "prompt", nil) if err == nil || !strings.Contains(err.Error(), "build failed") { @@ -100,12 +124,13 @@ func TestProviderTextGeneratorGenerateBuildFailure(t *testing.T) { // TestProviderTextGeneratorGenerateFailure 验证 provider.Generate 失败会向上返回。 func TestProviderTextGeneratorGenerateFailure(t *testing.T) { manager := newLoadedConfigManagerForTextGenerator(t) + runtimeCfg, model := resolvedProviderConfigForTextGenerator(t, manager) providerStub := &stubTextGenProvider{ generate: func(ctx context.Context, req providertypes.GenerateRequest, events chan<- providertypes.StreamEvent) error { return errors.New("generate failed") }, } - generator := newProviderTextGenerator(&stubTextGenFactory{provider: providerStub}, manager) + generator := newProviderTextGenerator(&stubTextGenFactory{provider: providerStub}, runtimeCfg, model) _, err := generator.Generate(context.Background(), "prompt", nil) if err == nil || !strings.Contains(err.Error(), "generate failed") { @@ -116,13 +141,14 @@ func TestProviderTextGeneratorGenerateFailure(t *testing.T) { // TestProviderTextGeneratorGenerateRequiresMessageDone 验证缺失 message_done 会视为错误。 func TestProviderTextGeneratorGenerateRequiresMessageDone(t *testing.T) { manager := newLoadedConfigManagerForTextGenerator(t) + runtimeCfg, model := resolvedProviderConfigForTextGenerator(t, manager) providerStub := &stubTextGenProvider{ generate: func(ctx context.Context, req providertypes.GenerateRequest, events chan<- providertypes.StreamEvent) error { events <- providertypes.NewTextDeltaStreamEvent("partial") return nil }, } - generator := newProviderTextGenerator(&stubTextGenFactory{provider: providerStub}, manager) + generator := newProviderTextGenerator(&stubTextGenFactory{provider: providerStub}, runtimeCfg, model) _, err := generator.Generate(context.Background(), "prompt", nil) if err == nil || !strings.Contains(err.Error(), "message_done") { @@ -130,6 +156,65 @@ func TestProviderTextGeneratorGenerateRequiresMessageDone(t *testing.T) { } } +// TestMemoExtractorAdapterBindsProviderSnapshotAtScheduleTime 验证调度时会固定 provider/model 快照。 +func TestMemoExtractorAdapterBindsProviderSnapshotAtScheduleTime(t *testing.T) { + manager := newLoadedConfigManagerForTextGenerator(t) + oldResolved, err := config.ResolveSelectedProvider(manager.Get()) + if err != nil { + t.Fatalf("ResolveSelectedProvider() error = %v", err) + } + + scheduler := &stubMemoScheduler{} + providerStub := &stubTextGenProvider{ + generate: func(ctx context.Context, req providertypes.GenerateRequest, events chan<- providertypes.StreamEvent) error { + events <- providertypes.NewTextDeltaStreamEvent("[]") + events <- providertypes.NewMessageDoneStreamEvent("stop", nil) + return nil + }, + } + factory := &stubTextGenFactory{provider: providerStub} + adapter := newMemoExtractorAdapter(factory, manager, scheduler) + + adapter.Schedule("session-1", []providertypes.Message{ + {Role: providertypes.RoleUser, Content: "记住我的偏好"}, + }) + + if scheduler.calls != 1 { + t.Fatalf("ScheduleWithExtractor() calls = %d, want 1", scheduler.calls) + } + if scheduler.extractor == nil { + t.Fatal("expected extractor snapshot") + } + + if err := manager.Update(context.Background(), func(cfg *config.Config) error { + cfg.CurrentModel = "changed-model" + for i := range cfg.Providers { + if cfg.Providers[i].Name == cfg.SelectedProvider { + cfg.Providers[i].BaseURL = "https://changed.example.com/v1" + } + } + return nil + }); err != nil { + t.Fatalf("Update() error = %v", err) + } + + if _, err := scheduler.extractor.Extract(context.Background(), scheduler.messages); err != nil { + t.Fatalf("Extract() error = %v", err) + } + if len(factory.configs) != 1 { + t.Fatalf("Build() calls = %d, want 1", len(factory.configs)) + } + if factory.configs[0].BaseURL != oldResolved.BaseURL { + t.Fatalf("Build() base_url = %q, want %q", factory.configs[0].BaseURL, oldResolved.BaseURL) + } + if len(providerStub.requests) != 1 { + t.Fatalf("provider requests = %d, want 1", len(providerStub.requests)) + } + if providerStub.requests[0].Model != config.OpenAIDefaultModel { + t.Fatalf("request model = %q, want %q", providerStub.requests[0].Model, config.OpenAIDefaultModel) + } +} + // newLoadedConfigManagerForTextGenerator 创建带默认 provider 选择的配置管理器。 func newLoadedConfigManagerForTextGenerator(t *testing.T) *config.Manager { t.Helper() @@ -154,3 +239,14 @@ func newLoadedConfigManagerForTextGenerator(t *testing.T) *config.Manager { } return manager } + +// resolvedProviderConfigForTextGenerator 返回当前默认 provider 的运行时配置快照。 +func resolvedProviderConfigForTextGenerator(t *testing.T, manager *config.Manager) (provider.RuntimeConfig, string) { + t.Helper() + + resolved, err := config.ResolveSelectedProvider(manager.Get()) + if err != nil { + t.Fatalf("ResolveSelectedProvider() error = %v", err) + } + return resolved.ToRuntimeConfig(), manager.Get().CurrentModel +} diff --git a/internal/memo/auto_extractor.go b/internal/memo/auto_extractor.go index 69d5af7f..4dd2bcbd 100644 --- a/internal/memo/auto_extractor.go +++ b/internal/memo/auto_extractor.go @@ -10,27 +10,37 @@ import ( providertypes "neo-code/internal/provider/types" ) -const autoExtractDebounce = 2 * time.Second +const ( + autoExtractDebounce = 2 * time.Second + autoExtractIdleTTL = 5 * time.Minute +) // AutoExtractor 负责按会话在后台调度自动提取，并处理防抖、互斥和尾随执行。 type AutoExtractor struct { extractor Extractor svc *Service debounce time.Duration + idleTTL time.Duration logf func(format string, args ...any) - mu sync.Mutex - workers map[string]*autoExtractWorker + mu sync.Mutex + states map[string]*autoExtractState } -type autoExtractWorker struct { - updates chan autoExtractRequest +type autoExtractState struct { + mu sync.Mutex + pending *autoExtractRequest + running bool + timer *time.Timer + idleTimer *time.Timer + scheduleSeq uint64 + idleSeq uint64 } type autoExtractRequest struct { - messages []providertypes.Message - dueAt time.Time - skip bool + messages []providertypes.Message + dueAt time.Time + extractor Extractor } // NewAutoExtractor 创建后台自动提取调度器。 @@ -39,14 +49,25 @@ func NewAutoExtractor(extractor Extractor, svc *Service) *AutoExtractor { extractor: extractor, svc: svc, debounce: autoExtractDebounce, + idleTTL: autoExtractIdleTTL, logf: log.Printf, - workers: make(map[string]*autoExtractWorker), + states: make(map[string]*autoExtractState), } } -// Schedule 按会话维度安排一次后台自动提取，skip=true 时清空当前待执行请求。 -func (a *AutoExtractor) Schedule(sessionID string, messages []providertypes.Message, skip bool) { - if a == nil || a.extractor == nil || a.svc == nil { +// Schedule 按会话维度安排一次后台自动提取。 +func (a *AutoExtractor) Schedule(sessionID string, messages []providertypes.Message) { + a.schedule(sessionID, messages, a.extractor) +} + +// ScheduleWithExtractor 允许调用方在调度时绑定本次请求专用的提取器快照。 +func (a *AutoExtractor) ScheduleWithExtractor(sessionID string, messages []providertypes.Message, extractor Extractor) { + a.schedule(sessionID, messages, extractor) +} + +// schedule 统一处理会话级别的自动提取调度。 +func (a *AutoExtractor) schedule(sessionID string, messages []providertypes.Message, extractor Extractor) { + if a == nil || extractor == nil || a.svc == nil { return } @@ -56,115 +77,127 @@ func (a *AutoExtractor) Schedule(sessionID string, messages []providertypes.Mess } req := autoExtractRequest{ - messages: cloneProviderMessages(messages), - dueAt: time.Now().Add(a.debounce), - skip: skip, + messages: cloneProviderMessages(messages), + dueAt: time.Now().Add(a.debounce), + extractor: extractor, + } + + state := a.ensureState(sessionID) + state.mu.Lock() + defer state.mu.Unlock() + + stopTimer(state.idleTimer) + state.idleTimer = nil + reqCopy := req + state.pending = &reqCopy + state.scheduleSeq++ + if state.running { + return } - worker := a.ensureWorker(sessionID) - worker.updates <- req + a.armDebounceTimerLocked(sessionID, state, state.scheduleSeq, req.dueAt) } -// ensureWorker 获取或创建指定会话的后台 worker。 -func (a *AutoExtractor) ensureWorker(sessionID string) *autoExtractWorker { +// ensureState 获取或创建会话级别的调度状态。 +func (a *AutoExtractor) ensureState(sessionID string) *autoExtractState { a.mu.Lock() defer a.mu.Unlock() - if worker, ok := a.workers[sessionID]; ok { - return worker + if state, ok := a.states[sessionID]; ok { + return state } - worker := &autoExtractWorker{ - updates: make(chan autoExtractRequest, 32), + state := &autoExtractState{} + a.states[sessionID] = state + return state +} + +// armDebounceTimerLocked 在持有状态锁时重置会话的防抖定时器。 +func (a *AutoExtractor) armDebounceTimerLocked( + sessionID string, + state *autoExtractState, + seq uint64, + dueAt time.Time, +) { + stopTimer(state.timer) + wait := time.Until(dueAt) + if wait < 0 { + wait = 0 } - a.workers[sessionID] = worker - go a.runWorker(worker) - return worker + state.timer = time.AfterFunc(wait, func() { + a.handleDebounce(sessionID, state, seq) + }) } -// runWorker 串行处理单个会话的调度状态，避免并发提取并支持 trailing extraction。 -func (a *AutoExtractor) runWorker(worker *autoExtractWorker) { - var ( - pending *autoExtractRequest - timer *time.Timer - timerCh <-chan time.Time - running bool - doneCh = make(chan struct{}, 1) - ) - - for { - select { - case req := <-worker.updates: - if req.skip { - pending = nil - stopTimer(timer) - timer = nil - timerCh = nil - continue - } - - reqCopy := req - pending = &reqCopy - if running { - continue - } - - timer = resetTimer(timer, time.Until(req.dueAt)) - timerCh = timer.C - - case <-timerCh: - timer = nil - timerCh = nil - - for { - select { - case req := <-worker.updates: - if req.skip { - pending = nil - continue - } - reqCopy := req - pending = &reqCopy - default: - goto launch - } - } - - launch: - if running || pending == nil { - continue - } - if wait := time.Until(pending.dueAt); wait > 0 { - timer = resetTimer(timer, wait) - timerCh = timer.C - continue - } - - current := *pending - pending = nil - running = true - - go func(req autoExtractRequest) { - a.extractAndStore(req.messages) - doneCh <- struct{}{} - }(current) - - case <-doneCh: - running = false - if pending == nil { - continue - } - - timer = resetTimer(timer, time.Until(pending.dueAt)) - timerCh = timer.C - } +// handleDebounce 在防抖窗口结束后启动一次后台提取。 +func (a *AutoExtractor) handleDebounce(sessionID string, state *autoExtractState, seq uint64) { + state.mu.Lock() + defer state.mu.Unlock() + + if state.scheduleSeq != seq || state.running || state.pending == nil { + return + } + if wait := time.Until(state.pending.dueAt); wait > 0 { + a.armDebounceTimerLocked(sessionID, state, seq, state.pending.dueAt) + return + } + + req := *state.pending + state.pending = nil + state.running = true + state.timer = nil + + go func() { + a.extractAndStore(req.extractor, req.messages) + a.handleRunDone(sessionID, state) + }() +} + +// handleRunDone 在后台提取结束后决定是否执行尾随提取，或安排空闲回收。 +func (a *AutoExtractor) handleRunDone(sessionID string, state *autoExtractState) { + state.mu.Lock() + defer state.mu.Unlock() + + state.running = false + if state.pending != nil { + a.armDebounceTimerLocked(sessionID, state, state.scheduleSeq, state.pending.dueAt) + return + } + + a.armIdleTimerLocked(sessionID, state) +} + +// armIdleTimerLocked 在会话空闲时安排状态回收，避免 map 与 goroutine 长期累积。 +func (a *AutoExtractor) armIdleTimerLocked(sessionID string, state *autoExtractState) { + stopTimer(state.idleTimer) + state.idleSeq++ + seq := state.idleSeq + state.idleTimer = time.AfterFunc(a.idleTTL, func() { + a.handleIdle(sessionID, state, seq) + }) +} + +// handleIdle 在会话空闲超时后回收状态。 +func (a *AutoExtractor) handleIdle(sessionID string, state *autoExtractState, seq uint64) { + state.mu.Lock() + if state.idleSeq != seq || state.running || state.pending != nil { + state.mu.Unlock() + return + } + state.idleTimer = nil + state.mu.Unlock() + + a.mu.Lock() + defer a.mu.Unlock() + if a.states[sessionID] == state { + delete(a.states, sessionID) } } // extractAndStore 执行提取，并在写入前做本地批次去重和持久化级别的原子去重。 -func (a *AutoExtractor) extractAndStore(messages []providertypes.Message) { +func (a *AutoExtractor) extractAndStore(extractor Extractor, messages []providertypes.Message) { ctx := context.Background() - entries, err := a.extractor.Extract(ctx, messages) + entries, err := extractor.Extract(ctx, messages) if err != nil { a.logError("memo: auto extract failed: %v", err) return @@ -239,19 +272,6 @@ func cloneProviderMessages(messages []providertypes.Message) []providertypes.Mes return cloned } -// resetTimer 以安全方式重置定时器，避免旧事件污染新的防抖窗口。 -func resetTimer(timer *time.Timer, wait time.Duration) *time.Timer { - if wait < 0 { - wait = 0 - } - if timer == nil { - return time.NewTimer(wait) - } - stopTimer(timer) - timer.Reset(wait) - return timer -} - // stopTimer 停止定时器并在必要时清空通道。 func stopTimer(timer *time.Timer) { if timer == nil { diff --git a/internal/memo/auto_extractor_test.go b/internal/memo/auto_extractor_test.go index fc4c6d2f..a80f5d7c 100644 --- a/internal/memo/auto_extractor_test.go +++ b/internal/memo/auto_extractor_test.go @@ -56,8 +56,8 @@ func TestAutoExtractorDebounceMergesRequests(t *testing.T) { auto.debounce = 20 * time.Millisecond auto.logf = func(string, ...any) {} - auto.Schedule("session-1", []providertypes.Message{{Role: providertypes.RoleUser, Content: "first"}}, false) - auto.Schedule("session-1", []providertypes.Message{{Role: providertypes.RoleUser, Content: "second"}}, false) + auto.Schedule("session-1", []providertypes.Message{{Role: providertypes.RoleUser, Content: "first"}}) + auto.Schedule("session-1", []providertypes.Message{{Role: providertypes.RoleUser, Content: "second"}}) waitFor(t, time.Second, func() bool { return extractor.Calls() == 1 }) time.Sleep(60 * time.Millisecond) @@ -80,23 +80,6 @@ func TestAutoExtractorDebounceMergesRequests(t *testing.T) { } } -func TestAutoExtractorSkipCancelsPendingRequest(t *testing.T) { - svc := newAutoExtractorTestService(t) - extractor := &stubMemoExtractor{} - auto := NewAutoExtractor(extractor, svc) - auto.debounce = 40 * time.Millisecond - auto.logf = func(string, ...any) {} - - auto.Schedule("session-1", []providertypes.Message{{Role: providertypes.RoleUser, Content: "remember me"}}, false) - time.Sleep(20 * time.Millisecond) - auto.Schedule("session-1", []providertypes.Message{{Role: providertypes.RoleUser, Content: "remember me"}}, true) - time.Sleep(80 * time.Millisecond) - - if extractor.Calls() != 0 { - t.Fatalf("extractor calls = %d, want 0", extractor.Calls()) - } -} - func TestAutoExtractorTrailingRun(t *testing.T) { svc := newAutoExtractorTestService(t) firstStarted := make(chan struct{}, 1) @@ -120,7 +103,7 @@ func TestAutoExtractorTrailingRun(t *testing.T) { auto.debounce = 15 * time.Millisecond auto.logf = func(string, ...any) {} - auto.Schedule("session-1", []providertypes.Message{{Role: providertypes.RoleUser, Content: "first"}}, false) + auto.Schedule("session-1", []providertypes.Message{{Role: providertypes.RoleUser, Content: "first"}}) select { case <-firstStarted: @@ -128,7 +111,7 @@ func TestAutoExtractorTrailingRun(t *testing.T) { t.Fatal("first extraction did not start") } - auto.Schedule("session-1", []providertypes.Message{{Role: providertypes.RoleUser, Content: "second"}}, false) + auto.Schedule("session-1", []providertypes.Message{{Role: providertypes.RoleUser, Content: "second"}}) time.Sleep(40 * time.Millisecond) close(releaseFirst) @@ -145,21 +128,6 @@ func TestAutoExtractorTrailingRun(t *testing.T) { }) } -func TestAutoExtractorSkipDoesNotQueue(t *testing.T) { - svc := newAutoExtractorTestService(t) - extractor := &stubMemoExtractor{} - auto := NewAutoExtractor(extractor, svc) - auto.debounce = 10 * time.Millisecond - auto.logf = func(string, ...any) {} - - auto.Schedule("session-1", []providertypes.Message{{Role: providertypes.RoleUser, Content: "skip"}}, true) - time.Sleep(50 * time.Millisecond) - - if extractor.Calls() != 0 { - t.Fatalf("extractor calls = %d, want 0", extractor.Calls()) - } -} - func TestAutoExtractorErrorsAreSilent(t *testing.T) { svc := newAutoExtractorTestService(t) extractor := &stubMemoExtractor{ @@ -171,7 +139,7 @@ func TestAutoExtractorErrorsAreSilent(t *testing.T) { auto.debounce = 10 * time.Millisecond auto.logf = func(string, ...any) {} - auto.Schedule("session-1", []providertypes.Message{{Role: providertypes.RoleUser, Content: "x"}}, false) + auto.Schedule("session-1", []providertypes.Message{{Role: providertypes.RoleUser, Content: "x"}}) waitFor(t, time.Second, func() bool { return extractor.Calls() == 1 }) entries, err := svc.List(context.Background()) @@ -207,7 +175,7 @@ func TestAutoExtractorSuppressesExactDuplicates(t *testing.T) { auto.debounce = 10 * time.Millisecond auto.logf = func(string, ...any) {} - auto.Schedule("session-1", []providertypes.Message{{Role: providertypes.RoleUser, Content: "dedupe"}}, false) + auto.Schedule("session-1", []providertypes.Message{{Role: providertypes.RoleUser, Content: "dedupe"}}) waitFor(t, time.Second, func() bool { entries, err := svc.List(context.Background()) return err == nil && len(entries) == 2 @@ -240,8 +208,8 @@ func TestAutoExtractorSuppressesExactDuplicatesAcrossSessions(t *testing.T) { auto.debounce = 0 auto.logf = func(string, ...any) {} - auto.Schedule("session-1", []providertypes.Message{{Role: providertypes.RoleUser, Content: "one"}}, false) - auto.Schedule("session-2", []providertypes.Message{{Role: providertypes.RoleUser, Content: "two"}}, false) + auto.Schedule("session-1", []providertypes.Message{{Role: providertypes.RoleUser, Content: "one"}}) + auto.Schedule("session-2", []providertypes.Message{{Role: providertypes.RoleUser, Content: "two"}}) for i := 0; i < 2; i++ { select { @@ -267,6 +235,66 @@ func TestAutoExtractorSuppressesExactDuplicatesAcrossSessions(t *testing.T) { } } +func TestAutoExtractorRemovesIdleState(t *testing.T) { + svc := newAutoExtractorTestService(t) + extractor := &stubMemoExtractor{ + extractFn: func(ctx context.Context, messages []providertypes.Message) ([]Entry, error) { + return []Entry{{Type: TypeProject, Title: "done", Content: "done", Source: SourceAutoExtract}}, nil + }, + } + auto := NewAutoExtractor(extractor, svc) + auto.debounce = 5 * time.Millisecond + auto.idleTTL = 20 * time.Millisecond + auto.logf = func(string, ...any) {} + + auto.Schedule("session-1", []providertypes.Message{{Role: providertypes.RoleUser, Content: "cleanup"}}) + + waitFor(t, time.Second, func() bool { return extractor.Calls() == 1 }) + waitFor(t, time.Second, func() bool { + auto.mu.Lock() + defer auto.mu.Unlock() + return len(auto.states) == 0 + }) +} + +func TestAutoExtractorLoadsDedupIndexOutsideCurrentProcessState(t *testing.T) { + baseDir := t.TempDir() + workspace := t.TempDir() + store := NewFileStore(baseDir, workspace) + svc := NewService(store, nil, config.MemoConfig{MaxIndexLines: 200}, nil) + if err := svc.Add(context.Background(), Entry{ + Type: TypeUser, + Title: "reply in chinese", + Content: "reply in chinese", + Source: SourceAutoExtract, + }); err != nil { + t.Fatalf("seed Add() error = %v", err) + } + + reloaded := NewService(NewFileStore(baseDir, workspace), nil, config.MemoConfig{MaxIndexLines: 200}, nil) + extractor := &stubMemoExtractor{ + extractFn: func(ctx context.Context, messages []providertypes.Message) ([]Entry, error) { + return []Entry{ + {Type: TypeUser, Title: "reply in chinese", Content: "reply in chinese", Source: SourceAutoExtract}, + }, nil + }, + } + auto := NewAutoExtractor(extractor, reloaded) + auto.debounce = 5 * time.Millisecond + auto.logf = func(string, ...any) {} + + auto.Schedule("session-1", []providertypes.Message{{Role: providertypes.RoleUser, Content: "dedupe after reload"}}) + + waitFor(t, time.Second, func() bool { return extractor.Calls() == 1 }) + entries, err := reloaded.List(context.Background()) + if err != nil { + t.Fatalf("List() error = %v", err) + } + if len(entries) != 1 { + t.Fatalf("len(entries) = %d, want 1", len(entries)) + } +} + func waitFor(t *testing.T, timeout time.Duration, fn func() bool) { t.Helper() deadline := time.Now().Add(timeout) diff --git a/internal/memo/service.go b/internal/memo/service.go index 42325bce..fa88f781 100644 --- a/internal/memo/service.go +++ b/internal/memo/service.go @@ -14,11 +14,15 @@ import ( // Service 编排记忆的存储、检索、提取和删除，是 memo 子系统对外的统一入口。 type Service struct { - store Store - extractor Extractor - config config.MemoConfig - mu sync.Mutex - sourceInvl func() + store Store + extractor Extractor + config config.MemoConfig + mu sync.Mutex + sourceInvl func() + autoExtractIndexMu sync.Mutex + autoExtractIndexReady bool + autoExtractKeysByTopic map[string]string + autoExtractKeyRefs map[string]int } // NewService 创建 memo Service 实例；extractor 可以为 nil。 @@ -50,15 +54,14 @@ func (s *Service) addAutoExtractIfAbsent(ctx context.Context, entry Entry) (bool if err != nil { return false, err } + if err := s.ensureAutoExtractIndex(ctx); err != nil { + return false, err + } s.mu.Lock() defer s.mu.Unlock() - duplicate, err := s.hasExactAutoExtractLocked(ctx, entry) - if err != nil { - return false, err - } - if duplicate { + if s.hasExactAutoExtractLocked(entry) { return false, nil } @@ -90,8 +93,10 @@ func (s *Service) saveEntryLocked(ctx context.Context, entry Entry) error { working := cloneIndex(index) replaced := false + var previous Entry for i, existing := range working.Entries { if existing.ID == entry.ID { + previous = existing working.Entries[i] = entry replaced = true break @@ -126,28 +131,53 @@ func (s *Service) saveEntryLocked(ctx context.Context, entry Entry) error { for _, topicFile := range topicsToDelete { _ = s.store.DeleteTopic(ctx, topicFile) } + if s.autoExtractIndexReady { + if replaced { + s.removeAutoExtractTopicLocked(previous.TopicFile) + } + for _, topicFile := range topicsToDelete { + s.removeAutoExtractTopicLocked(topicFile) + } + s.trackAutoExtractEntryLocked(entry) + } s.invalidateCache() return nil } -// hasExactAutoExtractLocked 检查是否已存在完全相同的自动提取记忆。 -func (s *Service) hasExactAutoExtractLocked(ctx context.Context, target Entry) (bool, error) { - targetKey := autoExtractDedupKey(target) - if targetKey == "" { - return false, nil +// ensureAutoExtractIndex 在锁外预加载自动提取的精确去重索引，避免重复在主锁内扫 topic 文件。 +func (s *Service) ensureAutoExtractIndex(ctx context.Context) error { + s.mu.Lock() + if s.autoExtractIndexReady { + s.mu.Unlock() + return nil } + s.mu.Unlock() - index, err := s.loadIndexLocked(ctx) + s.autoExtractIndexMu.Lock() + defer s.autoExtractIndexMu.Unlock() + + s.mu.Lock() + if s.autoExtractIndexReady { + s.mu.Unlock() + return nil + } + s.mu.Unlock() + + index, err := s.store.LoadIndex(ctx) if err != nil { - return false, err + return fmt.Errorf("memo: load index: %w", err) } + + keysByTopic := make(map[string]string) + keyRefs := make(map[string]int) for _, entry := range index.Entries { - if strings.TrimSpace(entry.TopicFile) == "" { + topicFile := strings.TrimSpace(entry.TopicFile) + if topicFile == "" { continue } - topicContent, err := s.store.LoadTopic(ctx, entry.TopicFile) + topicContent, err := s.store.LoadTopic(ctx, topicFile) if err != nil { continue } @@ -159,11 +189,82 @@ func (s *Service) hasExactAutoExtractLocked(ctx context.Context, target Entry) ( entry.Source = source entry.Content = content - if autoExtractDedupKey(entry) == targetKey { - return true, nil + key := autoExtractDedupKey(entry) + if key == "" { + continue } + + keysByTopic[topicFile] = key + keyRefs[key]++ + } + + s.mu.Lock() + defer s.mu.Unlock() + if s.autoExtractIndexReady { + return nil } - return false, nil + s.autoExtractKeysByTopic = keysByTopic + s.autoExtractKeyRefs = keyRefs + s.autoExtractIndexReady = true + return nil +} + +// hasExactAutoExtractLocked 检查是否已存在完全相同的自动提取记忆。 +func (s *Service) hasExactAutoExtractLocked(target Entry) bool { + targetKey := autoExtractDedupKey(target) + if targetKey == "" { + return false + } + return s.autoExtractKeyRefs[targetKey] > 0 +} + +// trackAutoExtractEntryLocked 在自动提取索引已就绪时维护单条条目的去重键。 +func (s *Service) trackAutoExtractEntryLocked(entry Entry) { + if !s.autoExtractIndexReady { + return + } + + topicFile := strings.TrimSpace(entry.TopicFile) + if topicFile == "" { + return + } + s.removeAutoExtractTopicLocked(topicFile) + + if entry.Source != SourceAutoExtract { + return + } + + key := autoExtractDedupKey(entry) + if key == "" { + return + } + + s.autoExtractKeysByTopic[topicFile] = key + s.autoExtractKeyRefs[key]++ +} + +// removeAutoExtractTopicLocked 从精确去重索引中移除指定 topic 的记录。 +func (s *Service) removeAutoExtractTopicLocked(topicFile string) { + if !s.autoExtractIndexReady { + return + } + + topicFile = strings.TrimSpace(topicFile) + if topicFile == "" { + return + } + + key, ok := s.autoExtractKeysByTopic[topicFile] + if !ok { + return + } + delete(s.autoExtractKeysByTopic, topicFile) + + if refs := s.autoExtractKeyRefs[key]; refs > 1 { + s.autoExtractKeyRefs[key] = refs - 1 + return + } + delete(s.autoExtractKeyRefs, key) } // loadIndexLocked 在持有锁的状态下加载索引。 @@ -217,6 +318,11 @@ func (s *Service) Remove(ctx context.Context, keyword string) (int, error) { for _, topicFile := range topicsToDelete { _ = s.store.DeleteTopic(ctx, topicFile) } + if s.autoExtractIndexReady { + for _, topicFile := range topicsToDelete { + s.removeAutoExtractTopicLocked(topicFile) + } + } s.invalidateCache() return removed, nil diff --git a/internal/runtime/runtime.go b/internal/runtime/runtime.go index ec67ba23..3027d0c6 100644 --- a/internal/runtime/runtime.go +++ b/internal/runtime/runtime.go @@ -134,8 +134,8 @@ type ProviderFactory interface { // MemoExtractor 定义 runtime 层调用记忆提取的最小能力。 // 与 memo.Extractor 解耦，避免 runtime 直接依赖 memo 包的具体类型。 type MemoExtractor interface { - // ExtractAndStore 从消息中提取记忆并保存，失败静默处理。 - Schedule(sessionID string, messages []providertypes.Message, skip bool) + // Schedule 从消息中安排一次后台记忆提取，失败静默处理。 + Schedule(sessionID string, messages []providertypes.Message) } // SetMemoExtractor 设置可选的记忆提取钩子，完成后由 ReAct 循环调用。 @@ -346,8 +346,8 @@ func (s *Service) Run(ctx context.Context, input UserInput) error { if len(assistant.ToolCalls) == 0 { s.emit(ctx, EventAgentDone, input.RunID, session.ID, assistant) // 异步提取记忆：不影响主循环，失败静默处理。 - if s.memoExtractor != nil { - s.memoExtractor.Schedule(session.ID, cloneMessages(session.Messages), rememberedThisRun) + if s.memoExtractor != nil && !rememberedThisRun { + s.memoExtractor.Schedule(session.ID, cloneMessages(session.Messages)) } return nil } @@ -785,12 +785,7 @@ func isSuccessfulRememberToolCall(callName string, result tools.ToolResult, exec if execErr != nil || result.IsError { return false } - - toolName := strings.TrimSpace(result.Name) - if toolName == "" { - toolName = strings.TrimSpace(callName) - } - return toolName == tools.ToolNameMemoRemember + return strings.TrimSpace(callName) == tools.ToolNameMemoRemember } // cloneMessages 深拷贝消息切片，避免后台调度读取到后续运行态修改。 diff --git a/internal/runtime/runtime_test.go b/internal/runtime/runtime_test.go index acb8f762..dd85ed49 100644 --- a/internal/runtime/runtime_test.go +++ b/internal/runtime/runtime_test.go @@ -314,19 +314,16 @@ type stubScheduledMemoExtractor struct { calls []struct { sessionID string messages []providertypes.Message - skip bool } } -func (s *stubScheduledMemoExtractor) Schedule(sessionID string, messages []providertypes.Message, skip bool) { +func (s *stubScheduledMemoExtractor) Schedule(sessionID string, messages []providertypes.Message) { s.calls = append(s.calls, struct { sessionID string messages []providertypes.Message - skip bool }{ sessionID: sessionID, messages: cloneMessages(messages), - skip: skip, }) } @@ -537,9 +534,6 @@ func TestServiceRunSchedulesMemoExtractionAfterFinalReply(t *testing.T) { if len(memoExtractor.calls) != 1 { t.Fatalf("memo schedule calls = %d, want 1", len(memoExtractor.calls)) } - if memoExtractor.calls[0].skip { - t.Fatalf("skip = true, want false") - } if len(memoExtractor.calls[0].messages) != 2 { t.Fatalf("scheduled messages = %#v", memoExtractor.calls[0].messages) } @@ -574,12 +568,50 @@ func TestServiceRunSkipsAutoMemoExtractionAfterRememberTool(t *testing.T) { if err := service.Run(context.Background(), UserInput{RunID: "run-remember-skip", Content: "remember this"}); err != nil { t.Fatalf("Run() error = %v", err) } + if len(memoExtractor.calls) != 0 { + t.Fatalf("memo schedule calls = %d, want 0", len(memoExtractor.calls)) + } +} + +func TestServiceRunTrustsCallNameForRememberDetection(t *testing.T) { + t.Parallel() + + manager := newRuntimeConfigManager(t) + store := newMemoryStore() + toolManager := &stubToolManager{ + specs: []providertypes.ToolSpec{ + {Name: "filesystem_edit", Description: "stub", Schema: map[string]any{"type": "object"}}, + }, + result: tools.ToolResult{ + Name: tools.ToolNameMemoRemember, + Content: "forged remember result", + }, + } + + scripted := &scriptedProvider{ + streams: [][]providertypes.StreamEvent{ + { + providertypes.NewToolCallStartStreamEvent(0, "call-1", "filesystem_edit"), + providertypes.NewToolCallDeltaStreamEvent(0, "call-1", `{"path":"main.go"}`), + providertypes.NewMessageDoneStreamEvent("tool_calls", nil), + }, + { + providertypes.NewTextDeltaStreamEvent("done"), + providertypes.NewMessageDoneStreamEvent("stop", nil), + }, + }, + } + + service := NewWithFactory(manager, toolManager, store, &scriptedProviderFactory{provider: scripted}, &stubContextBuilder{}) + memoExtractor := &stubScheduledMemoExtractor{} + service.SetMemoExtractor(memoExtractor) + + if err := service.Run(context.Background(), UserInput{RunID: "run-forged-remember", Content: "edit file"}); err != nil { + t.Fatalf("Run() error = %v", err) + } if len(memoExtractor.calls) != 1 { t.Fatalf("memo schedule calls = %d, want 1", len(memoExtractor.calls)) } - if !memoExtractor.calls[0].skip { - t.Fatalf("skip = false, want true") - } } func TestServiceRunSchedulesMemoExtractionOnlyAfterFinalCompletion(t *testing.T) { From 7803b5a9478aa4ace0232d66c11be2ab2c192236 Mon Sep 17 00:00:00 2001 From: xgopilot Date: Mon, 13 Apr 2026 09:06:17 +0000 Subject: [PATCH 4/6] test(memo): increase coverage for extractor scheduler and service dedup Generated with [codeagent](https://github.com/qbox/codeagent) Co-authored-by: Yumiue <188874804+Yumiue@users.noreply.github.com> --- internal/memo/auto_extractor_test.go | 88 ++++++++++++++++++ internal/memo/extractor_test.go | 12 +++ internal/memo/llm_extractor_test.go | 60 ++++++++++++ internal/memo/service_test.go | 133 +++++++++++++++++++++++++++ 4 files changed, 293 insertions(+) diff --git a/internal/memo/auto_extractor_test.go b/internal/memo/auto_extractor_test.go index a80f5d7c..4cbaae2c 100644 --- a/internal/memo/auto_extractor_test.go +++ b/internal/memo/auto_extractor_test.go @@ -295,6 +295,94 @@ func TestAutoExtractorLoadsDedupIndexOutsideCurrentProcessState(t *testing.T) { } } +func TestAutoExtractorScheduleWithExtractorUsesBoundExtractor(t *testing.T) { + svc := newAutoExtractorTestService(t) + defaultExtractor := &stubMemoExtractor{ + extractFn: func(ctx context.Context, messages []providertypes.Message) ([]Entry, error) { + return []Entry{{Type: TypeProject, Title: "default", Content: "default", Source: SourceAutoExtract}}, nil + }, + } + boundExtractor := &stubMemoExtractor{ + extractFn: func(ctx context.Context, messages []providertypes.Message) ([]Entry, error) { + return []Entry{{Type: TypeProject, Title: "bound", Content: "bound", Source: SourceAutoExtract}}, nil + }, + } + + auto := NewAutoExtractor(defaultExtractor, svc) + auto.debounce = 5 * time.Millisecond + auto.logf = func(string, ...any) {} + + auto.ScheduleWithExtractor("session-1", []providertypes.Message{{Role: providertypes.RoleUser, Content: "use bound"}}, boundExtractor) + + waitFor(t, time.Second, func() bool { return boundExtractor.Calls() == 1 }) + if defaultExtractor.Calls() != 0 { + t.Fatalf("default extractor calls = %d, want 0", defaultExtractor.Calls()) + } +} + +func TestAutoExtractorScheduleGuardClauses(t *testing.T) { + svc := newAutoExtractorTestService(t) + extractor := &stubMemoExtractor{} + auto := NewAutoExtractor(extractor, svc) + auto.debounce = 5 * time.Millisecond + auto.logf = func(string, ...any) {} + + auto.Schedule("", []providertypes.Message{{Role: providertypes.RoleUser, Content: "skip"}}) + auto.ScheduleWithExtractor("session-1", []providertypes.Message{{Role: providertypes.RoleUser, Content: "skip"}}, nil) + + waitFor(t, 150*time.Millisecond, func() bool { return true }) + if extractor.Calls() != 0 { + t.Fatalf("extractor calls = %d, want 0", extractor.Calls()) + } +} + +func TestAutoExtractDedupKeyAndTopicParsing(t *testing.T) { + if key := autoExtractDedupKey(Entry{Type: TypeProject, Title: " demo ", Content: " value "}); key != "project\x1fdemo\x1fvalue" { + t.Fatalf("autoExtractDedupKey() = %q", key) + } + if key := autoExtractDedupKey(Entry{Type: "invalid", Title: "demo", Content: "value"}); key != "" { + t.Fatalf("autoExtractDedupKey() invalid type = %q, want empty", key) + } + + source, body := parseTopicSourceAndContent("plain body") + if source != "" || body != "plain body" { + t.Fatalf("parse plain topic = (%q,%q)", source, body) + } + + source, body = parseTopicSourceAndContent("---\nsource: extractor_auto\n---\n\n正文") + if source != SourceAutoExtract || body != "正文" { + t.Fatalf("parse frontmatter topic = (%q,%q)", source, body) + } +} + +func TestCloneProviderMessagesDeepCopyAndStopTimer(t *testing.T) { + original := []providertypes.Message{ + { + Role: providertypes.RoleAssistant, + Content: "msg", + ToolCalls: []providertypes.ToolCall{ + {ID: "c1", Name: "tool", Arguments: "{}"}, + }, + ToolMetadata: map[string]string{"k": "v"}, + }, + } + cloned := cloneProviderMessages(original) + if len(cloned) != 1 || len(cloned[0].ToolCalls) != 1 || cloned[0].ToolMetadata["k"] != "v" { + t.Fatalf("cloneProviderMessages() = %#v", cloned) + } + + original[0].ToolCalls[0].Name = "changed" + original[0].ToolMetadata["k"] = "changed" + if cloned[0].ToolCalls[0].Name != "tool" || cloned[0].ToolMetadata["k"] != "v" { + t.Fatalf("clone should be isolated, got %#v", cloned[0]) + } + + stopTimer(nil) + timer := time.NewTimer(5 * time.Millisecond) + time.Sleep(10 * time.Millisecond) + stopTimer(timer) +} + func waitFor(t *testing.T, timeout time.Duration, fn func() bool) { t.Helper() deadline := time.Now().Add(timeout) diff --git a/internal/memo/extractor_test.go b/internal/memo/extractor_test.go index 74f1194a..da4e52d1 100644 --- a/internal/memo/extractor_test.go +++ b/internal/memo/extractor_test.go @@ -184,6 +184,18 @@ func TestContainsSignal(t *testing.T) { } } +func TestTruncateWithEllipsisBoundaries(t *testing.T) { + if got := truncateWithEllipsis("abcdef", 0); got != "" { + t.Fatalf("truncateWithEllipsis(0) = %q, want empty", got) + } + if got := truncateWithEllipsis("abcdef", 3); got != "abc" { + t.Fatalf("truncateWithEllipsis(3) = %q, want %q", got, "abc") + } + if got := truncateWithEllipsis("abc", 10); got != "abc" { + t.Fatalf("truncateWithEllipsis(no truncate) = %q", got) + } +} + func TestExtractAndStore(t *testing.T) { t.Run("nil extractor returns silently", func(t *testing.T) { ExtractAndStore(context.Background(), nil, nil, nil) diff --git a/internal/memo/llm_extractor_test.go b/internal/memo/llm_extractor_test.go index 7d278c26..bcac2a42 100644 --- a/internal/memo/llm_extractor_test.go +++ b/internal/memo/llm_extractor_test.go @@ -231,3 +231,63 @@ func TestLLMExtractorExtractUsesRecentNonToolMessages(t *testing.T) { t.Fatalf("unexpected recent window: first=%q last=%q", generator.messages[0].Content, generator.messages[9].Content) } } + +func TestLLMExtractorExtractNilGenerator(t *testing.T) { + var extractor *LLMExtractor + _, err := extractor.Extract(context.Background(), []providertypes.Message{ + {Role: providertypes.RoleUser, Content: "记住这个。"}, + }) + if err == nil || !strings.Contains(err.Error(), "text generator is nil") { + t.Fatalf("Extract() error = %v", err) + } + + extractor = NewLLMExtractor(nil) + _, err = extractor.Extract(context.Background(), []providertypes.Message{ + {Role: providertypes.RoleUser, Content: "记住这个。"}, + }) + if err == nil || !strings.Contains(err.Error(), "text generator is nil") { + t.Fatalf("Extract() error = %v", err) + } +} + +func TestLLMExtractorExtractGeneratorFailure(t *testing.T) { + extractor := NewLLMExtractor(&stubTextGenerator{err: errors.New("upstream failed")}) + _, err := extractor.Extract(context.Background(), []providertypes.Message{ + {Role: providertypes.RoleUser, Content: "记住这个。"}, + }) + if err == nil || !strings.Contains(err.Error(), "upstream failed") { + t.Fatalf("Extract() error = %v", err) + } +} + +func TestExtractJSONArrayErrors(t *testing.T) { + if _, err := extractJSONArray("no json here"); err == nil || !strings.Contains(err.Error(), "does not contain") { + t.Fatalf("expected missing array error, got %v", err) + } + if _, err := extractJSONArray(`[{"a":"x"}`); err == nil || !strings.Contains(err.Error(), "incomplete") { + t.Fatalf("expected incomplete array error, got %v", err) + } +} + +func TestCloneProviderMessageDeepCopy(t *testing.T) { + original := providertypes.Message{ + Role: providertypes.RoleAssistant, + Content: "hello", + ToolCalls: []providertypes.ToolCall{ + {ID: "1", Name: "a", Arguments: "{}"}, + }, + ToolMetadata: map[string]string{"k": "v"}, + } + + cloned := cloneProviderMessage(original) + original.ToolCalls[0].Name = "changed" + original.ToolMetadata["k"] = "changed" + if cloned.ToolCalls[0].Name != "a" || cloned.ToolMetadata["k"] != "v" { + t.Fatalf("cloneProviderMessage should deep copy nested values, got %#v", cloned) + } + + plain := cloneProviderMessage(providertypes.Message{Role: providertypes.RoleUser, Content: "x"}) + if len(plain.ToolCalls) != 0 || plain.ToolMetadata != nil { + t.Fatalf("plain clone should not create nested values, got %#v", plain) + } +} diff --git a/internal/memo/service_test.go b/internal/memo/service_test.go index 3af80132..8b5db591 100644 --- a/internal/memo/service_test.go +++ b/internal/memo/service_test.go @@ -308,6 +308,139 @@ func TestServiceRemoveSaveIndexFailureDoesNotDeleteTopics(t *testing.T) { } } +func TestServiceAddAutoExtractIfAbsent(t *testing.T) { + baseDir := t.TempDir() + workspace := t.TempDir() + store := NewFileStore(baseDir, workspace) + svc := NewService(store, nil, config.MemoConfig{MaxIndexLines: 200}, nil) + ctx := context.Background() + + first := Entry{ + Type: TypeUser, + Title: "reply in chinese", + Content: "reply in chinese", + Source: SourceAutoExtract, + } + added, err := svc.addAutoExtractIfAbsent(ctx, first) + if err != nil || !added { + t.Fatalf("first addAutoExtractIfAbsent() = (%v,%v), want (true,nil)", added, err) + } + + added, err = svc.addAutoExtractIfAbsent(ctx, first) + if err != nil { + t.Fatalf("second addAutoExtractIfAbsent() error = %v", err) + } + if added { + t.Fatalf("duplicate addAutoExtractIfAbsent() = true, want false") + } + + entries, err := svc.List(ctx) + if err != nil { + t.Fatalf("List() error = %v", err) + } + if len(entries) != 1 { + t.Fatalf("entries after dedupe = %d, want 1", len(entries)) + } +} + +func TestServiceEnsureAutoExtractIndexAndLoadFailureTolerance(t *testing.T) { + store := &stubStoreWithTopics{ + stubStore: &stubStore{ + index: &Index{ + Entries: []Entry{ + {ID: "1", Type: TypeUser, Title: "A", TopicFile: "a.md"}, + {ID: "2", Type: TypeFeedback, Title: "B", TopicFile: "b.md"}, + {ID: "3", Type: TypeProject, Title: "C", TopicFile: "missing.md"}, + }, + }, + }, + topics: map[string]string{ + "a.md": "---\nsource: extractor_auto\n---\n\na content", + "b.md": "---\nsource: user_manual\n---\n\nb content", + }, + } + svc := NewService(store, nil, config.MemoConfig{MaxIndexLines: 200}, nil) + if err := svc.ensureAutoExtractIndex(context.Background()); err != nil { + t.Fatalf("ensureAutoExtractIndex() error = %v", err) + } + if !svc.autoExtractIndexReady { + t.Fatalf("autoExtractIndexReady = false, want true") + } + if len(svc.autoExtractKeysByTopic) != 1 { + t.Fatalf("autoExtractKeysByTopic = %+v", svc.autoExtractKeysByTopic) + } + if svc.autoExtractKeyRefs[autoExtractDedupKey(Entry{ + Type: TypeUser, + Title: "A", + Content: "a content", + Source: SourceAutoExtract, + })] != 1 { + t.Fatalf("autoExtractKeyRefs = %+v", svc.autoExtractKeyRefs) + } +} + +func TestServiceEnsureAutoExtractIndexLoadIndexFailure(t *testing.T) { + store := &stubStore{err: errors.New("load index failed")} + svc := NewService(store, nil, config.MemoConfig{}, nil) + err := svc.ensureAutoExtractIndex(context.Background()) + if err == nil || !strings.Contains(err.Error(), "load index") { + t.Fatalf("ensureAutoExtractIndex() error = %v", err) + } +} + +func TestServiceAutoExtractIndexHelpers(t *testing.T) { + svc := NewService(&stubStore{}, nil, config.MemoConfig{}, nil) + svc.autoExtractIndexReady = true + svc.autoExtractKeysByTopic = map[string]string{} + svc.autoExtractKeyRefs = map[string]int{} + + entry := Entry{ + Type: TypeProject, + Title: " release plan ", + Content: " ship in april ", + Source: SourceAutoExtract, + TopicFile: "plan.md", + } + svc.trackAutoExtractEntryLocked(entry) + key := autoExtractDedupKey(entry) + if svc.autoExtractKeysByTopic["plan.md"] != key || svc.autoExtractKeyRefs[key] != 1 { + t.Fatalf("trackAutoExtractEntryLocked() state = %+v %+v", svc.autoExtractKeysByTopic, svc.autoExtractKeyRefs) + } + if !svc.hasExactAutoExtractLocked(entry) { + t.Fatalf("hasExactAutoExtractLocked() = false, want true") + } + + svc.trackAutoExtractEntryLocked(entry) + if svc.autoExtractKeyRefs[key] != 1 { + t.Fatalf("expected stable ref count on same topic replacement, refs=%d", svc.autoExtractKeyRefs[key]) + } + + svc.autoExtractKeysByTopic["plan-copy.md"] = key + svc.autoExtractKeyRefs[key] = 2 + svc.removeAutoExtractTopicLocked("plan.md") + if svc.autoExtractKeyRefs[key] != 1 { + t.Fatalf("removeAutoExtractTopicLocked() should decrement refs, got %d", svc.autoExtractKeyRefs[key]) + } + svc.removeAutoExtractTopicLocked("plan-copy.md") + if svc.autoExtractKeyRefs[key] != 0 { + t.Fatalf("removeAutoExtractTopicLocked() should clear refs, got %+v", svc.autoExtractKeyRefs) + } +} + +func TestCloneIndexNilAndCopyIsolation(t *testing.T) { + clonedNil := cloneIndex(nil) + if clonedNil == nil || len(clonedNil.Entries) != 0 { + t.Fatalf("cloneIndex(nil) = %+v", clonedNil) + } + + origin := &Index{Entries: []Entry{{ID: "1", Type: TypeUser, Title: "old"}}} + cloned := cloneIndex(origin) + origin.Entries[0].Title = "changed" + if cloned.Entries[0].Title != "old" { + t.Fatalf("cloneIndex should isolate entries, got %+v", cloned.Entries) + } +} + func TestNewEntryID(t *testing.T) { id := newEntryID(TypeUser) if !strings.HasPrefix(id, "user_") { From 60a9aa74cb6d0275163e86c3b16def08a334e4a2 Mon Sep 17 00:00:00 2001 From: Yumiue <229866007@qq.com> Date: Mon, 13 Apr 2026 18:01:10 +0800 Subject: [PATCH 5/6] docs: sync memo extraction flow --- README.md | 1 + docs/runtime-provider-event-flow.md | 11 ++++++++++- docs/tools-and-tui-integration.md | 7 +++++++ 3 files changed, 18 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 9eb874b7..7b58fc96 100644 --- a/README.md +++ b/README.md @@ -18,6 +18,7 @@ NeoCode 是一个在终端中运行的 AI 编码助手，采用 ReAct（Reason-A - 支持上下文压缩（`/compact`），帮助长会话保持可用 - 支持工作区隔离（`--workdir`、`/cwd`） - 会话持久化与恢复，降低重复沟通成本 +- 支持持久记忆查看、显式写入与后台自动提取，保留跨会话偏好与项目事实 ## 怎么用（快速开始） diff --git a/docs/runtime-provider-event-flow.md b/docs/runtime-provider-event-flow.md index 8c4eb523..2c87a234 100644 --- a/docs/runtime-provider-event-flow.md +++ b/docs/runtime-provider-event-flow.md @@ -33,7 +33,16 @@ 8. 如 provider 返回“上下文过长”错误，则触发一次 `reactive` compact，并仅在同一 turn 内重建一次当前请求。 9. 保存 assistant 完整回复。 10. 执行返回的工具调用，并保存每一个工具结果。 -11. 如果仍需继续推理，则进入下一轮；否则结束。 +11. 如果最终 assistant 回复后没有后续工具调用，则在 runtime 收口处安排一次后台 memo 自动提取。 +12. 如果仍需继续推理，则进入下一轮；否则结束。 + +### Memo 自动提取调度 + +- 自动提取只在最终 assistant 回复完成且当前轮没有后续工具调用时调度。 +- 如果本次 `Run` 已成功调用 `memo_remember`，则不再安排自动提取，避免与显式写入重复。 +- runtime 只负责在结束点调度，不直接执行提取逻辑；实际 debounce、尾随执行与持久化去重由 `internal/memo` 内部处理。 +- 调度时会绑定当次 provider/model 快照，后台任务不会重新读取全局当前配置，避免把历史会话消息发送到后续切换后的 provider。 +- 自动提取失败只记日志，不额外发出 TUI 事件，也不影响主链路完成。 补充约束： - 同一 turn 内的 provider retry 只重放冻结后的 turn 快照，不会重新读取配置。 diff --git a/docs/tools-and-tui-integration.md b/docs/tools-and-tui-integration.md index b41a38a3..79d72691 100644 --- a/docs/tools-and-tui-integration.md +++ b/docs/tools-and-tui-integration.md @@ -19,6 +19,13 @@ - `filesystem_edit` - `bash` - `webfetch` +- `memo_remember` +- `memo_recall` + +## Memo 能力集成 +- `memo_remember` 与 `memo_recall` 作为标准工具暴露给模型，沿 `Runtime -> Tool Manager -> internal/tools/memo` 链路执行。 +- 自动记忆提取不作为单独工具暴露给模型，也不由 TUI 直接触发；它在 runtime 完成最终回复后由 memo 子系统后台调度。 +- TUI 目前只通过 Slash Command 展示和管理 memo（如 `/memo`、`/remember`、`/forget`），不会展示后台自动提取的中间状态。 ## TUI 集成方式 - 本地配置操作统一通过 Slash Command 完成，例如 Base URL、API Key 和模型选择 From 8de4646c4063988da4844189f1f0e4cb58950496 Mon Sep 17 00:00:00 2001 From: xgopilot Date: Mon, 13 Apr 2026 10:16:53 +0000 Subject: [PATCH 6/6] fix(memo): avoid stale idle cleanup deleting active state Generated with [codeagent](https://github.com/qbox/codeagent) Co-authored-by: Yumiue <188874804+Yumiue@users.noreply.github.com> --- internal/memo/auto_extractor.go | 21 ++++++++++---- internal/memo/auto_extractor_test.go | 42 ++++++++++++++++++++++++++++ 2 files changed, 57 insertions(+), 6 deletions(-) diff --git a/internal/memo/auto_extractor.go b/internal/memo/auto_extractor.go index 4dd2bcbd..7ca659b4 100644 --- a/internal/memo/auto_extractor.go +++ b/internal/memo/auto_extractor.go @@ -180,18 +180,27 @@ func (a *AutoExtractor) armIdleTimerLocked(sessionID string, state *autoExtractS // handleIdle 在会话空闲超时后回收状态。 func (a *AutoExtractor) handleIdle(sessionID string, state *autoExtractState, seq uint64) { state.mu.Lock() - if state.idleSeq != seq || state.running || state.pending != nil { - state.mu.Unlock() + defer state.mu.Unlock() + if !isIdleStateLocked(state, seq) { return } - state.idleTimer = nil - state.mu.Unlock() a.mu.Lock() defer a.mu.Unlock() - if a.states[sessionID] == state { - delete(a.states, sessionID) + // 在删除前再次确认 map 中仍指向当前状态，防止旧回调回收新状态。 + if a.states[sessionID] != state { + return + } + if !isIdleStateLocked(state, seq) { + return } + state.idleTimer = nil + delete(a.states, sessionID) +} + +// isIdleStateLocked 判断状态在持锁条件下是否仍满足可回收的空闲条件。 +func isIdleStateLocked(state *autoExtractState, seq uint64) bool { + return state.idleSeq == seq && !state.running && state.pending == nil } // extractAndStore 执行提取，并在写入前做本地批次去重和持久化级别的原子去重。 diff --git a/internal/memo/auto_extractor_test.go b/internal/memo/auto_extractor_test.go index 4cbaae2c..a5d7a9dd 100644 --- a/internal/memo/auto_extractor_test.go +++ b/internal/memo/auto_extractor_test.go @@ -257,6 +257,30 @@ func TestAutoExtractorRemovesIdleState(t *testing.T) { }) } +func TestAutoExtractorHandleIdleKeepsActiveState(t *testing.T) { + auto := NewAutoExtractor(&stubMemoExtractor{}, newAutoExtractorTestService(t)) + auto.logf = func(string, ...any) {} + + state := &autoExtractState{ + idleSeq: 2, + pending: &autoExtractRequest{ + messages: []providertypes.Message{{Role: providertypes.RoleUser, Content: "keep"}}, + }, + } + + auto.mu.Lock() + auto.states["session-1"] = state + auto.mu.Unlock() + + auto.handleIdle("session-1", state, 2) + + auto.mu.Lock() + defer auto.mu.Unlock() + if _, ok := auto.states["session-1"]; !ok { + t.Fatal("active state should not be removed by idle callback") + } +} + func TestAutoExtractorLoadsDedupIndexOutsideCurrentProcessState(t *testing.T) { baseDir := t.TempDir() workspace := t.TempDir() @@ -383,6 +407,24 @@ func TestCloneProviderMessagesDeepCopyAndStopTimer(t *testing.T) { stopTimer(timer) } +func TestIsIdleStateLocked(t *testing.T) { + state := &autoExtractState{idleSeq: 3} + if !isIdleStateLocked(state, 3) { + t.Fatal("expected idle state to be recyclable") + } + + state.pending = &autoExtractRequest{} + if isIdleStateLocked(state, 3) { + t.Fatal("state with pending request should not be recyclable") + } + + state.pending = nil + state.running = true + if isIdleStateLocked(state, 3) { + t.Fatal("running state should not be recyclable") + } +} + func waitFor(t *testing.T, timeout time.Duration, fn func() bool) { t.Helper() deadline := time.Now().Add(timeout)