From b744b0f444aa94454f1c8bee6c4f3a03255b86b9 Mon Sep 17 00:00:00 2001 From: Kujtim Hoxha Date: Mon, 9 Jun 2025 11:49:54 +0200 Subject: [PATCH] refactor: new compact dialog --- internal/llm/agent/agent.go | 26 +- internal/llm/provider/anthropic.go | 3 +- .../components/dialogs/commands/commands.go | 45 +-- .../tui/components/dialogs/compact/compact.go | 266 ++++++++++++++++++ .../tui/components/dialogs/compact/keys.go | 61 ++++ internal/tui/components/dialogs/init/init.go | 42 ++- internal/tui/tui.go | 62 ++-- 7 files changed, 443 insertions(+), 62 deletions(-) create mode 100644 internal/tui/components/dialogs/compact/compact.go create mode 100644 internal/tui/components/dialogs/compact/keys.go diff --git a/internal/llm/agent/agent.go b/internal/llm/agent/agent.go index 9120c76aff8d5efa7161b4fab73577d31991e07a..c19451e1d0ef46597b8e3f9d56f9e0ebdf4362cb 100644 --- a/internal/llm/agent/agent.go +++ b/internal/llm/agent/agent.go @@ -586,22 +586,26 @@ func (a *agent) Summarize(ctx context.Context, sessionID string) error { a.Publish(pubsub.CreatedEvent, event) // Send the messages to the summarize provider - response, err := a.summarizeProvider.SendMessages( + response := a.summarizeProvider.StreamResponse( summarizeCtx, msgsWithPrompt, make([]tools.BaseTool, 0), ) - if err != nil { - event = AgentEvent{ - Type: AgentEventTypeError, - Error: fmt.Errorf("failed to summarize: %w", err), - Done: true, + var finalResponse *provider.ProviderResponse + for r := range response { + if r.Error != nil { + event = AgentEvent{ + Type: AgentEventTypeError, + Error: fmt.Errorf("failed to summarize: %w", err), + Done: true, + } + a.Publish(pubsub.CreatedEvent, event) + return } - a.Publish(pubsub.CreatedEvent, event) - return + finalResponse = r.Response } - summary := strings.TrimSpace(response.Content) + summary := strings.TrimSpace(finalResponse.Content) if summary == "" { event = AgentEvent{ Type: AgentEventTypeError, @@ -651,10 +655,10 @@ func (a *agent) Summarize(ctx context.Context, sessionID string) error { return } oldSession.SummaryMessageID = msg.ID - oldSession.CompletionTokens = response.Usage.OutputTokens + oldSession.CompletionTokens = finalResponse.Usage.OutputTokens oldSession.PromptTokens = 0 model := a.summarizeProvider.Model() - usage := response.Usage + usage := finalResponse.Usage cost := model.CostPer1MInCached/1e6*float64(usage.CacheCreationTokens) + model.CostPer1MOutCached/1e6*float64(usage.CacheReadTokens) + model.CostPer1MIn/1e6*float64(usage.InputTokens) + diff --git a/internal/llm/provider/anthropic.go b/internal/llm/provider/anthropic.go index 77edc8e0519e6f82b0c807626dfebbcd5c09d3a4..f5f627c228f5708307980efdcaf9e35a8a9f48c8 100644 --- a/internal/llm/provider/anthropic.go +++ b/internal/llm/provider/anthropic.go @@ -195,7 +195,7 @@ func (a *anthropicClient) preparedMessages(messages []anthropic.MessageParam, to } } -func (a *anthropicClient) send(ctx context.Context, messages []message.Message, tools []tools.BaseTool) (resposne *ProviderResponse, err error) { +func (a *anthropicClient) send(ctx context.Context, messages []message.Message, tools []tools.BaseTool) (response *ProviderResponse, err error) { preparedMessages := a.preparedMessages(a.convertMessages(messages), a.convertTools(tools)) cfg := config.Get() if cfg.Debug { @@ -339,6 +339,7 @@ func (a *anthropicClient) stream(ctx context.Context, messages []message.Message Usage: a.usage(accumulatedMessage), FinishReason: a.finishReason(string(accumulatedMessage.StopReason)), }, + Content: content, } } } diff --git a/internal/tui/components/dialogs/commands/commands.go b/internal/tui/components/dialogs/commands/commands.go index 823ad2ab72d84ac89e8b10ee686ae20dd8ad17d3..b140fc1246d36e806836359f5b17030e5b383a1b 100644 --- a/internal/tui/components/dialogs/commands/commands.go +++ b/internal/tui/components/dialogs/commands/commands.go @@ -50,14 +50,18 @@ type commandDialogCmp struct { help help.Model commandType int // SystemCommands or UserCommands userCommands []Command // User-defined commands + sessionID string // Current session ID } type ( SwitchSessionsMsg struct{} SwitchModelMsg struct{} + CompactMsg struct { + SessionID string + } ) -func NewCommandDialog() CommandsDialog { +func NewCommandDialog(sessionID string) CommandsDialog { listKeyMap := list.DefaultKeyMap() keyMap := DefaultCommandsDialogKeyMap() @@ -87,6 +91,7 @@ func NewCommandDialog() CommandsDialog { keyMap: DefaultCommandsDialogKeyMap(), help: help, commandType: SystemCommands, + sessionID: sessionID, } } @@ -222,7 +227,7 @@ func (c *commandDialogCmp) Position() (int, int) { } func (c *commandDialogCmp) defaultCommands() []Command { - return []Command{ + commands := []Command{ { ID: "init", Title: "Initialize Project", @@ -235,33 +240,35 @@ func (c *commandDialogCmp) defaultCommands() []Command { The file you create will be given to agentic coding agents (such as yourself) that operate in this repository. Make it about 20 lines long. If there's already a crush.md, improve it. If there are Cursor rules (in .cursor/rules/ or .cursorrules) or Copilot rules (in .github/copilot-instructions.md), make sure to include them.` - return tea.Batch( - util.CmdHandler(chat.SendMsg{ - Text: prompt, - }), - ) + return util.CmdHandler(chat.SendMsg{ + Text: prompt, + }) }, }, - { + } + + // Only show compact command if there's an active session + if c.sessionID != "" { + commands = append(commands, Command{ ID: "compact", Title: "Compact Session", Description: "Summarize the current session and create a new one with the summary", Handler: func(cmd Command) tea.Cmd { - return func() tea.Msg { - // TODO: implement compact message - return "" - } + return util.CmdHandler(CompactMsg{ + SessionID: c.sessionID, + }) }, - }, + }) + } + + return append(commands, []Command{ { ID: "switch_session", Title: "Switch Session", Description: "Switch to a different session", Shortcut: "ctrl+s", Handler: func(cmd Command) tea.Cmd { - return func() tea.Msg { - return SwitchSessionsMsg{} - } + return util.CmdHandler(SwitchSessionsMsg{}) }, }, { @@ -269,12 +276,10 @@ func (c *commandDialogCmp) defaultCommands() []Command { Title: "Switch Model", Description: "Switch to a different model", Handler: func(cmd Command) tea.Cmd { - return func() tea.Msg { - return SwitchModelMsg{} - } + return util.CmdHandler(SwitchModelMsg{}) }, }, - } + }...) } func (c *commandDialogCmp) ID() dialogs.DialogID { diff --git a/internal/tui/components/dialogs/compact/compact.go b/internal/tui/components/dialogs/compact/compact.go new file mode 100644 index 0000000000000000000000000000000000000000..afa7c8945009fb7b76d979e466cae290757f3f27 --- /dev/null +++ b/internal/tui/components/dialogs/compact/compact.go @@ -0,0 +1,266 @@ +package compact + +import ( + "context" + + "github.com/charmbracelet/bubbles/v2/key" + tea "github.com/charmbracelet/bubbletea/v2" + "github.com/charmbracelet/lipgloss/v2" + + "github.com/charmbracelet/crush/internal/llm/agent" + "github.com/charmbracelet/crush/internal/tui/components/core" + "github.com/charmbracelet/crush/internal/tui/components/dialogs" + "github.com/charmbracelet/crush/internal/tui/styles" + "github.com/charmbracelet/crush/internal/tui/util" +) + +const CompactDialogID dialogs.DialogID = "compact" + +// CompactDialog interface for the session compact dialog +type CompactDialog interface { + dialogs.DialogModel +} + +type compactDialogCmp struct { + wWidth, wHeight int + width, height int + selected int + keyMap KeyMap + sessionID string + state compactState + progress string + agent agent.Service + noAsk bool // If true, skip confirmation dialog +} + +type compactState int + +const ( + stateConfirm compactState = iota + stateCompacting + stateError +) + +// NewCompactDialogCmp creates a new session compact dialog +func NewCompactDialogCmp(agent agent.Service, sessionID string, noAsk bool) CompactDialog { + return &compactDialogCmp{ + sessionID: sessionID, + keyMap: DefaultKeyMap(), + state: stateConfirm, + selected: 0, + agent: agent, + noAsk: noAsk, + } +} + +func (c *compactDialogCmp) Init() tea.Cmd { + if c.noAsk { + // If noAsk is true, skip confirmation and start compaction immediately + return c.startCompaction() + } + return nil +} + +func (c *compactDialogCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) { + switch msg := msg.(type) { + case tea.WindowSizeMsg: + c.wWidth = msg.Width + c.wHeight = msg.Height + cmd := c.SetSize() + return c, cmd + + case tea.KeyPressMsg: + switch c.state { + case stateConfirm: + switch { + case key.Matches(msg, c.keyMap.ChangeSelection): + c.selected = (c.selected + 1) % 2 + return c, nil + case key.Matches(msg, c.keyMap.Select): + if c.selected == 0 { + return c, c.startCompaction() + } else { + return c, util.CmdHandler(dialogs.CloseDialogMsg{}) + } + case key.Matches(msg, c.keyMap.Y): + return c, c.startCompaction() + case key.Matches(msg, c.keyMap.N): + return c, util.CmdHandler(dialogs.CloseDialogMsg{}) + case key.Matches(msg, c.keyMap.Close): + return c, util.CmdHandler(dialogs.CloseDialogMsg{}) + } + case stateCompacting: + switch { + case key.Matches(msg, c.keyMap.Close): + return c, util.CmdHandler(dialogs.CloseDialogMsg{}) + } + case stateError: + switch { + case key.Matches(msg, c.keyMap.Select): + return c, util.CmdHandler(dialogs.CloseDialogMsg{}) + case key.Matches(msg, c.keyMap.Close): + return c, util.CmdHandler(dialogs.CloseDialogMsg{}) + } + } + + case agent.AgentEvent: + if msg.Type == agent.AgentEventTypeSummarize { + if msg.Error != nil { + c.state = stateError + c.progress = "Error: " + msg.Error.Error() + } else if msg.Done { + return c, util.CmdHandler( + dialogs.CloseDialogMsg{}, + ) + } else { + c.progress = msg.Progress + } + } + return c, nil + } + + return c, nil +} + +func (c *compactDialogCmp) startCompaction() tea.Cmd { + c.state = stateCompacting + c.progress = "Starting summarization..." + return func() tea.Msg { + err := c.agent.Summarize(context.Background(), c.sessionID) + if err != nil { + c.state = stateError + c.progress = "Error: " + err.Error() + } + return nil + } +} + +func (c *compactDialogCmp) renderButtons() string { + t := styles.CurrentTheme() + baseStyle := t.S().Base + + buttons := []core.ButtonOpts{ + { + Text: "Yes", + UnderlineIndex: 0, // "Y" + Selected: c.selected == 0, + }, + { + Text: "No", + UnderlineIndex: 0, // "N" + Selected: c.selected == 1, + }, + } + + content := core.SelectableButtons(buttons, " ") + + return baseStyle.AlignHorizontal(lipgloss.Right).Width(c.width - 4).Render(content) +} + +func (c *compactDialogCmp) renderContent() string { + t := styles.CurrentTheme() + baseStyle := t.S().Base + + switch c.state { + case stateConfirm: + explanation := t.S().Text. + Width(c.width - 4). + Render("This will summarize the current session and reset the context. The conversation history will be condensed into a summary to free up context space while preserving important information.") + + question := t.S().Text. + Width(c.width - 4). + Render("Do you want to continue?") + + return baseStyle.Render(lipgloss.JoinVertical( + lipgloss.Left, + explanation, + "", + question, + )) + case stateCompacting: + return baseStyle.Render(lipgloss.JoinVertical( + lipgloss.Left, + c.progress, + "", + "Please wait...", + )) + case stateError: + return baseStyle.Render(lipgloss.JoinVertical( + lipgloss.Left, + c.progress, + "", + "Press Enter to close", + )) + } + return "" +} + +func (c *compactDialogCmp) render() string { + t := styles.CurrentTheme() + baseStyle := t.S().Base + + var title string + switch c.state { + case stateConfirm: + title = "Compact Session" + case stateCompacting: + title = "Compacting Session" + case stateError: + title = "Compact Failed" + } + + titleView := core.Title(title, c.width-4) + content := c.renderContent() + + var dialogContent string + if c.state == stateConfirm { + buttons := c.renderButtons() + dialogContent = lipgloss.JoinVertical( + lipgloss.Top, + titleView, + "", + content, + "", + buttons, + "", + ) + } else { + dialogContent = lipgloss.JoinVertical( + lipgloss.Top, + titleView, + "", + content, + "", + ) + } + + return baseStyle. + Padding(0, 1). + Border(lipgloss.RoundedBorder()). + BorderForeground(t.BorderFocus). + Width(c.width). + Render(dialogContent) +} + +func (c *compactDialogCmp) View() tea.View { + return tea.NewView(c.render()) +} + +// SetSize sets the size of the component. +func (c *compactDialogCmp) SetSize() tea.Cmd { + c.width = min(90, c.wWidth) + c.height = min(15, c.wHeight) + return nil +} + +func (c *compactDialogCmp) Position() (int, int) { + row := (c.wHeight / 2) - (c.height / 2) + col := (c.wWidth / 2) - (c.width / 2) + return row, col +} + +// ID implements CompactDialog. +func (c *compactDialogCmp) ID() dialogs.DialogID { + return CompactDialogID +} + diff --git a/internal/tui/components/dialogs/compact/keys.go b/internal/tui/components/dialogs/compact/keys.go new file mode 100644 index 0000000000000000000000000000000000000000..0f176927a173ec44db9ec85a9f476723f0cb4b94 --- /dev/null +++ b/internal/tui/components/dialogs/compact/keys.go @@ -0,0 +1,61 @@ +package compact + +import ( + "github.com/charmbracelet/bubbles/v2/key" + "github.com/charmbracelet/crush/internal/tui/layout" +) + +// KeyMap defines the key bindings for the compact dialog. +type KeyMap struct { + ChangeSelection key.Binding + Select key.Binding + Y key.Binding + N key.Binding + Close key.Binding +} + +// DefaultKeyMap returns the default key bindings for the compact dialog. +func DefaultKeyMap() KeyMap { + return KeyMap{ + ChangeSelection: key.NewBinding( + key.WithKeys("tab", "left", "right", "h", "l"), + key.WithHelp("tab/←/→", "toggle selection"), + ), + Select: key.NewBinding( + key.WithKeys("enter"), + key.WithHelp("enter", "confirm"), + ), + Y: key.NewBinding( + key.WithKeys("y"), + key.WithHelp("y", "yes"), + ), + N: key.NewBinding( + key.WithKeys("n"), + key.WithHelp("n", "no"), + ), + Close: key.NewBinding( + key.WithKeys("esc"), + key.WithHelp("esc", "cancel"), + ), + } +} + +// FullHelp implements help.KeyMap. +func (k KeyMap) FullHelp() [][]key.Binding { + m := [][]key.Binding{} + slice := layout.KeyMapToSlice(k) + for i := 0; i < len(slice); i += 4 { + end := min(i+4, len(slice)) + m = append(m, slice[i:end]) + } + return m +} + +// ShortHelp implements help.KeyMap. +func (k KeyMap) ShortHelp() []key.Binding { + return []key.Binding{ + k.ChangeSelection, + k.Select, + k.Close, + } +} \ No newline at end of file diff --git a/internal/tui/components/dialogs/init/init.go b/internal/tui/components/dialogs/init/init.go index da792695e984454bc2439fd47fe940e655d843b2..ff4cbfb4d7b6933523cc873019758c9203ff8657 100644 --- a/internal/tui/components/dialogs/init/init.go +++ b/internal/tui/components/dialogs/init/init.go @@ -5,6 +5,8 @@ import ( tea "github.com/charmbracelet/bubbletea/v2" "github.com/charmbracelet/lipgloss/v2" + "github.com/charmbracelet/crush/internal/config" + cmpChat "github.com/charmbracelet/crush/internal/tui/components/chat" "github.com/charmbracelet/crush/internal/tui/components/core" "github.com/charmbracelet/crush/internal/tui/components/dialogs" "github.com/charmbracelet/crush/internal/tui/styles" @@ -51,7 +53,7 @@ func (m *initDialogCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) { case key.Matches(msg, m.keyMap.Close): return m, tea.Batch( util.CmdHandler(dialogs.CloseDialogMsg{}), - util.CmdHandler(CloseInitDialogMsg{Initialize: false}), + m.handleInitialization(false), ) case key.Matches(msg, m.keyMap.ChangeSelection): m.selected = (m.selected + 1) % 2 @@ -59,17 +61,17 @@ func (m *initDialogCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) { case key.Matches(msg, m.keyMap.Select): return m, tea.Batch( util.CmdHandler(dialogs.CloseDialogMsg{}), - util.CmdHandler(CloseInitDialogMsg{Initialize: m.selected == 0}), + m.handleInitialization(m.selected == 0), ) case key.Matches(msg, m.keyMap.Y): return m, tea.Batch( util.CmdHandler(dialogs.CloseDialogMsg{}), - util.CmdHandler(CloseInitDialogMsg{Initialize: true}), + m.handleInitialization(true), ) case key.Matches(msg, m.keyMap.N): return m, tea.Batch( util.CmdHandler(dialogs.CloseDialogMsg{}), - util.CmdHandler(CloseInitDialogMsg{Initialize: false}), + m.handleInitialization(false), ) } } @@ -168,6 +170,38 @@ func (m *initDialogCmp) Position() (int, int) { return row, col } +// handleInitialization handles the initialization logic when the dialog is closed. +func (m *initDialogCmp) handleInitialization(initialize bool) tea.Cmd { + if initialize { + // Run the initialization command + prompt := `Please analyze this codebase and create a Crush.md file containing: +1. Build/lint/test commands - especially for running a single test +2. Code style guidelines including imports, formatting, types, naming conventions, error handling, etc. + +The file you create will be given to agentic coding agents (such as yourself) that operate in this repository. Make it about 20 lines long. +If there's already a crush.md, improve it. +If there are Cursor rules (in .cursor/rules/ or .cursorrules) or Copilot rules (in .github/copilot-instructions.md), make sure to include them.` + + // Mark the project as initialized + if err := config.MarkProjectInitialized(); err != nil { + return util.ReportError(err) + } + + return tea.Sequence( + util.CmdHandler(cmpChat.SessionClearedMsg{}), + util.CmdHandler(cmpChat.SendMsg{ + Text: prompt, + }), + ) + } else { + // Mark the project as initialized without running the command + if err := config.MarkProjectInitialized(); err != nil { + return util.ReportError(err) + } + } + return nil +} + // CloseInitDialogMsg is a message that is sent when the init dialog is closed. type CloseInitDialogMsg struct { Initialize bool diff --git a/internal/tui/tui.go b/internal/tui/tui.go index a41b72bdbe2d0ad991ff9f2376948f202c9004b4..87f140838f224368fbabe59af47d1933b15312be 100644 --- a/internal/tui/tui.go +++ b/internal/tui/tui.go @@ -7,6 +7,7 @@ import ( tea "github.com/charmbracelet/bubbletea/v2" "github.com/charmbracelet/crush/internal/app" "github.com/charmbracelet/crush/internal/config" + "github.com/charmbracelet/crush/internal/llm/agent" "github.com/charmbracelet/crush/internal/llm/tools" "github.com/charmbracelet/crush/internal/logging" "github.com/charmbracelet/crush/internal/permission" @@ -16,6 +17,7 @@ import ( "github.com/charmbracelet/crush/internal/tui/components/core/status" "github.com/charmbracelet/crush/internal/tui/components/dialogs" "github.com/charmbracelet/crush/internal/tui/components/dialogs/commands" + "github.com/charmbracelet/crush/internal/tui/components/dialogs/compact" "github.com/charmbracelet/crush/internal/tui/components/dialogs/filepicker" initDialog "github.com/charmbracelet/crush/internal/tui/components/dialogs/init" "github.com/charmbracelet/crush/internal/tui/components/dialogs/models" @@ -157,6 +159,12 @@ func (a *appModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { Model: models.NewModelDialogCmp(), }, ) + // Compact + case commands.CompactMsg: + return a, util.CmdHandler(dialogs.OpenDialogMsg{ + Model: compact.NewCompactDialogCmp(a.app.CoderAgent, msg.SessionID, true), + }) + // File Picker case chat.OpenFilePickerMsg: if a.dialog.ActiveDialogId() == filepicker.FilePickerID { @@ -181,33 +189,35 @@ func (a *appModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { a.app.Permissions.Deny(msg.Permission) } return a, nil - // Init Dialog - case initDialog.CloseInitDialogMsg: - if msg.Initialize { - // Run the initialization command - prompt := `Please analyze this codebase and create a Crush.md file containing: -1. Build/lint/test commands - especially for running a single test -2. Code style guidelines including imports, formatting, types, naming conventions, error handling, etc. - -The file you create will be given to agentic coding agents (such as yourself) that operate in this repository. Make it about 20 lines long. -If there's already a crush.md, improve it. -If there are Cursor rules (in .cursor/rules/ or .cursorrules) or Copilot rules (in .github/copilot-instructions.md), make sure to include them.` - - // Mark the project as initialized - if err := config.MarkProjectInitialized(); err != nil { - return a, util.ReportError(err) - } - - return a, util.CmdHandler(cmpChat.SendMsg{ - Text: prompt, - }) - } else { - // Mark the project as initialized without running the command - if err := config.MarkProjectInitialized(); err != nil { - return a, util.ReportError(err) + // Agent Events + case pubsub.Event[agent.AgentEvent]: + payload := msg.Payload + + // Forward agent events to dialogs + if a.dialog.HasDialogs() && a.dialog.ActiveDialogId() == compact.CompactDialogID { + u, dialogCmd := a.dialog.Update(payload) + a.dialog = u.(dialogs.DialogCmp) + cmds = append(cmds, dialogCmd) + } + + // Handle auto-compact logic + if payload.Done && payload.Type == agent.AgentEventTypeResponse && a.selectedSessionID != "" { + // Get current session to check token usage + session, err := a.app.Sessions.Get(context.Background(), a.selectedSessionID) + if err == nil { + model := a.app.CoderAgent.Model() + contextWindow := model.ContextWindow + tokens := session.CompletionTokens + session.PromptTokens + if (tokens >= int64(float64(contextWindow)*0.95)) && config.Get().AutoCompact { + // Show compact confirmation dialog + cmds = append(cmds, util.CmdHandler(dialogs.OpenDialogMsg{ + Model: compact.NewCompactDialogCmp(a.app.CoderAgent, a.selectedSessionID, false), + })) + } } } - return a, nil + + return a, tea.Batch(cmds...) // Key Press Messages case tea.KeyPressMsg: if msg.String() == "ctrl+t" { @@ -296,7 +306,7 @@ func (a *appModel) handleKeyPressMsg(msg tea.KeyPressMsg) tea.Cmd { return util.CmdHandler(dialogs.CloseDialogMsg{}) } return util.CmdHandler(dialogs.OpenDialogMsg{ - Model: commands.NewCommandDialog(), + Model: commands.NewCommandDialog(a.selectedSessionID), }) case key.Matches(msg, a.keyMap.Sessions): if a.dialog.ActiveDialogId() == sessions.SessionsDialogID {