mirror of
https://github.com/mark3labs/kit.git
synced 2026-06-13 19:20:06 +00:00
fix(ui): execute slash commands in TUI instead of sending them to the LLM
The Bubble Tea refactor only wired /quit, /clear, and /clear-queue in InputComponent; the remaining commands (/help, /tools, /servers, /usage, /reset-usage) fell through as submitMsg and were forwarded to app.Run() as regular prompts. Intercept all recognized slash commands in AppModel.Update before they reach the app layer, and add print helpers that emit formatted output via tea.Println. Also create a UsageTracker for interactive mode so /usage and /reset-usage work correctly.
This commit is contained in:
+25
-14
@@ -714,19 +714,27 @@ func runNormalMode(ctx context.Context) error {
|
||||
CompactMode: viper.GetBool("compact"),
|
||||
}
|
||||
|
||||
// Wire usage tracker from the CLI (non-interactive non-quiet mode only).
|
||||
// In quiet mode or interactive mode, cli is nil and no tracker is attached.
|
||||
// Create a usage tracker that is shared between the app layer (for recording
|
||||
// usage after each step) and the TUI (for /usage display). For non-interactive
|
||||
// mode the tracker comes from the CLI factory; for interactive mode we create
|
||||
// one directly.
|
||||
var usageTracker *ui.UsageTracker
|
||||
if cli != nil {
|
||||
if tracker := cli.GetUsageTracker(); tracker != nil {
|
||||
appOpts.UsageTracker = tracker
|
||||
}
|
||||
usageTracker = cli.GetUsageTracker()
|
||||
} else {
|
||||
// Interactive mode: create a tracker using the same logic as SetupCLI.
|
||||
usageTracker = ui.CreateUsageTracker(modelString, viper.GetString("provider-api-key"))
|
||||
}
|
||||
if usageTracker != nil {
|
||||
appOpts.UsageTracker = usageTracker
|
||||
}
|
||||
|
||||
appInstance := app.New(appOpts, messages)
|
||||
defer appInstance.Close()
|
||||
|
||||
// Check if running in non-interactive mode
|
||||
if promptFlag != "" {
|
||||
return runNonInteractiveModeApp(ctx, appInstance, promptFlag, quietFlag, noExitFlag, modelName)
|
||||
return runNonInteractiveModeApp(ctx, appInstance, promptFlag, quietFlag, noExitFlag, modelName, serverNames, toolNames, usageTracker)
|
||||
}
|
||||
|
||||
// Quiet mode is not allowed in interactive mode
|
||||
@@ -734,7 +742,7 @@ func runNormalMode(ctx context.Context) error {
|
||||
return fmt.Errorf("--quiet flag can only be used with --prompt/-p")
|
||||
}
|
||||
|
||||
return runInteractiveModeBubbleTea(ctx, appInstance, modelName)
|
||||
return runInteractiveModeBubbleTea(ctx, appInstance, modelName, serverNames, toolNames, usageTracker)
|
||||
}
|
||||
|
||||
// runNonInteractiveModeApp executes a single prompt via the app layer and exits,
|
||||
@@ -747,14 +755,14 @@ func runNormalMode(ctx context.Context) error {
|
||||
//
|
||||
// When --no-exit is set, after RunOnce completes the interactive BubbleTea TUI
|
||||
// is started so the user can continue the conversation.
|
||||
func runNonInteractiveModeApp(ctx context.Context, appInstance *app.App, prompt string, _, noExit bool, modelName string) error {
|
||||
func runNonInteractiveModeApp(ctx context.Context, appInstance *app.App, prompt string, _, noExit bool, modelName string, serverNames, toolNames []string, usageTracker *ui.UsageTracker) error {
|
||||
if err := appInstance.RunOnce(ctx, prompt, os.Stdout); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// If --no-exit was requested, hand off to the interactive TUI.
|
||||
if noExit {
|
||||
return runInteractiveModeBubbleTea(ctx, appInstance, modelName)
|
||||
return runInteractiveModeBubbleTea(ctx, appInstance, modelName, serverNames, toolNames, usageTracker)
|
||||
}
|
||||
|
||||
return nil
|
||||
@@ -771,7 +779,7 @@ func runNonInteractiveModeApp(ctx context.Context, appInstance *app.App, prompt
|
||||
// 4. Calls program.Run() which blocks until the user quits (Ctrl+C or /quit).
|
||||
//
|
||||
// SetupCLI is not used for interactive mode; the TUI (AppModel) handles its own rendering.
|
||||
func runInteractiveModeBubbleTea(_ context.Context, appInstance *app.App, modelName string) error {
|
||||
func runInteractiveModeBubbleTea(_ context.Context, appInstance *app.App, modelName string, serverNames, toolNames []string, usageTracker *ui.UsageTracker) error {
|
||||
// Determine terminal size; fall back gracefully.
|
||||
termWidth, termHeight, err := term.GetSize(int(os.Stdout.Fd()))
|
||||
if err != nil || termWidth == 0 {
|
||||
@@ -780,10 +788,13 @@ func runInteractiveModeBubbleTea(_ context.Context, appInstance *app.App, modelN
|
||||
}
|
||||
|
||||
appModel := ui.NewAppModel(appInstance, ui.AppModelOptions{
|
||||
CompactMode: viper.GetBool("compact"),
|
||||
ModelName: modelName,
|
||||
Width: termWidth,
|
||||
Height: termHeight,
|
||||
CompactMode: viper.GetBool("compact"),
|
||||
ModelName: modelName,
|
||||
Width: termWidth,
|
||||
Height: termHeight,
|
||||
ServerNames: serverNames,
|
||||
ToolNames: toolNames,
|
||||
UsageTracker: usageTracker,
|
||||
})
|
||||
|
||||
program := tea.NewProgram(appModel)
|
||||
|
||||
@@ -38,6 +38,33 @@ func parseModelName(modelString string) (provider, model string) {
|
||||
return p, m
|
||||
}
|
||||
|
||||
// CreateUsageTracker creates a UsageTracker for the given model string and
|
||||
// provider API key. It returns nil when usage tracking is unavailable (e.g.
|
||||
// ollama or unrecognised models). This is used by the interactive TUI path
|
||||
// which doesn't go through SetupCLI.
|
||||
func CreateUsageTracker(modelString, providerAPIKey string) *UsageTracker {
|
||||
provider, model := parseModelName(modelString)
|
||||
if provider == "unknown" || model == "unknown" || provider == "ollama" {
|
||||
return nil
|
||||
}
|
||||
|
||||
registry := models.GetGlobalRegistry()
|
||||
modelInfo, err := registry.ValidateModel(provider, model)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
isOAuth := false
|
||||
if provider == "anthropic" {
|
||||
_, source, err := auth.GetAnthropicAPIKey(providerAPIKey)
|
||||
if err == nil && strings.HasPrefix(source, "stored OAuth") {
|
||||
isOAuth = true
|
||||
}
|
||||
}
|
||||
|
||||
return NewUsageTracker(modelInfo, provider, 80, isOAuth)
|
||||
}
|
||||
|
||||
// SetupCLI creates, configures, and initializes a CLI instance with the provided
|
||||
// options. It sets up model display, usage tracking for supported providers, and
|
||||
// shows initial loading information. Returns nil in quiet mode or an initialized
|
||||
|
||||
+159
-13
@@ -54,6 +54,16 @@ type AppModelOptions struct {
|
||||
|
||||
// Height is the initial terminal height in rows.
|
||||
Height int
|
||||
|
||||
// ServerNames holds loaded MCP server names for the /servers command.
|
||||
ServerNames []string
|
||||
|
||||
// ToolNames holds available tool names for the /tools command.
|
||||
ToolNames []string
|
||||
|
||||
// UsageTracker provides token usage statistics for /usage and /reset-usage.
|
||||
// May be nil if usage tracking is unavailable for the current model.
|
||||
UsageTracker *UsageTracker
|
||||
}
|
||||
|
||||
// AppModel is the root Bubble Tea model for the interactive TUI. It owns the
|
||||
@@ -105,6 +115,14 @@ type AppModel struct {
|
||||
// A second ESC within 2 seconds will cancel the current step.
|
||||
canceling bool
|
||||
|
||||
// serverNames, toolNames are used by /servers and /tools commands.
|
||||
serverNames []string
|
||||
toolNames []string
|
||||
|
||||
// usageTracker provides token usage stats for /usage and /reset-usage.
|
||||
// May be nil when usage tracking is unavailable.
|
||||
usageTracker *UsageTracker
|
||||
|
||||
// width and height track the terminal dimensions.
|
||||
width int
|
||||
height int
|
||||
@@ -156,14 +174,17 @@ func NewAppModel(appCtrl AppController, opts AppModelOptions) *AppModel {
|
||||
}
|
||||
|
||||
m := &AppModel{
|
||||
state: stateInput,
|
||||
appCtrl: appCtrl,
|
||||
renderer: NewMessageRenderer(width, false),
|
||||
compactRdr: NewCompactRenderer(width, false),
|
||||
compactMode: opts.CompactMode,
|
||||
modelName: opts.ModelName,
|
||||
width: width,
|
||||
height: height,
|
||||
state: stateInput,
|
||||
appCtrl: appCtrl,
|
||||
renderer: NewMessageRenderer(width, false),
|
||||
compactRdr: NewCompactRenderer(width, false),
|
||||
compactMode: opts.CompactMode,
|
||||
modelName: opts.ModelName,
|
||||
serverNames: opts.ServerNames,
|
||||
toolNames: opts.ToolNames,
|
||||
usageTracker: opts.UsageTracker,
|
||||
width: width,
|
||||
height: height,
|
||||
}
|
||||
|
||||
// Wire up child components now that we have the concrete implementations.
|
||||
@@ -255,12 +276,15 @@ func (m *AppModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
|
||||
|
||||
// ── Input submitted ──────────────────────────────────────────────────────
|
||||
case submitMsg:
|
||||
// Handle /quit (and its aliases) before sending to the app layer:
|
||||
// look up the command and check if it resolves to "/quit".
|
||||
if cmd := GetCommandByName(msg.Text); cmd != nil && cmd.Name == "/quit" {
|
||||
return m, tea.Quit
|
||||
// Handle slash commands locally — they should never reach app.Run().
|
||||
if sc := GetCommandByName(msg.Text); sc != nil {
|
||||
if cmd := m.handleSlashCommand(sc); cmd != nil {
|
||||
cmds = append(cmds, cmd)
|
||||
}
|
||||
return m, tea.Batch(cmds...)
|
||||
}
|
||||
// Print user message immediately to scrollback.
|
||||
|
||||
// Regular prompt — print user message and forward to the app layer.
|
||||
cmds = append(cmds, m.printUserMessage(msg.Text))
|
||||
if m.appCtrl != nil {
|
||||
// app.Run() handles queueing internally if a step is in progress.
|
||||
@@ -519,6 +543,128 @@ func (m *AppModel) printErrorResponse(evt app.StepErrorEvent) tea.Cmd {
|
||||
return tea.Println(rendered)
|
||||
}
|
||||
|
||||
// --------------------------------------------------------------------------
|
||||
// Slash command handlers
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
// handleSlashCommand executes a recognized slash command and returns a tea.Cmd
|
||||
// that emits the appropriate output to scrollback. Returns tea.Quit for /quit,
|
||||
// nil for commands with no visible output, or a tea.Println cmd for display.
|
||||
func (m *AppModel) handleSlashCommand(sc *SlashCommand) tea.Cmd {
|
||||
switch sc.Name {
|
||||
case "/quit":
|
||||
return tea.Quit
|
||||
case "/help":
|
||||
return m.printHelpMessage()
|
||||
case "/tools":
|
||||
return m.printToolsMessage()
|
||||
case "/servers":
|
||||
return m.printServersMessage()
|
||||
case "/usage":
|
||||
return m.printUsageMessage()
|
||||
case "/reset-usage":
|
||||
return m.printResetUsage()
|
||||
case "/clear":
|
||||
if m.appCtrl != nil {
|
||||
m.appCtrl.ClearMessages()
|
||||
}
|
||||
return m.printSystemMessage("Conversation cleared. Starting fresh.")
|
||||
case "/clear-queue":
|
||||
if m.appCtrl != nil {
|
||||
m.appCtrl.ClearQueue()
|
||||
}
|
||||
return nil
|
||||
default:
|
||||
return m.printSystemMessage(fmt.Sprintf("Unknown command: %s", sc.Name))
|
||||
}
|
||||
}
|
||||
|
||||
// printSystemMessage renders a system-level message and emits it above the BT region.
|
||||
func (m *AppModel) printSystemMessage(text string) tea.Cmd {
|
||||
var rendered string
|
||||
if m.compactMode {
|
||||
msg := m.compactRdr.RenderSystemMessage(text, time.Now())
|
||||
rendered = msg.Content
|
||||
} else {
|
||||
msg := m.renderer.RenderSystemMessage(text, time.Now())
|
||||
rendered = msg.Content
|
||||
}
|
||||
return tea.Println(rendered)
|
||||
}
|
||||
|
||||
// printHelpMessage renders the help text listing all available slash commands.
|
||||
func (m *AppModel) printHelpMessage() tea.Cmd {
|
||||
help := "## Available Commands\n\n" +
|
||||
"- `/help`: Show this help message\n" +
|
||||
"- `/tools`: List all available tools\n" +
|
||||
"- `/servers`: List configured MCP servers\n" +
|
||||
"- `/usage`: Show token usage and cost statistics\n" +
|
||||
"- `/reset-usage`: Reset usage statistics\n" +
|
||||
"- `/clear`: Clear message history\n" +
|
||||
"- `/quit`: Exit the application\n" +
|
||||
"- `Ctrl+C`: Exit at any time\n" +
|
||||
"- `ESC` (x2): Cancel ongoing LLM generation\n\n" +
|
||||
"You can also just type your message to chat with the AI assistant."
|
||||
return m.printSystemMessage(help)
|
||||
}
|
||||
|
||||
// printToolsMessage renders the list of available tools.
|
||||
func (m *AppModel) printToolsMessage() tea.Cmd {
|
||||
var content string
|
||||
content = "## Available Tools\n\n"
|
||||
if len(m.toolNames) == 0 {
|
||||
content += "No tools are currently available."
|
||||
} else {
|
||||
for i, tool := range m.toolNames {
|
||||
content += fmt.Sprintf("%d. `%s`\n", i+1, tool)
|
||||
}
|
||||
}
|
||||
return m.printSystemMessage(content)
|
||||
}
|
||||
|
||||
// printServersMessage renders the list of configured MCP servers.
|
||||
func (m *AppModel) printServersMessage() tea.Cmd {
|
||||
var content string
|
||||
content = "## Configured MCP Servers\n\n"
|
||||
if len(m.serverNames) == 0 {
|
||||
content += "No MCP servers are currently configured."
|
||||
} else {
|
||||
for i, server := range m.serverNames {
|
||||
content += fmt.Sprintf("%d. `%s`\n", i+1, server)
|
||||
}
|
||||
}
|
||||
return m.printSystemMessage(content)
|
||||
}
|
||||
|
||||
// printUsageMessage renders token usage statistics.
|
||||
func (m *AppModel) printUsageMessage() tea.Cmd {
|
||||
if m.usageTracker == nil {
|
||||
return m.printSystemMessage("Usage tracking is not available for this model.")
|
||||
}
|
||||
|
||||
sessionStats := m.usageTracker.GetSessionStats()
|
||||
lastStats := m.usageTracker.GetLastRequestStats()
|
||||
|
||||
content := "## Usage Statistics\n\n"
|
||||
if lastStats != nil {
|
||||
content += fmt.Sprintf("**Last Request:** %d input + %d output tokens = $%.6f\n",
|
||||
lastStats.InputTokens, lastStats.OutputTokens, lastStats.TotalCost)
|
||||
}
|
||||
content += fmt.Sprintf("**Session Total:** %d input + %d output tokens = $%.6f (%d requests)\n",
|
||||
sessionStats.TotalInputTokens, sessionStats.TotalOutputTokens, sessionStats.TotalCost, sessionStats.RequestCount)
|
||||
|
||||
return m.printSystemMessage(content)
|
||||
}
|
||||
|
||||
// printResetUsage resets usage statistics and prints a confirmation.
|
||||
func (m *AppModel) printResetUsage() tea.Cmd {
|
||||
if m.usageTracker == nil {
|
||||
return m.printSystemMessage("Usage tracking is not available for this model.")
|
||||
}
|
||||
m.usageTracker.Reset()
|
||||
return m.printSystemMessage("Usage statistics have been reset.")
|
||||
}
|
||||
|
||||
// flushStreamContent gets the rendered content from the stream component,
|
||||
// emits it above the BT region via tea.Println, and resets the stream. This
|
||||
// is called before printing tool calls (streaming completes before tools fire)
|
||||
|
||||
Reference in New Issue
Block a user