From 6100e8b3a81abea23c27a25e0330ccad6121f009 Mon Sep 17 00:00:00 2001 From: Ed Zynda Date: Sun, 7 Jun 2026 18:05:20 +0300 Subject: [PATCH] feat(ui): add /retry slash command for resubmitting last user message MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add PopLastUserMessage() on *App: walks the current tree branch back to the parent of the most recent user message, syncs the in-memory store, and returns the prompt + image parts for resubmission. - Register /retry (alias /rt) and wire handleRetryCommand which rebuilds the visible ScrollList from the truncated branch before resubmitting via Run/RunWithFiles. Mirrors SubmitMsg display path (badges, pending prints, stateWorking transition). - Recovers from transient provider errors (overloaded, timeout) without duplicating the user message in context — the failed turn's entries become orphaned off-branch rather than being re-sent to the LLM. - Update help text, AppController interface, and stub controller. - Add unit tests covering busy/closed/no-session guards, the happy-path truncation, and the empty-branch error case. --- internal/app/app.go | 85 ++++++++++++++++++ internal/app/app_test.go | 146 +++++++++++++++++++++++++++++++ internal/ui/commands/commands.go | 6 ++ internal/ui/model.go | 84 ++++++++++++++++++ internal/ui/model_test.go | 5 ++ 5 files changed, 326 insertions(+) diff --git a/internal/app/app.go b/internal/app/app.go index e1efb5d8..14759a00 100644 --- a/internal/app/app.go +++ b/internal/app/app.go @@ -13,6 +13,7 @@ import ( "charm.land/fantasy" "github.com/mark3labs/kit/internal/extensions" + "github.com/mark3labs/kit/internal/message" "github.com/mark3labs/kit/internal/session" kit "github.com/mark3labs/kit/pkg/kit" ) @@ -343,6 +344,90 @@ func (a *App) SwitchTreeSession(ts *session.TreeManager) { } } +// PopLastUserMessage truncates the tree session back to the parent of the +// most recent user message on the current branch, syncs the in-memory +// message store, and returns the user prompt text plus any image file +// parts so the caller can resubmit via Run/RunWithFiles. +// +// This is the building block for /retry: the user message and any orphaned +// assistant/tool entries produced by a failed turn become unreachable on +// the current branch (they remain in the session file under a different +// leaf) and are excluded from the next LLM context. +// +// Returns an error when: +// - the agent is currently working (busy) +// - the app has been closed +// - no tree session is active (sessions disabled via --no-session) +// - no user message exists on the current branch +// +// Satisfies ui.AppController. +func (a *App) PopLastUserMessage() (string, []kit.LLMFilePart, error) { + a.mu.Lock() + if a.closed { + a.mu.Unlock() + return "", nil, fmt.Errorf("app is closed") + } + if a.busy { + a.mu.Unlock() + return "", nil, fmt.Errorf("cannot retry while the agent is working") + } + a.mu.Unlock() + + ts := a.opts.TreeSession + if ts == nil { + return "", nil, fmt.Errorf("no tree session active; /retry requires a session") + } + + // Walk the current branch backwards to find the most recent user message. + branch := ts.GetBranch("") + var target *session.MessageEntry + for i := len(branch) - 1; i >= 0; i-- { + me, ok := branch[i].(*session.MessageEntry) + if !ok { + continue + } + if me.Role == string(message.RoleUser) { + target = me + break + } + } + if target == nil { + return "", nil, fmt.Errorf("no user message to retry") + } + + // Extract the prompt text and any image parts from the target entry. + msg, err := target.ToMessage() + if err != nil { + return "", nil, fmt.Errorf("decode user message: %w", err) + } + prompt := msg.Content() + var files []kit.LLMFilePart + for _, part := range msg.Parts { + if ic, ok := part.(message.ImageContent); ok { + files = append(files, kit.LLMFilePart{ + Data: ic.Data, + MediaType: ic.MediaType, + }) + } + } + + // Move the leaf to the parent of the user message. The failed turn's + // entries (user message + any partial assistant/tool entries) are still + // in the tree file but no longer on the active branch, so they will not + // be re-sent to the LLM. runTurn() will append a fresh user message on + // the next call. + if err := ts.Branch(target.ParentID); err != nil { + return "", nil, fmt.Errorf("branch to parent: %w", err) + } + + // Sync the in-memory store with the new branch position so subsequent + // reads (and ReloadMessagesFromTree() consumers) see the truncated view. + a.store.Clear() + a.store.Replace(ts.GetLLMMessages()) + + return prompt, files, nil +} + // AddContextMessage adds a user-role message to the conversation history // without triggering an LLM response. Used by the ! shell command prefix // to inject command output into context so the LLM can reference it in diff --git a/internal/app/app_test.go b/internal/app/app_test.go index 645dfe7b..bf214a8f 100644 --- a/internal/app/app_test.go +++ b/internal/app/app_test.go @@ -9,7 +9,10 @@ import ( "time" tea "charm.land/bubbletea/v2" + "charm.land/fantasy" kit "github.com/mark3labs/kit/pkg/kit" + + "github.com/mark3labs/kit/internal/session" ) // -------------------------------------------------------------------------- @@ -969,3 +972,146 @@ func TestReleaseBusyAfterCompact_dropsQueueWhenClosed(t *testing.T) { t.Fatalf("expected 0 PromptFunc calls on closed app, got %d", n) } } + +// -------------------------------------------------------------------------- +// PopLastUserMessage (/retry building block) +// -------------------------------------------------------------------------- + +// TestPopLastUserMessage_NoTreeSession verifies that PopLastUserMessage +// returns an error when no tree session is active. +func TestPopLastUserMessage_NoTreeSession(t *testing.T) { + app := newTestApp(newStub()) + defer app.Close() + + prompt, files, err := app.PopLastUserMessage() + if err == nil { + t.Fatal("expected error when no tree session is active") + } + if prompt != "" || files != nil { + t.Fatalf("expected zero values on error, got prompt=%q files=%v", prompt, files) + } +} + +// TestPopLastUserMessage_WhileBusy verifies that PopLastUserMessage +// refuses to truncate while the agent is busy (would race with executeBatch). +func TestPopLastUserMessage_WhileBusy(t *testing.T) { + app := newTestApp(newStub()) + defer app.Close() + + app.mu.Lock() + app.busy = true + app.mu.Unlock() + + _, _, err := app.PopLastUserMessage() + if err == nil { + t.Fatal("expected error when agent is busy") + } + if !strings.Contains(err.Error(), "working") { + t.Fatalf("expected error mentioning busy/working, got %q", err.Error()) + } +} + +// TestPopLastUserMessage_WhenClosed verifies that PopLastUserMessage +// returns an error after Close(). +func TestPopLastUserMessage_WhenClosed(t *testing.T) { + app := newTestApp(newStub()) + app.Close() + + _, _, err := app.PopLastUserMessage() + if err == nil { + t.Fatal("expected error on closed app") + } +} + +// TestPopLastUserMessage_TruncatesAndReturnsPrompt verifies the happy path: +// a real tree session with user→assistant→user→assistant entries is +// truncated back to before the most recent user message, and that user's +// text is returned. +func TestPopLastUserMessage_TruncatesAndReturnsPrompt(t *testing.T) { + dir := t.TempDir() + ts, err := session.CreateTreeSession(dir) + if err != nil { + t.Fatalf("create tree session: %v", err) + } + defer func() { _ = ts.Close() }() + + // Build history: user "first" → assistant "ack 1" → user "second" → assistant "ack 2". + if _, err := ts.AppendLLMMessage(fantasy.NewUserMessage("first")); err != nil { + t.Fatal(err) + } + if _, err := ts.AppendLLMMessage(fantasy.Message{ + Role: fantasy.MessageRoleAssistant, + Content: []fantasy.MessagePart{fantasy.TextPart{Text: "ack 1"}}, + }); err != nil { + t.Fatal(err) + } + if _, err := ts.AppendLLMMessage(fantasy.NewUserMessage("second")); err != nil { + t.Fatal(err) + } + if _, err := ts.AppendLLMMessage(fantasy.Message{ + Role: fantasy.MessageRoleAssistant, + Content: []fantasy.MessagePart{fantasy.TextPart{Text: "ack 2"}}, + }); err != nil { + t.Fatal(err) + } + + app := New(Options{TreeSession: ts, PromptFunc: newStub().fn}, nil) + defer app.Close() + + prompt, files, err := app.PopLastUserMessage() + if err != nil { + t.Fatalf("PopLastUserMessage: %v", err) + } + if prompt != "second" { + t.Fatalf("expected prompt=%q, got %q", "second", prompt) + } + if files != nil { + t.Fatalf("expected no files, got %v", files) + } + + // After truncation the branch should only contain the first user + // message and its assistant response (the "second" turn is orphaned). + msgs := ts.GetLLMMessages() + if len(msgs) != 2 { + t.Fatalf("expected 2 messages on truncated branch, got %d", len(msgs)) + } + if got := messageText(msgs[0]); got != "first" { + t.Fatalf("expected first message %q, got %q", "first", got) + } + if got := messageText(msgs[1]); got != "ack 1" { + t.Fatalf("expected second message %q, got %q", "ack 1", got) + } +} + +// messageText extracts concatenated TextPart content from a fantasy.Message. +func messageText(m fantasy.Message) string { + var out strings.Builder + for _, p := range m.Content { + if tp, ok := p.(fantasy.TextPart); ok { + out.WriteString(tp.Text) + } + } + return out.String() +} + +// TestPopLastUserMessage_NoUserOnBranch verifies that an empty tree (no +// user messages at all) returns a friendly error rather than panicking. +func TestPopLastUserMessage_NoUserOnBranch(t *testing.T) { + dir := t.TempDir() + ts, err := session.CreateTreeSession(dir) + if err != nil { + t.Fatalf("create tree session: %v", err) + } + defer func() { _ = ts.Close() }() + + app := New(Options{TreeSession: ts, PromptFunc: newStub().fn}, nil) + defer app.Close() + + _, _, err = app.PopLastUserMessage() + if err == nil { + t.Fatal("expected error when no user message exists on branch") + } + if !strings.Contains(err.Error(), "no user message") { + t.Fatalf("expected error mentioning missing user message, got %q", err.Error()) + } +} diff --git a/internal/ui/commands/commands.go b/internal/ui/commands/commands.go index 61f2d873..84febbaf 100644 --- a/internal/ui/commands/commands.go +++ b/internal/ui/commands/commands.go @@ -167,6 +167,12 @@ var SlashCommands = []SlashCommand{ Category: "System", Aliases: []string{"/cp"}, }, + { + Name: "/retry", + Description: "Resubmit the last user message (e.g. after a provider error)", + Category: "System", + Aliases: []string{"/rt"}, + }, { Name: "/edit", Description: "Open a file in $EDITOR (fuzzy-find a path, then edit)", diff --git a/internal/ui/model.go b/internal/ui/model.go index ce4d1f54..a9a870d4 100644 --- a/internal/ui/model.go +++ b/internal/ui/model.go @@ -126,6 +126,14 @@ type AppController interface { // attachments (e.g. pasted images) into the currently running agent // turn. Behaves like Steer but includes file parts alongside the text. SteerWithFiles(prompt string, files []kit.LLMFilePart) int + // PopLastUserMessage truncates the tree session at the parent of the + // most recent user message on the current branch, syncs the in-memory + // message store, and returns that user prompt (plus any image file + // parts) so the caller can resubmit it. Used by /retry to recover from + // provider errors (overloaded, timeout) without duplicating the user + // message in context. Returns an error if the agent is busy, no tree + // session is active, or no user message exists on the current branch. + PopLastUserMessage() (string, []kit.LLMFilePart, error) } // SkillItem holds display metadata about a loaded skill for the startup @@ -3288,6 +3296,8 @@ func (m *AppModel) handleSlashCommand(sc *commands.SlashCommand, args string) te return m.handleExportCommand(args) case "/copy": return m.handleCopyCommand() + case "/retry": + return m.handleRetryCommand() case "/edit": return m.handleEditCommand(args) case "/share": @@ -3715,6 +3725,7 @@ func (m *AppModel) printHelpMessage() { "- `/compact [instructions]`: Summarise older messages to free context space\n" + "- `/clear`: Clear message history\n" + "- `/copy`: Copy the last message to the system clipboard\n" + + "- `/retry`: Resubmit the last user message (e.g. after a provider error)\n" + "- `/edit [path]`: Open a file in `$EDITOR` (fuzzy-find from cwd)\n" + "- `/export [path]`: Export session as JSONL\n" + "- `/import `: Import session from JSONL file\n" + @@ -4566,6 +4577,79 @@ func (m *AppModel) handleCopyCommand() tea.Cmd { return clipboard.CopyToClipboard(text) } +// handleRetryCommand resubmits the most recent user message on the current +// branch. Used to recover from transient provider errors (overloaded, +// timeout) without users having to retype — and without the duplicate-user- +// message bloat that retyping creates. +// +// Flow: +// 1. App.PopLastUserMessage() truncates the tree at the parent of the last +// user message and returns its text + any image parts. The failed turn's +// entries become orphaned (still on disk, off-branch) so they will not +// be re-sent to the LLM. +// 2. The visible message list is rebuilt from the truncated branch so the +// prior user message + any partial assistant + error rendering vanish. +// 3. The prompt is resubmitted via Run/RunWithFiles, mirroring the normal +// SubmitMsg display path (badge formatting, pending-prints flush, +// stateWorking transition). +func (m *AppModel) handleRetryCommand() tea.Cmd { + if m.appCtrl == nil { + m.printSystemMessage("App controller unavailable.") + return nil + } + + prompt, files, err := m.appCtrl.PopLastUserMessage() + if err != nil { + m.printSystemMessage(fmt.Sprintf("Cannot retry: %v", err)) + return nil + } + + // Rebuild the visible ScrollList from the truncated branch so the failed + // turn's user message and any partial assistant/error rendering disappear + // before the resubmit prints a fresh user message. + m.messages = []MessageItem{} + m.renderSessionHistory() + + // Mirror SubmitMsg's badge formatting for the display text. + var imageCount, fileOnlyCount int + for _, f := range files { + if strings.HasPrefix(f.MediaType, "image/") { + imageCount++ + } else { + fileOnlyCount++ + } + } + displayText := prompt + if imageCount > 0 || fileOnlyCount > 0 { + var badges []string + if imageCount > 0 { + badges = append(badges, fmt.Sprintf("%d image(s) pasted", imageCount)) + } + if fileOnlyCount > 0 { + badges = append(badges, fmt.Sprintf("%d file(s) attached", fileOnlyCount)) + } + displayText = fmt.Sprintf("%s\n[%s]", prompt, strings.Join(badges, ", ")) + } + + var qLen int + if len(files) > 0 { + qLen = m.appCtrl.RunWithFiles(prompt, files) + } else { + qLen = m.appCtrl.Run(prompt) + } + if qLen > 0 { + m.queuedMessages = append(m.queuedMessages, displayText) + m.layoutDirty = true + } else { + m.pendingUserPrints = append(m.pendingUserPrints, displayText) + m.flushStreamAndPendingUserMessages() + } + if m.state != stateWorking { + m.state = stateWorking + } + return nil +} + // handleEditCommand opens the supplied path in $EDITOR via tea.ExecProcess, // pausing the TUI for the duration of the editor session. The path is // resolved relative to cwd; ~/ and absolute paths are honoured. Non-existent diff --git a/internal/ui/model_test.go b/internal/ui/model_test.go index dae272b6..9d810850 100644 --- a/internal/ui/model_test.go +++ b/internal/ui/model_test.go @@ -2,6 +2,7 @@ package ui import ( "errors" + "fmt" "strings" "testing" @@ -87,6 +88,10 @@ func (s *stubAppController) SteerWithFiles(prompt string, _ []kit.LLMFilePart) i return s.queueLen } +func (s *stubAppController) PopLastUserMessage() (string, []kit.LLMFilePart, error) { + return "", nil, fmt.Errorf("no user message to retry") +} + // -------------------------------------------------------------------------- // Stub child components // --------------------------------------------------------------------------