diff --git a/cmd/root.go b/cmd/root.go index e9af179e..99e15f0e 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -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) diff --git a/internal/ui/factory.go b/internal/ui/factory.go index 9ee9c778..c566df98 100644 --- a/internal/ui/factory.go +++ b/internal/ui/factory.go @@ -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 diff --git a/internal/ui/model.go b/internal/ui/model.go index adab7020..6d9865db 100644 --- a/internal/ui/model.go +++ b/internal/ui/model.go @@ -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)