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

wt/agent/
model.rs
