use std::borrow::Cow; use std::fmt::Write; use std::path::PathBuf; use convert_case::{Case, Casing}; use derive_setters::Setters; use forge_api::{AgentId, Effort, ModelId, Usage}; use nu_ansi_term::{Color, Style}; use crate::display_constants::markers; use crate::utils::humanize_number; // Nerd font symbols — left prompt const DIR_SYMBOL: &str = "\u{ea83}"; // 󪃃 folder icon const BRANCH_SYMBOL: &str = "\u{f418}"; // branch icon const SUCCESS_SYMBOL: &str = "\u{f013e}"; // 󰄾 chevron // Nerd font symbols — right prompt (ZSH rprompt) const AGENT_SYMBOL: &str = "\u{f167a}"; const MODEL_SYMBOL: &str = "\u{ec19}"; /// Terminal width at which the reasoning effort label switches from the /// compact three-letter form (e.g. `MED`) to the full uppercase label /// (e.g. `MEDIUM`). Matches [`crate::zsh::rprompt`] so the CLI and zsh /// integration render identically on equivalent terminals. const WIDE_TERMINAL_THRESHOLD: usize = 100; /// Very Specialized Prompt for the Agent Chat #[derive(Clone, Setters)] #[setters(strip_option, borrow_self)] pub struct ForgePrompt { pub cwd: PathBuf, pub usage: Option, pub agent_id: AgentId, pub model: Option, /// Currently configured reasoning effort level for the active model, /// rendered to the right of the model when set. `Effort::None` is /// suppressed (see [`ForgePrompt::render_prompt_right`]). pub reasoning_effort: Option, pub git_branch: Option, } impl ForgePrompt { /// Creates a new `ForgePrompt`, resolving the git branch once at /// construction time. pub fn new(cwd: PathBuf, agent_id: AgentId) -> Self { let git_branch = get_git_branch(); Self { cwd, usage: None, agent_id, model: None, reasoning_effort: None, git_branch, } } pub fn refresh(&mut self) -> &mut Self { let git_branch = get_git_branch(); self.git_branch = git_branch; self } pub fn render_prompt_left(&self) -> Cow<'_, str> { // Left prompt layout: // // AGENT_NAME 󪃃 dir branch // 󰄾 // // Colors: // agent → bold white (identifies the active agent) // dir → bold cyan // branch → bold green // chevron → bold green let dir_style = Style::new().fg(Color::Cyan).bold(); let branch_style = Style::new().fg(Color::LightGreen).bold(); let chevron_style = Style::new().fg(Color::LightGreen).bold(); let current_dir = self .cwd .file_name() .and_then(|name| name.to_str()) .map(String::from) .unwrap_or_else(|| markers::EMPTY.to_string()); let mut result = String::with_capacity(80); // Directory — folder icon + name, bold cyan write!( result, "{}", dir_style.paint(format!("{DIR_SYMBOL} {current_dir}")) ) .unwrap(); // Git branch — branch icon + name, bold green (only when present and // different from the directory name, matching existing behaviour) if let Some(branch) = self.git_branch.as_deref() && branch != current_dir { write!( result, " {}", branch_style.paint(format!("{BRANCH_SYMBOL} {branch}")) ) .unwrap(); } // Second line: success chevron write!(result, "\n{} ", chevron_style.paint(SUCCESS_SYMBOL)).unwrap(); Cow::Owned(result) } pub fn render_prompt_right(&self) -> Cow<'_, str> { // Right prompt layout: agent · tokens · cost · model // Active (tokens > 0): bright white for agent/tokens, green for cost // Inactive (no tokens): all segments dimmed let total_tokens = self.usage.as_ref().map(|u| u.total_tokens); let active = total_tokens.map(|t| *t > 0).unwrap_or(false); let agent_color = if active { Color::LightGray } else { Color::DarkGray }; let mut result = String::with_capacity(64); // Agent name with nerd font symbol let agent_str = format!( "{AGENT_SYMBOL} {}", self.agent_id.as_str().to_case(Case::UpperSnake) ); write!( result, " {}", Style::new().bold().fg(agent_color).paint(&agent_str) ) .unwrap(); // Token count (only shown when active) if let Some(tokens) = total_tokens && active { let prefix = match tokens { forge_api::TokenCount::Actual(_) => "", forge_api::TokenCount::Approx(_) => "~", }; let count_str = format!("{}{}", prefix, humanize_number(*tokens)); write!( result, " {}", Style::new().bold().fg(Color::LightGray).paint(&count_str) ) .unwrap(); } // Cost (only shown when active) if let Some(cost) = self.usage.as_ref().and_then(|u| u.cost) && active { let cost_str = format!("\u{f155}{cost:.2}"); write!( result, " {}", Style::new().bold().fg(Color::Green).paint(&cost_str) ) .unwrap(); } // Model with nerd font symbol if let Some(model) = self.model.as_ref() { let model_str = model.to_string(); let short_model = model_str.split('/').next_back().unwrap_or(model.as_str()); let model_label = format!("{MODEL_SYMBOL} {short_model}"); let color = if active { Color::LightMagenta } else { Color::DarkGray }; write!(result, " {}", Style::new().fg(color).paint(&model_label)).unwrap(); } // Reasoning effort — rendered to the right of the model, matching the // ZSH rprompt. `Effort::None` is suppressed (see zsh/rprompt.rs). On // narrow terminals the label collapses to its first three characters // so the prompt stays compact. if let Some(ref effort) = self.reasoning_effort && !matches!(effort, Effort::None) { let effort_label = effort_label(effort, term_width()); let color = if active { Color::Yellow } else { Color::DarkGray }; write!(result, " {}", Style::new().fg(color).paint(&effort_label)).unwrap(); } Cow::Owned(result) } pub fn render_prompt_indicator(&self) -> Cow<'_, str> { Cow::Borrowed("") } } /// Gets the current git branch name if available fn get_git_branch() -> Option { let repo = gix::discover(".").ok()?; let head = repo.head().ok()?; head.referent_name().map(|r| r.shorten().to_string()) } /// Returns the current terminal width in columns, falling back to 80 when /// the size cannot be detected. fn term_width() -> usize { terminal_size::terminal_size() .map(|(w, _)| w.0 as usize) .unwrap_or(80) } /// Formats an [`Effort`] as its uppercase label, collapsing to the first three /// characters on narrow terminals (< [`WIDE_TERMINAL_THRESHOLD`] columns). fn effort_label(effort: &Effort, width: usize) -> String { let full = effort.to_string().to_uppercase(); if width >= WIDE_TERMINAL_THRESHOLD { full } else { // `chars().take(3)` rather than `&full[..3]` to satisfy the // `clippy::string_slice` lint denied in CI. full.chars().take(3).collect() } } #[cfg(test)] mod tests { use nu_ansi_term::Style; use pretty_assertions::assert_eq; use super::*; impl Default for ForgePrompt { fn default() -> Self { ForgePrompt { cwd: PathBuf::from("."), usage: None, agent_id: AgentId::default(), model: None, reasoning_effort: None, git_branch: None, } } } enum PromptHistorySearchStatus { Passing, Failing, } struct PromptHistorySearch { status: PromptHistorySearchStatus, term: String, } fn render_prompt_history_search_indicator( history_search: PromptHistorySearch, ) -> Cow<'static, str> { let prefix = match history_search.status { PromptHistorySearchStatus::Passing => "", PromptHistorySearchStatus::Failing => "failing ", }; let mut result = String::with_capacity(32); if history_search.term.is_empty() { write!(result, "({prefix}reverse-search) ").unwrap(); } else { write!( result, "({}reverse-search: {}) ", prefix, history_search.term ) .unwrap(); } Cow::Owned(Style::new().fg(Color::White).paint(&result).to_string()) } #[test] fn test_render_prompt_left() { let prompt = ForgePrompt::default(); let actual = prompt.render_prompt_left(); // Starship directory icon present assert!(actual.contains(DIR_SYMBOL)); // Starship success chevron present assert!(actual.contains(SUCCESS_SYMBOL)); } #[test] fn test_render_prompt_left_with_branch() { let prompt = ForgePrompt { git_branch: Some("main".to_string()), ..Default::default() }; let actual = prompt.render_prompt_left(); // Agent name is on the right prompt, not the left // Branch icon and name present assert!(actual.contains(BRANCH_SYMBOL)); assert!(actual.contains("main")); } #[test] fn test_render_prompt_right_inactive() { // No tokens → dimmed agent + model, no token/cost segments let mut prompt = ForgePrompt::default(); let _ = prompt.model(ModelId::new("gpt-4")); let actual = prompt.render_prompt_right(); // Agent symbol and name present assert!(actual.contains(AGENT_SYMBOL)); assert!(actual.contains("FORGE")); // Model symbol and name present assert!(actual.contains(MODEL_SYMBOL)); assert!(actual.contains("gpt-4")); // No token count text in inactive state (no humanized number segment) assert!(!actual.contains("1k") && !actual.contains("~")); } #[test] fn test_render_prompt_right_active_with_tokens() { // Tokens > 0 → active colours; approx tokens show "~" prefix let usage = Usage { prompt_tokens: forge_api::TokenCount::Actual(10), completion_tokens: forge_api::TokenCount::Actual(20), total_tokens: forge_api::TokenCount::Approx(30), ..Default::default() }; let mut prompt = ForgePrompt::default(); let _ = prompt.usage(usage); let actual = prompt.render_prompt_right(); assert!(actual.contains("~30")); assert!(actual.contains(AGENT_SYMBOL)); } #[test] fn test_render_prompt_history_search_indicator_passing() { let history_search = PromptHistorySearch { status: PromptHistorySearchStatus::Passing, term: "test".to_string(), }; let actual = render_prompt_history_search_indicator(history_search); let expected = Style::new() .fg(Color::White) .paint("(reverse-search: test) ") .to_string(); assert_eq!(actual, expected); } #[test] fn test_render_prompt_history_search_indicator_failing() { let history_search = PromptHistorySearch { status: PromptHistorySearchStatus::Failing, term: "test".to_string(), }; let actual = render_prompt_history_search_indicator(history_search); let expected = Style::new() .fg(Color::White) .paint("(failing reverse-search: test) ") .to_string(); assert_eq!(actual, expected); } #[test] fn test_render_prompt_history_search_indicator_empty_term() { let history_search = PromptHistorySearch { status: PromptHistorySearchStatus::Passing, term: "".to_string(), }; let actual = render_prompt_history_search_indicator(history_search); let expected = Style::new() .fg(Color::White) .paint("(reverse-search) ") .to_string(); assert_eq!(actual, expected); } #[test] fn test_render_prompt_right_strips_provider_prefix() { // Model ID like "anthropic/claude-3" should show only "claude-3" let usage = Usage { prompt_tokens: forge_api::TokenCount::Actual(10), completion_tokens: forge_api::TokenCount::Actual(20), total_tokens: forge_api::TokenCount::Actual(30), ..Default::default() }; let mut prompt = ForgePrompt::default(); let _ = prompt.usage(usage); let _ = prompt.model(ModelId::new("anthropic/claude-3")); let actual = prompt.render_prompt_right(); assert!(actual.contains("claude-3")); assert!(!actual.contains("anthropic/claude-3")); assert!(actual.contains("30")); } #[test] fn test_render_prompt_right_with_cost() { // Cost shown when active let usage = Usage { total_tokens: forge_api::TokenCount::Actual(1500), cost: Some(0.01), ..Default::default() }; let mut prompt = ForgePrompt::default(); let _ = prompt.usage(usage); let actual = prompt.render_prompt_right(); assert!(actual.contains("0.01")); assert!(actual.contains("1.5k")); } #[test] fn test_render_prompt_right_with_reasoning_effort() { // When reasoning effort is set, its uppercase label appears after the // model segment. let mut prompt = ForgePrompt::default(); let _ = prompt.model(ModelId::new("gpt-4")); let _ = prompt.reasoning_effort(Effort::High); let actual = prompt.render_prompt_right(); assert!(actual.contains("HIGH") || actual.contains("HIG")); } #[test] fn test_render_prompt_right_hides_effort_none() { // `Effort::None` carries no useful info — it must not be rendered. let mut prompt = ForgePrompt::default(); let _ = prompt.model(ModelId::new("gpt-4")); let _ = prompt.reasoning_effort(Effort::None); let actual = prompt.render_prompt_right(); assert!(!actual.to_uppercase().contains("NONE")); } #[test] fn test_effort_label_narrow_vs_wide() { assert_eq!(effort_label(&Effort::Medium, 80), "MED"); assert_eq!( effort_label(&Effort::Medium, WIDE_TERMINAL_THRESHOLD), "MEDIUM" ); } }