use std::collections::{BTreeMap, BTreeSet};
use std::path::{Path, PathBuf};
use serde::{Deserialize, Serialize};
use crate::config::Config;
const NUDGE_CLASSIFIER_FILE: &str = "nudge.toml";
const BUNDLED_NUDGE_CLASSIFIER: &str = include_str!("classifiers/nudge.toml");
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum NudgeClass {
PendingAction,
PlanUpdate,
FinalAnswer,
Unknown,
}
#[derive(Debug, Clone, PartialEq)]
pub struct NudgeClassification {
pub class: NudgeClass,
pub score: f32,
}
impl NudgeClassification {
pub fn is_pending_action(&self) -> bool {
matches!(
self.class,
NudgeClass::PendingAction | NudgeClass::PlanUpdate
)
}
pub fn is_plan_update(&self) -> bool {
self.class == NudgeClass::PlanUpdate
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct NudgeClassifierConfig {
#[serde(default = "default_classifier_version")]
pub version: u32,
#[serde(default = "default_min_score")]
pub min_score: f32,
#[serde(default = "default_min_margin")]
pub min_margin: f32,
#[serde(default)]
pub classes: BTreeMap<String, NudgeClassConfig>,
}
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
pub struct NudgeClassConfig {
#[serde(default)]
pub matchers: Vec<String>,
#[serde(default)]
pub nudge: String,
}
impl Default for NudgeClassifierConfig {
fn default() -> Self {
bundled_nudge_config()
}
}
fn default_classifier_version() -> u32 {
1
}
fn default_min_score() -> f32 {
0.28
}
fn default_min_margin() -> f32 {
0.03
}
fn bundled_nudge_config() -> NudgeClassifierConfig {
toml::from_str(BUNDLED_NUDGE_CLASSIFIER).expect("bundled nudge classifier template is valid")
}
fn bundled_nudge_classes() -> BTreeMap<String, NudgeClassConfig> {
bundled_nudge_config().classes
}
impl NudgeClassifierConfig {
pub fn load_file(path: &Path) -> anyhow::Result<Self> {
match std::fs::read_to_string(path) {
Ok(text) => {
let cfg: Self = toml::from_str(&text)?;
Ok(cfg.with_builtin_fallbacks())
}
Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(Self::default()),
Err(e) => Err(e.into()),
}
}
pub fn load_from_dir(dir: &Path) -> anyhow::Result<Self> {
Self::load_file(&dir.join(NUDGE_CLASSIFIER_FILE))
}
fn with_builtin_fallbacks(mut self) -> Self {
let builtins = bundled_nudge_classes();
for (class, bundled) in builtins {
self.classes
.entry(class)
.and_modify(|configured| {
if configured.matchers.is_empty() {
configured.matchers = bundled.matchers.clone();
}
if configured.nudge.trim().is_empty() {
configured.nudge = bundled.nudge.clone();
}
})
.or_insert(bundled);
}
if self.min_score <= 0.0 {
self.min_score = default_min_score();
}
if self.min_margin < 0.0 {
self.min_margin = default_min_margin();
}
self
}
}
#[derive(Debug, Clone)]
pub struct NudgeClassifier {
cfg: NudgeClassifierConfig,
}
impl Default for NudgeClassifier {
fn default() -> Self {
Self::builtin()
}
}
impl NudgeClassifier {
pub fn builtin() -> Self {
Self {
cfg: NudgeClassifierConfig::default(),
}
}
pub fn from_config(cfg: NudgeClassifierConfig) -> Self {
Self {
cfg: cfg.with_builtin_fallbacks(),
}
}
pub fn load_from_dir(dir: &Path) -> anyhow::Result<Self> {
Ok(Self::from_config(NudgeClassifierConfig::load_from_dir(
dir,
)?))
}
pub fn load_default() -> Self {
#[cfg(test)]
{
Self::builtin()
}
#[cfg(not(test))]
{
let Some(dir) = classifier_config_dir() else {
return Self::builtin();
};
match Self::load_from_dir(&dir) {
Ok(classifier) => classifier,
Err(e) => {
tracing::warn!(
path = %dir.join(NUDGE_CLASSIFIER_FILE).display(),
error = %e,
"failed to load nudge classifier config; using built-ins"
);
Self::builtin()
}
}
}
}
pub fn classify(&self, text: &str) -> NudgeClassification {
let query = tokens(text);
if query.is_empty() {
return NudgeClassification {
class: NudgeClass::Unknown,
score: 0.0,
};
}
let mut scored: Vec<(NudgeClass, f32)> = self
.cfg
.classes
.iter()
.filter_map(|(class, class_cfg)| {
let class = parse_nudge_class(class)?;
let best = class_cfg
.matchers
.iter()
.map(|example| prototype_similarity(&query, &tokens(example)))
.fold(0.0_f32, f32::max);
Some((class, best))
})
.collect();
scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
let Some((class, score)) = scored.first().copied() else {
return NudgeClassification {
class: NudgeClass::Unknown,
score: 0.0,
};
};
let runner_up = scored.get(1).map(|(_, s)| *s).unwrap_or(0.0);
if score < self.cfg.min_score || (score - runner_up) < self.cfg.min_margin {
return NudgeClassification {
class: NudgeClass::Unknown,
score,
};
}
NudgeClassification { class, score }
}
pub fn is_pending_action(&self, text: &str) -> bool {
self.classify(text).is_pending_action()
}
pub fn direction_for(&self, class: NudgeClass) -> Option<&str> {
self.cfg
.classes
.get(nudge_class_key(class))
.map(|class_cfg| class_cfg.nudge.as_str())
.map(str::trim)
.filter(|direction| !direction.is_empty())
}
}
pub fn classifier_config_dir() -> Option<PathBuf> {
Config::user_config_dir().map(|dir| dir.join("classifiers"))
}
fn parse_nudge_class(s: &str) -> Option<NudgeClass> {
match s.trim().to_ascii_lowercase().replace('-', "_").as_str() {
"pending_action" | "continue" | "continuation" => Some(NudgeClass::PendingAction),
"plan_update" | "update_plan" | "stale_plan" | "replan" => Some(NudgeClass::PlanUpdate),
"final_answer" | "final" | "done" => Some(NudgeClass::FinalAnswer),
_ => None,
}
}
fn nudge_class_key(class: NudgeClass) -> &'static str {
match class {
NudgeClass::PendingAction => "pending_action",
NudgeClass::PlanUpdate => "plan_update",
NudgeClass::FinalAnswer => "final_answer",
NudgeClass::Unknown => "unknown",
}
}
fn tokens(s: &str) -> BTreeSet<String> {
s.split(|c: char| !c.is_alphanumeric())
.filter_map(|raw| {
let t = raw.trim().to_ascii_lowercase();
(t.len() >= 3).then_some(t)
})
.collect()
}
fn jaccard(a: &BTreeSet<String>, b: &BTreeSet<String>) -> f32 {
if a.is_empty() || b.is_empty() {
return 0.0;
}
let intersection = a.intersection(b).count() as f32;
let union = a.union(b).count() as f32;
if union == 0.0 {
0.0
} else {
intersection / union
}
}
fn prototype_similarity(query: &BTreeSet<String>, prototype: &BTreeSet<String>) -> f32 {
if query.is_empty() || prototype.is_empty() {
return 0.0;
}
let overlap = query.intersection(prototype).count() as f32;
let prototype_recall = overlap / prototype.len() as f32;
jaccard(query, prototype).max(prototype_recall)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn nudge_classifier_builtin_matches_known_phrases() {
let classifier = NudgeClassifier::builtin();
assert!(classifier.is_pending_action(
"Plan is current — no update needed. Continuing with step 2: inserting the progressive dispatch into lib.rs."
));
assert!(classifier.is_pending_action(
"\
Summary
Progress so far: Step 2 is in progress.
Current blocker: The for loop needs to iterate over the slice directly.
The fix is to use for entry in page.entries.iter().
Next steps needed:
1. Fix the iteration type error in lib.rs.
2. Add tests.
3. Run just check."
));
assert!(classifier.is_pending_action("Now I'll add the --home flag to the Cli struct."));
assert!(classifier.is_pending_action("Let me keep editing now."));
assert!(classifier.is_pending_action(
"I found the issue - there's an extra closing brace } on line 809 of help_sections.rs that's causing a syntax error. I need to remove this stray brace."
));
assert!(classifier.is_pending_action(
"I have two issues: duplicate topic_has_rollups and a stray brace. Let me fix both — read around 490 to see what needs removing, then verify with a build check."
));
assert!(!classifier.is_pending_action("The capital of France is Paris."));
assert!(!classifier.is_pending_action("Done. Let me know if you want any further changes."));
assert!(!classifier.is_pending_action(
"The duplicate helper definitions and stray brace were removed, and the build check passed."
));
assert!(
classifier
.direction_for(NudgeClass::PendingAction)
.is_some_and(|d| d.contains("emit the tool call now")),
"bundled classifier should carry output direction text"
);
}
#[test]
fn nudge_classifier_separates_plan_update_summaries_from_final_summaries() {
let classifier = NudgeClassifier::builtin();
let findings = "\
Summary of Findings
Across the tool calls, I observed two issues in newt-tui/src/help_sections.rs:
1. Duplicate function
2. Stray closing brace
Current Status
The build is broken due to these syntax errors. The plan was at step 2, but we need to fix the immediate compilation issues first.
Next Steps Required
To continue, I would need to remove the duplicate function using edit_file, remove the stray brace using edit_file, verify cargo check, then proceed with step 2 of the plan.
However, I've reached the tool-call limit and cannot make these edits now.";
let classified = classifier.classify(findings);
assert_eq!(classified.class, NudgeClass::PlanUpdate);
assert!(classified.is_pending_action());
assert!(classified.is_plan_update());
let resume_handoff = "\
Summary
I reached the tool-call limit. Current state of newt-tui/src/help_sections.rs:
duplicate topic_has_rollups and rollup_page_for_topic definitions need to be removed.
Recommended next action if session resumes:
1. Fix the duplicate functions.
2. Clean up the broken test block.
3. Read lib.rs and wire the progressive dispatch.
The build is currently broken due to the duplicate definitions — that's the blocker for any further progress.";
assert_eq!(
classifier.classify(resume_handoff).class,
NudgeClass::PlanUpdate
);
let final_summary =
classifier.classify("Here is a summary of what I found across the tool calls.");
assert_eq!(final_summary.class, NudgeClass::FinalAnswer);
}
#[test]
fn nudge_classifier_loads_dropin_from_classifiers_dir() {
let dir = tempfile::tempdir().unwrap();
let classifiers = dir.path().join("classifiers");
std::fs::create_dir_all(&classifiers).unwrap();
std::fs::write(
classifiers.join(NUDGE_CLASSIFIER_FILE),
r#"
version = 1
min_score = 0.20
min_margin = 0.01
[classes.pending_action]
matchers = ["Proceeding with the patch by editing the target file."]
nudge = "Call the edit tool now."
[classes.final_answer]
matchers = ["No changes are needed."]
"#,
)
.unwrap();
let classifier = NudgeClassifier::load_from_dir(&classifiers).unwrap();
assert!(
classifier.is_pending_action("Proceeding with the patch by editing the target file.")
);
assert!(!classifier.is_pending_action("No changes are needed."));
assert_eq!(
classifier.direction_for(NudgeClass::PendingAction),
Some("Call the edit tool now.")
);
assert!(
classifier
.direction_for(NudgeClass::PlanUpdate)
.is_some_and(|d| d.contains("Update the plan first")),
"partial user configs still inherit bundled nudges"
);
}
#[test]
fn classifier_dir_is_under_user_config_dir() {
let dir = classifier_config_dir().unwrap_or_else(|| PathBuf::from(".newt/classifiers"));
assert!(dir.ends_with("classifiers"));
}
}