minor fixes

Kujtim Hoxha created 7 months ago

Change summary

internal/app/app.go              |  5 +--
internal/config/config.go        |  5 +++
internal/llm/agent/agent-tool.go |  8 ++++--
internal/llm/agent/agent.go      | 17 ++++++++++++++
internal/llm/agent/coder.go      |  2 
internal/llm/agent/task.go       | 39 +++++++++++++++++----------------
6 files changed, 49 insertions(+), 27 deletions(-)

Detailed changes

internal/app/app.go

@@ -48,9 +48,10 @@ func New(ctx context.Context, conn *sql.DB) (*App, error) {
 		LSPClients:  make(map[string]*lsp.Client),
 	}
 
+	app.initLSPClients(ctx)
+
 	var err error
 	app.CoderAgent, err = agent.NewCoderAgent(
-
 		app.Permissions,
 		app.Sessions,
 		app.Messages,
@@ -61,8 +62,6 @@ func New(ctx context.Context, conn *sql.DB) (*App, error) {
 		return nil, err
 	}
 
-	app.initLSPClients(ctx)
-
 	return app, nil
 }
 

internal/config/config.go

@@ -272,5 +272,8 @@ func Get() *Config {
 
 // WorkingDirectory returns the current working directory from the configuration.
 func WorkingDirectory() string {
-	return viper.GetString("wd")
+	if cfg == nil {
+		panic("config not loaded")
+	}
+	return cfg.WorkingDir
 }

internal/llm/agent/agent-tool.go

@@ -53,7 +53,7 @@ func (b *agentTool) Run(ctx context.Context, call tools.ToolCall) (tools.ToolRes
 		return tools.ToolResponse{}, fmt.Errorf("session_id and message_id are required")
 	}
 
-	agent, err := NewTaskAgent(b.lspClients)
+	agent, err := NewTaskAgent(b.messages, b.sessions, b.lspClients)
 	if err != nil {
 		return tools.ToolResponse{}, fmt.Errorf("error creating agent: %s", err)
 	}
@@ -105,9 +105,11 @@ func (b *agentTool) Run(ctx context.Context, call tools.ToolCall) (tools.ToolRes
 func NewAgentTool(
 	Sessions session.Service,
 	Messages message.Service,
+	LspClients map[string]*lsp.Client,
 ) tools.BaseTool {
 	return &agentTool{
-		sessions: Sessions,
-		messages: Messages,
+		sessions:   Sessions,
+		messages:   Messages,
+		lspClients: LspClients,
 	}
 }

internal/llm/agent/agent.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"os"
+	"runtime/debug"
 	"strings"
 	"sync"
 
@@ -88,6 +90,21 @@ func (a *agent) Generate(ctx context.Context, sessionID string, content string)
 		defer func() {
 			if r := recover(); r != nil {
 				logging.ErrorPersist(fmt.Sprintf("Panic in Generate: %v", r))
+
+				// dump stack trace into a file
+				file, err := os.Create("panic.log")
+				if err != nil {
+					logging.ErrorPersist(fmt.Sprintf("Failed to create panic log: %v", err))
+					return
+				}
+
+				defer file.Close()
+
+				stackTrace := debug.Stack()
+				if _, err := file.Write(stackTrace); err != nil {
+					logging.ErrorPersist(fmt.Sprintf("Failed to write panic log: %v", err))
+				}
+
 			}
 		}()
 		defer a.activeRequests.Delete(sessionID)

internal/llm/agent/coder.go

@@ -49,7 +49,7 @@ func NewCoderAgent(
 				tools.NewSourcegraphTool(),
 				tools.NewViewTool(lspClients),
 				tools.NewWriteTool(lspClients, permissions),
-				NewAgentTool(sessions, messages),
+				NewAgentTool(sessions, messages, lspClients),
 			}, otherTools...,
 		),
 	)

internal/llm/agent/task.go

@@ -8,39 +8,40 @@ import (
 	"github.com/kujtimiihoxha/termai/internal/llm/models"
 	"github.com/kujtimiihoxha/termai/internal/llm/tools"
 	"github.com/kujtimiihoxha/termai/internal/lsp"
+	"github.com/kujtimiihoxha/termai/internal/message"
+	"github.com/kujtimiihoxha/termai/internal/session"
 )
 
 type taskAgent struct {
-	*agent
+	Service
 }
 
-func (c *taskAgent) Generate(ctx context.Context, sessionID string, content string) error {
-	return c.generate(ctx, sessionID, content)
-}
-
-func NewTaskAgent(lspClients map[string]*lsp.Client) (Service, error) {
+func NewTaskAgent(messages message.Service, sessions session.Service, lspClients map[string]*lsp.Client) (Service, error) {
 	model, ok := models.SupportedModels[config.Get().Model.Coder]
 	if !ok {
 		return nil, errors.New("model not supported")
 	}
 
 	ctx := context.Background()
-	agentProvider, titleGenerator, err := getAgentProviders(ctx, model)
+
+	agent, err := NewAgent(
+		ctx,
+		sessions,
+		messages,
+		model,
+		[]tools.BaseTool{
+			tools.NewGlobTool(),
+			tools.NewGrepTool(),
+			tools.NewLsTool(),
+			tools.NewSourcegraphTool(),
+			tools.NewViewTool(lspClients),
+		},
+	)
 	if err != nil {
 		return nil, err
 	}
+
 	return &taskAgent{
-		agent: &agent{
-			tools: []tools.BaseTool{
-				tools.NewGlobTool(),
-				tools.NewGrepTool(),
-				tools.NewLsTool(),
-				tools.NewSourcegraphTool(),
-				tools.NewViewTool(lspClients),
-			},
-			model:          model,
-			agent:          agentProvider,
-			titleGenerator: titleGenerator,
-		},
+		agent,
 	}, nil
 }