wip: refactoring

3 files changed, 15 insertions, 10 deletions

diff --git a/internal/llm/agent/agent-tool.go b/internal/llm/agent/agent-tool.go
index 713b0690d..d451eafed 100644
--- a/internal/llm/agent/agent-tool.go
+++ b/internal/llm/agent/agent-tool.go
@@ -91,7 +91,7 @@ func (b *agentTool) Run(ctx context.Context, call tools.ToolCall) (tools.ToolRes
 	parentSession.PromptTokens += updatedSession.PromptTokens
 	parentSession.CompletionTokens += updatedSession.CompletionTokens
 
-	_, err = b.sessions.Save(ctx, parentSession)
+	_, err = b.sessions.Update(ctx, parentSession)
 	if err != nil {
 		return tools.ToolResponse{}, fmt.Errorf("error saving parent session: %s", err)
 	}
diff --git a/internal/llm/agent/agent.go b/internal/llm/agent/agent.go
index 295ac4654..f6f816256 100644
--- a/internal/llm/agent/agent.go
+++ b/internal/llm/agent/agent.go
@@ -156,7 +156,7 @@ func (a *agent) generateTitle(ctx context.Context, sessionID string, content str
 	}
 
 	session.Title = title
-	_, err = a.sessions.Save(ctx, session)
+	_, err = a.sessions.Update(ctx, session)
 	return err
 }
 
@@ -459,7 +459,7 @@ out:
 
 func (a *agent) finishMessage(ctx context.Context, msg *message.Message, finishReson message.FinishReason) {
 	msg.AddFinish(finishReson)
-	_ = a.messages.Update(ctx, *msg)
+	_, _ = a.messages.Update(ctx, *msg)
 }
 
 func (a *agent) processEvent(ctx context.Context, sessionID string, assistantMsg *message.Message, event provider.ProviderEvent) error {
@@ -477,13 +477,16 @@ func (a *agent) processEvent(ctx context.Context, sessionID string, assistantMsg
 	switch event.Type {
 	case provider.EventThinkingDelta:
 		assistantMsg.AppendReasoningContent(event.Content)
-		return a.messages.Update(ctx, *assistantMsg)
+		_, err := a.messages.Update(ctx, *assistantMsg)
+		return err
 	case provider.EventContentDelta:
 		assistantMsg.AppendContent(event.Content)
-		return a.messages.Update(ctx, *assistantMsg)
+		_, err := a.messages.Update(ctx, *assistantMsg)
+		return err
 	case provider.EventToolUseStart:
 		assistantMsg.AddToolCall(*event.ToolCall)
-		return a.messages.Update(ctx, *assistantMsg)
+		_, err := a.messages.Update(ctx, *assistantMsg)
+		return err
 	// TODO: see how to handle this
 	// case provider.EventToolUseDelta:
 	// 	tm := time.Unix(assistantMsg.UpdatedAt, 0)
@@ -495,7 +498,8 @@ func (a *agent) processEvent(ctx context.Context, sessionID string, assistantMsg
 	// 	}
 	case provider.EventToolUseStop:
 		assistantMsg.FinishToolCall(event.ToolCall.ID)
-		return a.messages.Update(ctx, *assistantMsg)
+		_, err := a.messages.Update(ctx, *assistantMsg)
+		return err
 	case provider.EventError:
 		if errors.Is(event.Error, context.Canceled) {
 			status.Info(fmt.Sprintf("Event processing canceled for session: %s", sessionID))
@@ -506,7 +510,7 @@ func (a *agent) processEvent(ctx context.Context, sessionID string, assistantMsg
 	case provider.EventComplete:
 		assistantMsg.SetToolCalls(event.Response.ToolCalls)
 		assistantMsg.AddFinish(event.Response.FinishReason)
-		if err := a.messages.Update(ctx, *assistantMsg); err != nil {
+		if _, err := a.messages.Update(ctx, *assistantMsg); err != nil {
 			return fmt.Errorf("failed to update message: %w", err)
 		}
 		return a.TrackUsage(ctx, sessionID, a.provider.Model(), event.Response.Usage)
@@ -540,7 +544,7 @@ func (a *agent) TrackUsage(ctx context.Context, sessionID string, model models.M
 	sess.CompletionTokens += usage.OutputTokens
 	sess.PromptTokens += usage.InputTokens
 
-	_, err = a.sessions.Save(ctx, sess)
+	_, err = a.sessions.Update(ctx, sess)
 	if err != nil {
 		return fmt.Errorf("failed to save session: %w", err)
 	}
@@ -691,7 +695,7 @@ func (a *agent) CompactSession(ctx context.Context, sessionID string) error {
 	session.SummarizedAt = currentTime
 
 	// Save the updated session
-	_, err = a.sessions.Save(ctx, session)
+	_, err = a.sessions.Update(ctx, session)
 	if err != nil {
 		return fmt.Errorf("failed to save session with summary: %w", err)
 	}
diff --git a/internal/llm/agent/mcp-tools.go b/internal/llm/agent/mcp-tools.go
index 9966b99d9..a6a21d028 100644
--- a/internal/llm/agent/mcp-tools.go
+++ b/internal/llm/agent/mcp-tools.go
@@ -86,6 +86,7 @@ func (b *mcpTool) Run(ctx context.Context, params tools.ToolCall) (tools.ToolRes
 	}
 	permissionDescription := fmt.Sprintf("execute %s with the following parameters: %s", b.Info().Name, params.Input)
 	p := b.permissions.Request(
+		ctx,
 		permission.CreatePermissionRequest{
 			SessionID:   sessionID,
 			Path:        config.WorkingDirectory(),