fix logs and add cancellation

5 files changed, 146 insertions, 24 deletions

diff --git a/internal/llm/agent/agent-tool.go b/internal/llm/agent/agent-tool.go
index bf5e31f8f..deb6aed60 100644
--- a/internal/llm/agent/agent-tool.go
+++ b/internal/llm/agent/agent-tool.go
@@ -56,7 +56,7 @@ func (b *agentTool) Run(ctx context.Context, call tools.ToolCall) (tools.ToolRes
 		return tools.NewTextErrorResponse(fmt.Sprintf("error creating session: %s", err)), nil
 	}
 
-	err = agent.Generate(session.ID, params.Prompt)
+	err = agent.Generate(ctx, session.ID, params.Prompt)
 	if err != nil {
 		return tools.NewTextErrorResponse(fmt.Sprintf("error generating agent: %s", err)), nil
 	}
diff --git a/internal/llm/agent/agent.go b/internal/llm/agent/agent.go
index 43ba0fc26..998dc1551 100644
--- a/internal/llm/agent/agent.go
+++ b/internal/llm/agent/agent.go
@@ -13,11 +13,12 @@ import (
 	"github.com/kujtimiihoxha/termai/internal/llm/prompt"
 	"github.com/kujtimiihoxha/termai/internal/llm/provider"
 	"github.com/kujtimiihoxha/termai/internal/llm/tools"
+	"github.com/kujtimiihoxha/termai/internal/logging"
 	"github.com/kujtimiihoxha/termai/internal/message"
 )
 
 type Agent interface {
-	Generate(sessionID string, content string) error
+	Generate(ctx context.Context, sessionID string, content string) error
 }
 
 type agent struct {
@@ -28,9 +29,9 @@ type agent struct {
 	titleGenerator provider.Provider
 }
 
-func (c *agent) handleTitleGeneration(sessionID, content string) {
+func (c *agent) handleTitleGeneration(ctx context.Context, sessionID, content string) {
 	response, err := c.titleGenerator.SendMessages(
-		c.Context,
+		ctx,
 		[]message.Message{
 			{
 				Role: message.User,
@@ -91,13 +92,16 @@ func (c *agent) processEvent(
 		assistantMsg.AppendContent(event.Content)
 		return c.Messages.Update(*assistantMsg)
 	case provider.EventError:
-		c.App.Logger.PersistError(event.Error.Error())
+		if errors.Is(event.Error, context.Canceled) {
+			return nil
+		}
+		logging.ErrorPersist(event.Error.Error())
 		return event.Error
 	case provider.EventWarning:
-		c.App.Logger.PersistWarn(event.Info)
+		logging.WarnPersist(event.Info)
 		return nil
 	case provider.EventInfo:
-		c.App.Logger.PersistInfo(event.Info)
+		logging.InfoPersist(event.Info)
 	case provider.EventComplete:
 		assistantMsg.SetToolCalls(event.Response.ToolCalls)
 		assistantMsg.AddFinish(event.Response.FinishReason)
@@ -115,12 +119,37 @@ func (c *agent) ExecuteTools(ctx context.Context, toolCalls []message.ToolCall,
 	var wg sync.WaitGroup
 	toolResults := make([]message.ToolResult, len(toolCalls))
 	mutex := &sync.Mutex{}
+	errChan := make(chan error, 1)
+
+	// Create a child context that can be canceled
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
 
 	for i, tc := range toolCalls {
 		wg.Add(1)
 		go func(index int, toolCall message.ToolCall) {
 			defer wg.Done()
 
+			// Check if context is already canceled
+			select {
+			case <-ctx.Done():
+				mutex.Lock()
+				toolResults[index] = message.ToolResult{
+					ToolCallID: toolCall.ID,
+					Content:    "Tool execution canceled",
+					IsError:    true,
+				}
+				mutex.Unlock()
+
+				// Send cancellation error to error channel if it's empty
+				select {
+				case errChan <- ctx.Err():
+				default:
+				}
+				return
+			default:
+			}
+
 			response := ""
 			isError := false
 			found := false
@@ -133,8 +162,19 @@ func (c *agent) ExecuteTools(ctx context.Context, toolCalls []message.ToolCall,
 						Name:  toolCall.Name,
 						Input: toolCall.Input,
 					})
+
 					if toolErr != nil {
-						response = fmt.Sprintf("error running tool: %s", toolErr)
+						if errors.Is(toolErr, context.Canceled) {
+							response = "Tool execution canceled"
+
+							// Send cancellation error to error channel if it's empty
+							select {
+							case errChan <- ctx.Err():
+							default:
+							}
+						} else {
+							response = fmt.Sprintf("error running tool: %s", toolErr)
+						}
 						isError = true
 					} else {
 						response = toolResult.Content
@@ -160,7 +200,24 @@ func (c *agent) ExecuteTools(ctx context.Context, toolCalls []message.ToolCall,
 		}(i, tc)
 	}
 
-	wg.Wait()
+	// Wait for all goroutines to finish or context to be canceled
+	done := make(chan struct{})
+	go func() {
+		wg.Wait()
+		close(done)
+	}()
+
+	select {
+	case <-done:
+		// All tools completed successfully
+	case err := <-errChan:
+		// One of the tools encountered a cancellation
+		return toolResults, err
+	case <-ctx.Done():
+		// Context was canceled externally
+		return toolResults, ctx.Err()
+	}
+
 	return toolResults, nil
 }
 
@@ -188,14 +245,14 @@ func (c *agent) handleToolExecution(
 	return &msg, err
 }
 
-func (c *agent) generate(sessionID string, content string) error {
+func (c *agent) generate(ctx context.Context, sessionID string, content string) error {
 	messages, err := c.Messages.List(sessionID)
 	if err != nil {
 		return err
 	}
 
 	if len(messages) == 0 {
-		go c.handleTitleGeneration(sessionID, content)
+		go c.handleTitleGeneration(ctx, sessionID, content)
 	}
 
 	userMsg, err := c.Messages.Create(sessionID, message.CreateMessageParams{
@@ -212,9 +269,36 @@ func (c *agent) generate(sessionID string, content string) error {
 
 	messages = append(messages, userMsg)
 	for {
+		select {
+		case <-ctx.Done():
+			assistantMsg, err := c.Messages.Create(sessionID, message.CreateMessageParams{
+				Role:  message.Assistant,
+				Parts: []message.ContentPart{},
+			})
+			if err != nil {
+				return err
+			}
+			assistantMsg.AddFinish("canceled")
+			c.Messages.Update(assistantMsg)
+			return context.Canceled
+		default:
+			// Continue processing
+		}
 
-		eventChan, err := c.agent.StreamResponse(c.Context, messages, c.tools)
+		eventChan, err := c.agent.StreamResponse(ctx, messages, c.tools)
 		if err != nil {
+			if errors.Is(err, context.Canceled) {
+				assistantMsg, err := c.Messages.Create(sessionID, message.CreateMessageParams{
+					Role:  message.Assistant,
+					Parts: []message.ContentPart{},
+				})
+				if err != nil {
+					return err
+				}
+				assistantMsg.AddFinish("canceled")
+				c.Messages.Update(assistantMsg)
+				return context.Canceled
+			}
 			return err
 		}
 
@@ -228,19 +312,47 @@ func (c *agent) generate(sessionID string, content string) error {
 		for event := range eventChan {
 			err = c.processEvent(sessionID, &assistantMsg, event)
 			if err != nil {
+				if errors.Is(err, context.Canceled) {
+					assistantMsg.AddFinish("canceled")
+					c.Messages.Update(assistantMsg)
+					return context.Canceled
+				}
 				assistantMsg.AddFinish("error:" + err.Error())
 				c.Messages.Update(assistantMsg)
 				return err
 			}
+
+			select {
+			case <-ctx.Done():
+				assistantMsg.AddFinish("canceled")
+				c.Messages.Update(assistantMsg)
+				return context.Canceled
+			default:
+			}
 		}
 
-		msg, err := c.handleToolExecution(c.Context, assistantMsg)
+		// Check for context cancellation before tool execution
+		select {
+		case <-ctx.Done():
+			assistantMsg.AddFinish("canceled")
+			c.Messages.Update(assistantMsg)
+			return context.Canceled
+		default:
+			// Continue processing
+		}
 
-		c.Messages.Update(assistantMsg)
+		msg, err := c.handleToolExecution(ctx, assistantMsg)
 		if err != nil {
+			if errors.Is(err, context.Canceled) {
+				assistantMsg.AddFinish("canceled")
+				c.Messages.Update(assistantMsg)
+				return context.Canceled
+			}
 			return err
 		}
 
+		c.Messages.Update(assistantMsg)
+
 		if len(assistantMsg.ToolCalls()) == 0 {
 			break
 		}
@@ -249,6 +361,16 @@ func (c *agent) generate(sessionID string, content string) error {
 		if msg != nil {
 			messages = append(messages, *msg)
 		}
+
+		// Check for context cancellation after tool execution
+		select {
+		case <-ctx.Done():
+			assistantMsg.AddFinish("canceled")
+			c.Messages.Update(assistantMsg)
+			return context.Canceled
+		default:
+			// Continue processing
+		}
 	}
 	return nil
 }
diff --git a/internal/llm/agent/coder.go b/internal/llm/agent/coder.go
index 8ff3c61aa..5deff05a8 100644
--- a/internal/llm/agent/coder.go
+++ b/internal/llm/agent/coder.go
@@ -1,6 +1,7 @@
 package agent
 
 import (
+	"context"
 	"errors"
 
 	"github.com/kujtimiihoxha/termai/internal/app"
@@ -28,9 +29,9 @@ func (c *coderAgent) setAgentTool(sessionID string) {
 	}
 }
 
-func (c *coderAgent) Generate(sessionID string, content string) error {
+func (c *coderAgent) Generate(ctx context.Context, sessionID string, content string) error {
 	c.setAgentTool(sessionID)
-	return c.generate(sessionID, content)
+	return c.generate(ctx, sessionID, content)
 }
 
 func NewCoderAgent(app *app.App) (Agent, error) {
diff --git a/internal/llm/agent/mcp-tools.go b/internal/llm/agent/mcp-tools.go
index dcf880fe7..b1c97b512 100644
--- a/internal/llm/agent/mcp-tools.go
+++ b/internal/llm/agent/mcp-tools.go
@@ -22,8 +22,6 @@ type mcpTool struct {
 	permissions permission.Service
 }
 
-var logger = logging.Get()
-
 type MCPClient interface {
 	Initialize(
 		ctx context.Context,
@@ -143,13 +141,13 @@ func getTools(ctx context.Context, name string, m config.MCPServer, permissions
 
 	_, err := c.Initialize(ctx, initRequest)
 	if err != nil {
-		logger.Error("error initializing mcp client", "error", err)
+		logging.Error("error initializing mcp client", "error", err)
 		return stdioTools
 	}
 	toolsRequest := mcp.ListToolsRequest{}
 	tools, err := c.ListTools(ctx, toolsRequest)
 	if err != nil {
-		logger.Error("error listing tools", "error", err)
+		logging.Error("error listing tools", "error", err)
 		return stdioTools
 	}
 	for _, t := range tools.Tools {
@@ -172,7 +170,7 @@ func GetMcpTools(ctx context.Context, permissions permission.Service) []tools.Ba
 				m.Args...,
 			)
 			if err != nil {
-				logger.Error("error creating mcp client", "error", err)
+				logging.Error("error creating mcp client", "error", err)
 				continue
 			}
 
@@ -183,7 +181,7 @@ func GetMcpTools(ctx context.Context, permissions permission.Service) []tools.Ba
 				client.WithHeaders(m.Headers),
 			)
 			if err != nil {
-				logger.Error("error creating mcp client", "error", err)
+				logging.Error("error creating mcp client", "error", err)
 				continue
 			}
 			mcpTools = append(mcpTools, getTools(ctx, name, m, permissions, c)...)
diff --git a/internal/llm/agent/task.go b/internal/llm/agent/task.go
index 9737d41b8..034e93460 100644
--- a/internal/llm/agent/task.go
+++ b/internal/llm/agent/task.go
@@ -1,6 +1,7 @@
 package agent
 
 import (
+	"context"
 	"errors"
 
 	"github.com/kujtimiihoxha/termai/internal/app"
@@ -13,8 +14,8 @@ type taskAgent struct {
 	*agent
 }
 
-func (c *taskAgent) Generate(sessionID string, content string) error {
-	return c.generate(sessionID, content)
+func (c *taskAgent) Generate(ctx context.Context, sessionID string, content string) error {
+	return c.generate(ctx, sessionID, content)
 }
 
 func NewTaskAgent(app *app.App) (Agent, error) {