use std::collections::HashMap;
use serde::{Deserialize, Serialize};
use crate::error::TuiError;
use crate::security::validate::validate_argument;
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct HarnessCapability {
pub model: bool,
pub reasoning_effort: bool,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct ModelAssignment {
pub agent_id: String,
pub harness: String,
pub model: Option<String>,
#[serde(default)]
pub model_provider: Option<String>,
#[serde(default)]
pub model_fallback_from: Option<String>,
#[serde(default)]
pub model_warning: Option<String>,
pub reasoning_effort: Option<String>,
pub model_source: String,
pub reasoning_source: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct ModelAssignments {
pub manifest_version: u32,
pub generated_by: String,
pub policy_sha256: String,
pub capabilities: HashMap<String, HarnessCapability>,
pub assignments: Vec<ModelAssignment>,
}
impl ModelAssignments {
pub fn for_agent(&self, agent_id: &str) -> Vec<&ModelAssignment> {
self.assignments
.iter()
.filter(|a| a.agent_id == agent_id)
.collect()
}
}
pub const CAPABLE_HARNESSES: &[&str] = &["codex", "claude-code", "cursor"];
pub const REASONING_EFFORTS: &[&str] = &[
"auto", "none", "minimal", "low", "medium", "high", "xhigh", "max",
];
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ModelScope {
All,
Provider(String),
Role(String),
Agent(String),
}
impl ModelScope {
pub fn to_flag(&self) -> String {
match self {
ModelScope::All => "all".to_string(),
ModelScope::Provider(id) => format!("provider={id}"),
ModelScope::Role(id) => format!("role={id}"),
ModelScope::Agent(id) => format!("agent={id}"),
}
}
pub fn display(&self) -> String {
match self {
ModelScope::All => "All agents".to_string(),
ModelScope::Provider(id) => format!("Provider: {id}"),
ModelScope::Role(id) => format!("Role: {id}"),
ModelScope::Agent(id) => format!("Agent: {id}"),
}
}
}
#[derive(Debug, Clone)]
pub struct ModelPolicyCommand {
pub scope: ModelScope,
pub harness: String,
pub model: Option<String>,
pub reasoning: Option<String>,
pub dry_run: bool,
}
impl ModelPolicyCommand {
pub fn validate(&self) -> Result<(), TuiError> {
match &self.scope {
ModelScope::All => {}
ModelScope::Provider(id) | ModelScope::Role(id) | ModelScope::Agent(id) => {
validate_argument(id)?;
}
}
validate_argument(&self.harness)?;
if let Some(model) = &self.model {
validate_argument(model)?;
}
if let Some(reasoning) = &self.reasoning {
validate_argument(reasoning)?;
}
if self.model.is_none() && self.reasoning.is_none() {
return Err(TuiError::ValidationRejected {
value: "<empty>".to_string(),
rule: "set a model and/or a reasoning effort".to_string(),
});
}
Ok(())
}
pub fn to_args(&self) -> Vec<String> {
let mut args = vec![
"set".to_string(),
"--scope".to_string(),
self.scope.to_flag(),
"--harness".to_string(),
self.harness.clone(),
];
if let Some(model) = &self.model {
args.push("--model".to_string());
args.push(model.clone());
}
if let Some(reasoning) = &self.reasoning {
args.push("--reasoning".to_string());
args.push(reasoning.clone());
}
if self.dry_run {
args.push("--dry-run".to_string());
}
args
}
pub fn display_command(&self) -> String {
format!("node scripts/model-policy.mjs {}", self.to_args().join(" "))
}
}
#[cfg(test)]
mod tests {
use super::*;
fn cmd(scope: ModelScope) -> ModelPolicyCommand {
ModelPolicyCommand {
scope,
harness: "codex".to_string(),
model: Some("gpt-5.4".to_string()),
reasoning: Some("high".to_string()),
dry_run: true,
}
}
#[test]
fn to_args_all_scope() {
let args = cmd(ModelScope::All).to_args();
assert_eq!(
args,
vec![
"set",
"--scope",
"all",
"--harness",
"codex",
"--model",
"gpt-5.4",
"--reasoning",
"high",
"--dry-run",
]
);
}
#[test]
fn to_args_role_scope_without_dry_run() {
let mut c = cmd(ModelScope::Role("cloud-security-engineer".to_string()));
c.dry_run = false;
c.reasoning = None;
let args = c.to_args();
assert!(args.contains(&"role=cloud-security-engineer".to_string()));
assert!(!args.contains(&"--dry-run".to_string()));
assert!(!args.contains(&"--reasoning".to_string()));
}
#[test]
fn to_args_provider_and_agent_scopes() {
assert!(cmd(ModelScope::Provider("gcp".to_string()))
.to_args()
.contains(&"provider=gcp".to_string()));
assert!(cmd(ModelScope::Agent(
"gcp-firebase-developer-agent".to_string()
))
.to_args()
.contains(&"agent=gcp-firebase-developer-agent".to_string()));
}
#[test]
fn validate_accepts_safe_command() {
assert!(cmd(ModelScope::All).validate().is_ok());
}
#[test]
fn validate_rejects_metachar_in_model() {
let mut c = cmd(ModelScope::All);
c.model = Some("gpt-5.4; rm -rf /".to_string());
assert!(c.validate().is_err());
}
#[test]
fn validate_rejects_metachar_in_scope_id() {
let c = cmd(ModelScope::Role("role$(cmd)".to_string()));
assert!(c.validate().is_err());
}
#[test]
fn validate_rejects_empty_field_set() {
let mut c = cmd(ModelScope::All);
c.model = None;
c.reasoning = None;
assert!(c.validate().is_err());
}
#[test]
fn display_command_format() {
let display = cmd(ModelScope::All).display_command();
assert!(display.starts_with("node scripts/model-policy.mjs set"));
assert!(display.contains("--scope all"));
}
#[test]
fn assignments_for_agent_filters() {
let doc = ModelAssignments {
manifest_version: 1,
generated_by: "test".to_string(),
policy_sha256: "0".repeat(64),
capabilities: HashMap::new(),
assignments: vec![
ModelAssignment {
agent_id: "a".to_string(),
harness: "codex".to_string(),
model: Some("gpt-5.4".to_string()),
model_provider: None,
model_fallback_from: None,
model_warning: None,
reasoning_effort: Some("high".to_string()),
model_source: "all".to_string(),
reasoning_source: "all".to_string(),
},
ModelAssignment {
agent_id: "b".to_string(),
harness: "codex".to_string(),
model: None,
model_provider: None,
model_fallback_from: None,
model_warning: None,
reasoning_effort: None,
model_source: "default".to_string(),
reasoning_source: "default".to_string(),
},
],
};
let for_a = doc.for_agent("a");
assert_eq!(for_a.len(), 1);
assert_eq!(for_a[0].model.as_deref(), Some("gpt-5.4"));
assert!(doc.for_agent("missing").is_empty());
}
#[test]
fn parses_generated_assignments_shape() {
let json = r#"{
"manifest_version": 1,
"generated_by": "scripts/model-policy.mjs",
"policy_sha256": "abc",
"capabilities": {
"codex": { "model": true, "reasoning_effort": true }
},
"assignments": [
{
"agent_id": "x",
"harness": "codex",
"model": null,
"model_provider": null,
"reasoning_effort": "high",
"model_source": "default",
"reasoning_source": "all"
}
]
}"#;
let parsed: ModelAssignments = serde_json::from_str(json).expect("parse");
assert_eq!(parsed.assignments.len(), 1);
assert!(parsed.assignments[0].model.is_none());
assert!(parsed.assignments[0].model_fallback_from.is_none());
assert!(parsed.assignments[0].model_warning.is_none());
}
#[test]
fn parses_assignment_with_lifecycle_warning() {
let json = r#"{
"manifest_version": 1,
"generated_by": "scripts/model-policy.mjs",
"policy_sha256": "abc",
"capabilities": {
"codex": { "model": true, "reasoning_effort": true }
},
"assignments": [
{
"agent_id": "x",
"harness": "codex",
"model": "gpt-5.5",
"model_provider": null,
"model_fallback_from": "gpt-5-2025-08-07",
"model_warning": "model \"gpt-5-2025-08-07\" was retired by the provider — projecting documented successor \"gpt-5.5\"; migrate the policy rule",
"reasoning_effort": "high",
"model_source": "agent:x",
"reasoning_source": "agent:x"
}
]
}"#;
let parsed: ModelAssignments = serde_json::from_str(json).expect("parse");
let a = &parsed.assignments[0];
assert_eq!(a.model.as_deref(), Some("gpt-5.5"));
assert_eq!(a.model_fallback_from.as_deref(), Some("gpt-5-2025-08-07"));
assert!(a.model_warning.as_deref().unwrap().contains("was retired"));
}
}