basic undo feature (#1268)

Co-authored-by: adamdotdevin <[email protected]>
Co-authored-by: Jay V <[email protected]>
Co-authored-by: Aiden Cline <[email protected]>
Co-authored-by: Andrew Joslin <[email protected]>
Co-authored-by: GitHub Action <[email protected]>
Co-authored-by: Tobias Walle <[email protected]>

2 files changed, 275 insertions, 15 deletions

diff --git a/packages/tui/internal/components/chat/editor.go b/packages/tui/internal/components/chat/editor.go
index 4a5212786..4311a4ca7 100644
--- a/packages/tui/internal/components/chat/editor.go
+++ b/packages/tui/internal/components/chat/editor.go
@@ -21,6 +21,7 @@ import (
 	"github.com/sst/opencode/internal/commands"
 	"github.com/sst/opencode/internal/components/dialog"
 	"github.com/sst/opencode/internal/components/textarea"
+	"github.com/sst/opencode/internal/components/toast"
 	"github.com/sst/opencode/internal/styles"
 	"github.com/sst/opencode/internal/theme"
 	"github.com/sst/opencode/internal/util"
@@ -57,6 +58,7 @@ type editorComponent struct {
 	historyIndex           int    // -1 means current (not in history)
 	currentText            string // Store current text when navigating history
 	pasteCounter           int
+	reverted               bool
 }
 
 func (m *editorComponent) Init() tea.Cmd {
@@ -122,10 +124,34 @@ func (m *editorComponent) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 		}
 		// Maximize editor responsiveness for printable characters
 		if msg.Text != "" {
+			m.reverted = false
 			m.textarea, cmd = m.textarea.Update(msg)
 			cmds = append(cmds, cmd)
 			return m, tea.Batch(cmds...)
 		}
+	case app.MessageRevertedMsg:
+		if msg.Session.ID == m.app.Session.ID {
+			switch msg.Message.Info.(type) {
+			case opencode.UserMessage:
+				prompt, err := msg.Message.ToPrompt()
+				if err != nil {
+					return m, toast.NewErrorToast("Failed to revert message")
+				}
+				m.RestoreFromPrompt(*prompt)
+				m.textarea.MoveToEnd()
+				m.reverted = true
+				return m, nil
+			}
+		}
+	case app.SessionUnrevertedMsg:
+		if msg.Session.ID == m.app.Session.ID {
+			if m.reverted {
+				updated, cmd := m.Clear()
+				m = updated.(*editorComponent)
+				return m, cmd
+			}
+			return m, nil
+		}
 	case tea.PasteMsg:
 		text := string(msg)
 
@@ -646,21 +672,14 @@ func NewEditorComponent(app *app.App) EditorComponent {
 	return m
 }
 
-// RestoreFromHistory restores a message from history at the given index
-func (m *editorComponent) RestoreFromHistory(index int) {
-	if index < 0 || index >= len(m.app.State.MessageHistory) {
-		return
-	}
-
-	entry := m.app.State.MessageHistory[index]
-
+func (m *editorComponent) RestoreFromPrompt(prompt app.Prompt) {
 	m.textarea.Reset()
-	m.textarea.SetValue(entry.Text)
+	m.textarea.SetValue(prompt.Text)
 
 	// Sort attachments by start index in reverse order (process from end to beginning)
 	// This prevents index shifting issues
-	attachmentsCopy := make([]*attachment.Attachment, len(entry.Attachments))
-	copy(attachmentsCopy, entry.Attachments)
+	attachmentsCopy := make([]*attachment.Attachment, len(prompt.Attachments))
+	copy(attachmentsCopy, prompt.Attachments)
 
 	for i := 0; i < len(attachmentsCopy)-1; i++ {
 		for j := i + 1; j < len(attachmentsCopy); j++ {
@@ -677,6 +696,15 @@ func (m *editorComponent) RestoreFromHistory(index int) {
 	}
 }
 
+// RestoreFromHistory restores a message from history at the given index
+func (m *editorComponent) RestoreFromHistory(index int) {
+	if index < 0 || index >= len(m.app.State.MessageHistory) {
+		return
+	}
+	entry := m.app.State.MessageHistory[index]
+	m.RestoreFromPrompt(entry)
+}
+
 func getMediaTypeFromExtension(ext string) string {
 	switch strings.ToLower(ext) {
 	case ".jpg":
diff --git a/packages/tui/internal/components/chat/messages.go b/packages/tui/internal/components/chat/messages.go
index cbea349ca..e498de9f0 100644
--- a/packages/tui/internal/components/chat/messages.go
+++ b/packages/tui/internal/components/chat/messages.go
@@ -1,6 +1,7 @@
 package chat
 
 import (
+	"context"
 	"fmt"
 	"log/slog"
 	"slices"
@@ -11,6 +12,7 @@ import (
 	"github.com/charmbracelet/x/ansi"
 	"github.com/sst/opencode-sdk-go"
 	"github.com/sst/opencode/internal/app"
+	"github.com/sst/opencode/internal/commands"
 	"github.com/sst/opencode/internal/components/dialog"
 	"github.com/sst/opencode/internal/components/toast"
 	"github.com/sst/opencode/internal/layout"
@@ -31,6 +33,8 @@ type MessagesComponent interface {
 	GotoTop() (tea.Model, tea.Cmd)
 	GotoBottom() (tea.Model, tea.Cmd)
 	CopyLastMessage() (tea.Model, tea.Cmd)
+	UndoLastMessage() (tea.Model, tea.Cmd)
+	RedoLastMessage() (tea.Model, tea.Cmd)
 }
 
 type messagesComponent struct {
@@ -161,10 +165,22 @@ func (m *messagesComponent) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 		m.tail = true
 		m.loading = true
 		return m, m.renderView()
+	case app.SessionUnrevertedMsg:
+		if msg.Session.ID == m.app.Session.ID {
+			m.cache.Clear()
+			m.tail = true
+			return m, m.renderView()
+		}
+	case app.MessageRevertedMsg:
+		if msg.Session.ID == m.app.Session.ID {
+			m.cache.Clear()
+			m.tail = true
+			return m, m.renderView()
+		}
 
 	case opencode.EventListResponseEventSessionUpdated:
 		if msg.Properties.Info.ID == m.app.Session.ID {
-			m.header = m.renderHeader()
+			cmds = append(cmds, m.renderView())
 		}
 	case opencode.EventListResponseEventMessageUpdated:
 		if msg.Properties.Info.SessionID == m.app.Session.ID {
@@ -205,7 +221,6 @@ type renderCompleteMsg struct {
 }
 
 func (m *messagesComponent) renderView() tea.Cmd {
-
 	if m.rendering {
 		slog.Debug("pending render, skipping")
 		m.dirty = true
@@ -233,6 +248,9 @@ func (m *messagesComponent) renderView() tea.Cmd {
 
 		width := m.width // always use full width
 
+		reverted := false
+		revertedMessageCount := 0
+		revertedToolCount := 0
 		lastAssistantMessage := "zzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzz"
 		for _, msg := range slices.Backward(m.app.Messages) {
 			if assistant, ok := msg.Info.(opencode.AssistantMessage); ok {
@@ -246,6 +264,17 @@ func (m *messagesComponent) renderView() tea.Cmd {
 
 			switch casted := message.Info.(type) {
 			case opencode.UserMessage:
+				if casted.ID == m.app.Session.Revert.MessageID {
+					reverted = true
+					revertedMessageCount = 1
+					revertedToolCount = 0
+					continue
+				}
+				if reverted {
+					revertedMessageCount++
+					continue
+				}
+
 				for partIndex, part := range message.Parts {
 					switch part := part.(type) {
 					case opencode.TextPart:
@@ -324,10 +353,18 @@ func (m *messagesComponent) renderView() tea.Cmd {
 				}
 
 			case opencode.AssistantMessage:
+				if casted.ID == m.app.Session.Revert.MessageID {
+					reverted = true
+					revertedMessageCount = 1
+					revertedToolCount = 0
+				}
 				hasTextPart := false
 				for partIndex, p := range message.Parts {
 					switch part := p.(type) {
 					case opencode.TextPart:
+						if reverted {
+							continue
+						}
 						hasTextPart = true
 						finished := part.Time.End > 0
 						remainingParts := message.Parts[partIndex+1:]
@@ -406,6 +443,10 @@ func (m *messagesComponent) renderView() tea.Cmd {
 							blocks = append(blocks, content)
 						}
 					case opencode.ToolPart:
+						if reverted {
+							revertedToolCount++
+							continue
+						}
 						if !m.showToolDetails {
 							if !hasTextPart {
 								orphanedToolCalls = append(orphanedToolCalls, part)
@@ -472,7 +513,7 @@ func (m *messagesComponent) renderView() tea.Cmd {
 				}
 			}
 
-			if error != "" {
+			if error != "" && !reverted {
 				error = styles.NewStyle().Width(width - 6).Render(error)
 				error = renderContentBlock(
 					m.app,
@@ -491,6 +532,44 @@ func (m *messagesComponent) renderView() tea.Cmd {
 			}
 		}
 
+		if revertedMessageCount > 0 || revertedToolCount > 0 {
+			messagePlural := ""
+			toolPlural := ""
+			if revertedMessageCount != 1 {
+				messagePlural = "s"
+			}
+			if revertedToolCount != 1 {
+				toolPlural = "s"
+			}
+			revertedStyle := styles.NewStyle().
+				Background(t.BackgroundPanel()).
+				Foreground(t.TextMuted())
+
+			content := revertedStyle.Render(fmt.Sprintf(
+				"%d message%s reverted, %d tool call%s reverted",
+				revertedMessageCount,
+				messagePlural,
+				revertedToolCount,
+				toolPlural,
+			))
+			hintStyle := styles.NewStyle().Background(t.BackgroundPanel()).Foreground(t.Text())
+			hint := hintStyle.Render(m.app.Keybind(commands.MessagesRedoCommand))
+			hint += revertedStyle.Render(" (or /redo) to restore")
+
+			content += "\n" + hint
+			content = styles.NewStyle().
+				Background(t.BackgroundPanel()).
+				Width(width - 6).
+				Render(content)
+			content = renderContentBlock(
+				m.app,
+				content,
+				width,
+				WithBorderColor(t.BackgroundPanel()),
+			)
+			blocks = append(blocks, content)
+		}
+
 		final := []string{}
 		clipboard := []string{}
 		var selection *selection
@@ -522,7 +601,11 @@ func (m *messagesComponent) renderView() tea.Cmd {
 					middle := strings.TrimRight(ansi.Strip(ansi.Cut(line, left, right)), " ")
 					suffix := ansi.Cut(line, left+ansi.StringWidth(middle), width)
 					clipboard = append(clipboard, middle)
-					line = prefix + styles.NewStyle().Background(t.Accent()).Foreground(t.BackgroundPanel()).Render(middle) + suffix
+					line = prefix + styles.NewStyle().
+						Background(t.Accent()).
+						Foreground(t.BackgroundPanel()).
+						Render(ansi.Strip(middle)) +
+						suffix
 				}
 				final = append(final, line)
 			}
@@ -773,6 +856,155 @@ func (m *messagesComponent) CopyLastMessage() (tea.Model, tea.Cmd) {
 	return m, tea.Batch(cmds...)
 }
 
+func (m *messagesComponent) UndoLastMessage() (tea.Model, tea.Cmd) {
+	after := float64(0)
+	var revertedMessage app.Message
+	reversedMessages := []app.Message{}
+	for i := len(m.app.Messages) - 1; i >= 0; i-- {
+		reversedMessages = append(reversedMessages, m.app.Messages[i])
+		switch casted := m.app.Messages[i].Info.(type) {
+		case opencode.UserMessage:
+			if casted.ID == m.app.Session.Revert.MessageID {
+				after = casted.Time.Created
+			}
+		case opencode.AssistantMessage:
+			if casted.ID == m.app.Session.Revert.MessageID {
+				after = casted.Time.Created
+			}
+		}
+		if m.app.Session.Revert.PartID != "" {
+			for _, part := range m.app.Messages[i].Parts {
+				switch casted := part.(type) {
+				case opencode.TextPart:
+					if casted.ID == m.app.Session.Revert.PartID {
+						after = casted.Time.Start
+					}
+				case opencode.ToolPart:
+					// TODO: handle tool parts
+				}
+			}
+		}
+	}
+
+	messageID := ""
+	for _, msg := range reversedMessages {
+		switch casted := msg.Info.(type) {
+		case opencode.UserMessage:
+			if after > 0 && casted.Time.Created >= after {
+				continue
+			}
+			messageID = casted.ID
+			revertedMessage = msg
+		}
+		if messageID != "" {
+			break
+		}
+	}
+
+	if messageID == "" {
+		return m, nil
+	}
+
+	return m, func() tea.Msg {
+		response, err := m.app.Client.Session.Revert(
+			context.Background(),
+			m.app.Session.ID,
+			opencode.SessionRevertParams{
+				MessageID: opencode.F(messageID),
+			},
+		)
+		if err != nil {
+			slog.Error("Failed to undo message", "error", err)
+			return toast.NewErrorToast("Failed to undo message")
+		}
+		if response == nil {
+			return toast.NewErrorToast("Failed to undo message")
+		}
+		return app.MessageRevertedMsg{Session: *response, Message: revertedMessage}
+	}
+}
+
+func (m *messagesComponent) RedoLastMessage() (tea.Model, tea.Cmd) {
+	before := float64(0)
+	var revertedMessage app.Message
+	for _, message := range m.app.Messages {
+		switch casted := message.Info.(type) {
+		case opencode.UserMessage:
+			if casted.ID == m.app.Session.Revert.MessageID {
+				before = casted.Time.Created
+			}
+		case opencode.AssistantMessage:
+			if casted.ID == m.app.Session.Revert.MessageID {
+				before = casted.Time.Created
+			}
+		}
+		if m.app.Session.Revert.PartID != "" {
+			for _, part := range message.Parts {
+				switch casted := part.(type) {
+				case opencode.TextPart:
+					if casted.ID == m.app.Session.Revert.PartID {
+						before = casted.Time.Start
+					}
+				case opencode.ToolPart:
+					// TODO: handle tool parts
+				}
+			}
+		}
+	}
+
+	messageID := ""
+	for _, msg := range m.app.Messages {
+		switch casted := msg.Info.(type) {
+		case opencode.UserMessage:
+			if casted.Time.Created <= before {
+				continue
+			}
+			messageID = casted.ID
+			revertedMessage = msg
+		}
+		if messageID != "" {
+			break
+		}
+	}
+
+	if messageID == "" {
+		return m, func() tea.Msg {
+			// unrevert back to original state
+			response, err := m.app.Client.Session.Unrevert(
+				context.Background(),
+				m.app.Session.ID,
+			)
+			if err != nil {
+				slog.Error("Failed to unrevert session", "error", err)
+				return toast.NewErrorToast("Failed to redo message")
+			}
+			if response == nil {
+				return toast.NewErrorToast("Failed to redo message")
+			}
+			return app.SessionUnrevertedMsg{Session: *response}
+		}
+	}
+
+	return m, func() tea.Msg {
+		// calling revert on a "later" message is like a redo
+		response, err := m.app.Client.Session.Revert(
+			context.Background(),
+			m.app.Session.ID,
+			opencode.SessionRevertParams{
+				MessageID: opencode.F(messageID),
+			},
+		)
+		if err != nil {
+			slog.Error("Failed to redo message", "error", err)
+			return toast.NewErrorToast("Failed to redo message")
+		}
+		if response == nil {
+			return toast.NewErrorToast("Failed to redo message")
+		}
+		return app.MessageRevertedMsg{Session: *response, Message: revertedMessage}
+	}
+}
+
 func NewMessagesComponent(app *app.App) MessagesComponent {
 	vp := viewport.New()
 	vp.KeyMap = viewport.KeyMap{}