From 380d524ec1f9c77dbff62298adb05c40d7f4b937 Mon Sep 17 00:00:00 2001 From: Carlos Alexandro Becker Date: Tue, 22 Jul 2025 12:10:26 -0300 Subject: [PATCH] perf: init tools in background --- internal/app/app.go | 2 +- internal/cmd/root.go | 4 ++ internal/llm/agent/agent.go | 114 ++++++++++++++++++++---------------- 3 files changed, 70 insertions(+), 50 deletions(-) diff --git a/internal/app/app.go b/internal/app/app.go index d63c90c6e2599f63e3a65cd8069b53638f45cc5f..170debad340e2cb33fcb7a1c9fe814c184573c9b 100644 --- a/internal/app/app.go +++ b/internal/app/app.go @@ -80,7 +80,7 @@ func New(ctx context.Context, conn *sql.DB, cfg *config.Config) (*App, error) { app.setupEvents() // Initialize LSP clients in the background. - go app.initLSPClients(ctx) + app.initLSPClients(ctx) // TODO: remove the concept of agent config, most likely. if cfg.IsConfigured() { diff --git a/internal/cmd/root.go b/internal/cmd/root.go index d63160992141da26b6a26610b06f1b601213e00d..e782898f42e38bc4b06f890b76a2deebb96b9df1 100644 --- a/internal/cmd/root.go +++ b/internal/cmd/root.go @@ -6,6 +6,7 @@ import ( "io" "log/slog" "os" + "time" tea "github.com/charmbracelet/bubbletea/v2" "github.com/charmbracelet/crush/internal/app" @@ -83,12 +84,15 @@ to assist developers in writing, debugging, and understanding code directly from return err } + slog.Info("Initing...") + now := time.Now() app, err := app.New(ctx, conn, cfg) if err != nil { slog.Error(fmt.Sprintf("Failed to create app instance: %v", err)) return err } defer app.Shutdown() + slog.Info("Init done", "took", time.Since(now).String()) prompt, err = maybePrependStdin(prompt) if err != nil { diff --git a/internal/llm/agent/agent.go b/internal/llm/agent/agent.go index 907a2348f838aa2f2ba6792db9b768eb656904a8..291528d36594a5178f72ae64bb565f18e36c1cae 100644 --- a/internal/llm/agent/agent.go +++ b/internal/llm/agent/agent.go @@ -8,6 +8,7 @@ import ( "slices" "strings" "sync" + "sync/atomic" "time" "github.com/charmbracelet/crush/internal/config" @@ -67,7 +68,9 @@ type agent struct { sessions session.Service messages message.Service - tools []tools.BaseTool + toolsDone atomic.Bool + tools []tools.BaseTool + provider provider.Provider providerID string @@ -94,46 +97,7 @@ func NewAgent( ) (Service, error) { ctx := context.Background() cfg := config.Get() - otherTools := GetMCPTools(ctx, permissions, cfg) - if len(lspClients) > 0 { - otherTools = append(otherTools, tools.NewDiagnosticsTool(lspClients)) - } - cwd := cfg.WorkingDir() - allTools := []tools.BaseTool{ - tools.NewBashTool(permissions, cwd), - tools.NewDownloadTool(permissions, cwd), - tools.NewEditTool(lspClients, permissions, history, cwd), - tools.NewFetchTool(permissions, cwd), - tools.NewGlobTool(cwd), - tools.NewGrepTool(cwd), - tools.NewLsTool(cwd), - tools.NewSourcegraphTool(), - tools.NewViewTool(lspClients, cwd), - tools.NewWriteTool(lspClients, permissions, history, cwd), - } - - if agentCfg.ID == "coder" { - taskAgentCfg := config.Get().Agents["task"] - if taskAgentCfg.ID == "" { - return nil, fmt.Errorf("task agent not found in config") - } - taskAgent, err := NewAgent(taskAgentCfg, permissions, sessions, messages, history, lspClients) - if err != nil { - return nil, fmt.Errorf("failed to create task agent: %w", err) - } - - allTools = append( - allTools, - NewAgentTool( - taskAgent, - sessions, - messages, - ), - ) - } - - allTools = append(allTools, otherTools...) providerCfg := config.Get().GetProviderForModel(agentCfg.Model) if providerCfg == nil { return nil, fmt.Errorf("provider for agent %s not found in config", agentCfg.Name) @@ -190,15 +154,22 @@ func NewAgent( return nil, err } - agentTools := []tools.BaseTool{} - if agentCfg.AllowedTools == nil { - agentTools = allTools - } else { - for _, tool := range allTools { - if slices.Contains(agentCfg.AllowedTools, tool.Name()) { - agentTools = append(agentTools, tool) - } + var agentTool tools.BaseTool + if agentCfg.ID == "coder" { + taskAgentCfg := config.Get().Agents["task"] + if taskAgentCfg.ID == "" { + return nil, fmt.Errorf("task agent not found in config") } + taskAgent, err := NewAgent(taskAgentCfg, permissions, sessions, messages, history, lspClients) + if err != nil { + return nil, fmt.Errorf("failed to create task agent: %w", err) + } + + agentTool = NewAgentTool( + taskAgent, + sessions, + messages, + ) } agent := &agent{ @@ -208,13 +179,55 @@ func NewAgent( providerID: string(providerCfg.ID), messages: messages, sessions: sessions, - tools: agentTools, titleProvider: titleProvider, summarizeProvider: summarizeProvider, summarizeProviderID: string(smallModelProviderCfg.ID), activeRequests: sync.Map{}, } + go func() { + slog.Info("Initializing agent tools", "agent", agentCfg.ID) + + cwd := cfg.WorkingDir() + allTools := []tools.BaseTool{ + tools.NewBashTool(permissions, cwd), + tools.NewDownloadTool(permissions, cwd), + tools.NewEditTool(lspClients, permissions, history, cwd), + tools.NewFetchTool(permissions, cwd), + tools.NewGlobTool(cwd), + tools.NewGrepTool(cwd), + tools.NewLsTool(cwd), + tools.NewSourcegraphTool(), + tools.NewViewTool(lspClients, cwd), + tools.NewWriteTool(lspClients, permissions, history, cwd), + } + + mcpTools := GetMCPTools(ctx, permissions, cfg) + if len(lspClients) > 0 { + mcpTools = append(mcpTools, tools.NewDiagnosticsTool(lspClients)) + } + allTools = append(allTools, mcpTools...) + + if agentTool != nil { + allTools = append(allTools, agentTool) + } + + agentTools := []tools.BaseTool{} + if agentCfg.AllowedTools == nil { + agentTools = allTools + } else { + for _, tool := range allTools { + if slices.Contains(agentCfg.AllowedTools, tool.Name()) { + agentTools = append(agentTools, tool) + } + } + } + + slog.Info("Initialized agent tools", "agent", agentCfg.ID) + agent.tools = agentTools + agent.toolsDone.Store(true) + }() + return agent, nil } @@ -437,6 +450,9 @@ func (a *agent) createUserMessage(ctx context.Context, sessionID, content string func (a *agent) streamAndHandleEvents(ctx context.Context, sessionID string, msgHistory []message.Message) (message.Message, *message.Message, error) { ctx = context.WithValue(ctx, tools.SessionIDContextKey, sessionID) + if !a.toolsDone.Load() { + return message.Message{}, nil, fmt.Errorf("tools not initialized yet") + } eventChan := a.provider.StreamResponse(ctx, msgHistory, a.tools) assistantMsg, err := a.messages.Create(ctx, sessionID, message.CreateMessageParams{