use std::time::Duration;
#[derive(Debug, Clone, PartialEq)]
pub struct ShellConfig {
pub suggest_timeout_ms: u64,
pub max_model_size: usize,
pub max_history_size: usize,
pub max_suggestions: usize,
pub max_prefix_length: usize,
pub min_prefix_length: usize,
pub min_quality_score: f32,
}
impl Default for ShellConfig {
fn default() -> Self {
Self {
suggest_timeout_ms: 100,
max_model_size: 100 * 1024 * 1024, max_history_size: 500 * 1024 * 1024, max_suggestions: 10,
max_prefix_length: 500,
min_prefix_length: 2,
min_quality_score: 0.3,
}
}
}
impl ShellConfig {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn with_suggest_timeout_ms(mut self, timeout: u64) -> Self {
self.suggest_timeout_ms = timeout;
self
}
#[must_use]
pub fn with_max_model_size(mut self, size: usize) -> Self {
self.max_model_size = size;
self
}
#[must_use]
pub fn with_max_history_size(mut self, size: usize) -> Self {
self.max_history_size = size;
self
}
#[must_use]
pub fn with_max_suggestions(mut self, count: usize) -> Self {
self.max_suggestions = count;
self
}
#[must_use]
pub fn with_max_prefix_length(mut self, length: usize) -> Self {
self.max_prefix_length = length;
self
}
#[must_use]
pub fn with_min_prefix_length(mut self, length: usize) -> Self {
self.min_prefix_length = length;
self
}
#[must_use]
pub fn with_min_quality_score(mut self, score: f32) -> Self {
self.min_quality_score = score.clamp(0.0, 1.0);
self
}
#[must_use]
pub fn suggest_timeout(&self) -> Duration {
Duration::from_millis(self.suggest_timeout_ms)
}
#[must_use]
pub fn is_model_size_valid(&self, size: usize) -> bool {
size <= self.max_model_size
}
#[must_use]
pub fn is_history_size_valid(&self, size: usize) -> bool {
size <= self.max_history_size
}
#[must_use]
pub fn is_prefix_valid(&self, prefix: &str) -> bool {
let len = prefix.len();
len >= self.min_prefix_length && len <= self.max_prefix_length
}
#[must_use]
pub fn truncate_prefix<'a>(&self, prefix: &'a str) -> &'a str {
if prefix.len() > self.max_prefix_length {
let mut end = self.max_prefix_length;
while end > 0 && !prefix.is_char_boundary(end) {
end -= 1;
}
&prefix[..end]
} else {
prefix
}
}
}
impl ShellConfig {
#[must_use]
pub fn fast() -> Self {
Self {
suggest_timeout_ms: 50,
max_model_size: 50 * 1024 * 1024, max_history_size: 100 * 1024 * 1024, max_suggestions: 5,
max_prefix_length: 200,
min_prefix_length: 2,
min_quality_score: 0.5,
}
}
#[must_use]
pub fn thorough() -> Self {
Self {
suggest_timeout_ms: 500,
max_model_size: 500 * 1024 * 1024, max_history_size: 1024 * 1024 * 1024, max_suggestions: 20,
max_prefix_length: 1000,
min_prefix_length: 1,
min_quality_score: 0.1,
}
}
}
use crate::model::MarkovModel;
use crate::quality::suggestion_quality_score;
use crate::security::is_sensitive_command;
use std::time::Instant;
pub fn suggest_with_fallback(
prefix: &str,
model: Option<&MarkovModel>,
config: &ShellConfig,
) -> Vec<(String, f32)> {
let model = match model {
Some(m) => m,
None => return vec![],
};
if !is_prefix_processable(prefix, config) {
return vec![];
}
let prefix = config.truncate_prefix(prefix);
let raw_suggestions = model.suggest(prefix, config.max_suggestions * 2);
filter_suggestions(raw_suggestions, config)
}
fn is_prefix_processable(prefix: &str, config: &ShellConfig) -> bool {
prefix.len() >= config.min_prefix_length
}
fn filter_suggestions(
raw_suggestions: Vec<(String, f32)>,
config: &ShellConfig,
) -> Vec<(String, f32)> {
let deadline = Instant::now() + config.suggest_timeout();
let mut results = Vec::with_capacity(config.max_suggestions);
for (suggestion, score) in raw_suggestions {
if should_stop_filtering(&results, &deadline, config) {
break;
}
if let Some(scored) = process_suggestion(&suggestion, score, config) {
results.push(scored);
}
}
results
}
fn should_stop_filtering(
results: &[(String, f32)],
deadline: &Instant,
config: &ShellConfig,
) -> bool {
Instant::now() > *deadline || results.len() >= config.max_suggestions
}
fn process_suggestion(suggestion: &str, score: f32, config: &ShellConfig) -> Option<(String, f32)> {
if is_sensitive_command(suggestion) {
return None;
}
let quality = suggestion_quality_score(suggestion);
if quality < config.min_quality_score {
return None;
}
Some((suggestion.to_string(), score * quality))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_default_values() {
let config = ShellConfig::default();
assert_eq!(config.suggest_timeout_ms, 100);
assert_eq!(config.max_model_size, 100 * 1024 * 1024);
assert_eq!(config.max_history_size, 500 * 1024 * 1024);
assert_eq!(config.max_suggestions, 10);
assert_eq!(config.max_prefix_length, 500);
assert_eq!(config.min_prefix_length, 2);
}
#[test]
fn test_builder_pattern() {
let config = ShellConfig::new()
.with_suggest_timeout_ms(50)
.with_max_suggestions(5)
.with_min_quality_score(0.5);
assert_eq!(config.suggest_timeout_ms, 50);
assert_eq!(config.max_suggestions, 5);
assert!((config.min_quality_score - 0.5).abs() < f32::EPSILON);
}
#[test]
fn test_quality_score_clamped() {
let config = ShellConfig::new().with_min_quality_score(1.5);
assert!((config.min_quality_score - 1.0).abs() < f32::EPSILON);
let config = ShellConfig::new().with_min_quality_score(-0.5);
assert!((config.min_quality_score - 0.0).abs() < f32::EPSILON);
}
#[test]
fn test_suggest_timeout_duration() {
let config = ShellConfig::new().with_suggest_timeout_ms(100);
assert_eq!(config.suggest_timeout(), Duration::from_millis(100));
}
#[test]
fn test_model_size_validation() {
let config = ShellConfig::new().with_max_model_size(1024);
assert!(config.is_model_size_valid(512));
assert!(config.is_model_size_valid(1024));
assert!(!config.is_model_size_valid(2048));
}
#[test]
fn test_history_size_validation() {
let config = ShellConfig::new().with_max_history_size(1024);
assert!(config.is_history_size_valid(512));
assert!(!config.is_history_size_valid(2048));
}
#[test]
fn test_prefix_validation() {
let config = ShellConfig::new()
.with_min_prefix_length(2)
.with_max_prefix_length(10);
assert!(!config.is_prefix_valid("a")); assert!(config.is_prefix_valid("ab")); assert!(config.is_prefix_valid("hello")); assert!(config.is_prefix_valid("0123456789")); assert!(!config.is_prefix_valid("01234567890")); }
#[test]
fn test_truncate_prefix() {
let config = ShellConfig::new().with_max_prefix_length(5);
assert_eq!(config.truncate_prefix("abc"), "abc");
assert_eq!(config.truncate_prefix("abcde"), "abcde");
assert_eq!(config.truncate_prefix("abcdefgh"), "abcde");
}
#[test]
fn test_truncate_prefix_utf8_boundary() {
let config = ShellConfig::new().with_max_prefix_length(5);
let jp = "日本";
let truncated = config.truncate_prefix(jp);
assert!(truncated.len() <= 5);
assert!(truncated.is_char_boundary(truncated.len()));
}
#[test]
fn test_fast_preset() {
let config = ShellConfig::fast();
assert_eq!(config.suggest_timeout_ms, 50);
assert_eq!(config.max_suggestions, 5);
}
#[test]
fn test_thorough_preset() {
let config = ShellConfig::thorough();
assert_eq!(config.suggest_timeout_ms, 500);
assert_eq!(config.max_suggestions, 20);
}
#[test]
fn test_clone_and_eq() {
let config1 = ShellConfig::default();
let config2 = config1.clone();
assert_eq!(config1, config2);
}
#[test]
fn test_suggest_without_model_returns_empty() {
let config = ShellConfig::default();
let suggestions = suggest_with_fallback("git ", None, &config);
assert!(suggestions.is_empty());
}
#[test]
fn test_suggest_with_short_prefix_returns_empty() {
let config = ShellConfig::default().with_min_prefix_length(3);
let mut model = MarkovModel::new(3);
model.train(&["git status".to_string()]);
let suggestions = suggest_with_fallback("g", Some(&model), &config);
assert!(suggestions.is_empty());
}
#[test]
fn test_suggest_with_model_returns_results() {
let config = ShellConfig::default();
let mut model = MarkovModel::new(3);
model.train(&[
"git status".to_string(),
"git commit".to_string(),
"git push".to_string(),
]);
let suggestions = suggest_with_fallback("git ", Some(&model), &config);
assert!(!suggestions.is_empty());
}
#[test]
fn test_suggest_respects_max_suggestions() {
let config = ShellConfig::default().with_max_suggestions(2);
let mut model = MarkovModel::new(3);
model.train(&[
"git status".to_string(),
"git commit".to_string(),
"git push".to_string(),
"git pull".to_string(),
"git fetch".to_string(),
]);
let suggestions = suggest_with_fallback("git ", Some(&model), &config);
assert!(suggestions.len() <= 2);
}
#[test]
fn test_suggest_filters_sensitive_commands() {
let config = ShellConfig::default().with_min_quality_score(0.0);
let mut model = MarkovModel::new(3);
model.train(&[
"git status".to_string(),
"export SECRET=abc".to_string(),
"curl -u admin:pass http://localhost".to_string(),
]);
let suggestions = suggest_with_fallback("export ", Some(&model), &config);
for (suggestion, _) in &suggestions {
assert!(!suggestion.contains("SECRET="));
}
}
}