use std::time::Duration;
use colored::Colorize;
use inquire::MultiSelect;
use crate::{config::Config, git, jj, openai};
const OUTPUT_AND_REASONING_RESERVE_TOKENS: usize = 25_000;
fn input_token_budget(context: usize) -> usize {
context.saturating_sub(OUTPUT_AND_REASONING_RESERVE_TOKENS)
}
fn ensure_diff_fits(used_tokens: usize, diff_tokens: usize, context: usize) -> anyhow::Result<()> {
let budget = input_token_budget(context);
if used_tokens.saturating_add(diff_tokens) > budget {
return Err(anyhow::anyhow!(
"Selected diff is still too large: about {} input tokens for a {}-token budget. Select fewer files.",
used_tokens.saturating_add(diff_tokens),
budget
));
}
Ok(())
}
pub fn decide_diff(
repo: &git2::Repository,
used_tokens: usize,
context: usize,
always_select_files: bool,
) -> anyhow::Result<(String, usize)> {
let staged_files = git::staged_files(&repo)?;
let mut diff = git::diff(&repo, &staged_files)?;
let mut diff_tokens = openai::count_token(&diff)?;
let input_budget = input_token_budget(context);
if diff_tokens == 0 {
println!(
"{} {}",
"No staged files.".red(),
"Please stage the files you want to commit.".bright_black()
);
std::process::exit(1);
}
if always_select_files || used_tokens.saturating_add(diff_tokens) > input_budget {
if always_select_files {
println!(
"{} {}",
"File selection mode:".blue(),
"Select the files you want to include in the commit.".bright_black()
);
} else {
println!(
"{} {}",
"The request is too long!".red(),
format!(
"The request is ~{} input tokens long, while the input budget is {} ({} tokens are reserved for reasoning and output).",
used_tokens.saturating_add(diff_tokens),
input_budget,
OUTPUT_AND_REASONING_RESERVE_TOKENS
)
.bright_black()
);
}
let selected_files = MultiSelect::new(
"Select the files you want to include in the diff:",
staged_files.clone(),
)
.prompt()?;
diff = git::diff(&repo, &selected_files)?;
diff_tokens = openai::count_token(&diff)?;
}
ensure_diff_fits(used_tokens, diff_tokens, context)?;
Ok((diff, diff_tokens))
}
pub fn decide_diff_jj(
used_tokens: usize,
context: usize,
always_select_files: bool,
revision: Option<&str>,
) -> anyhow::Result<(String, usize)> {
let modified_files = jj::get_jj_modified_files()?;
let mut diff = jj::get_jj_diff(revision)?;
let mut diff_tokens = openai::count_token(&diff)?;
let input_budget = input_token_budget(context);
if diff_tokens == 0 {
let revision_msg = if let Some(rev) = revision {
format!("No changes detected for revision '{}'.", rev)
} else {
"No changes detected.".to_string()
};
println!(
"{} {}",
revision_msg.red(),
"Please make some changes before running turbocommit.".bright_black()
);
std::process::exit(1);
}
if always_select_files || used_tokens.saturating_add(diff_tokens) > input_budget {
if always_select_files {
println!(
"{} {}",
"File selection mode:".blue(),
"Select the files you want to include in the diff.".bright_black()
);
} else {
println!(
"{} {}",
"The request is too long!".red(),
format!(
"The request is ~{} input tokens long, while the input budget is {} ({} tokens are reserved for reasoning and output).",
used_tokens.saturating_add(diff_tokens),
input_budget,
OUTPUT_AND_REASONING_RESERVE_TOKENS
)
.bright_black()
);
}
let selected_files = MultiSelect::new(
"Select the files you want to include in the diff:",
modified_files.clone(),
)
.prompt()?;
diff = jj::get_jj_diff_for_files(revision, &selected_files)?;
diff_tokens = openai::count_token(&diff)?;
}
ensure_diff_fits(used_tokens, diff_tokens, context)?;
Ok((diff, diff_tokens))
}
pub fn check_config_age(max_age: Duration) -> bool {
let path = Config::path();
let metadata = match std::fs::metadata(&path) {
Ok(metadata) => metadata,
Err(_) => {
return false;
}
};
let last_modified = metadata.modified().unwrap();
let now = std::time::SystemTime::now();
match now.duration_since(last_modified) {
Ok(duration) => duration > max_age,
Err(_) => false,
}
}
pub fn is_system_prompt_same_as_default(system_msg: &str) -> bool {
let default = Config::default().system_msg;
system_msg == default
}
pub async fn check_version() {
let client = match crates_io_api::AsyncClient::new(
"turbocommit latest version checker",
Duration::from_millis(1000),
) {
Ok(client) => client,
Err(_) => {
return;
}
};
let turbo = match client.get_crate("turbocommit").await {
Ok(turbo) => turbo,
Err(_) => {
return;
}
};
let newest_version = turbo.versions[0].num.clone();
let current_version = env!("CARGO_PKG_VERSION");
if current_version != newest_version {
println!(
"\n{} {}",
"New version available!".yellow(),
format!("v{}", newest_version).purple()
);
println!(
"To update, run\n{}",
"cargo install --force turbocommit".purple()
);
}
}
pub fn choose_message(choices: Vec<String>) -> Option<String> {
if choices.len() == 1 {
return Some(process_response(&choices[0]));
}
let max_index = choices.len();
let commit_index = match inquire::CustomType::<usize>::new(&format!(
"Which commit message do you want to use? {}",
"<ESC> to cancel".bright_black()
))
.with_validator(move |i: &usize| {
if *i >= max_index {
Err(inquire::CustomUserError::from("Invalid index"))
} else {
Ok(inquire::validator::Validation::Valid)
}
})
.prompt()
{
Ok(index) => index,
Err(_) => {
return None;
}
};
Some(process_response(&choices[commit_index]))
}
pub fn format_token_count(tokens: usize) -> String {
format!("{:.2}k", tokens as f64 / 1000.0)
}
fn process_response(response: &str) -> String {
if let Some(think_start) = response.find("<think>") {
if let Some(think_end) = response.find("</think>") {
let thinking = &response[think_start + 7..think_end];
let message_part = response[think_end + 8..].trim_matches(|c: char| c.is_whitespace());
println!("\n{}", "AI's Thought Process:".blue().bold());
println!("{}", thinking.bright_black());
println!("\n{}", "Generated Commit Message:".blue().bold());
println!("[0] {}", "=".repeat(76));
println!("{}", message_part);
return message_part.to_string();
}
}
response
.trim_matches(|c: char| c.is_whitespace())
.to_string()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn reserves_context_for_reasoning_and_output() {
assert_eq!(input_token_budget(1_050_000), 1_025_000);
assert!(ensure_diff_fits(500, 1_024_500, 1_050_000).is_ok());
assert!(ensure_diff_fits(501, 1_024_500, 1_050_000).is_err());
}
#[test]
fn small_contexts_saturate_safely() {
assert_eq!(input_token_budget(10_000), 0);
assert!(ensure_diff_fits(0, 1, 10_000).is_err());
}
}