From 6f646c9151a63d2c609f349847ef670420dddc4e Mon Sep 17 00:00:00 2001 From: Carlos Alexandro Becker Date: Mon, 12 Jan 2026 10:37:05 -0300 Subject: [PATCH] perf: improve startup and shutdown speed (#1829) Signed-off-by: Carlos Alexandro Becker --- internal/agent/agent.go | 5 +++++ internal/agent/coordinator.go | 17 +++++++++++------ internal/agent/tools/mcp/init.go | 29 ++++++++++++++--------------- internal/app/app.go | 11 +++++++---- internal/lsp/client.go | 4 ---- internal/shell/background.go | 26 +++++++++++++++++++------- internal/skills/skills.go | 2 +- 7 files changed, 57 insertions(+), 37 deletions(-) diff --git a/internal/agent/agent.go b/internal/agent/agent.go index dae219a53f5387e45dcb084fa6ac5ae4a7165a4b..c64a4c7838b8545eed8bbaa3e32f33bab437f8d6 100644 --- a/internal/agent/agent.go +++ b/internal/agent/agent.go @@ -68,6 +68,7 @@ type SessionAgent interface { Run(context.Context, SessionAgentCall) (*fantasy.AgentResult, error) SetModels(large Model, small Model) SetTools(tools []fantasy.AgentTool) + SetSystemPrompt(systemPrompt string) Cancel(sessionID string) CancelAll() IsSessionBusy(sessionID string) bool @@ -964,6 +965,10 @@ func (a *sessionAgent) SetTools(tools []fantasy.AgentTool) { a.tools = tools } +func (a *sessionAgent) SetSystemPrompt(systemPrompt string) { + a.systemPrompt = systemPrompt +} + func (a *sessionAgent) Model() Model { return a.largeModel } diff --git a/internal/agent/coordinator.go b/internal/agent/coordinator.go index b13603bb131090c86eaff3f6ea9527cb1b9dacf6..d48c074ff5d9c02db1427aa44ee13aed4c9af7e2 100644 --- a/internal/agent/coordinator.go +++ b/internal/agent/coordinator.go @@ -322,17 +322,12 @@ func (c *coordinator) buildAgent(ctx context.Context, prompt *prompt.Prompt, age return nil, err } - systemPrompt, err := prompt.Build(ctx, large.Model.Provider(), large.Model.Model(), *c.cfg) - if err != nil { - return nil, err - } - largeProviderCfg, _ := c.cfg.Providers.Get(large.ModelCfg.Provider) result := NewSessionAgent(SessionAgentOptions{ large, small, largeProviderCfg.SystemPromptPrefix, - systemPrompt, + "", isSubAgent, c.cfg.Options.DisableAutoSummarize, c.permissions.SkipRequests(), @@ -340,6 +335,16 @@ func (c *coordinator) buildAgent(ctx context.Context, prompt *prompt.Prompt, age c.messages, nil, }) + + c.readyWg.Go(func() error { + systemPrompt, err := prompt.Build(ctx, large.Model.Provider(), large.Model.Model(), *c.cfg) + if err != nil { + return err + } + result.SetSystemPrompt(systemPrompt) + return nil + }) + c.readyWg.Go(func() error { tools, err := c.buildTools(ctx, agent) if err != nil { diff --git a/internal/agent/tools/mcp/init.go b/internal/agent/tools/mcp/init.go index c28da6c1722413d276bbb47b3dbf3e9f66826263..bb43f7f157dc1cf2d094354a4e709e0beb1f52b6 100644 --- a/internal/agent/tools/mcp/init.go +++ b/internal/agent/tools/mcp/init.go @@ -107,29 +107,28 @@ func GetState(name string) (ClientInfo, bool) { // Close closes all MCP clients. This should be called during application shutdown. func Close() error { - var errs []error var wg sync.WaitGroup - for name, session := range sessions.Seq2() { - wg.Go(func() { - done := make(chan bool, 1) - go func() { + done := make(chan struct{}, 1) + go func() { + for name, session := range sessions.Seq2() { + wg.Go(func() { if err := session.Close(); err != nil && !errors.Is(err, io.EOF) && !errors.Is(err, context.Canceled) && err.Error() != "signal: killed" { - errs = append(errs, fmt.Errorf("close mcp: %s: %w", name, err)) + slog.Warn("Failed to shutdown MCP client", "name", name, "error", err) } - done <- true - }() - select { - case <-done: - case <-time.After(time.Millisecond * 250): - } - }) + }) + } + wg.Wait() + done <- struct{}{} + }() + select { + case <-done: + case <-time.After(5 * time.Second): } - wg.Wait() broker.Shutdown() - return errors.Join(errs...) + return nil } // Initialize initializes MCP clients based on the provided configuration. diff --git a/internal/app/app.go b/internal/app/app.go index 436579d0b9593a1f9fd36606ae1e1b81fd89e737..24f1500cc707a16964f00a6eac1a62c0ac094850 100644 --- a/internal/app/app.go +++ b/internal/app/app.go @@ -405,12 +405,15 @@ func (app *App) Shutdown() { }) // Shutdown all LSP clients. + shutdownCtx, cancel := context.WithTimeout(app.globalCtx, 5*time.Second) + defer cancel() for name, client := range app.LSPClients.Seq2() { wg.Go(func() { - shutdownCtx, cancel := context.WithTimeout(app.globalCtx, 5*time.Second) - defer cancel() - if err := client.Close(shutdownCtx); err != nil { - slog.Error("Failed to shutdown LSP client", "name", name, "error", err) + if err := client.Close(shutdownCtx); err != nil && + !errors.Is(err, io.EOF) && + !errors.Is(err, context.Canceled) && + err.Error() != "signal: killed" { + slog.Warn("Failed to shutdown LSP client", "name", name, "error", err) } }) } diff --git a/internal/lsp/client.go b/internal/lsp/client.go index 7dba52fdf48a2205bc9fca8436390326be2a2d39..79220cc1f315fec30a1bee2aa0dcd106bc311a02 100644 --- a/internal/lsp/client.go +++ b/internal/lsp/client.go @@ -153,10 +153,6 @@ func (c *Client) Initialize(ctx context.Context, workspaceDir string) (*protocol // Close closes the LSP client. func (c *Client) Close(ctx context.Context) error { - // Try to close all open files first - ctx, cancel := context.WithTimeout(ctx, 5*time.Second) - defer cancel() - c.CloseAllFiles(ctx) // Shutdown and exit the client diff --git a/internal/shell/background.go b/internal/shell/background.go index 37dd5d909adb0ef2230a4a84ab394851dd8167a4..bc81369ec877586c92fa9bc701d8b78b669f23d5 100644 --- a/internal/shell/background.go +++ b/internal/shell/background.go @@ -4,6 +4,7 @@ import ( "bytes" "context" "fmt" + "slices" "sync" "sync/atomic" "time" @@ -163,15 +164,26 @@ func (m *BackgroundShellManager) Cleanup() int { // KillAll terminates all background shells. func (m *BackgroundShellManager) KillAll() { - shells := make([]*BackgroundShell, 0, m.shells.Len()) - for shell := range m.shells.Seq() { - shells = append(shells, shell) - } + shells := slices.Collect(m.shells.Seq()) m.shells.Reset(map[string]*BackgroundShell{}) + done := make(chan struct{}, 1) + go func() { + var wg sync.WaitGroup + for _, shell := range shells { + wg.Go(func() { + shell.cancel() + <-shell.done + }) + } + wg.Wait() + done <- struct{}{} + }() - for _, shell := range shells { - shell.cancel() - <-shell.done + select { + case <-done: + return + case <-time.After(time.Second * 5): + return } } diff --git a/internal/skills/skills.go b/internal/skills/skills.go index cba0b994d188e184fdcb55cf98bd080764e34327..e0488cff4a8c42ceee68a6a89608bd90acafdb2e 100644 --- a/internal/skills/skills.go +++ b/internal/skills/skills.go @@ -147,7 +147,7 @@ func Discover(paths []string) []*Skill { slog.Warn("Skill validation failed", "path", path, "error", err) return nil } - slog.Info("Successfully loaded skill", "name", skill.Name, "path", path) + slog.Debug("Successfully loaded skill", "name", skill.Name, "path", path) mu.Lock() skills = append(skills, skill) mu.Unlock()