zeph_commands/handlers/
think_tokens.rs

1// SPDX-FileCopyrightText: 2026 Andrei G <bug-ops>
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4//! `/think-tokens` command handler — runtime Claude/Gemini thinking-token budget control.
5
6use std::future::Future;
7use std::pin::Pin;
8
9use crate::context::CommandContext;
10use crate::{CommandError, CommandHandler, CommandOutput, SlashCategory};
11
12/// Parse a `/think-tokens` argument into a token budget.
13///
14/// Accepts a bare integer, or an integer with a case-insensitive `k` (×1000) or `M`
15/// (×`1_000_000`) suffix. One decimal place is allowed on the numeric part (e.g. `10.5k`),
16/// rounded to the nearest integer. `0` and `off` (case-insensitive) both mean "disable" and
17/// parse to `Ok(None)`. Negative numbers and malformed input return a descriptive `Err`.
18///
19/// # Examples
20///
21/// ```
22/// use zeph_commands::handlers::think_tokens::parse_token_budget;
23///
24/// assert_eq!(parse_token_budget("8k"), Ok(Some(8_000)));
25/// assert_eq!(parse_token_budget("10.5k"), Ok(Some(10_500)));
26/// assert_eq!(parse_token_budget("1M"), Ok(Some(1_000_000)));
27/// assert_eq!(parse_token_budget("off"), Ok(None));
28/// assert_eq!(parse_token_budget("0"), Ok(None));
29/// assert!(parse_token_budget("-1").is_err());
30/// ```
31///
32/// # Errors
33///
34/// Returns `Err(String)` with a descriptive message when `arg` is empty, negative, or does
35/// not parse as a number with an optional `k`/`M` suffix.
36pub fn parse_token_budget(arg: &str) -> Result<Option<u32>, String> {
37    let trimmed = arg.trim();
38    if trimmed.is_empty() {
39        return Err("empty token budget — expected a number (e.g. 8k, 1M, 0, off)".to_owned());
40    }
41    if trimmed.eq_ignore_ascii_case("off") {
42        return Ok(None);
43    }
44
45    let (numeric, multiplier) = match trimmed.chars().last() {
46        Some(c) if c.eq_ignore_ascii_case(&'k') => (&trimmed[..trimmed.len() - 1], 1_000.0),
47        Some(c) if c.eq_ignore_ascii_case(&'m') => (&trimmed[..trimmed.len() - 1], 1_000_000.0),
48        _ => (trimmed, 1.0),
49    };
50
51    if numeric.is_empty() {
52        return Err(format!(
53            "'{trimmed}' is missing a numeric value before the suffix"
54        ));
55    }
56
57    let value: f64 = numeric
58        .parse()
59        .map_err(|_| format!("'{trimmed}' is not a valid token budget"))?;
60    if value.is_sign_negative() {
61        return Err(format!("token budget must not be negative: '{trimmed}'"));
62    }
63
64    let scaled = value * multiplier;
65    if !scaled.is_finite() || scaled > f64::from(u32::MAX) {
66        return Err(format!("'{trimmed}' is too large for a token budget"));
67    }
68
69    #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
70    let rounded = scaled.round() as u32;
71    if rounded == 0 {
72        return Ok(None);
73    }
74    Ok(Some(rounded))
75}
76
77/// Show or set the active provider's runtime thinking-token budget.
78///
79/// - `/think-tokens` — display the current budget (or "off").
80/// - `/think-tokens 8k` / `/think-tokens 8000` — set an explicit budget.
81/// - `/think-tokens off` / `/think-tokens 0` — disable thinking.
82///
83/// Session-only: never persisted across restarts or `/provider` switches. Only Claude and
84/// Gemini support a thinking-token budget; other providers return an explicit "not supported"
85/// message.
86pub struct ThinkTokensCommand;
87
88impl CommandHandler<CommandContext<'_>> for ThinkTokensCommand {
89    fn name(&self) -> &'static str {
90        "/think-tokens"
91    }
92
93    fn description(&self) -> &'static str {
94        "Show or set the active provider's runtime thinking-token budget"
95    }
96
97    fn args_hint(&self) -> &'static str {
98        "[N|Nk|NM|off]"
99    }
100
101    fn category(&self) -> SlashCategory {
102        SlashCategory::Configuration
103    }
104
105    fn requires_auth(&self) -> bool {
106        true
107    }
108
109    fn handle<'a>(
110        &'a self,
111        ctx: &'a mut CommandContext<'_>,
112        args: &'a str,
113    ) -> Pin<Box<dyn Future<Output = Result<CommandOutput, CommandError>> + Send + 'a>> {
114        use tracing::Instrument as _;
115        let span = tracing::info_span!("commands.think_tokens.handle");
116        Box::pin(
117            async move {
118                let result = ctx.agent.handle_think_tokens(args).await;
119                Ok(CommandOutput::message_or_silent(result))
120            }
121            .instrument(span),
122        )
123    }
124}
125
126#[cfg(test)]
127mod tests {
128    use super::*;
129    use crate::handlers::test_helpers::{MockDebug, MockMessages, MockSession, make_ctx};
130    use crate::sink::NullSink;
131    use std::assert_matches;
132
133    #[test]
134    fn think_tokens_name_and_description() {
135        assert_eq!(ThinkTokensCommand.name(), "/think-tokens");
136        assert!(!ThinkTokensCommand.description().is_empty());
137    }
138
139    #[tokio::test]
140    async fn think_tokens_returns_silent_when_agent_returns_empty() {
141        let mut sink = NullSink;
142        let mut debug = MockDebug;
143        let mut messages = MockMessages;
144        let session = MockSession;
145        let mut agent = crate::NullAgent;
146        let mut ctx = make_ctx(&mut sink, &mut debug, &mut messages, &session, &mut agent);
147        let out = ThinkTokensCommand.handle(&mut ctx, "").await.unwrap();
148        assert_matches!(out, CommandOutput::Silent);
149    }
150
151    // ── parse_token_budget ───────────────────────────────────────────────
152
153    #[test]
154    fn parse_token_budget_empty_is_error() {
155        assert!(parse_token_budget("").is_err());
156        assert!(parse_token_budget("   ").is_err());
157    }
158
159    #[test]
160    fn parse_token_budget_bare_k_is_error() {
161        assert!(parse_token_budget("k").is_err());
162    }
163
164    #[test]
165    fn parse_token_budget_negative_is_error() {
166        assert!(parse_token_budget("-1").is_err());
167    }
168
169    #[test]
170    fn parse_token_budget_malformed_compound_is_error() {
171        assert!(parse_token_budget("1.2.3k").is_err());
172    }
173
174    #[test]
175    fn parse_token_budget_off_disables() {
176        assert_eq!(parse_token_budget("off"), Ok(None));
177        assert_eq!(parse_token_budget("OFF"), Ok(None));
178        assert_eq!(parse_token_budget("Off"), Ok(None));
179    }
180
181    #[test]
182    fn parse_token_budget_zero_disables() {
183        assert_eq!(parse_token_budget("0"), Ok(None));
184    }
185
186    #[test]
187    fn parse_token_budget_k_suffix() {
188        assert_eq!(parse_token_budget("8k"), Ok(Some(8_000)));
189        assert_eq!(parse_token_budget("8K"), Ok(Some(8_000)));
190    }
191
192    #[test]
193    fn parse_token_budget_decimal_k_suffix_rounds() {
194        assert_eq!(parse_token_budget("10.5k"), Ok(Some(10_500)));
195    }
196
197    #[test]
198    fn parse_token_budget_m_suffix() {
199        assert_eq!(parse_token_budget("1M"), Ok(Some(1_000_000)));
200        assert_eq!(parse_token_budget("1m"), Ok(Some(1_000_000)));
201    }
202
203    #[test]
204    fn parse_token_budget_bare_integer() {
205        assert_eq!(parse_token_budget("1024"), Ok(Some(1_024)));
206    }
207
208    #[test]
209    fn parse_token_budget_overflow_is_error() {
210        assert!(parse_token_budget("999999999999999M").is_err());
211    }
212
213    #[test]
214    fn parse_token_budget_trims_whitespace() {
215        assert_eq!(parse_token_budget("  8k  "), Ok(Some(8_000)));
216    }
217}
zeph_commands/handlers/think_tokens.rs

zeph_commands/handlers/
think_tokens.rs
