wip: refactoring

8 files changed, 38 insertions, 24 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(),
diff --git a/internal/llm/tools/bash.go b/internal/llm/tools/bash.go
index 7231e1d2a..eb27026e5 100644
--- a/internal/llm/tools/bash.go
+++ b/internal/llm/tools/bash.go
@@ -268,6 +268,7 @@ func (b *bashTool) Run(ctx context.Context, call ToolCall) (ToolResponse, error)
 	}
 	if !isSafeReadOnly {
 		p := b.permissions.Request(
+			ctx,
 			permission.CreatePermissionRequest{
 				SessionID:   sessionID,
 				Path:        config.WorkingDirectory(),
diff --git a/internal/llm/tools/edit.go b/internal/llm/tools/edit.go
index 44787f525..df5c94290 100644
--- a/internal/llm/tools/edit.go
+++ b/internal/llm/tools/edit.go
@@ -37,7 +37,7 @@ type EditResponseMetadata struct {
 type editTool struct {
 	lspClients  map[string]*lsp.Client
 	permissions permission.Service
-	files       history.Service
+	history     history.Service
 }
 
 const (
@@ -95,7 +95,7 @@ func NewEditTool(lspClients map[string]*lsp.Client, permissions permission.Servi
 	return &editTool{
 		lspClients:  lspClients,
 		permissions: permissions,
-		files:       files,
+		history:     files,
 	}
 }
 
@@ -202,6 +202,7 @@ func (e *editTool) createNewFile(ctx context.Context, filePath, content string)
 		permissionPath = rootDir
 	}
 	p := e.permissions.Request(
+		ctx,
 		permission.CreatePermissionRequest{
 			SessionID:   sessionID,
 			Path:        permissionPath,
@@ -224,14 +225,14 @@ func (e *editTool) createNewFile(ctx context.Context, filePath, content string)
 	}
 
 	// File can't be in the history so we create a new file history
-	_, err = e.files.Create(ctx, sessionID, filePath, "")
+	_, err = e.history.Create(ctx, sessionID, filePath, "")
 	if err != nil {
 		// Log error but don't fail the operation
 		return ToolResponse{}, fmt.Errorf("error creating file history: %w", err)
 	}
 
 	// Add the new content to the file history
-	_, err = e.files.CreateVersion(ctx, sessionID, filePath, content)
+	_, err = e.history.CreateVersion(ctx, sessionID, filePath, content)
 	if err != nil {
 		// Log error but don't fail the operation
 		slog.Debug("Error creating file history version", "error", err)
@@ -313,6 +314,7 @@ func (e *editTool) deleteContent(ctx context.Context, filePath, oldString string
 		permissionPath = rootDir
 	}
 	p := e.permissions.Request(
+		ctx,
 		permission.CreatePermissionRequest{
 			SessionID:   sessionID,
 			Path:        permissionPath,
@@ -335,9 +337,9 @@ func (e *editTool) deleteContent(ctx context.Context, filePath, oldString string
 	}
 
 	// Check if file exists in history
-	file, err := e.files.GetByPathAndSession(ctx, filePath, sessionID)
+	file, err := e.history.GetLatestByPathAndSession(ctx, filePath, sessionID)
 	if err != nil {
-		_, err = e.files.Create(ctx, sessionID, filePath, oldContent)
+		_, err = e.history.Create(ctx, sessionID, filePath, oldContent)
 		if err != nil {
 			// Log error but don't fail the operation
 			return ToolResponse{}, fmt.Errorf("error creating file history: %w", err)
@@ -345,13 +347,13 @@ func (e *editTool) deleteContent(ctx context.Context, filePath, oldString string
 	}
 	if file.Content != oldContent {
 		// User Manually changed the content store an intermediate version
-		_, err = e.files.CreateVersion(ctx, sessionID, filePath, oldContent)
+		_, err = e.history.CreateVersion(ctx, sessionID, filePath, oldContent)
 		if err != nil {
 			slog.Debug("Error creating file history version", "error", err)
 		}
 	}
 	// Store the new version
-	_, err = e.files.CreateVersion(ctx, sessionID, filePath, "")
+	_, err = e.history.CreateVersion(ctx, sessionID, filePath, "")
 	if err != nil {
 		slog.Debug("Error creating file history version", "error", err)
 	}
@@ -433,6 +435,7 @@ func (e *editTool) replaceContent(ctx context.Context, filePath, oldString, newS
 		permissionPath = rootDir
 	}
 	p := e.permissions.Request(
+		ctx,
 		permission.CreatePermissionRequest{
 			SessionID:   sessionID,
 			Path:        permissionPath,
@@ -455,9 +458,9 @@ func (e *editTool) replaceContent(ctx context.Context, filePath, oldString, newS
 	}
 
 	// Check if file exists in history
-	file, err := e.files.GetByPathAndSession(ctx, filePath, sessionID)
+	file, err := e.history.GetLatestByPathAndSession(ctx, filePath, sessionID)
 	if err != nil {
-		_, err = e.files.Create(ctx, sessionID, filePath, oldContent)
+		_, err = e.history.Create(ctx, sessionID, filePath, oldContent)
 		if err != nil {
 			// Log error but don't fail the operation
 			return ToolResponse{}, fmt.Errorf("error creating file history: %w", err)
@@ -465,13 +468,13 @@ func (e *editTool) replaceContent(ctx context.Context, filePath, oldString, newS
 	}
 	if file.Content != oldContent {
 		// User Manually changed the content store an intermediate version
-		_, err = e.files.CreateVersion(ctx, sessionID, filePath, oldContent)
+		_, err = e.history.CreateVersion(ctx, sessionID, filePath, oldContent)
 		if err != nil {
 			slog.Debug("Error creating file history version", "error", err)
 		}
 	}
 	// Store the new version
-	_, err = e.files.CreateVersion(ctx, sessionID, filePath, newContent)
+	_, err = e.history.CreateVersion(ctx, sessionID, filePath, newContent)
 	if err != nil {
 		slog.Debug("Error creating file history version", "error", err)
 	}
diff --git a/internal/llm/tools/fetch.go b/internal/llm/tools/fetch.go
index 863532a0b..d2ab148b9 100644
--- a/internal/llm/tools/fetch.go
+++ b/internal/llm/tools/fetch.go
@@ -122,6 +122,7 @@ func (t *fetchTool) Run(ctx context.Context, call ToolCall) (ToolResponse, error
 	}
 
 	p := t.permissions.Request(
+		ctx,
 		permission.CreatePermissionRequest{
 			SessionID:   sessionID,
 			Path:        config.WorkingDirectory(),
diff --git a/internal/llm/tools/patch.go b/internal/llm/tools/patch.go
index e0c0bf5bc..d70e6a451 100644
--- a/internal/llm/tools/patch.go
+++ b/internal/llm/tools/patch.go
@@ -193,6 +193,7 @@ func (p *patchTool) Run(ctx context.Context, call ToolCall) (ToolResponse, error
 			dir := filepath.Dir(path)
 			patchDiff, _, _ := diff.GenerateDiff("", *change.NewContent, path)
 			p := p.permissions.Request(
+				ctx,
 				permission.CreatePermissionRequest{
 					SessionID:   sessionID,
 					Path:        dir,
@@ -220,6 +221,7 @@ func (p *patchTool) Run(ctx context.Context, call ToolCall) (ToolResponse, error
 			patchDiff, _, _ := diff.GenerateDiff(currentContent, newContent, path)
 			dir := filepath.Dir(path)
 			p := p.permissions.Request(
+				ctx,
 				permission.CreatePermissionRequest{
 					SessionID:   sessionID,
 					Path:        dir,
@@ -239,6 +241,7 @@ func (p *patchTool) Run(ctx context.Context, call ToolCall) (ToolResponse, error
 			dir := filepath.Dir(path)
 			patchDiff, _, _ := diff.GenerateDiff(*change.OldContent, "", path)
 			p := p.permissions.Request(
+				ctx,
 				permission.CreatePermissionRequest{
 					SessionID:   sessionID,
 					Path:        dir,
@@ -313,7 +316,7 @@ func (p *patchTool) Run(ctx context.Context, call ToolCall) (ToolResponse, error
 		totalRemovals += removals
 
 		// Update history
-		file, err := p.files.GetByPathAndSession(ctx, absPath, sessionID)
+		file, err := p.files.GetLatestByPathAndSession(ctx, absPath, sessionID)
 		if err != nil && change.Type != diff.ActionAdd {
 			// If not adding a file, create history entry for existing file
 			_, err = p.files.Create(ctx, sessionID, absPath, oldContent)
diff --git a/internal/llm/tools/write.go b/internal/llm/tools/write.go
index 617d69c29..d826b3a49 100644
--- a/internal/llm/tools/write.go
+++ b/internal/llm/tools/write.go
@@ -167,6 +167,7 @@ func (w *writeTool) Run(ctx context.Context, call ToolCall) (ToolResponse, error
 		permissionPath = rootDir
 	}
 	p := w.permissions.Request(
+		ctx,
 		permission.CreatePermissionRequest{
 			SessionID:   sessionID,
 			Path:        permissionPath,
@@ -189,7 +190,7 @@ func (w *writeTool) Run(ctx context.Context, call ToolCall) (ToolResponse, error
 	}
 
 	// Check if file exists in history
-	file, err := w.files.GetByPathAndSession(ctx, filePath, sessionID)
+	file, err := w.files.GetLatestByPathAndSession(ctx, filePath, sessionID)
 	if err != nil {
 		_, err = w.files.Create(ctx, sessionID, filePath, oldContent)
 		if err != nil {