perf: improve startup and shutdown speed (#1829)

Carlos Alexandro Becker created 3 weeks ago

Signed-off-by: Carlos Alexandro Becker <[e-mail]>

Change summary

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(-)

Detailed changes

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
 }

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 {

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.

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)
 			}
 		})
 	}

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

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
 	}
 }
 

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()