feat: delete messages below

Carlos Alexandro Becker created 1 month ago

Signed-off-by: Carlos Alexandro Becker <[e-mail]>

Change summary

internal/message/message.go           | 24 ++++++++++++++++
internal/session/session.go           |  2 
internal/session/session_fork_test.go |  2 
internal/ui/dialog/actions.go         |  1 
internal/ui/dialog/commands.go        | 10 ++++++
internal/ui/model/ui.go               | 42 ++++++++++++++++++++++++++++
6 files changed, 77 insertions(+), 4 deletions(-)

Detailed changes

internal/message/message.go

@@ -31,6 +31,7 @@ type Service interface {
 	Copy(ctx context.Context, sessionID string, message Message) (Message, error)
 	Delete(ctx context.Context, id string) error
 	DeleteSessionMessages(ctx context.Context, sessionID string) error
+	DeleteMessagesFrom(ctx context.Context, sessionID, messageID string) error
 }
 
 type service struct {
@@ -123,6 +124,29 @@ func (s *service) DeleteSessionMessages(ctx context.Context, sessionID string) e
 	return nil
 }
 
+func (s *service) DeleteMessagesFrom(ctx context.Context, sessionID, messageID string) error {
+	allMessages, err := s.List(ctx, sessionID)
+	if err != nil {
+		return err
+	}
+	targetIndex := -1
+	for i, msg := range allMessages {
+		if msg.ID == messageID {
+			targetIndex = i
+			break
+		}
+	}
+	if targetIndex == -1 {
+		return fmt.Errorf("message not found: %s", messageID)
+	}
+	for i := len(allMessages) - 1; i >= targetIndex; i-- {
+		if err := s.Delete(ctx, allMessages[i].ID); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 func (s *service) Update(ctx context.Context, message Message) error {
 	parts, err := marshalParts(message.Parts)
 	if err != nil {

internal/session/session.go

@@ -173,7 +173,7 @@ func (s *service) Fork(ctx context.Context, sourceSessionID, upToMessageID strin
 		return Session{}, fmt.Errorf("creating session: %w", err)
 	}
 
-	for i := 0; i <= targetIndex; i++ {
+	for i := 0; i < targetIndex; i++ {
 		_, err = messageSvc.Copy(ctx, newSession.ID, messages[i])
 		if err != nil {
 			_ = s.Delete(ctx, newSession.ID)

internal/session/session_fork_test.go

@@ -53,7 +53,7 @@ func TestFork(t *testing.T) {
 	require.Contains(t, newSession.Title, sourceSession.Title)
 
 	forkedMessages := getMessages(newSession.ID)
-	require.Len(t, forkedMessages, 3)
+	require.Len(t, forkedMessages, 2)
 
 	for i, msg := range forkedMessages {
 		require.Equal(t, sourceMessages[i].Role, msg.Role)

internal/ui/dialog/actions.go

@@ -45,6 +45,7 @@ type ActionSelectModel struct {
 type (
 	ActionNewSession        struct{}
 	ActionForkConversation  struct{}
+	ActionDeleteMessages    struct{}
 	ActionToggleHelp        struct{}
 	ActionToggleCompactMode struct{}
 	ActionToggleThinking    struct{}

internal/ui/dialog/commands.go

@@ -66,12 +66,13 @@ type Commands struct {
 	customCommands []commands.CustomCommand
 	mcpPrompts     []commands.MCPPrompt
 	canFork        bool // whether a user message is selected for forking
+	canDelete      bool // whether a user message is selected for deletion
 }
 
 var _ Dialog = (*Commands)(nil)
 
 // NewCommands creates a new commands dialog.
-func NewCommands(com *common.Common, sessionID string, customCommands []commands.CustomCommand, mcpPrompts []commands.MCPPrompt, canFork bool) (*Commands, error) {
+func NewCommands(com *common.Common, sessionID string, customCommands []commands.CustomCommand, mcpPrompts []commands.MCPPrompt, canFork, canDelete bool) (*Commands, error) {
 	c := &Commands{
 		com:            com,
 		selected:       SystemCommands,
@@ -79,6 +80,7 @@ func NewCommands(com *common.Common, sessionID string, customCommands []commands
 		customCommands: customCommands,
 		mcpPrompts:     mcpPrompts,
 		canFork:        canFork,
+		canDelete:      canDelete,
 	}
 
 	help := help.New()
@@ -396,6 +398,12 @@ func (c *Commands) defaultCommands() []*CommandItem {
 		)
 	}
 
+	if c.canDelete {
+		commands = append(commands,
+			NewCommandItem(c.com.Styles, "delete_messages", "Delete Messages (including this)", "", ActionDeleteMessages{}),
+		)
+	}
+
 	commands = append(commands,
 		NewCommandItem(c.com.Styles, "switch_session", "Sessions", "ctrl+s", ActionOpenDialog{SessionsID}),
 		NewCommandItem(c.com.Styles, "switch_model", "Switch Model", "ctrl+l", ActionOpenDialog{ModelsID}),

internal/ui/model/ui.go

@@ -1193,6 +1193,44 @@ func (m *UI) handleDialogMsg(msg tea.Msg) tea.Cmd {
 			}
 		})
 		m.dialog.CloseDialog(dialog.CommandsID)
+		m.chat.list.Blur()
+		m.focus = uiFocusEditor
+		cmds = append(cmds, m.textarea.Focus())
+
+	case dialog.ActionDeleteMessages:
+		if m.isAgentBusy() {
+			cmds = append(cmds, uiutil.ReportWarn("Agent is busy, please wait..."))
+			break
+		}
+		if m.session == nil {
+			cmds = append(cmds, uiutil.ReportWarn("No session to delete from..."))
+			break
+		}
+		selectedItem := m.chat.SelectedItem()
+		if selectedItem == nil {
+			cmds = append(cmds, uiutil.ReportWarn("No message selected..."))
+			break
+		}
+		if _, ok := selectedItem.(*chat.UserMessageItem); !ok {
+			cmds = append(cmds, uiutil.ReportWarn("Can only delete from user messages..."))
+			break
+		}
+		messageID, ok := m.getMessageIDFromItem(selectedItem)
+		if !ok {
+			cmds = append(cmds, uiutil.ReportWarn("Cannot get message ID from selected item..."))
+			break
+		}
+		cmds = append(cmds, func() tea.Msg {
+			err := m.com.App.Messages.DeleteMessagesFrom(context.Background(), m.session.ID, messageID)
+			if err != nil {
+				return uiutil.ReportError(err)()
+			}
+			return m.loadSession(m.session.ID)()
+		})
+		m.dialog.CloseDialog(dialog.CommandsID)
+		m.chat.list.Blur()
+		m.focus = uiFocusEditor
+		cmds = append(cmds, m.textarea.Focus())
 
 	case dialog.ActionSelectModel:
 		if m.isAgentBusy() {
@@ -2752,10 +2790,12 @@ func (m *UI) openModelsDialog() tea.Cmd {
 // openCommandsDialog opens the commands dialog.
 func (m *UI) openCommandsDialog() tea.Cmd {
 	canFork := false
+	canDelete := false
 	selectedItem := m.chat.SelectedItem()
 	if selectedItem != nil {
 		if _, ok := selectedItem.(*chat.UserMessageItem); ok {
 			canFork = true
+			canDelete = true
 		}
 	}
 
@@ -2769,7 +2809,7 @@ func (m *UI) openCommandsDialog() tea.Cmd {
 		sessionID = m.session.ID
 	}
 
-	commands, err := dialog.NewCommands(m.com, sessionID, m.customCommands, m.mcpPrompts, canFork)
+	commands, err := dialog.NewCommands(m.com, sessionID, m.customCommands, m.mcpPrompts, canFork, canDelete)
 	if err != nil {
 		return uiutil.ReportError(err)
 	}