Skip to content

Commit e2b31ec

Browse files
committed
fix: integrate Chutes TEE-only model validation
- Make validate_chutes_model() case-insensitive and trim whitespace - Fix ProviderManager.set_provider() to actually set the provider name - Add validation call in set_model() for Chutes provider - Fix test assertion to expect 2 providers (cortex + chutes) - Add unit tests for validation functions
1 parent 378cbc3 commit e2b31ec

3 files changed

Lines changed: 49 additions & 5 deletions

File tree

src/cortex-common/src/model_presets/presets.rs

Lines changed: 37 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -836,7 +836,8 @@ pub fn get_models_for_provider(provider: &str) -> Vec<&'static ModelPreset> {
836836
/// Chutes only allows TEE (Trusted Execution Environment) models for security.
837837
/// Returns Ok(()) if valid, Err with message if invalid.
838838
pub fn validate_chutes_model(model: &str) -> Result<(), String> {
839-
if !model.ends_with("-TEE") {
839+
let model = model.trim();
840+
if !model.to_uppercase().ends_with("-TEE") {
840841
return Err(format!(
841842
"Chutes provider only allows TEE models (models ending with '-TEE'). \
842843
Model '{}' is not a TEE model. Available TEE models: moonshotai/Kimi-K2.5-TEE",
@@ -852,3 +853,38 @@ pub fn provider_allows_custom_models(provider: &str) -> bool {
852853
// Chutes does NOT allow custom models - only predefined TEE models
853854
provider != "chutes"
854855
}
856+
857+
#[cfg(test)]
858+
mod tests {
859+
use super::*;
860+
861+
#[test]
862+
fn test_validate_chutes_model_valid() {
863+
// Valid TEE models
864+
assert!(validate_chutes_model("moonshotai/Kimi-K2.5-TEE").is_ok());
865+
assert!(validate_chutes_model("some-model-TEE").is_ok());
866+
// Case insensitive
867+
assert!(validate_chutes_model("model-tee").is_ok());
868+
assert!(validate_chutes_model("model-Tee").is_ok());
869+
// Whitespace handling
870+
assert!(validate_chutes_model(" model-TEE ").is_ok());
871+
}
872+
873+
#[test]
874+
fn test_validate_chutes_model_invalid() {
875+
assert!(validate_chutes_model("gpt-4").is_err());
876+
assert!(validate_chutes_model("claude-3").is_err());
877+
assert!(validate_chutes_model("model-TEE-v2").is_err());
878+
assert!(validate_chutes_model("").is_err());
879+
}
880+
881+
#[test]
882+
fn test_provider_allows_custom_models() {
883+
// Chutes does NOT allow custom models
884+
assert!(!provider_allows_custom_models("chutes"));
885+
// Other providers allow custom models
886+
assert!(provider_allows_custom_models("cortex"));
887+
assert!(provider_allows_custom_models("openai"));
888+
assert!(provider_allows_custom_models("anthropic"));
889+
}
890+
}

src/cortex-tui/src/modal/providers.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -398,7 +398,7 @@ mod tests {
398398
let modal = ProvidersModal::with_known_providers(Some("cortex".to_string()));
399399

400400
assert_eq!(modal.title(), "Providers");
401-
assert_eq!(modal.providers.len(), 1);
401+
assert_eq!(modal.providers.len(), 2);
402402
}
403403

404404
#[test]

src/cortex-tui/src/providers/manager.rs

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ use cortex_engine::client::{
1111

1212
use super::config::CortexConfig;
1313
use super::models::{ModelInfo, get_models_for_provider, get_popular_models};
14+
use cortex_common::model_presets::validate_chutes_model;
1415

1516
// ============================================================
1617
// PROVIDER MANAGER
@@ -208,16 +209,23 @@ impl ProviderManager {
208209
|| cortex_login::has_valid_auth()
209210
}
210211

211-
/// Sets the current provider (for compatibility - always sets to "cortex").
212-
pub fn set_provider(&mut self, _provider: &str) -> Result<()> {
213-
self.current_provider = "cortex".to_string();
212+
/// Sets the current provider.
213+
pub fn set_provider(&mut self, provider: &str) -> Result<()> {
214+
self.current_provider = provider.to_string();
214215
self.client = None;
215216
Ok(())
216217
}
217218

218219
/// Sets the current model.
219220
pub fn set_model(&mut self, model: &str) -> Result<()> {
220221
let resolved = self.config.resolve_alias(model);
222+
223+
// Validate model for Chutes provider (TEE-only security requirement)
224+
if self.current_provider == "chutes" {
225+
validate_chutes_model(&resolved)
226+
.map_err(|e| anyhow::anyhow!(e))?;
227+
}
228+
221229
self.current_model = resolved;
222230
self.client = None;
223231
Ok(())

0 commit comments

Comments
 (0)