Skip to main content

wt/agent/
model.rs

1//! Model and effort selection for code-agent runs (the AI PR auto-fill path).
2//!
3//! [`AgentModel`] is a small, curated set of selectable model tiers; [`Effort`]
4//! is how hard the agent should work; [`AgentOptions`] bundles the two for one
5//! run. All three are pure data with `parse`/`id`/`label`/`next` helpers so they
6//! drive the config layer, the CLI flags, and the TUI's live cycle keys without
7//! any process or I/O.
8
9use serde::Serialize;
10
11/// A selectable model tier for a code agent. The variants currently encode the
12/// Claude tiers (the only supported agent); the `id` doubles as the CLI
13/// `--model` value — a stable alias that resolves to the latest model of that
14/// tier — so labels can track the current family without breaking selection.
15#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Default)]
16#[serde(rename_all = "lowercase")]
17pub enum AgentModel {
18    /// Most capable, highest latency (Claude Opus).
19    Opus,
20    /// Balanced capability and speed (Claude Sonnet) — the default.
21    #[default]
22    Sonnet,
23    /// Fastest and lightest (Claude Haiku).
24    Haiku,
25}
26
27impl AgentModel {
28    /// Every selectable model, in display and cycle order.
29    pub fn all() -> &'static [AgentModel] {
30        &[AgentModel::Opus, AgentModel::Sonnet, AgentModel::Haiku]
31    }
32
33    /// The stable lowercase identifier, used both in config/flags and as the
34    /// agent CLI's `--model` value (e.g. `"sonnet"`).
35    pub fn id(self) -> &'static str {
36        match self {
37            AgentModel::Opus => "opus",
38            AgentModel::Sonnet => "sonnet",
39            AgentModel::Haiku => "haiku",
40        }
41    }
42
43    /// A human-readable label for the status display; tracks the current model
44    /// family (the `id` alias always selects the latest of that tier).
45    pub fn label(self) -> &'static str {
46        match self {
47            AgentModel::Opus => "Opus 4.8",
48            AgentModel::Sonnet => "Sonnet 4.6",
49            AgentModel::Haiku => "Haiku 4.5",
50        }
51    }
52
53    /// Parses a model identifier (case-insensitive: `opus`/`sonnet`/`haiku`),
54    /// returning `None` if unknown.
55    pub fn parse(s: &str) -> Option<AgentModel> {
56        match s.trim().to_ascii_lowercase().as_str() {
57            "opus" => Some(AgentModel::Opus),
58            "sonnet" => Some(AgentModel::Sonnet),
59            "haiku" => Some(AgentModel::Haiku),
60            _ => None,
61        }
62    }
63
64    /// The next model in cycle order (wraps), for the TUI's `Ctrl-M` picker.
65    pub fn next(self) -> AgentModel {
66        match self {
67            AgentModel::Opus => AgentModel::Sonnet,
68            AgentModel::Sonnet => AgentModel::Haiku,
69            AgentModel::Haiku => AgentModel::Opus,
70        }
71    }
72
73    /// The previous model in cycle order (wraps), for navigating the TUI's
74    /// model dropdown upward (`↑`).
75    pub fn prev(self) -> AgentModel {
76        match self {
77            AgentModel::Opus => AgentModel::Haiku,
78            AgentModel::Sonnet => AgentModel::Opus,
79            AgentModel::Haiku => AgentModel::Sonnet,
80        }
81    }
82}
83
84/// How much effort the agent should spend on a draft. Claude has no native
85/// headless effort flag, so `wt` conveys effort as a one-line directive
86/// prepended to the prompt (see [`Effort::directive`]) — a safe, never-failing
87/// lever that shapes the model's deliberation and can map to native reasoning
88/// controls per agent in the future.
89#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Default)]
90#[serde(rename_all = "lowercase")]
91pub enum Effort {
92    /// Quick, minimal deliberation.
93    Low,
94    /// Balanced effort — the default (no directive).
95    #[default]
96    Medium,
97    /// Maximum deliberation and care.
98    High,
99}
100
101impl Effort {
102    /// Every effort level, in display and cycle order.
103    pub fn all() -> &'static [Effort] {
104        &[Effort::Low, Effort::Medium, Effort::High]
105    }
106
107    /// The stable lowercase identifier, used in config and `--effort`.
108    pub fn id(self) -> &'static str {
109        match self {
110            Effort::Low => "low",
111            Effort::Medium => "medium",
112            Effort::High => "high",
113        }
114    }
115
116    /// A human-readable label (currently identical to [`Effort::id`]).
117    pub fn label(self) -> &'static str {
118        self.id()
119    }
120
121    /// Parses an effort identifier (case-insensitive: `low`, `medium`/`med`,
122    /// `high`), returning `None` if unknown.
123    pub fn parse(s: &str) -> Option<Effort> {
124        match s.trim().to_ascii_lowercase().as_str() {
125            "low" => Some(Effort::Low),
126            "medium" | "med" => Some(Effort::Medium),
127            "high" => Some(Effort::High),
128            _ => None,
129        }
130    }
131
132    /// The next effort level in cycle order (wraps), for the TUI's `Ctrl-E` key.
133    pub fn next(self) -> Effort {
134        match self {
135            Effort::Low => Effort::Medium,
136            Effort::Medium => Effort::High,
137            Effort::High => Effort::Low,
138        }
139    }
140
141    /// The previous effort level in cycle order (wraps), for navigating the TUI's
142    /// effort dropdown upward (`↑`).
143    pub fn prev(self) -> Effort {
144        match self {
145            Effort::Low => Effort::High,
146            Effort::Medium => Effort::Low,
147            Effort::High => Effort::Medium,
148        }
149    }
150
151    /// A one-line instruction conveying this effort to the agent, prepended to
152    /// the prompt; `None` for the balanced baseline (medium).
153    pub fn directive(self) -> Option<&'static str> {
154        match self {
155            Effort::Low => Some("Work quickly and keep your reasoning brief."),
156            Effort::Medium => None,
157            Effort::High => Some("Think carefully and review the diff thoroughly before writing."),
158        }
159    }
160}
161
162/// The model and effort selected for a single agent run.
163#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
164pub struct AgentOptions {
165    /// The model tier to drive.
166    pub model: AgentModel,
167    /// How much effort to spend.
168    pub effort: Effort,
169}
170
171#[cfg(test)]
172mod tests {
173    use super::*;
174
175    #[test]
176    fn model_parse_roundtrips_and_rejects_unknown() {
177        for &m in AgentModel::all() {
178            assert_eq!(AgentModel::parse(m.id()), Some(m));
179        }
180        assert_eq!(AgentModel::parse("OPUS"), Some(AgentModel::Opus));
181        assert_eq!(AgentModel::parse(" sonnet "), Some(AgentModel::Sonnet));
182        assert_eq!(AgentModel::parse("gpt"), None);
183    }
184
185    #[test]
186    fn model_cycle_visits_every_variant() {
187        let mut seen = vec![AgentModel::Opus];
188        let mut cur = AgentModel::Opus;
189        for _ in 0..AgentModel::all().len() - 1 {
190            cur = cur.next();
191            seen.push(cur);
192        }
193        assert_eq!(cur.next(), AgentModel::Opus); // wraps
194        assert_eq!(seen.len(), AgentModel::all().len());
195    }
196
197    #[test]
198    fn model_serializes_lowercase() {
199        assert_eq!(
200            serde_json::to_string(&AgentModel::Sonnet).unwrap(),
201            "\"sonnet\""
202        );
203    }
204
205    #[test]
206    fn effort_parse_accepts_aliases() {
207        assert_eq!(Effort::parse("low"), Some(Effort::Low));
208        assert_eq!(Effort::parse("MED"), Some(Effort::Medium));
209        assert_eq!(Effort::parse("medium"), Some(Effort::Medium));
210        assert_eq!(Effort::parse("High"), Some(Effort::High));
211        assert_eq!(Effort::parse("max"), None);
212    }
213
214    #[test]
215    fn effort_directive_only_for_non_baseline() {
216        assert!(Effort::Low.directive().is_some());
217        assert!(Effort::Medium.directive().is_none());
218        assert!(Effort::High.directive().is_some());
219    }
220
221    #[test]
222    fn effort_cycle_wraps() {
223        assert_eq!(Effort::Low.next(), Effort::Medium);
224        assert_eq!(Effort::Medium.next(), Effort::High);
225        assert_eq!(Effort::High.next(), Effort::Low);
226    }
227
228    #[test]
229    fn prev_is_the_inverse_of_next() {
230        for &m in AgentModel::all() {
231            assert_eq!(m.next().prev(), m);
232            assert_eq!(m.prev().next(), m);
233        }
234        for &e in Effort::all() {
235            assert_eq!(e.next().prev(), e);
236            assert_eq!(e.prev().next(), e);
237        }
238    }
239
240    #[test]
241    fn defaults_are_sonnet_and_medium() {
242        let opts = AgentOptions::default();
243        assert_eq!(opts.model, AgentModel::Sonnet);
244        assert_eq!(opts.effort, Effort::Medium);
245    }
246}