diff --git a/server/src/cursor/model_catalog.rs b/server/src/cursor/model_catalog.rs index 2a7364e..aefe280 100644 --- a/server/src/cursor/model_catalog.rs +++ b/server/src/cursor/model_catalog.rs @@ -14,7 +14,7 @@ use crate::{ proxy::{self, CursorProxy}, CursorSessionRegistry, }, - model::ProviderModel, + model::{format_token_count, parse_token_count, ProviderModel}, Error, Result, }; @@ -214,7 +214,7 @@ fn context_options(model: &ProviderModel) -> Vec<(String, String)> { let value = tokens.to_string(); let duplicate = contexts .iter() - .any(|(existing, _)| parse_context_tokens(existing) == Some(tokens)); + .any(|(existing, _)| parse_token_count(existing) == Some(tokens)); if !duplicate { contexts.push((value, format!("{} (Custom)", format_token_count(tokens)))); } @@ -222,26 +222,6 @@ fn context_options(model: &ProviderModel) -> Vec<(String, String)> { contexts } -fn parse_context_tokens(value: &str) -> Option { - let value = value.trim().to_ascii_lowercase(); - let (number, multiplier) = match value.chars().last()? { - 'k' => (&value[..value.len() - 1], 1_000), - 'm' => (&value[..value.len() - 1], 1_000_000), - _ => (value.as_str(), 1), - }; - number.parse::().ok()?.checked_mul(multiplier) -} - -fn format_token_count(tokens: u64) -> String { - if tokens >= 1_000_000 && tokens.is_multiple_of(1_000_000) { - format!("{}M", tokens / 1_000_000) - } else if tokens >= 1_000 && tokens.is_multiple_of(1_000) { - format!("{}K", tokens / 1_000) - } else { - tokens.to_string() - } -} - pub async fn available_models( State(registry): State, Extension(proxy): Extension, diff --git a/server/src/cursor/request/model.rs b/server/src/cursor/request/model.rs index 16f0972..2af1442 100644 --- a/server/src/cursor/request/model.rs +++ b/server/src/cursor/request/model.rs @@ -1,6 +1,9 @@ use crate::{ cursor::proto::agent::v1 as pb, - model::{ModelLatency, ModelSpec, ReasoningSpec, SubagentKind, SubagentModelOverride}, + model::{ + parse_token_count, ModelLatency, ModelSpec, ReasoningSpec, SubagentKind, + SubagentModelOverride, + }, Error, Result, }; @@ -136,16 +139,6 @@ fn parse_bool(parameter: &pb::requested_model::ModelParameterValue) -> Result Option { - let value = value.trim().to_ascii_lowercase(); - let (number, multiplier) = match value.chars().last()? { - 'k' => (&value[..value.len() - 1], 1_000), - 'm' => (&value[..value.len() - 1], 1_000_000), - _ => (value.as_str(), 1), - }; - number.parse::().ok()?.checked_mul(multiplier) -} - #[cfg(test)] mod tests { use super::*; diff --git a/server/src/model/mod.rs b/server/src/model/mod.rs index 3ca7fe9..ae3637a 100644 --- a/server/src/model/mod.rs +++ b/server/src/model/mod.rs @@ -9,6 +9,7 @@ mod projection; mod provider; mod run; mod runtime_tag; +mod token_count; mod tool; mod usage; @@ -23,5 +24,6 @@ pub use projection::*; pub use provider::*; pub use run::*; pub use runtime_tag::*; +pub(crate) use token_count::*; pub use tool::*; pub use usage::*; diff --git a/server/src/model/token_count.rs b/server/src/model/token_count.rs new file mode 100644 index 0000000..126b344 --- /dev/null +++ b/server/src/model/token_count.rs @@ -0,0 +1,39 @@ +pub(crate) fn parse_token_count(value: &str) -> Option { + let value = value.trim().to_ascii_lowercase(); + let (number, multiplier) = match value.chars().last()? { + 'k' => (&value[..value.len() - 1], 1_000), + 'm' => (&value[..value.len() - 1], 1_000_000), + _ => (value.as_str(), 1), + }; + number.parse::().ok()?.checked_mul(multiplier) +} + +pub(crate) fn format_token_count(tokens: u64) -> String { + if tokens >= 1_000_000 && tokens.is_multiple_of(1_000_000) { + format!("{}M", tokens / 1_000_000) + } else if tokens >= 1_000 && tokens.is_multiple_of(1_000) { + format!("{}K", tokens / 1_000) + } else { + tokens.to_string() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn token_counts_parse_plain_and_abbreviated_values() { + assert_eq!(parse_token_count("272000"), Some(272_000)); + assert_eq!(parse_token_count("272K"), Some(272_000)); + assert_eq!(parse_token_count("1m"), Some(1_000_000)); + assert_eq!(parse_token_count("invalid"), None); + } + + #[test] + fn token_counts_format_exact_thousands_and_millions() { + assert_eq!(format_token_count(272_000), "272K"); + assert_eq!(format_token_count(1_000_000), "1M"); + assert_eq!(format_token_count(272_001), "272001"); + } +}