Skip to main content

lean_ctx/core/gain/
task_classifier.rs

1use serde::{Deserialize, Serialize};
2
3#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
4pub enum TaskCategory {
5    Coding,
6    Debugging,
7    Refactoring,
8    Testing,
9    Exploration,
10    Planning,
11    Delegation,
12    Git,
13    BuildDeploy,
14    Knowledge,
15    Architecture,
16    Review,
17    General,
18}
19
20impl TaskCategory {
21    pub fn label(&self) -> &'static str {
22        match self {
23            TaskCategory::Coding => "Coding",
24            TaskCategory::Debugging => "Debugging",
25            TaskCategory::Refactoring => "Refactoring",
26            TaskCategory::Testing => "Testing",
27            TaskCategory::Exploration => "Exploration",
28            TaskCategory::Planning => "Planning",
29            TaskCategory::Delegation => "Delegation",
30            TaskCategory::Git => "Git",
31            TaskCategory::BuildDeploy => "Build/Deploy",
32            TaskCategory::Knowledge => "Knowledge",
33            TaskCategory::Architecture => "Architecture",
34            TaskCategory::Review => "Review",
35            TaskCategory::General => "General",
36        }
37    }
38}
39
40pub struct TaskClassifier;
41
42impl TaskClassifier {
43    pub fn classify_tool(tool_name: &str) -> TaskCategory {
44        let t = normalize(tool_name);
45        match t.as_str() {
46            "ctx_edit" | "ctx_fill" => TaskCategory::Refactoring,
47            "ctx_read" | "ctx_multi_read" | "ctx_smart_read" | "ctx_delta" | "ctx_tree"
48            | "ctx_search" | "ctx_outline" | "ctx_graph" | "ctx_callgraph" => {
49                TaskCategory::Exploration
50            }
51            "ctx_semantic_search" | "ctx_architecture" | "ctx_impact" => TaskCategory::Architecture,
52            "ctx_overview" | "ctx_preload" | "ctx_task" | "ctx_intent" | "ctx_workflow" => {
53                TaskCategory::Planning
54            }
55            "ctx_handoff" | "ctx_agent" | "ctx_share" => TaskCategory::Delegation,
56            "ctx_session" | "ctx_knowledge" | "ctx_compress_memory" => TaskCategory::Knowledge,
57            "ctx_cost" | "ctx_gain" | "ctx_metrics" | "ctx_heatmap" => TaskCategory::Review,
58            "ctx_shell" | "ctx_execute" => TaskCategory::Debugging,
59            _ => TaskCategory::General,
60        }
61    }
62
63    pub fn classify_command_key(cmd_key: &str) -> TaskCategory {
64        let k = normalize(cmd_key);
65        if k.is_empty() {
66            return TaskCategory::General;
67        }
68
69        // lean-ctx records its own activity as "commands" too: MCP tool calls
70        // (`ctx_*`) and CLI read-mode / hook keys (`cli_*`). Route those through the
71        // tool/mode classifier so reads land in Exploration, edits in Refactoring, etc.
72        // — instead of collapsing everything into General (the shell heuristics below
73        // only understand real shell commands like git/cargo/grep).
74        if k.starts_with("ctx_") {
75            return Self::classify_tool(&k);
76        }
77        if let Some(mode) = k.strip_prefix("cli_") {
78            return Self::classify_cli_mode(mode);
79        }
80
81        if k.starts_with("git ") || k == "git" {
82            return TaskCategory::Git;
83        }
84
85        if k.starts_with("cargo ") {
86            let sub = k.trim_start_matches("cargo ").trim();
87            if matches!(sub, "test" | "nextest" | "llvm-cov" | "tarpaulin") {
88                return TaskCategory::Testing;
89            }
90            if matches!(sub, "build" | "check" | "clippy" | "fmt" | "run" | "doc") {
91                return TaskCategory::BuildDeploy;
92            }
93            return TaskCategory::BuildDeploy;
94        }
95
96        if k.contains("test") || k.contains("pytest") || k.contains("jest") || k.contains("vitest")
97        {
98            return TaskCategory::Testing;
99        }
100        if k.contains("build")
101            || k.contains("deploy")
102            || k.contains("docker")
103            || k.contains("compose")
104            || k.contains("kubectl")
105            || k.contains("helm")
106            || k.contains("terraform")
107        {
108            return TaskCategory::BuildDeploy;
109        }
110        if k.contains("lint") || k.contains("clippy") || k.contains("fmt") || k.contains("format") {
111            return TaskCategory::BuildDeploy;
112        }
113        if k.contains("grep") || k.contains("rg") || k.contains("ripgrep") {
114            return TaskCategory::Exploration;
115        }
116
117        TaskCategory::General
118    }
119
120    /// Classifies a CLI command key with the `cli_` prefix already stripped. These are the
121    /// shell-hook compression modes: read/inspect modes and search map to Exploration; the
122    /// catch-all `shell` bucket aggregates arbitrary shell commands, so it stays General.
123    fn classify_cli_mode(mode: &str) -> TaskCategory {
124        match mode {
125            "grep" | "rg" | "ripgrep" | "search" | "find" | "ls" | "tree" => {
126                TaskCategory::Exploration
127            }
128            "full" | "map" | "signatures" | "aggressive" | "entropy" | "diff" | "lines"
129            | "reference" | "task" | "auto" | "outline" | "read" => TaskCategory::Exploration,
130            _ => TaskCategory::General,
131        }
132    }
133}
134
135fn normalize(s: &str) -> String {
136    s.trim().to_lowercase()
137}
138
139#[cfg(test)]
140mod tests {
141    use super::*;
142
143    #[test]
144    fn classify_git() {
145        assert_eq!(
146            TaskClassifier::classify_command_key("git status"),
147            TaskCategory::Git
148        );
149    }
150
151    #[test]
152    fn classify_tools() {
153        assert_eq!(
154            TaskClassifier::classify_tool("ctx_semantic_search"),
155            TaskCategory::Architecture
156        );
157    }
158
159    #[test]
160    fn ctx_command_keys_route_through_tool_classifier() {
161        // The regression behind the "everything is General" task breakdown: lean-ctx's own
162        // tool/mode command keys must not collapse into General.
163        assert_eq!(
164            TaskClassifier::classify_command_key("ctx_search"),
165            TaskCategory::Exploration
166        );
167        assert_eq!(
168            TaskClassifier::classify_command_key("ctx_read"),
169            TaskCategory::Exploration
170        );
171        assert_eq!(
172            TaskClassifier::classify_command_key("ctx_edit"),
173            TaskCategory::Refactoring
174        );
175        assert_eq!(
176            TaskClassifier::classify_command_key("ctx_knowledge"),
177            TaskCategory::Knowledge
178        );
179    }
180
181    #[test]
182    fn cli_mode_keys_are_classified() {
183        assert_eq!(
184            TaskClassifier::classify_command_key("cli_grep"),
185            TaskCategory::Exploration
186        );
187        assert_eq!(
188            TaskClassifier::classify_command_key("cli_full"),
189            TaskCategory::Exploration
190        );
191        assert_eq!(
192            TaskClassifier::classify_command_key("cli_signatures"),
193            TaskCategory::Exploration
194        );
195        // The mixed shell bucket stays General (it aggregates arbitrary commands).
196        assert_eq!(
197            TaskClassifier::classify_command_key("cli_shell"),
198            TaskCategory::General
199        );
200    }
201
202    #[test]
203    fn real_shell_commands_still_classify() {
204        assert_eq!(
205            TaskClassifier::classify_command_key("cargo build"),
206            TaskCategory::BuildDeploy
207        );
208        assert_eq!(
209            TaskClassifier::classify_command_key("git status"),
210            TaskCategory::Git
211        );
212    }
213}