fix: improvements

Carlos Alexandro Becker created 3 months ago

Change summary

internal/cmd/root.go        |  4 ----
internal/llm/agent/agent.go | 29 ++++++++++++++++-------------
2 files changed, 16 insertions(+), 17 deletions(-)

Detailed changes

internal/cmd/root.go

@@ -6,7 +6,6 @@ import (
 	"io"
 	"log/slog"
 	"os"
-	"time"
 
 	tea "github.com/charmbracelet/bubbletea/v2"
 	"github.com/charmbracelet/crush/internal/app"
@@ -84,15 +83,12 @@ 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 {

internal/llm/agent/agent.go

@@ -187,6 +187,10 @@ func NewAgent(
 
 	go func() {
 		slog.Info("Initializing agent tools", "agent", agentCfg.ID)
+		defer func() {
+			slog.Info("Initialized agent tools", "agent", agentCfg.ID)
+			agent.toolsDone.Store(true)
+		}()
 
 		cwd := cfg.WorkingDir()
 		allTools := []tools.BaseTool{
@@ -203,29 +207,28 @@ func NewAgent(
 		}
 
 		mcpTools := GetMCPTools(ctx, permissions, cfg)
+		allTools = append(allTools, mcpTools...)
+
 		if len(lspClients) > 0 {
-			mcpTools = append(mcpTools, tools.NewDiagnosticsTool(lspClients))
+			allTools = append(allTools, 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)
-				}
-			}
+			agent.tools = allTools
+			return
 		}
 
-		slog.Info("Initialized agent tools", "agent", agentCfg.ID)
-		agent.tools = agentTools
-		agent.toolsDone.Store(true)
+		var filteredTools []tools.BaseTool
+		for _, tool := range allTools {
+			if slices.Contains(agentCfg.AllowedTools, tool.Name()) {
+				filteredTools = append(filteredTools, tool)
+			}
+		}
+		agent.tools = filteredTools
 	}()
 
 	return agent, nil