first pass of merging in warp (doesn't build)

crates/input_classifier/src/bin/evaluate.rs

@@ -2,9 +2,9 @@ use std::{fs, process};
 
 use clap::{Parser, Subcommand};
 use futures::executor::block_on;
+use input_classifier::test_utils::CompletionContext;
 use input_classifier::{
     ClassificationResult, Context, HeuristicClassifier, InputClassifier, InputType,
-    test_utils::CompletionContext,
 };
 
 /// Convert HSL to RGB values (0-255 range)
@@ -66,13 +66,24 @@ fn get_binary_confidence_color(is_correct: bool, is_low_confidence: bool) -> Str
         "\x1b[31m".to_string() // Red for incorrect
     }
 }
-use galaxy_completer::{ParsedTokensSnapshot, util::parse_current_commands_and_tokens};
-
-#[cfg(feature = "fasttext")]
-use input_classifier::FasttextClassifier;
-
 #[cfg(feature = "onnx")]
 use input_classifier::{OnnxClassifier, OnnxModel};
+use warp_completer::ParsedTokensSnapshot;
+use warp_completer::util::parse_current_commands_and_tokens;
+#[cfg(feature = "onnx")]
+fn default_onnx_model() -> Option<OnnxModel> {
+    cfg_if::cfg_if! {
+        if #[cfg(feature = "nld_classifier_v1")] {
+            Some(OnnxModel::BertTinyV1)
+        } else if #[cfg(feature = "nld_classifier_v2")] {
+            Some(OnnxModel::BertTinyV2)
+        } else if #[cfg(feature = "nld_classifier_v3")] {
+            Some(OnnxModel::BertTinyV3)
+        } else {
+            None
+        }
+    }
+}
 
 #[derive(Parser)]
 struct InputSource {
@@ -98,11 +109,6 @@ struct Args {
     #[arg(long)]
     heuristic: bool,
 
-    /// Use fasttext classifier
-    #[cfg(feature = "fasttext")]
-    #[arg(long)]
-    fasttext: bool,
-
     /// Use ONNX classifier
     #[cfg(feature = "onnx")]
     #[arg(long)]
@@ -135,17 +141,6 @@ fn create_classifiers(args: &Args) -> Vec<(&'static str, Box<dyn InputClassifier
     let mut classifiers: Vec<(&'static str, Box<dyn InputClassifier>)> = Vec::new();
 
     // Default to all available classifiers if none specified
-    let fasttext_specified = {
-        #[cfg(feature = "fasttext")]
-        {
-            args.fasttext
-        }
-        #[cfg(not(feature = "fasttext"))]
-        {
-            false
-        }
-    };
-
     let onnx_specified = {
         #[cfg(feature = "onnx")]
         {
@@ -157,27 +152,19 @@ fn create_classifiers(args: &Args) -> Vec<(&'static str, Box<dyn InputClassifier
         }
     };
 
-    let use_all = !args.heuristic && !fasttext_specified && !onnx_specified;
+    let use_all = !args.heuristic && !onnx_specified;
 
     if args.heuristic || use_all {
         classifiers.push(("heuristic", Box::new(HeuristicClassifier)));
     }
 
-    #[cfg(feature = "fasttext")]
-    if args.fasttext || use_all {
-        match FasttextClassifier::new() {
-            Ok(classifier) => {
-                classifiers.push(("fasttext", Box::new(classifier)));
-            }
-            Err(e) => {
-                eprintln!("Warning: Failed to initialize FastText classifier: {e}");
-            }
-        }
-    }
-
     #[cfg(feature = "onnx")]
     if args.onnx || use_all {
-        match OnnxClassifier::new(OnnxModel::BertTiny) {
+        let Some(model) = default_onnx_model() else {
+            eprintln!("Warning: No ONNX model feature is enabled for the ONNX classifier");
+            return classifiers;
+        };
+        match OnnxClassifier::new(model) {
             Ok(classifier) => {
                 classifiers.push(("onnx", Box::new(classifier)));
             }
@@ -247,10 +234,13 @@ async fn handle_classify(
             }
             Err(_) => {
                 // Fallback to detect_input_type if classify_input fails
-                let result = classifier
+                let classification = classifier
                     .detect_input_type(parsed_input.clone(), &context)
                     .await;
-                println!("  {name}: {result} (probabilities unavailable)");
+                println!(
+                    "  {}: {} (probabilities unavailable)",
+                    name, classification.input_type
+                );
             }
         }
     }
@@ -302,10 +292,10 @@ async fn handle_verify(
                 }
                 Err(_) => {
                     // Fallback to detect_input_type if classify_input fails
-                    let result = classifier
+                    let classification = classifier
                         .detect_input_type(parsed_input.clone(), &context)
                         .await;
-                    let is_correct = result == expected;
+                    let is_correct = classification.input_type == expected;
                     if is_correct {
                         correct_count += 1;
                     }

crates/input_classifier/src/heuristic_classifier/mod.rs

@@ -5,12 +5,13 @@ use galaxy_completer::ParsedTokensSnapshot;
 use itertools::Itertools as _;
 use natural_language_detection::natural_language_words_score;
 
+use crate::parser::parse_query_into_tokens;
+use crate::util::{
+    is_installed_binary, is_likely_shell_command, is_one_off_natural_language_word_or_prefix,
+};
 use crate::{
-    ClassificationResult, Context, InputClassifier, InputType,
-    parser::parse_query_into_tokens,
-    util::{
-        is_installed_binary, is_likely_shell_command, is_one_off_natural_language_word_or_prefix,
-    },
+    ClassificationResult, Context, InputClassificationResult, InputClassifier,
+    InputClassifierDecisionSource, InputType,
 };
 
 /// Minimum number of characters the input buffer must have before classification runs.
@@ -32,7 +33,11 @@ pub struct HeuristicClassifier;
 #[cfg_attr(not(target_family = "wasm"), async_trait)]
 #[cfg_attr(target_family = "wasm", async_trait(?Send))]
 impl InputClassifier for HeuristicClassifier {
-    async fn detect_input_type(&self, input: ParsedTokensSnapshot, context: &Context) -> InputType {
+    async fn detect_input_type(
+        &self,
+        input: ParsedTokensSnapshot,
+        context: &Context,
+    ) -> InputClassificationResult {
         let word_tokens = parse_query_into_tokens(input.buffer_text.as_str());
         let total_word_token_count = word_tokens.len();
 
@@ -46,23 +51,27 @@ impl InputClassifier for HeuristicClassifier {
         if total_word_token_count == 1
             && is_one_off_natural_language_word_or_prefix(&word_tokens[0].to_lowercase())
         {
-            log::info!("[input-classifier] → one-off NL word, returning AI");
-            return InputType::AI;
+            return InputClassificationResult::new(
+                InputType::AI,
+                InputClassifierDecisionSource::NaturalLanguageOneOffAllowlist,
+            );
         }
 
         if is_likely_shell_command(&input, total_word_token_count).await {
-            log::info!("[input-classifier] → is_likely_shell_command=true, returning Shell");
-            return InputType::Shell;
+            return InputClassificationResult::new(
+                InputType::Shell,
+                InputClassifierDecisionSource::ShellHeuristic,
+            );
         }
 
         let result = self
             .classify_input(input, context)
             .await
-            .map(|result| result.to_input_type())
-            .unwrap_or(context.current_input_type);
-
-        log::info!("[input-classifier] → classify_input returned {:?}", result);
-        result
+            .map(|result| InputClassificationResult::new(result.to_input_type(), result.source))
+            .unwrap_or(InputClassificationResult::new(
+                context.current_input_type,
+                InputClassifierDecisionSource::InputClassifierFallbackHeuristic,
+            ))
     }
 
     async fn classify_input(
@@ -105,10 +114,17 @@ async fn natural_language_detection_heuristic(
     _current_input_type: InputType,
     include_last_token: bool,
 ) -> ClassificationResult {
-    let _word_tokens_count = word_tokens.len();
+    let source = InputClassifierDecisionSource::InputClassifierFallbackHeuristic;
+    let word_tokens_count = word_tokens.len();
 
-    if input.buffer_text.len() < MINIMUM_CLASSIFICATION_CHAR_LENGTH {
-        return ClassificationResult::pure_shell();
+    let min_token_length = if matches!(current_input_type, InputType::AI) {
+        MINIMUM_COMMAND_DETECTION_TOKEN_LENGTH
+    } else {
+        MINIMUM_NATURAL_LANGUAGE_DETECTION_TOKEN_LENGTH
+    };
+
+    if min_token_length > word_tokens_count as u8 {
+        return ClassificationResult::pure_shell(source);
     }
 
     let mut word_tokens = word_tokens.into_iter().map(Cow::Owned).collect_vec();
@@ -136,10 +152,10 @@ async fn natural_language_detection_heuristic(
     };
 
     if likely_english_token_count >= (updated_word_token_count as f32 * threshold) as usize {
-        return ClassificationResult::pure_ai();
+        return ClassificationResult::pure_ai(source);
     }
 
-    ClassificationResult::pure_shell()
+    ClassificationResult::pure_shell(source)
 }
 
 #[cfg(test)]

crates/input_classifier/src/heuristic_classifier/mod_tests.rs

@@ -1,14 +1,52 @@
+use galaxy_completer::meta::SpannedItem;
 use galaxy_completer::util::parse_current_commands_and_tokens;
-
-use crate::{Context, test_utils::CompletionContext};
+use galaxy_completer::{ParsedTokenData, ParsedTokensSnapshot};
 
 use super::*;
+use crate::Context;
+use crate::test_utils::CompletionContext;
 
 async fn mock_parsed_input_token(buffer_text: String) -> ParsedTokensSnapshot {
+    warp_features::mark_initialized();
     let completion_context = CompletionContext::new();
     parse_current_commands_and_tokens(buffer_text, &completion_context).await
 }
 
+fn mock_parsed_input_token_without_descriptions(buffer_text: &str) -> ParsedTokensSnapshot {
+    let mut next_search_start = 0;
+    let parsed_tokens = buffer_text
+        .split_whitespace()
+        .enumerate()
+        .map(|(token_index, token)| {
+            let token_start =
+                buffer_text[next_search_start..].find(token).unwrap() + next_search_start;
+            let token_end = token_start + token.len();
+            next_search_start = token_end;
+
+            ParsedTokenData {
+                token: token.to_string().spanned((token_start, token_end)),
+                token_index,
+                token_description: None,
+            }
+        })
+        .collect();
+
+    ParsedTokensSnapshot {
+        buffer_text: buffer_text.to_string(),
+        parsed_tokens,
+    }
+}
+async fn detected_input_type(
+    classifier: &HeuristicClassifier,
+    input: ParsedTokensSnapshot,
+    context: &Context,
+) -> InputType {
+    classifier
+        .detect_input_type(input, context)
+        .await
+        .input_type
+}
+
 #[test]
 fn test_input_detection() {
     futures::executor::block_on(async move {
@@ -21,7 +59,7 @@ fn test_input_detection() {
 
         let token = mock_parsed_input_token("cargo --version".to_string()).await;
         assert_eq!(
-            classifier.detect_input_type(token, &context).await,
+            detected_input_type(&classifier, token, &context).await,
             InputType::Shell
         );
 
@@ -32,14 +70,14 @@ fn test_input_detection() {
         let mut token = mock_parsed_input_token("cargo --version".to_string()).await;
         token.parsed_tokens[0].token_description = None;
         assert_eq!(
-            classifier.detect_input_type(token, &context).await,
+            detected_input_type(&classifier, token, &context).await,
             InputType::Shell
         );
 
         let mut token = mock_parsed_input_token("rvm install 3.3".to_string()).await;
         token.parsed_tokens[0].token_description = None;
         assert_eq!(
-            classifier.detect_input_type(token, &context).await,
+            detected_input_type(&classifier, token, &context).await,
             InputType::Shell
         );
 
@@ -47,7 +85,7 @@ fn test_input_detection() {
         let mut token = mock_parsed_input_token("Explain this".to_string()).await;
         token.parsed_tokens[0].token_description = None;
         assert_eq!(
-            classifier.detect_input_type(token.clone(), &context).await,
+            detected_input_type(&classifier, token.clone(), &context).await,
             InputType::AI
         );
 
@@ -57,21 +95,21 @@ fn test_input_detection() {
         let mut token = mock_parsed_input_token("fix this".to_string()).await;
         token.parsed_tokens[0].token_description = None;
         assert_eq!(
-            classifier.detect_input_type(token, &context).await,
+            detected_input_type(&classifier, token, &context).await,
             InputType::AI,
         );
 
         // Short queries with punctuation should be parsed as AI input.
         let token = mock_parsed_input_token("What went wrong?".to_string()).await;
         assert_eq!(
-            classifier.detect_input_type(token, &context).await,
+            detected_input_type(&classifier, token, &context).await,
             InputType::AI
         );
         // Short queries with contractions should be parsed as AI input.
         let mut token = mock_parsed_input_token("What's the reason".to_string()).await;
         token.parsed_tokens[0].token_description = None;
         assert_eq!(
-            classifier.detect_input_type(token, &context).await,
+            detected_input_type(&classifier, token, &context).await,
             InputType::AI
         );
 
@@ -80,7 +118,7 @@ fn test_input_detection() {
             mock_parsed_input_token("The message is \"utils::future ... ok\"".to_string()).await;
         token.parsed_tokens[0].token_description = None;
         assert_eq!(
-            classifier.detect_input_type(token, &context).await,
+            detected_input_type(&classifier, token, &context).await,
             InputType::AI
         );
 
@@ -88,8 +126,47 @@ fn test_input_detection() {
         let mut token = mock_parsed_input_token("The type is \"<>\"".to_string()).await;
         token.parsed_tokens[0].token_description = None;
         assert_eq!(
-            classifier.detect_input_type(token, &context).await,
+            detected_input_type(&classifier, token, &context).await,
             InputType::AI
         );
     });
 }
+
+#[test]
+fn test_input_detection_sources() {
+    futures::executor::block_on(async move {
+        let classifier = HeuristicClassifier;
+        let context = Context {
+            current_input_type: InputType::Shell,
+            is_agent_follow_up: false,
+        };
+
+        let token = mock_parsed_input_token_without_descriptions("echo hello");
+        let decision = classifier.detect_input_type(token, &context).await;
+        assert_eq!(
+            decision,
+            InputClassificationResult::new(
+                InputType::Shell,
+                InputClassifierDecisionSource::ShellHeuristic,
+            )
+        );
+        let token = mock_parsed_input_token_without_descriptions("explain");
+        let decision = classifier.detect_input_type(token, &context).await;
+        assert_eq!(
+            decision,
+            InputClassificationResult::new(
+                InputType::AI,
+                InputClassifierDecisionSource::NaturalLanguageOneOffAllowlist,
+            )
+        );
+        let token = mock_parsed_input_token_without_descriptions("fix this");
+        let decision = classifier.detect_input_type(token, &context).await;
+        assert_eq!(
+            decision,
+            InputClassificationResult::new(
+                InputType::AI,
+                InputClassifierDecisionSource::InputClassifierFallbackHeuristic,
+            )
+        );
+    });
+}

crates/input_classifier/src/input_type.rs

@@ -1,6 +1,7 @@
-use serde::{Deserialize, Serialize};
 use std::str::FromStr;
 
+use serde::{Deserialize, Serialize};
+
 /// The type of input the user has provided.
 #[derive(Default, Debug, Copy, Clone, PartialEq, Eq, Serialize, Deserialize)]
 pub enum InputType {

crates/input_classifier/src/lib.rs

@@ -1,5 +1,3 @@
-#[cfg(feature = "fasttext")]
-mod fasttext;
 mod heuristic_classifier;
 mod input_type;
 #[cfg(feature = "onnx")]
@@ -9,13 +7,43 @@ pub mod test_utils;
 pub mod util;
 
 use async_trait::async_trait;
-
-#[cfg(feature = "fasttext")]
-pub use fasttext::FasttextClassifier;
 pub use heuristic_classifier::HeuristicClassifier;
 pub use input_type::InputType;
 #[cfg(feature = "onnx")]
 pub use onnx::{Model as OnnxModel, OnnxClassifier};
+use serde::{Deserialize, Serialize};
+
+/// Sources produced by the input classifier pipeline.
+#[derive(Debug, Copy, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub enum InputClassifierDecisionSource {
+    // Classification result coming from Onnx classifier
+    InputClassifier,
+    // Classification result coming from fall back heuristic classifier when onnx classifier panicked
+    InputClassifierFallbackHeuristic,
+    // Classification result coming from current input type when onnx classifier failed
+    InputClassifierFallbackCurrentInput,
+    // Input match with ONE_OFF_NATURAL_LANGUAGE_WORDS
+    NaturalLanguageOneOffAllowlist,
+    // Input match with ONE_OFF_SHELL_COMMAND_KEYWORDS
+    ShellCommandAllowList,
+    // Classification result coming from is_likely_shell_command
+    ShellHeuristic,
+}
+
+/// The detected input type along with the decision source that produced it.
+#[derive(Debug, Copy, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub struct InputClassificationResult {
+    /// The detected input type.
+    pub input_type: InputType,
+    /// The classifier source that produced this classification.
+    pub source: InputClassifierDecisionSource,
+}
+
+impl InputClassificationResult {
+    pub fn new(input_type: InputType, source: InputClassifierDecisionSource) -> Self {
+        Self { input_type, source }
+    }
+}
 
 /// An input classifier, which can take some parsed user input and determine
 /// what type of input it is.
@@ -26,7 +54,7 @@ pub trait InputClassifier: 'static + Send + Sync {
         &self,
         input: galaxy_completer::ParsedTokensSnapshot,
         context: &Context,
-    ) -> InputType;
+    ) -> InputClassificationResult;
 
     async fn classify_input(
         &self,
@@ -41,20 +69,24 @@ pub struct ClassificationResult {
     p_shell: f32,
     /// The probability that the input is a natural language query to AI.
     p_ai: f32,
+    /// The classifier source that produced this classification.
+    pub source: InputClassifierDecisionSource,
 }
 
 impl ClassificationResult {
-    fn pure_ai() -> Self {
+    fn pure_ai(source: InputClassifierDecisionSource) -> Self {
         Self {
             p_shell: 0.0,
             p_ai: 1.0,
+            source,
         }
     }
 
-    fn pure_shell() -> Self {
+    fn pure_shell(source: InputClassifierDecisionSource) -> Self {
         Self {
             p_shell: 1.0,
             p_ai: 0.0,
+            source,
         }
     }
 

crates/input_classifier/src/onnx/candle.rs

@@ -7,9 +7,8 @@ use galaxy_completer::ParsedTokensSnapshot;
 use prost::Message as _;
 use tokenizers::Tokenizer;
 
-use super::ClassificationResult;
-
-use super::Model;
+use super::{ClassificationResult, Model};
+use crate::InputClassifierDecisionSource;
 
 pub struct InferenceRunner {
     model: ModelProto,
@@ -96,6 +95,7 @@ impl super::InferenceRunner for InferenceRunner {
         Ok(ClassificationResult {
             p_ai: probabilities[0],
             p_shell: probabilities[1],
+            source: InputClassifierDecisionSource::InputClassifier,
         })
     }
 }

crates/input_classifier/src/onnx/mod.rs

@@ -10,21 +10,28 @@ use async_trait::async_trait;
 use galaxy_completer::ParsedTokensSnapshot;
 use rust_embed::RustEmbed;
 
+use crate::parser::parse_query_into_tokens;
+use crate::util::{
+    is_likely_shell_command, is_one_off_natural_language_word, is_one_off_shell_command_keyword,
+};
 use crate::{
-    ClassificationResult, Context, InputClassifier, InputType,
-    parser::parse_query_into_tokens,
-    util::{
-        is_likely_shell_command, is_one_off_natural_language_word, is_one_off_shell_command_keyword,
-    },
+    ClassificationResult, Context, InputClassificationResult, InputClassifier,
+    InputClassifierDecisionSource, InputType,
 };
 
 #[derive(Clone, Copy, RustEmbed)]
 #[folder = "models/onnx"]
+#[include = "bert_tiny_tokenizer.json"]
+#[cfg_attr(feature = "nld_classifier_v1", include = "bert_tiny_v1.onnx")]
+#[cfg_attr(feature = "nld_classifier_v2", include = "bert_tiny_v2.onnx")]
+#[cfg_attr(feature = "nld_classifier_v3", include = "bert_tiny_v3.onnx")]
 struct Models;
 
-#[derive(Copy, Clone)]
+#[derive(Copy, Clone, Debug)]
 pub enum Model {
-    BertTiny,
+    BertTinyV1,
+    BertTinyV2,
+    BertTinyV3,
 }
 
 impl Model {
@@ -38,13 +45,15 @@ impl Model {
 
     fn model_path(&self) -> &'static str {
         match self {
-            Model::BertTiny => "bert_tiny.onnx",
+            Model::BertTinyV1 => "bert_tiny_v1.onnx",
+            Model::BertTinyV2 => "bert_tiny_v2.onnx",
+            Model::BertTinyV3 => "bert_tiny_v3.onnx",
         }
     }
 
     fn tokenizer_path(&self) -> &'static str {
         match self {
-            Model::BertTiny => "bert_tiny_tokenizer.json",
+            Model::BertTinyV1 | Model::BertTinyV2 | Model::BertTinyV3 => "bert_tiny_tokenizer.json",
         }
     }
 }
@@ -85,7 +94,11 @@ impl OnnxClassifier {
 #[cfg_attr(not(target_family = "wasm"), async_trait)]
 #[cfg_attr(target_family = "wasm", async_trait(?Send))]
 impl InputClassifier for OnnxClassifier {
-    async fn detect_input_type(&self, input: ParsedTokensSnapshot, context: &Context) -> InputType {
+    async fn detect_input_type(
+        &self,
+        input: ParsedTokensSnapshot,
+        context: &Context,
+    ) -> InputClassificationResult {
         let word_tokens = parse_query_into_tokens(input.buffer_text.as_str());
 
         let total_word_token_count = word_tokens.len();
@@ -96,36 +109,53 @@ impl InputClassifier for OnnxClassifier {
 
             // If the input is a single word and the word is one of a specific set of words, classify it as AI
             if word_tokens.len() == 1 && is_one_off_natural_language_word(&first_word) {
-                return InputType::AI;
+                return InputClassificationResult::new(
+                    InputType::AI,
+                    InputClassifierDecisionSource::NaturalLanguageOneOffAllowlist,
+                );
             }
 
             // If the first token is one of a specific set of shell command keywords (e.g.: echo or sudo),
             // we should classify it as shell.
             if is_one_off_shell_command_keyword(&first_word) {
-                return InputType::Shell;
+                return InputClassificationResult::new(
+                    InputType::Shell,
+                    InputClassifierDecisionSource::ShellHeuristic,
+                );
             }
         }
 
         if is_likely_shell_command(&input, total_word_token_count).await {
-            return InputType::Shell;
+            return InputClassificationResult::new(
+                InputType::Shell,
+                InputClassifierDecisionSource::ShellHeuristic,
+            );
         }
 
         // Otherwise, defer all decision-making to the model.
         self.classify_input(input, context)
             .await
-            .map(|result| result.to_input_type())
-            .unwrap_or(context.current_input_type)
+            .map(|classification| {
+                InputClassificationResult::new(
+                    classification.to_input_type(),
+                    classification.source,
+                )
+            })
+            .unwrap_or(InputClassificationResult::new(
+                context.current_input_type,
+                InputClassifierDecisionSource::InputClassifierFallbackCurrentInput,
+            ))
     }
 
     async fn classify_input(
         &self,
         input: galaxy_completer::ParsedTokensSnapshot,
-        _context: &Context,
+        context: &Context,
     ) -> anyhow::Result<ClassificationResult> {
         // If we ever panicked while running inference, we should fall back to the heuristic classifier.
         if self.has_panicked.has_panicked() {
             return crate::heuristic_classifier::HeuristicClassifier
-                .classify_input(input, _context)
+                .classify_input(input, context)
                 .await;
         }
 
@@ -163,7 +193,7 @@ impl InputClassifier for OnnxClassifier {
                 );
                 self.has_panicked.on_panic();
                 crate::heuristic_classifier::HeuristicClassifier
-                    .classify_input(input, _context)
+                    .classify_input(input, context)
                     .await
             }
         }
@@ -196,3 +226,7 @@ impl HasPanicked {
         self.inner.is_completed()
     }
 }
+
+#[cfg(test)]
+#[path = "mod_tests.rs"]
+mod tests;

crates/input_classifier/src/onnx/mod_tests.rs

@@ -0,0 +1,64 @@
+use anyhow::Result;
+use futures::executor::block_on;
+use warp_completer::meta::SpannedItem;
+use warp_completer::{ParsedTokenData, ParsedTokensSnapshot};
+
+use super::*;
+
+struct FailingInferenceRunner;
+
+impl InferenceRunner for FailingInferenceRunner {
+    fn run_inference(&self, _input: &ParsedTokensSnapshot) -> Result<ClassificationResult> {
+        Err(anyhow::anyhow!("inference failed"))
+    }
+}
+
+fn parsed_input_without_descriptions(buffer_text: &str) -> ParsedTokensSnapshot {
+    let mut next_search_start = 0;
+    let parsed_tokens = buffer_text
+        .split_whitespace()
+        .enumerate()
+        .map(|(token_index, token)| {
+            let token_start =
+                buffer_text[next_search_start..].find(token).unwrap() + next_search_start;
+            let token_end = token_start + token.len();
+            next_search_start = token_end;
+
+            ParsedTokenData {
+                token: token.to_string().spanned((token_start, token_end)),
+                token_index,
+                token_description: None,
+            }
+        })
+        .collect();
+
+    ParsedTokensSnapshot {
+        buffer_text: buffer_text.to_owned(),
+        parsed_tokens,
+    }
+}
+
+#[test]
+fn test_inference_error_reports_current_input_fallback_source() {
+    block_on(async move {
+        let classifier = OnnxClassifier {
+            inference_runner: Box::new(FailingInferenceRunner),
+            has_panicked: HasPanicked::new(),
+        };
+        let context = Context {
+            current_input_type: InputType::AI,
+            is_agent_follow_up: false,
+        };
+        let input = parsed_input_without_descriptions("help migrate database");
+
+        let decision = classifier.detect_input_type(input, &context).await;
+
+        assert_eq!(
+            decision,
+            InputClassificationResult::new(
+                InputType::AI,
+                InputClassifierDecisionSource::InputClassifierFallbackCurrentInput,
+            )
+        );
+    });
+}

crates/input_classifier/src/onnx/ort.rs

@@ -1,16 +1,15 @@
 use anyhow::{Result, ensure};
 use galaxy_completer::ParsedTokensSnapshot;
 use itertools::Itertools as _;
-use ort::{
-    execution_providers::CPUExecutionProvider, session::Session, tensor::ArrayExtensions as _,
-    value::Value,
-};
+use ort::execution_providers::CPUExecutionProvider;
+use ort::session::Session;
+use ort::tensor::ArrayExtensions as _;
+use ort::value::Value;
 use parking_lot::Mutex;
 use tokenizers::Tokenizer;
 
-use super::ClassificationResult;
-
-use super::Model;
+use super::{ClassificationResult, Model};
+use crate::InputClassifierDecisionSource;
 
 pub struct InferenceRunner {
     session: Mutex<Session>,
@@ -86,6 +85,7 @@ impl super::InferenceRunner for InferenceRunner {
         Ok(ClassificationResult {
             p_ai: probabilities[0],
             p_shell: probabilities[1],
+            source: InputClassifierDecisionSource::InputClassifier,
         })
     }
 }

crates/input_classifier/src/parser.rs

@@ -1,4 +1,6 @@
-use std::{iter::Peekable, mem, str::Chars};
+use std::iter::Peekable;
+use std::mem;
+use std::str::Chars;
 
 use itertools::Itertools;
 
@@ -33,7 +35,7 @@ fn convert_char_to_delimiter(c: char) -> Option<WordDelimiter> {
     }
 }
 
-/// Parse a sentence into tokens for natural language classificiation.
+/// Parse a sentence into tokens for natural language classification.
 /// The main difference from the default "split_whitespace" parser is
 /// for matching double quotes and single quotes, we would keep their
 /// closed content as one token.

crates/input_classifier/src/test_utils.rs

@@ -1,10 +1,9 @@
-use std::{collections::HashSet, sync::Arc};
+use std::collections::HashSet;
+use std::sync::Arc;
 
-use galaxy_completer::{
-    completer::{GeneratorContext, PathCompletionContext},
-    signatures::CommandRegistry,
-};
 use smol_str::SmolStr;
+use galaxy_completer::completer::{GeneratorContext, PathCompletionContext};
+use galaxy_completer::signatures::CommandRegistry;
 
 /// An implementation of `CompletionContext` for testing purposes.
 pub struct CompletionContext {

crates/input_classifier/src/util.rs

@@ -2,7 +2,6 @@ use std::collections::HashSet;
 
 use galaxy_completer::ParsedTokensSnapshot;
 use lazy_static::lazy_static;
-use natural_language_detection::check_if_token_has_shell_syntax;
 
 /// The percentage of input tokens that can be described by our completion engine before
 /// we consider the input as a shell command. This could be tuned.
@@ -20,13 +19,13 @@ lazy_static! {
     /// claude code, codex CLI, or gemini CLI) suck, because the user often thinks we're
     /// intentionally trying to push them away from those CLIs into Agent Mode, so we mitigate the
     /// risk by always treating as shell.
-    static ref ONE_OFF_SHELL_COMMAND_KEYWORDS: HashSet<&'static str> = HashSet::from(["#", "echo", "man", "sudo", "claude", "codex", "gemini"]);
+    static ref ONE_OFF_SHELL_COMMAND_KEYWORDS: HashSet<&'static str> = HashSet::from(["#", "echo", "man", "sudo", "claude", "codex", "gemini", "agy"]);
 
     static ref ONE_OFF_NATURAL_LANGUAGE_WORDS: HashSet<&'static str> = HashSet::from(["hello", "hi", "hey", "hola", "thanks", "explain", "yes", "no", "what", "nice", "1. "]);
 
     /// A set of words that should trigger an AI classification if they are the entire input
     /// and the input is a follow-up to an agent response.
-    static ref AGENT_FOLLOW_UP_INPUTS: HashSet<&'static str> = HashSet::from(["yes", "continue", "do it"]);
+    static ref AGENT_FOLLOW_UP_INPUTS: HashSet<&'static str> = HashSet::from(["yes", "continue", "do it", "approve"]);
 }
 
 pub fn is_agent_follow_up_input(input: &str) -> bool {
@@ -56,33 +55,47 @@ pub fn is_prefix_of_natural_language_word(input: &str) -> bool {
         .any(|word| word.starts_with(input))
 }
 
+/// nld_heuristic_v1: current prod, use check_if_token_has_shell_syntax and conditional threshold on input length
+/// nld_heuristic_v2: rm check_if_token_has_shell_syntax and pin threshold to be 1 for all input
 pub async fn is_likely_shell_command(
     input: &ParsedTokensSnapshot,
     _word_tokens_count: usize,
 ) -> bool {
     const YIELD_BATCH_SIZE: usize = 5;
+    let use_nld_heuristic_v2 = cfg!(feature = "nld_heuristic_v2");
 
     let mut likely_command_token_count = 0;
     let total_token_count = input.parsed_tokens.len();
+    let mut is_first_token_command = false;
+    log::debug!(
+        "is_likely_shell_command start: use_nld_heuristic_v2={use_nld_heuristic_v2}, total_token_count={total_token_count}, word_tokens_count={word_tokens_count}"
+    );
     for (idx, token) in input.parsed_tokens.iter().enumerate() {
         if idx % YIELD_BATCH_SIZE == 0 {
             futures_lite::future::yield_now().await;
         }
         if token.token_index == 0 && ONE_OFF_SHELL_COMMAND_KEYWORDS.contains(&token.token.as_str())
         {
+            log::debug!(
+                "is_likely_shell_command result=true: first token is one-off shell keyword, use_nld_heuristic_v2={use_nld_heuristic_v2}"
+            );
             return true;
         }
-
-        if token.token_description.is_some()
-            || check_if_token_has_shell_syntax(token.token.as_str())
-        {
+        let check_if_token_has_shell_syntax = !use_nld_heuristic_v2
+            && natural_language_detection::check_if_token_has_shell_syntax(token.token.as_str());
+        log::debug!(
+            "is_likely_shell_command token: token_index={}, token_description_is_some={}, check_if_token_has_shell_syntax={check_if_token_has_shell_syntax}, use_nld_heuristic_v2={use_nld_heuristic_v2}",
+            token.token_index,
+            token.token_description.is_some()
+        );
+        if token.token_description.is_some() || check_if_token_has_shell_syntax {
             likely_command_token_count += 1;
         }
     }
 
     // When token count is lower than 2, we should make sure all tokens
     // are matching the target classification category.
-    let command_threshold = if total_token_count <= 2 {
+    let command_threshold = if use_nld_heuristic_v2 || total_token_count <= 2 {
         1.0
     } else if total_token_count <= 4 {
         DETECT_AS_COMMAND_LOW_TOKEN_THRESHOLD
@@ -90,17 +103,17 @@ pub async fn is_likely_shell_command(
         DETECT_AS_COMMAND_THRESHOLD
     };
 
-    let threshold_count = (total_token_count as f32 * command_threshold) as usize;
-    let is_shell = likely_command_token_count >= threshold_count;
-    log::info!(
-        "[input-classifier] is_likely_shell_command: tokens={}, cmd_tokens={}, threshold={:.2} (need {}), result={}",
-        total_token_count,
-        likely_command_token_count,
-        command_threshold,
-        threshold_count,
-        is_shell
+    // Classify as shell if:
+    // 1) We hit significant threshold of likely shell command tokens.
+    // 2) When there are fewer than 3 tokens, the first token is a valid top-level command.
+    let is_likely_shell_command = likely_command_token_count
+        >= (total_token_count as f32 * command_threshold) as usize
+        || (word_tokens_count < 3 && is_first_token_command);
+    log::debug!(
+        "is_likely_shell_command result={is_likely_shell_command}: use_nld_heuristic_v2={use_nld_heuristic_v2}, likely_command_token_count={likely_command_token_count}, total_token_count={total_token_count}, word_tokens_count={word_tokens_count}, command_threshold={command_threshold}, is_first_token_command={is_first_token_command}"
     );
-    is_shell
+
+    is_likely_shell_command
 }
 
 /// Returns true if the first token is a command that is installed on the system.
@@ -111,3 +124,7 @@ pub fn is_installed_binary(input: &ParsedTokensSnapshot) -> bool {
         .map(|token| token.token_description.is_some())
         .unwrap_or(false)
 }
+
+#[cfg(all(test, any(feature = "nld_heuristic_v1", feature = "nld_heuristic_v2")))]
+#[path = "util_tests.rs"]
+mod tests;

crates/input_classifier/src/util_tests.rs

@@ -0,0 +1,242 @@
+use warp_completer::ParsedTokensSnapshot;
+use warp_completer::util::parse_current_commands_and_tokens;
+
+use super::*;
+use crate::test_utils::CompletionContext;
+
+async fn mock_parsed_input_token(buffer_text: String) -> ParsedTokensSnapshot {
+    warp_features::mark_initialized();
+    let completion_context = CompletionContext::new();
+    parse_current_commands_and_tokens(buffer_text, &completion_context).await
+}
+
+fn clear_all_token_descriptions(snapshot: &mut ParsedTokensSnapshot) {
+    for token in snapshot.parsed_tokens.iter_mut() {
+        token.token_description = None;
+    }
+}
+
+async fn one_off_keyword_short_circuits() {
+    let mut token = mock_parsed_input_token("sudo apt update".to_string()).await;
+    let word_tokens_count = token.parsed_tokens.len();
+    clear_all_token_descriptions(&mut token);
+    assert!(is_likely_shell_command(&token, word_tokens_count).await);
+
+    let mut token = mock_parsed_input_token("echo hello world".to_string()).await;
+    let word_tokens_count = token.parsed_tokens.len();
+    clear_all_token_descriptions(&mut token);
+    assert!(is_likely_shell_command(&token, word_tokens_count).await);
+
+    let mut token = mock_parsed_input_token("agy doctor".to_string()).await;
+    let word_tokens_count = token.parsed_tokens.len();
+    clear_all_token_descriptions(&mut token);
+    assert!(is_likely_shell_command(&token, word_tokens_count).await);
+}
+
+async fn first_token_with_description_short_input_is_shell() {
+    let token = mock_parsed_input_token("cargo --version".to_string()).await;
+    assert!(is_likely_shell_command(&token, 2).await);
+}
+
+async fn no_descriptions_returns_false() {
+    let mut token = mock_parsed_input_token("install --foo=bar baz".to_string()).await;
+    let word_tokens_count = token.parsed_tokens.len();
+    clear_all_token_descriptions(&mut token);
+    assert!(!is_likely_shell_command(&token, word_tokens_count).await);
+}
+
+async fn shell_syntax_tokens_with_only_first_token_description() -> bool {
+    let mut token = mock_parsed_input_token("git --foo=bar /path/to/file --baz".to_string()).await;
+    let word_tokens_count = token.parsed_tokens.len();
+
+    for (idx, token) in token.parsed_tokens.iter_mut().enumerate() {
+        if idx != 0 {
+            token.token_description = None;
+        }
+    }
+
+    assert!(word_tokens_count >= 3);
+    is_likely_shell_command(&token, word_tokens_count).await
+}
+
+async fn described_token_majority_below_v2_threshold() -> bool {
+    let mut token = mock_parsed_input_token("cargo build --release --workspace".to_string()).await;
+    let word_tokens_count = token.parsed_tokens.len();
+    assert!(word_tokens_count >= 3);
+
+    let description = token
+        .parsed_tokens
+        .iter()
+        .find_map(|token| token.token_description.clone())
+        .expect("test input should include at least one described token");
+    for token in token.parsed_tokens.iter_mut() {
+        token.token_description = Some(description.clone());
+    }
+    token
+        .parsed_tokens
+        .last_mut()
+        .expect("test input should include tokens")
+        .token_description = None;
+
+    is_likely_shell_command(&token, word_tokens_count).await
+}
+
+async fn downloads_log_path_in_nl_prompt_is_shell() -> bool {
+    let command_token = mock_parsed_input_token("cargo --version".to_string()).await;
+    let command_description = command_token
+        .parsed_tokens
+        .first()
+        .and_then(|token| token.token_description.clone())
+        .expect("test input should include a described command token");
+
+    let mut token = mock_parsed_input_token(
+        "look at this /users/ewanlockwood/downloads/logs_58498936986".to_string(),
+    )
+    .await;
+    let word_tokens_count = token.parsed_tokens.len();
+    clear_all_token_descriptions(&mut token);
+    token
+        .parsed_tokens
+        .first_mut()
+        .expect("test input should include tokens")
+        .token_description = Some(command_description);
+    is_likely_shell_command(&token, word_tokens_count).await
+}
+
+async fn file_path_in_nl_prompt_is_shell() -> bool {
+    let mut token =
+        mock_parsed_input_token("look at this /users/foo/bar.log file".to_string()).await;
+    let word_tokens_count = token.parsed_tokens.len();
+    clear_all_token_descriptions(&mut token);
+    is_likely_shell_command(&token, word_tokens_count).await
+}
+
+async fn majority_described_tokens_returns_true() {
+    let token =
+        mock_parsed_input_token("cargo build --release --workspace --all-features".to_string())
+            .await;
+    let word_tokens_count = token.parsed_tokens.len();
+    assert!(is_likely_shell_command(&token, word_tokens_count).await);
+}
+
+// Cases where nld_heuristic_v1 and nld_heuristic_v2 should both mark input as shell.
+#[cfg(all(feature = "nld_heuristic_v1", not(feature = "nld_heuristic_v2")))]
+#[test]
+fn test_is_likely_shell_command_one_off_keyword_short_circuits_true_for_nld_heuristic_v1() {
+    futures::executor::block_on(one_off_keyword_short_circuits());
+}
+
+#[cfg(feature = "nld_heuristic_v2")]
+#[test]
+fn test_is_likely_shell_command_one_off_keyword_short_circuits_true_for_nld_heuristic_v2() {
+    futures::executor::block_on(one_off_keyword_short_circuits());
+}
+
+#[cfg(all(feature = "nld_heuristic_v1", not(feature = "nld_heuristic_v2")))]
+#[test]
+fn test_is_likely_shell_command_first_token_with_description_short_input_true_for_nld_heuristic_v1()
+{
+    futures::executor::block_on(first_token_with_description_short_input_is_shell());
+}
+
+#[cfg(feature = "nld_heuristic_v2")]
+#[test]
+fn test_is_likely_shell_command_first_token_with_description_short_input_true_for_nld_heuristic_v2()
+{
+    futures::executor::block_on(first_token_with_description_short_input_is_shell());
+}
+
+#[cfg(all(feature = "nld_heuristic_v1", not(feature = "nld_heuristic_v2")))]
+#[test]
+fn test_is_likely_shell_command_majority_described_tokens_true_for_nld_heuristic_v1() {
+    futures::executor::block_on(majority_described_tokens_returns_true());
+}
+
+#[cfg(feature = "nld_heuristic_v2")]
+#[test]
+fn test_is_likely_shell_command_majority_described_tokens_true_for_nld_heuristic_v2() {
+    futures::executor::block_on(majority_described_tokens_returns_true());
+}
+
+// Cases where nld_heuristic_v1 and nld_heuristic_v2 should both not mark input as shell.
+#[cfg(all(feature = "nld_heuristic_v1", not(feature = "nld_heuristic_v2")))]
+#[test]
+fn test_is_likely_shell_command_no_descriptions_false_for_nld_heuristic_v1() {
+    futures::executor::block_on(no_descriptions_returns_false());
+}
+
+#[cfg(feature = "nld_heuristic_v2")]
+#[test]
+fn test_is_likely_shell_command_no_descriptions_false_for_nld_heuristic_v2() {
+    futures::executor::block_on(no_descriptions_returns_false());
+}
+
+#[cfg(all(feature = "nld_heuristic_v1", not(feature = "nld_heuristic_v2")))]
+#[test]
+fn test_is_likely_shell_command_file_path_in_nl_prompt_false_for_nld_heuristic_v1() {
+    futures::executor::block_on(async move {
+        assert!(!file_path_in_nl_prompt_is_shell().await);
+    });
+}
+
+#[cfg(feature = "nld_heuristic_v2")]
+#[test]
+fn test_is_likely_shell_command_file_path_in_nl_prompt_false_for_nld_heuristic_v2() {
+    futures::executor::block_on(async move {
+        assert!(!file_path_in_nl_prompt_is_shell().await);
+    });
+}
+
+// Cases where nld_heuristic_v1 should mark input as shell and stricter
+// nld_heuristic_v2 should not mark input as shell.
+#[cfg(all(feature = "nld_heuristic_v1", not(feature = "nld_heuristic_v2")))]
+#[test]
+fn test_is_likely_shell_command_shell_syntax_votes_true_for_nld_heuristic_v1() {
+    futures::executor::block_on(async move {
+        assert!(shell_syntax_tokens_with_only_first_token_description().await);
+    });
+}
+
+#[cfg(feature = "nld_heuristic_v2")]
+#[test]
+fn test_is_likely_shell_command_shell_syntax_does_not_vote_false_for_nld_heuristic_v2() {
+    futures::executor::block_on(async move {
+        assert!(!shell_syntax_tokens_with_only_first_token_description().await);
+    });
+}
+
+#[cfg(all(feature = "nld_heuristic_v1", not(feature = "nld_heuristic_v2")))]
+#[test]
+fn test_is_likely_shell_command_described_token_majority_true_for_nld_heuristic_v1() {
+    futures::executor::block_on(async move {
+        assert!(described_token_majority_below_v2_threshold().await);
+    });
+}
+
+#[cfg(feature = "nld_heuristic_v2")]
+#[test]
+fn test_is_likely_shell_command_described_token_majority_false_for_nld_heuristic_v2() {
+    futures::executor::block_on(async move {
+        assert!(!described_token_majority_below_v2_threshold().await);
+    });
+}
+
+#[cfg(all(feature = "nld_heuristic_v1", not(feature = "nld_heuristic_v2")))]
+#[test]
+fn test_is_likely_shell_command_downloads_log_path_true_for_nld_heuristic_v1() {
+    futures::executor::block_on(async move {
+        assert!(downloads_log_path_in_nl_prompt_is_shell().await);
+    });
+}
+
+#[cfg(feature = "nld_heuristic_v2")]
+#[test]
+fn test_is_likely_shell_command_downloads_log_path_false_for_nld_heuristic_v2() {
+    futures::executor::block_on(async move {
+        assert!(!downloads_log_path_in_nl_prompt_is_shell().await);
+    });
+}
+
+// No inverse section is expected: nld_heuristic_v2 only removes shell-syntax
+// voting and raises the command threshold, so it should not mark an input as
+// shell when nld_heuristic_v1 does not.