diff --git a/internal/app/bootstrap.go b/internal/app/bootstrap.go index 2a509906..a9df98a2 100644 --- a/internal/app/bootstrap.go +++ b/internal/app/bootstrap.go @@ -2,7 +2,6 @@ package app import ( "context" - "fmt" "log" "strings" "time" @@ -34,6 +33,8 @@ const utf8CodePage = 65001 var ( setConsoleOutputCodePage = platformSetConsoleOutputCodePage setConsoleInputCodePage = platformSetConsoleInputCodePage + buildToolManagerFunc = buildToolManager + newTUIWithMemo = tui.NewWithMemo ) // BootstrapOptions 描述应用启动时可注入的运行时选项。 @@ -41,6 +42,44 @@ type BootstrapOptions struct { Workdir string } +type memoExtractorScheduler interface { + ScheduleWithExtractor(sessionID string, messages []providertypes.Message, extractor memo.Extractor) +} + +func newMemoExtractorAdapter( + factory agentruntime.ProviderFactory, + cm *config.Manager, + scheduler memoExtractorScheduler, +) agentruntime.MemoExtractor { + return runtimeMemoExtractorFunc(func(sessionID string, messages []providertypes.Message) { + if scheduler == nil { + return + } + + cfg := cm.Get() + resolved, err := config.ResolveSelectedProvider(cfg) + if err != nil { + log.Printf("memo: resolve selected provider failed: %v", err) + return + } + + generator := textGenAdapter(func(ctx context.Context, prompt string, msgs []providertypes.Message) (string, error) { + p, err := factory.Build(ctx, resolved.ToRuntimeConfig()) + if err != nil { + return "", err + } + + return provider.GenerateText(ctx, p, providertypes.GenerateRequest{ + Model: cfg.CurrentModel, + SystemPrompt: prompt, + Messages: append([]providertypes.Message(nil), msgs...), + }) + }) + + scheduler.ScheduleWithExtractor(sessionID, messages, memo.NewLLMExtractor(generator)) + }) +} + // RuntimeBundle 聚合 CLI 与 TUI 共享的运行时依赖。 type RuntimeBundle struct { Config config.Config @@ -48,6 +87,7 @@ type RuntimeBundle struct { Runtime agentruntime.Runtime ProviderSelection *configstate.Service MemoService *memo.Service + Close func() error // 用于清理 bundle 运行期间拉起的系统资源 } // EnsureConsoleUTF8 负责在 Windows 控制台中尽量启用 UTF-8 编码。 @@ -83,11 +123,18 @@ func BuildRuntime(ctx context.Context, opts BootstrapOptions) (RuntimeBundle, er cfg := manager.Get() - toolRegistry, err := buildToolRegistry(cfg) + toolRegistry, toolsCleanup, err := buildToolRegistry(cfg) if err != nil { return RuntimeBundle{}, err } - toolManager, err := buildToolManager(toolRegistry) + needCleanup := true + defer func() { + if needCleanup && toolsCleanup != nil { + _ = toolsCleanup() + } + }() + + toolManager, err := buildToolManagerFunc(toolRegistry) if err != nil { return RuntimeBundle{}, err } @@ -127,6 +174,7 @@ func BuildRuntime(ctx context.Context, opts BootstrapOptions) (RuntimeBundle, er memo.NewAutoExtractor(nil, memoSvc), )) } + needCleanup = false return RuntimeBundle{ Config: cfg, @@ -134,25 +182,29 @@ func BuildRuntime(ctx context.Context, opts BootstrapOptions) (RuntimeBundle, er Runtime: runtimeSvc, ProviderSelection: providerSelection, MemoService: memoSvc, + Close: toolsCleanup, }, nil } -// NewProgram 基于共享运行时依赖构建并返回 TUI 程序。 -func NewProgram(ctx context.Context, opts BootstrapOptions) (*tea.Program, error) { +// NewProgram 基于共享运行时依赖构建并返回 TUI 程序，同时返回退出时应调用的资源清理函数。 +func NewProgram(ctx context.Context, opts BootstrapOptions) (*tea.Program, func() error, error) { bundle, err := BuildRuntime(ctx, opts) if err != nil { - return nil, err + return nil, nil, err } - tuiApp, err := tui.NewWithMemo(&bundle.Config, bundle.ConfigManager, bundle.Runtime, bundle.ProviderSelection, bundle.MemoService) + tuiApp, err := newTUIWithMemo(&bundle.Config, bundle.ConfigManager, bundle.Runtime, bundle.ProviderSelection, bundle.MemoService) if err != nil { - return nil, err + if bundle.Close != nil { + _ = bundle.Close() + } + return nil, nil, err } return tea.NewProgram( tuiApp, tea.WithAltScreen(), tea.WithMouseCellMotion(), - ), nil + ), bundle.Close, nil } // bootstrapDefaultConfig 负责计算本次启动应使用的默认配置快照。 @@ -176,7 +228,7 @@ func resolveBootstrapWorkdir(workdir string) (string, error) { return agentsession.ResolveExistingDir(workdir) } -func buildToolRegistry(cfg config.Config) (*tools.Registry, error) { +func buildToolRegistry(cfg config.Config) (*tools.Registry, func() error, error) { toolRegistry := tools.NewRegistry() toolRegistry.Register(filesystem.New(cfg.Workdir)) toolRegistry.Register(filesystem.NewWrite(cfg.Workdir)) @@ -191,7 +243,7 @@ func buildToolRegistry(cfg config.Config) (*tools.Registry, error) { })) mcpRegistry, err := buildMCPRegistry(cfg) if err != nil { - return nil, err + return nil, nil, err } if mcpRegistry != nil { toolRegistry.SetMCPRegistry(mcpRegistry) @@ -201,7 +253,10 @@ func buildToolRegistry(cfg config.Config) (*tools.Registry, error) { Agents: buildMCPAgentExposureRules(cfg.Tools.MCP.Exposure.Agents), })) } - return toolRegistry, nil + if mcpRegistry == nil { + return toolRegistry, nil, nil + } + return toolRegistry, mcpRegistry.Close, nil } // buildMCPAgentExposureRules 将配置层的 agent 过滤规则转换为 tools/mcp 层输入。 @@ -227,136 +282,14 @@ func buildToolManager(registry *tools.Registry) (tools.Manager, error) { return tools.NewManager(registry, engine, security.NewWorkspaceSandbox()) } -type memoExtractorScheduler interface { - ScheduleWithExtractor(sessionID string, messages []providertypes.Message, extractor memo.Extractor) -} +type runtimeMemoExtractorFunc func(sessionID string, messages []providertypes.Message) -// memoExtractorAdapter 在调度时绑定 provider 配置快照，避免后台任务读取全局可变配置。 -type memoExtractorAdapter struct { - factory agentruntime.ProviderFactory - configManager *config.Manager - scheduler memoExtractorScheduler +func (f runtimeMemoExtractorFunc) Schedule(sessionID string, messages []providertypes.Message) { + f(sessionID, messages) } -// 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 -} - -// 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) { - modelProvider, err := g.factory.Build(ctx, g.runtimeCfg) - 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 - }() +type textGenAdapter func(ctx context.Context, prompt string, msgs []providertypes.Message) (string, error) - err = modelProvider.Generate(ctx, providertypes.GenerateRequest{ - Model: g.model, - 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 +func (f textGenAdapter) Generate(ctx context.Context, prompt string, msgs []providertypes.Message) (string, error) { + return f(ctx, prompt, msgs) } diff --git a/internal/app/bootstrap_test.go b/internal/app/bootstrap_test.go index d7dea0ad..9059b7b8 100644 --- a/internal/app/bootstrap_test.go +++ b/internal/app/bootstrap_test.go @@ -19,8 +19,13 @@ import ( "neo-code/internal/config" configstate "neo-code/internal/config/state" + "neo-code/internal/memo" + "neo-code/internal/provider" + providertypes "neo-code/internal/provider/types" + agentruntime "neo-code/internal/runtime" "neo-code/internal/tools" "neo-code/internal/tools/mcp" + "neo-code/internal/tui" ) func TestNewProgram(t *testing.T) { @@ -30,13 +35,16 @@ func TestNewProgram(t *testing.T) { t.Setenv("HOME", home) t.Setenv("USERPROFILE", home) - program, err := NewProgram(context.Background(), BootstrapOptions{}) + program, cleanup, err := NewProgram(context.Background(), BootstrapOptions{}) if err != nil { t.Fatalf("NewProgram() error = %v", err) } if program == nil { t.Fatalf("expected tea program") } + if cleanup != nil { + defer cleanup() + } configPath := filepath.Join(home, ".neocode", "config.yaml") if _, err := os.Stat(configPath); err != nil { @@ -62,13 +70,16 @@ func TestNewProgramNormalizesInvalidCurrentModelOnStartup(t *testing.T) { t.Fatalf("write config: %v", err) } - program, err := NewProgram(context.Background(), BootstrapOptions{}) + program, cleanup, err := NewProgram(context.Background(), BootstrapOptions{}) if err != nil { t.Fatalf("NewProgram() error = %v", err) } if program == nil { t.Fatalf("expected tea program") } + if cleanup != nil { + defer cleanup() + } data, err := os.ReadFile(configPath) if err != nil { @@ -137,10 +148,13 @@ func TestBuildToolRegistryUsesWebFetchConfig(t *testing.T) { cfg.Workdir = t.TempDir() cfg.Tools.WebFetch.MaxResponseBytes = 4 - registry, err := buildToolRegistry(cfg) + registry, cleanup, err := buildToolRegistry(cfg) if err != nil { t.Fatalf("buildToolRegistry() error = %v", err) } + if cleanup != nil { + defer cleanup() + } tool, err := registry.Get("webfetch") if err != nil { t.Fatalf("registry.Get(webfetch) error = %v", err) @@ -348,10 +362,13 @@ func TestBuildToolRegistryIncludesMCPFromConfig(t *testing.T) { return registry.RefreshServerTools(context.Background(), server.ID) } - registry, err := buildToolRegistry(cfg) + registry, cleanup, err := buildToolRegistry(cfg) if err != nil { t.Fatalf("buildToolRegistry() error = %v", err) } + if cleanup != nil { + defer cleanup() + } specs, err := registry.ListAvailableSpecs(context.Background(), tools.SpecListInput{}) if err != nil { t.Fatalf("ListAvailableSpecs() error = %v", err) @@ -410,10 +427,13 @@ func TestBuildToolRegistryAppliesMCPExposureConfig(t *testing.T) { return registry.RefreshServerTools(context.Background(), server.ID) } - registry, err := buildToolRegistry(cfg) + registry, cleanup, err := buildToolRegistry(cfg) if err != nil { t.Fatalf("buildToolRegistry() error = %v", err) } + if cleanup != nil { + defer cleanup() + } specs, err := registry.ListAvailableSpecs(context.Background(), tools.SpecListInput{Agent: "planner"}) if err != nil { @@ -448,7 +468,10 @@ func TestBuildToolRegistryReturnsMCPSourceError(t *testing.T) { }, } - _, err := buildToolRegistry(cfg) + _, cleanup, err := buildToolRegistry(cfg) + if cleanup != nil { + defer cleanup() + } if err == nil || !strings.Contains(strings.ToLower(err.Error()), "unsupported mcp source") { t.Fatalf("expected unsupported mcp source error, got %v", err) } @@ -782,6 +805,119 @@ func TestBuildRuntimeRejectsUnsupportedMCPSource(t *testing.T) { } } +func TestBuildRuntimeCleansResourcesWhenToolManagerBuildFails(t *testing.T) { + disableBuiltinProviderAPIKeys(t) + + home := t.TempDir() + t.Setenv("HOME", home) + t.Setenv("USERPROFILE", home) + + configDir := filepath.Join(home, ".neocode") + if err := os.MkdirAll(configDir, 0o755); err != nil { + t.Fatalf("mkdir config dir: %v", err) + } + configPath := filepath.Join(configDir, "config.yaml") + raw := []byte(strings.Join([]string{ + "selected_provider: openai", + "current_model: " + config.OpenAIDefaultModel, + "shell: powershell", + "tools:", + " mcp:", + " servers:", + " - id: docs", + " enabled: true", + " source: stdio", + " stdio:", + " command: mock", + }, "\n") + "\n") + if err := os.WriteFile(configPath, raw, 0o644); err != nil { + t.Fatalf("write config: %v", err) + } + + closed := false + originalRegister := registerMCPStdioServer + t.Cleanup(func() { registerMCPStdioServer = originalRegister }) + registerMCPStdioServer = func(registry *mcp.Registry, cfg config.Config, server config.MCPServerConfig) error { + client := &closeableStubMCPServerClient{closed: &closed} + return registry.RegisterServer(server.ID, "stdio", server.Version, client) + } + + originalBuildToolManager := buildToolManagerFunc + t.Cleanup(func() { buildToolManagerFunc = originalBuildToolManager }) + buildToolManagerFunc = func(registry *tools.Registry) (tools.Manager, error) { + return nil, errors.New("build tool manager failed") + } + + _, err := BuildRuntime(context.Background(), BootstrapOptions{}) + if err == nil || !strings.Contains(err.Error(), "build tool manager failed") { + t.Fatalf("expected tool manager build error, got %v", err) + } + if !closed { + t.Fatalf("expected MCP resources to be closed on BuildRuntime failure") + } +} + +func TestNewProgramCleansResourcesWhenTUIBuildFails(t *testing.T) { + disableBuiltinProviderAPIKeys(t) + + home := t.TempDir() + t.Setenv("HOME", home) + t.Setenv("USERPROFILE", home) + + configDir := filepath.Join(home, ".neocode") + if err := os.MkdirAll(configDir, 0o755); err != nil { + t.Fatalf("mkdir config dir: %v", err) + } + configPath := filepath.Join(configDir, "config.yaml") + raw := []byte(strings.Join([]string{ + "selected_provider: openai", + "current_model: " + config.OpenAIDefaultModel, + "shell: powershell", + "tools:", + " mcp:", + " servers:", + " - id: docs", + " enabled: true", + " source: stdio", + " stdio:", + " command: mock", + }, "\n") + "\n") + if err := os.WriteFile(configPath, raw, 0o644); err != nil { + t.Fatalf("write config: %v", err) + } + + closed := false + originalRegister := registerMCPStdioServer + t.Cleanup(func() { registerMCPStdioServer = originalRegister }) + registerMCPStdioServer = func(registry *mcp.Registry, cfg config.Config, server config.MCPServerConfig) error { + client := &closeableStubMCPServerClient{closed: &closed} + return registry.RegisterServer(server.ID, "stdio", server.Version, client) + } + + originalNewTUIWithMemo := newTUIWithMemo + t.Cleanup(func() { newTUIWithMemo = originalNewTUIWithMemo }) + newTUIWithMemo = func( + cfg *config.Config, + configManager *config.Manager, + runtime agentruntime.Runtime, + providerSvc tui.ProviderController, + memoSvc *memo.Service, + ) (tui.App, error) { + return tui.App{}, errors.New("tui init failed") + } + + _, cleanup, err := NewProgram(context.Background(), BootstrapOptions{}) + if cleanup != nil { + t.Fatalf("expected nil cleanup on NewProgram failure") + } + if err == nil || !strings.Contains(err.Error(), "tui init failed") { + t.Fatalf("expected tui init error, got %v", err) + } + if !closed { + t.Fatalf("expected MCP resources to be closed when NewProgram fails") + } +} + func TestNewProgramRejectsInvalidWorkdirOverride(t *testing.T) { disableBuiltinProviderAPIKeys(t) @@ -789,7 +925,10 @@ func TestNewProgramRejectsInvalidWorkdirOverride(t *testing.T) { t.Setenv("HOME", home) t.Setenv("USERPROFILE", home) - _, err := NewProgram(context.Background(), BootstrapOptions{Workdir: filepath.Join(t.TempDir(), "missing", "中文")}) + _, cleanup, err := NewProgram(context.Background(), BootstrapOptions{Workdir: filepath.Join(t.TempDir(), "missing", "中文")}) + if cleanup != nil { + defer cleanup() + } if err == nil || !strings.Contains(strings.ToLower(err.Error()), "resolve workdir") { t.Fatalf("expected invalid workdir error, got %v", err) } @@ -866,6 +1005,213 @@ func TestEnsureConsoleUTF8SkipsInputWhenOutputFails(t *testing.T) { } } +func TestRuntimeMemoExtractorFuncSchedule(t *testing.T) { + t.Parallel() + + called := false + extractor := runtimeMemoExtractorFunc(func(sessionID string, messages []providertypes.Message) { + called = true + if sessionID != "session-1" { + t.Fatalf("unexpected session id %q", sessionID) + } + if len(messages) != 1 || messages[0].Content != "hi" { + t.Fatalf("unexpected messages %+v", messages) + } + }) + + extractor.Schedule("session-1", []providertypes.Message{ + {Role: providertypes.RoleUser, Content: "hi"}, + }) + + if !called { + t.Fatalf("expected schedule callback to be called") + } +} + +func TestTextGenAdapterGenerate(t *testing.T) { + t.Parallel() + + called := false + adapter := textGenAdapter(func(ctx context.Context, prompt string, msgs []providertypes.Message) (string, error) { + called = true + if prompt != "prompt" { + t.Fatalf("unexpected prompt %q", prompt) + } + if len(msgs) != 1 || msgs[0].Content != "msg" { + t.Fatalf("unexpected messages %+v", msgs) + } + return "ok", nil + }) + + result, err := adapter.Generate(context.Background(), "prompt", []providertypes.Message{ + {Role: providertypes.RoleUser, Content: "msg"}, + }) + if err != nil { + t.Fatalf("Generate() error = %v", err) + } + if result != "ok" { + t.Fatalf("unexpected result %q", result) + } + if !called { + t.Fatalf("expected adapter callback to be called") + } +} + +func TestNewMemoExtractorAdapterSkipsWhenSchedulerNil(t *testing.T) { + t.Parallel() + + cfg := config.StaticDefaults().Clone() + cfg.SelectedProvider = config.OpenAIName + manager := config.NewManager(config.NewLoader("", &cfg)) + extractor := newMemoExtractorAdapter(nil, manager, nil) + + extractor.Schedule("session-1", []providertypes.Message{ + {Role: providertypes.RoleUser, Content: "remember this"}, + }) +} + +func TestNewMemoExtractorAdapterSkipsWhenResolveProviderFails(t *testing.T) { + t.Parallel() + + cfg := config.StaticDefaults().Clone() + cfg.SelectedProvider = "" + manager := config.NewManager(config.NewLoader("", &cfg)) + scheduler := &stubMemoExtractorScheduler{} + extractor := newMemoExtractorAdapter(nil, manager, scheduler) + + extractor.Schedule("session-1", []providertypes.Message{ + {Role: providertypes.RoleUser, Content: "remember this"}, + }) + + if scheduler.called { + t.Fatalf("expected scheduler not to be called when provider resolution fails") + } +} + +func TestNewMemoExtractorAdapterSchedulesExtractorAndGenerates(t *testing.T) { + t.Setenv(config.OpenAIDefaultAPIKeyEnv, "token") + cfg := config.StaticDefaults().Clone() + cfg.SelectedProvider = config.OpenAIName + manager := config.NewManager(config.NewLoader("", &cfg)) + providerStub := &stubMemoProvider{ + generate: func(ctx context.Context, req providertypes.GenerateRequest, events chan<- providertypes.StreamEvent) error { + if req.Model != cfg.CurrentModel { + t.Fatalf("unexpected model %q", req.Model) + } + if strings.TrimSpace(req.SystemPrompt) == "" { + t.Fatalf("expected non-empty system prompt") + } + events <- providertypes.NewTextDeltaStreamEvent(`[{"type":"user","title":"偏好","content":"记住使用 Go","keywords":["go"]}]`) + events <- providertypes.NewMessageDoneStreamEvent("stop", nil) + return nil + }, + } + factory := &stubMemoProviderFactory{provider: providerStub} + scheduler := &stubMemoExtractorScheduler{} + extractor := newMemoExtractorAdapter(factory, manager, scheduler) + + inputMessages := []providertypes.Message{ + {Role: providertypes.RoleUser, Content: "remember 我偏好 Go"}, + } + extractor.Schedule("session-1", inputMessages) + + if !scheduler.called || scheduler.extractor == nil { + t.Fatalf("expected scheduler to receive extractor") + } + if scheduler.sessionID != "session-1" { + t.Fatalf("unexpected session id %q", scheduler.sessionID) + } + entries, err := scheduler.extractor.Extract(context.Background(), inputMessages) + if err != nil { + t.Fatalf("extractor.Extract() error = %v", err) + } + if len(entries) != 1 { + t.Fatalf("expected one extracted entry, got %+v", entries) + } + if !factory.called { + t.Fatalf("expected provider factory Build to be called") + } +} + +func TestNewMemoExtractorAdapterKeepsScheduledConfigSnapshot(t *testing.T) { + t.Setenv(config.OpenAIDefaultAPIKeyEnv, "openai-token") + t.Setenv(config.QiniuDefaultAPIKeyEnv, "qiniu-token") + home := t.TempDir() + t.Setenv("HOME", home) + t.Setenv("USERPROFILE", home) + + cfg := config.StaticDefaults().Clone() + cfg.SelectedProvider = config.OpenAIName + cfg.CurrentModel = "snapshot-model" + manager := config.NewManager(config.NewLoader("", &cfg)) + providerStub := &stubMemoProvider{ + generate: func(ctx context.Context, req providertypes.GenerateRequest, events chan<- providertypes.StreamEvent) error { + if req.Model != "snapshot-model" { + t.Fatalf("expected scheduled model snapshot, got %q", req.Model) + } + events <- providertypes.NewTextDeltaStreamEvent( + `[{"type":"user","title":"偏好","content":"snapshot","keywords":["snapshot"]}]`, + ) + events <- providertypes.NewMessageDoneStreamEvent("stop", nil) + return nil + }, + } + factory := &stubMemoProviderFactory{provider: providerStub} + scheduler := &stubMemoExtractorScheduler{} + extractor := newMemoExtractorAdapter(factory, manager, scheduler) + + messages := []providertypes.Message{{Role: providertypes.RoleUser, Content: "remember snapshot"}} + extractor.Schedule("session-1", messages) + if !scheduler.called || scheduler.extractor == nil { + t.Fatalf("expected scheduler to receive extractor") + } + + if err := manager.Update(context.Background(), func(next *config.Config) error { + next.SelectedProvider = config.QiniuName + next.CurrentModel = "drifted-model" + return nil + }); err != nil { + t.Fatalf("manager.Update() error = %v", err) + } + + entries, err := scheduler.extractor.Extract(context.Background(), messages) + if err != nil { + t.Fatalf("extractor.Extract() error = %v", err) + } + if len(entries) != 1 { + t.Fatalf("expected one extracted entry, got %+v", entries) + } + if !factory.called { + t.Fatalf("expected provider factory Build to be called") + } + if factory.cfg.Name != config.OpenAIName { + t.Fatalf("expected scheduled provider snapshot %q, got %q", config.OpenAIName, factory.cfg.Name) + } +} + +func TestNewMemoExtractorAdapterPropagatesFactoryBuildError(t *testing.T) { + t.Setenv(config.OpenAIDefaultAPIKeyEnv, "token") + cfg := config.StaticDefaults().Clone() + cfg.SelectedProvider = config.OpenAIName + manager := config.NewManager(config.NewLoader("", &cfg)) + factory := &stubMemoProviderFactory{buildErr: errors.New("build failed")} + scheduler := &stubMemoExtractorScheduler{} + extractor := newMemoExtractorAdapter(factory, manager, scheduler) + + inputMessages := []providertypes.Message{ + {Role: providertypes.RoleUser, Content: "remember coding style"}, + } + extractor.Schedule("session-1", inputMessages) + if !scheduler.called || scheduler.extractor == nil { + t.Fatalf("expected scheduler to receive extractor") + } + + _, err := scheduler.extractor.Extract(context.Background(), inputMessages) + if err == nil || !strings.Contains(err.Error(), "build failed") { + t.Fatalf("expected build failure, got %v", err) + } +} + type stubToolForBootstrap struct { name string content string @@ -932,6 +1278,59 @@ func (s *closeableStubMCPServerClient) Close() error { return nil } +type stubMemoExtractorScheduler struct { + called bool + sessionID string + messages []providertypes.Message + extractor memo.Extractor +} + +func (s *stubMemoExtractorScheduler) ScheduleWithExtractor( + sessionID string, + messages []providertypes.Message, + extractor memo.Extractor, +) { + s.called = true + s.sessionID = sessionID + s.messages = append([]providertypes.Message(nil), messages...) + s.extractor = extractor +} + +type stubMemoProviderFactory struct { + called bool + cfg provider.RuntimeConfig + provider provider.Provider + buildErr error +} + +func (s *stubMemoProviderFactory) Build(ctx context.Context, cfg provider.RuntimeConfig) (provider.Provider, error) { + s.called = true + s.cfg = cfg + if s.buildErr != nil { + return nil, s.buildErr + } + if s.provider != nil { + return s.provider, nil + } + return &stubMemoProvider{}, nil +} + +type stubMemoProvider struct { + generate func(ctx context.Context, req providertypes.GenerateRequest, events chan<- providertypes.StreamEvent) error +} + +func (s *stubMemoProvider) Generate( + ctx context.Context, + req providertypes.GenerateRequest, + events chan<- providertypes.StreamEvent, +) error { + if s.generate != nil { + return s.generate(ctx, req, events) + } + events <- providertypes.NewMessageDoneStreamEvent("stop", nil) + return nil +} + func TestHelperProcessAppMCPStdioServer(t *testing.T) { if os.Getenv("GO_WANT_APP_MCP_STDIO_HELPER") != "1" { return diff --git a/internal/app/provider_text_generator_test.go b/internal/app/provider_text_generator_test.go deleted file mode 100644 index b02e82bd..00000000 --- a/internal/app/provider_text_generator_test.go +++ /dev/null @@ -1,252 +0,0 @@ -package app - -import ( - "context" - "errors" - "strings" - "testing" - - "neo-code/internal/config" - "neo-code/internal/memo" - "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 -} - -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 ") - events <- providertypes.NewTextDeltaStreamEvent("world") - events <- providertypes.NewMessageDoneStreamEvent("stop", nil) - return nil - }, - } - factory := &stubTextGenFactory{provider: providerStub} - generator := newProviderTextGenerator(factory, runtimeCfg, model) - - 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) - } - 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, runtimeCfg, model) - - _, 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) - 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}, runtimeCfg, model) - - _, 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) - 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}, runtimeCfg, model) - - _, err := generator.Generate(context.Background(), "prompt", nil) - if err == nil || !strings.Contains(err.Error(), "message_done") { - t.Fatalf("Generate() error = %v", err) - } -} - -// 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() - 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 -} - -// 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/cli/root.go b/internal/cli/root.go index b48d8600..b5b3c5dd 100644 --- a/internal/cli/root.go +++ b/internal/cli/root.go @@ -2,6 +2,7 @@ package cli import ( "context" + "errors" "strings" "github.com/spf13/cobra" @@ -49,11 +50,24 @@ func NewRootCommand() *cobra.Command { } // defaultRootProgramLauncher 负责在默认根命令路径下启动 TUI。 -func defaultRootProgramLauncher(ctx context.Context, opts app.BootstrapOptions) error { - program, err := newRootProgram(ctx, opts) +func defaultRootProgramLauncher(ctx context.Context, opts app.BootstrapOptions) (err error) { + program, cleanup, err := newRootProgram(ctx, opts) if err != nil { return err } + if cleanup != nil { + defer func() { + cleanupErr := cleanup() + if cleanupErr == nil { + return + } + if err == nil { + err = cleanupErr + return + } + err = errors.Join(err, cleanupErr) + }() + } _, err = program.Run() return err } diff --git a/internal/cli/root_test.go b/internal/cli/root_test.go index 0ad9103a..68760522 100644 --- a/internal/cli/root_test.go +++ b/internal/cli/root_test.go @@ -96,14 +96,18 @@ func TestDefaultRootProgramLauncherRunsProgram(t *testing.T) { originalNewProgram := newRootProgram t.Cleanup(func() { newRootProgram = originalNewProgram }) - newRootProgram = func(ctx context.Context, opts app.BootstrapOptions) (*tea.Program, error) { + cleanedUp := false + newRootProgram = func(ctx context.Context, opts app.BootstrapOptions) (*tea.Program, func() error, error) { model := quitModel{} - return tea.NewProgram(model, tea.WithInput(nil), tea.WithOutput(io.Discard)), nil + return tea.NewProgram(model, tea.WithInput(nil), tea.WithOutput(io.Discard)), func() error { cleanedUp = true; return nil }, nil } if err := defaultRootProgramLauncher(context.Background(), app.BootstrapOptions{Workdir: `D:\项目\中文目录`}); err != nil { t.Fatalf("defaultRootProgramLauncher() error = %v", err) } + if !cleanedUp { + t.Fatalf("expected cleanup to be called") + } } func TestDefaultRootProgramLauncherReturnsNewProgramError(t *testing.T) { @@ -111,8 +115,8 @@ func TestDefaultRootProgramLauncherReturnsNewProgramError(t *testing.T) { t.Cleanup(func() { newRootProgram = originalNewProgram }) expected := errors.New("new program failed") - newRootProgram = func(ctx context.Context, opts app.BootstrapOptions) (*tea.Program, error) { - return nil, expected + newRootProgram = func(ctx context.Context, opts app.BootstrapOptions) (*tea.Program, func() error, error) { + return nil, nil, expected } err := defaultRootProgramLauncher(context.Background(), app.BootstrapOptions{}) @@ -121,6 +125,47 @@ func TestDefaultRootProgramLauncherReturnsNewProgramError(t *testing.T) { } } +func TestDefaultRootProgramLauncherReturnsCleanupErrorWhenRunSucceeds(t *testing.T) { + originalNewProgram := newRootProgram + t.Cleanup(func() { newRootProgram = originalNewProgram }) + + cleanupErr := errors.New("cleanup failed") + newRootProgram = func(ctx context.Context, opts app.BootstrapOptions) (*tea.Program, func() error, error) { + model := quitModel{} + return tea.NewProgram(model, tea.WithInput(nil), tea.WithOutput(io.Discard)), func() error { + return cleanupErr + }, nil + } + + err := defaultRootProgramLauncher(context.Background(), app.BootstrapOptions{}) + if !errors.Is(err, cleanupErr) { + t.Fatalf("expected cleanup error %v, got %v", cleanupErr, err) + } +} + +func TestDefaultRootProgramLauncherJoinsRunAndCleanupErrors(t *testing.T) { + originalNewProgram := newRootProgram + t.Cleanup(func() { newRootProgram = originalNewProgram }) + + runErr := context.Canceled + cleanupErr := errors.New("cleanup failed") + newRootProgram = func(ctx context.Context, opts app.BootstrapOptions) (*tea.Program, func() error, error) { + cancelledCtx, cancel := context.WithCancel(context.Background()) + cancel() + return tea.NewProgram(quitModel{}, tea.WithContext(cancelledCtx), tea.WithInput(nil), tea.WithOutput(io.Discard)), func() error { + return cleanupErr + }, nil + } + + err := defaultRootProgramLauncher(context.Background(), app.BootstrapOptions{}) + if !errors.Is(err, runErr) { + t.Fatalf("expected joined error to include run error %v, got %v", runErr, err) + } + if !errors.Is(err, cleanupErr) { + t.Fatalf("expected joined error to include cleanup error %v, got %v", cleanupErr, err) + } +} + type quitModel struct{} func (quitModel) Init() tea.Cmd { diff --git a/internal/provider/generate.go b/internal/provider/generate.go new file mode 100644 index 00000000..c1d2b544 --- /dev/null +++ b/internal/provider/generate.go @@ -0,0 +1,64 @@ +package provider + +import ( + "context" + "fmt" + "strings" + + providertypes "neo-code/internal/provider/types" +) + +// GenerateText 聚合并消费流式事件，直接返回完整字符串。消灭上层流处理样板代码。 +func GenerateText(ctx context.Context, p Provider, req providertypes.GenerateRequest) (string, error) { + events := make(chan providertypes.StreamEvent, 32) + done := make(chan error, 1) + var builder strings.Builder + + go func() { + var streamErr error + messageDone := false + 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("unexpected provider stream event %q", event.Type) + } + } + } + if streamErr == nil && !messageDone { + streamErr = fmt.Errorf("provider stream ended without message_done event") + } + done <- streamErr + }() + + err := p.Generate(ctx, req, events) + close(events) + + if streamErr := <-done; streamErr != nil { + if err != nil { + return "", fmt.Errorf("generate failed: %v: %w", streamErr, err) + } + return "", streamErr + } + if err != nil { + return "", err + } + return builder.String(), nil +} diff --git a/internal/provider/generate_test.go b/internal/provider/generate_test.go new file mode 100644 index 00000000..2636aaa5 --- /dev/null +++ b/internal/provider/generate_test.go @@ -0,0 +1,148 @@ +package provider_test + +import ( + "context" + "errors" + "strings" + "testing" + + "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 +} + +func TestGenerateTextSuccess(t *testing.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 + }, + } + + req := providertypes.GenerateRequest{ + Model: "test-model", + SystemPrompt: "test prompt", + Messages: []providertypes.Message{ + {Role: providertypes.RoleUser, Content: "test message"}, + }, + } + + text, err := provider.GenerateText(context.Background(), providerStub, req) + if err != nil { + t.Fatalf("GenerateText() error = %v", err) + } + if text != "hello world" { + t.Fatalf("text = %q, want %q", text, "hello world") + } + if len(providerStub.requests) != 1 { + t.Fatalf("provider requests = %d, want 1", len(providerStub.requests)) + } +} + +func TestGenerateTextProviderError(t *testing.T) { + providerStub := &stubTextGenProvider{ + generate: func(ctx context.Context, req providertypes.GenerateRequest, events chan<- providertypes.StreamEvent) error { + return errors.New("provider error") + }, + } + + _, err := provider.GenerateText(context.Background(), providerStub, providertypes.GenerateRequest{}) + if err == nil || !strings.Contains(err.Error(), "provider error") { + t.Fatalf("GenerateText() error = %v", err) + } +} + +func TestGenerateTextReturnsEmptyTextWhenProviderErrorsAfterStreaming(t *testing.T) { + providerStub := &stubTextGenProvider{ + generate: func(ctx context.Context, req providertypes.GenerateRequest, events chan<- providertypes.StreamEvent) error { + events <- providertypes.NewTextDeltaStreamEvent("partial") + events <- providertypes.NewMessageDoneStreamEvent("stop", nil) + return errors.New("provider error") + }, + } + + text, err := provider.GenerateText(context.Background(), providerStub, providertypes.GenerateRequest{}) + if err == nil || !strings.Contains(err.Error(), "provider error") { + t.Fatalf("GenerateText() error = %v", err) + } + if text != "" { + t.Fatalf("text = %q, want empty text on provider error", text) + } +} + +func TestGenerateTextRequiresMessageDone(t *testing.T) { + providerStub := &stubTextGenProvider{ + generate: func(ctx context.Context, req providertypes.GenerateRequest, events chan<- providertypes.StreamEvent) error { + events <- providertypes.NewTextDeltaStreamEvent("partial") + return nil + }, + } + + _, err := provider.GenerateText(context.Background(), providerStub, providertypes.GenerateRequest{}) + if err == nil || !strings.Contains(err.Error(), "message_done") { + t.Fatalf("GenerateText() error = %v", err) + } +} + +func TestGenerateTextRejectsUnexpectedEvent(t *testing.T) { + providerStub := &stubTextGenProvider{ + generate: func(ctx context.Context, req providertypes.GenerateRequest, events chan<- providertypes.StreamEvent) error { + events <- providertypes.StreamEvent{Type: "unexpected"} + events <- providertypes.NewMessageDoneStreamEvent("stop", nil) + return nil + }, + } + + _, err := provider.GenerateText(context.Background(), providerStub, providertypes.GenerateRequest{}) + if err == nil || !strings.Contains(err.Error(), "unexpected provider stream event") { + t.Fatalf("GenerateText() error = %v", err) + } +} + +func TestGenerateTextRejectsTextDeltaWithNilPayload(t *testing.T) { + providerStub := &stubTextGenProvider{ + generate: func(ctx context.Context, req providertypes.GenerateRequest, events chan<- providertypes.StreamEvent) error { + events <- providertypes.StreamEvent{Type: providertypes.StreamEventTextDelta} + events <- providertypes.NewMessageDoneStreamEvent("stop", nil) + return nil + }, + } + + _, err := provider.GenerateText(context.Background(), providerStub, providertypes.GenerateRequest{}) + if err == nil || !strings.Contains(err.Error(), "text_delta event payload is nil") { + t.Fatalf("GenerateText() error = %v", err) + } +} + +func TestGenerateTextRejectsMessageDoneWithNilPayload(t *testing.T) { + providerStub := &stubTextGenProvider{ + generate: func(ctx context.Context, req providertypes.GenerateRequest, events chan<- providertypes.StreamEvent) error { + events <- providertypes.NewTextDeltaStreamEvent("partial") + events <- providertypes.StreamEvent{Type: providertypes.StreamEventMessageDone} + return nil + }, + } + + _, err := provider.GenerateText(context.Background(), providerStub, providertypes.GenerateRequest{}) + if err == nil || !strings.Contains(err.Error(), "message_done event payload is nil") { + t.Fatalf("GenerateText() error = %v", err) + } +} diff --git a/internal/tools/mcp/registry.go b/internal/tools/mcp/registry.go index ea85e0cc..fa0b76ba 100644 --- a/internal/tools/mcp/registry.go +++ b/internal/tools/mcp/registry.go @@ -132,6 +132,22 @@ func (r *Registry) UnregisterServer(serverID string) bool { return true } +// Close 会卸载并关闭注册表中所有的 MCP client，通常用于在系统退出时清理所有的 stdio 子进程。 +func (r *Registry) Close() error { + if r == nil { + return nil + } + r.mu.Lock() + servers := r.servers + r.servers = make(map[string]*serverEntry) + r.mu.Unlock() + + for _, entry := range servers { + closeServerClient(entry.client) + } + return nil +} + // SetServerStatus 更新指定 server 的生命周期状态。 func (r *Registry) SetServerStatus(serverID string, status ServerStatus) error { if r == nil { diff --git a/internal/tools/mcp/registry_test.go b/internal/tools/mcp/registry_test.go index 13d34179..f58d2065 100644 --- a/internal/tools/mcp/registry_test.go +++ b/internal/tools/mcp/registry_test.go @@ -321,6 +321,39 @@ func TestRegistryUnregisterServerClosesClient(t *testing.T) { } } +func TestRegistryClose(t *testing.T) { + registry := NewRegistry() + client1 := &closableStubServerClient{} + client2 := &closableStubServerClient{} + + if err := registry.RegisterServer("srv-1", "src1", "v1", client1); err != nil { + t.Fatalf("RegisterServer() error = %v", err) + } + if err := registry.RegisterServer("srv-2", "src2", "v2", client2); err != nil { + t.Fatalf("RegisterServer() error = %v", err) + } + + if err := registry.Close(); err != nil { + t.Fatalf("Close() error = %v", err) + } + + if !client1.closed || !client2.closed { + t.Fatalf("expected all clients to be closed") + } + if len(registry.Snapshot()) != 0 { + t.Fatalf("expected registry to be empty after Close") + } +} + +func TestRegistryCloseNilReceiver(t *testing.T) { + t.Parallel() + + var nilRegistry *Registry + if err := nilRegistry.Close(); err != nil { + t.Fatalf("expected nil receiver close to return nil, got %v", err) + } +} + func TestCloseServerClientBoundaries(t *testing.T) { t.Parallel()