1use std::collections::HashMap;
12
13use serde::{Deserialize, Serialize};
14
15use crate::error::TuiError;
16use crate::security::validate::{validate_argument, validate_model_argument};
17
18#[derive(Debug, Clone, Serialize, Deserialize)]
20#[serde(deny_unknown_fields)]
21pub struct HarnessCapability {
22 pub model: bool,
23 pub reasoning_effort: bool,
24}
25
26#[derive(Debug, Clone, Serialize, Deserialize)]
46#[serde(deny_unknown_fields)]
47pub struct ModelAssignment {
48 pub agent_id: String,
49 pub harness: String,
50 pub model: Option<String>,
51 #[serde(default)]
52 pub model_provider: Option<String>,
53 #[serde(default)]
54 pub model_fallback_from: Option<String>,
55 #[serde(default)]
56 pub model_warning: Option<String>,
57 pub reasoning_effort: Option<String>,
58 pub model_source: String,
59 pub reasoning_source: String,
60}
61
62#[derive(Debug, Clone, Serialize, Deserialize)]
64#[serde(deny_unknown_fields)]
65pub struct ModelAssignments {
66 pub manifest_version: u32,
67 pub generated_by: String,
68 pub policy_sha256: String,
69 pub capabilities: HashMap<String, HarnessCapability>,
70 pub assignments: Vec<ModelAssignment>,
71}
72
73impl ModelAssignments {
74 pub fn for_agent(&self, agent_id: &str) -> Vec<&ModelAssignment> {
76 self.assignments
77 .iter()
78 .filter(|a| a.agent_id == agent_id)
79 .collect()
80 }
81}
82
83pub const CAPABLE_HARNESSES: &[&str] = &["codex", "claude-code", "cursor"];
87
88pub const REASONING_EFFORTS: &[&str] = &[
93 "auto", "none", "minimal", "low", "medium", "high", "xhigh", "max",
94];
95
96#[derive(Debug, Clone, PartialEq, Eq)]
98pub enum ModelScope {
99 All,
101 Provider(String),
103 Role(String),
105 Agent(String),
107}
108
109impl ModelScope {
110 pub fn to_flag(&self) -> String {
112 match self {
113 ModelScope::All => "all".to_string(),
114 ModelScope::Provider(id) => format!("provider={id}"),
115 ModelScope::Role(id) => format!("role={id}"),
116 ModelScope::Agent(id) => format!("agent={id}"),
117 }
118 }
119
120 pub fn display(&self) -> String {
122 match self {
123 ModelScope::All => "All agents".to_string(),
124 ModelScope::Provider(id) => format!("Provider: {id}"),
125 ModelScope::Role(id) => format!("Role: {id}"),
126 ModelScope::Agent(id) => format!("Agent: {id}"),
127 }
128 }
129}
130
131#[derive(Debug, Clone)]
134pub struct ModelPolicyCommand {
135 pub scope: ModelScope,
136 pub harness: String,
137 pub model: Option<String>,
139 pub reasoning: Option<String>,
141 pub dry_run: bool,
142}
143
144impl ModelPolicyCommand {
145 pub fn validate(&self) -> Result<(), TuiError> {
149 match &self.scope {
150 ModelScope::All => {}
151 ModelScope::Provider(id) | ModelScope::Role(id) | ModelScope::Agent(id) => {
152 validate_argument(id)?;
153 }
154 }
155 validate_argument(&self.harness)?;
156 if let Some(model) = &self.model {
157 validate_model_argument(model)?;
160 }
161 if let Some(reasoning) = &self.reasoning {
162 validate_argument(reasoning)?;
163 }
164 if self.model.is_none() && self.reasoning.is_none() {
165 return Err(TuiError::ValidationRejected {
166 value: "<empty>".to_string(),
167 rule: "set a model and/or a reasoning effort".to_string(),
168 });
169 }
170 Ok(())
171 }
172
173 pub fn to_args(&self) -> Vec<String> {
175 let mut args = vec![
176 "set".to_string(),
177 "--scope".to_string(),
178 self.scope.to_flag(),
179 "--harness".to_string(),
180 self.harness.clone(),
181 ];
182 if let Some(model) = &self.model {
183 args.push("--model".to_string());
184 args.push(model.clone());
185 }
186 if let Some(reasoning) = &self.reasoning {
187 args.push("--reasoning".to_string());
188 args.push(reasoning.clone());
189 }
190 if self.dry_run {
191 args.push("--dry-run".to_string());
192 }
193 args
194 }
195
196 pub fn display_command(&self) -> String {
198 format!("node scripts/model-policy.mjs {}", self.to_args().join(" "))
199 }
200}
201
202#[cfg(test)]
203mod tests {
204 use super::*;
205
206 fn cmd(scope: ModelScope) -> ModelPolicyCommand {
207 ModelPolicyCommand {
208 scope,
209 harness: "codex".to_string(),
210 model: Some("gpt-5.4".to_string()),
211 reasoning: Some("high".to_string()),
212 dry_run: true,
213 }
214 }
215
216 #[test]
217 fn to_args_all_scope() {
218 let args = cmd(ModelScope::All).to_args();
219 assert_eq!(
220 args,
221 vec![
222 "set",
223 "--scope",
224 "all",
225 "--harness",
226 "codex",
227 "--model",
228 "gpt-5.4",
229 "--reasoning",
230 "high",
231 "--dry-run",
232 ]
233 );
234 }
235
236 #[test]
237 fn to_args_role_scope_without_dry_run() {
238 let mut c = cmd(ModelScope::Role("cloud-security-engineer".to_string()));
239 c.dry_run = false;
240 c.reasoning = None;
241 let args = c.to_args();
242 assert!(args.contains(&"role=cloud-security-engineer".to_string()));
243 assert!(!args.contains(&"--dry-run".to_string()));
244 assert!(!args.contains(&"--reasoning".to_string()));
245 }
246
247 #[test]
248 fn to_args_provider_and_agent_scopes() {
249 assert!(cmd(ModelScope::Provider("gcp".to_string()))
250 .to_args()
251 .contains(&"provider=gcp".to_string()));
252 assert!(cmd(ModelScope::Agent(
253 "gcp-firebase-developer-agent".to_string()
254 ))
255 .to_args()
256 .contains(&"agent=gcp-firebase-developer-agent".to_string()));
257 }
258
259 #[test]
260 fn validate_accepts_safe_command() {
261 assert!(cmd(ModelScope::All).validate().is_ok());
262 }
263
264 #[test]
265 fn validate_rejects_metachar_in_model() {
266 let mut c = cmd(ModelScope::All);
267 c.model = Some("gpt-5.4; rm -rf /".to_string());
268 assert!(c.validate().is_err());
269 }
270
271 #[test]
272 fn validate_rejects_metachar_in_scope_id() {
273 let c = cmd(ModelScope::Role("role$(cmd)".to_string()));
274 assert!(c.validate().is_err());
275 }
276
277 #[test]
278 fn validate_rejects_empty_field_set() {
279 let mut c = cmd(ModelScope::All);
280 c.model = None;
281 c.reasoning = None;
282 assert!(c.validate().is_err());
283 }
284
285 #[test]
286 fn display_command_format() {
287 let display = cmd(ModelScope::All).display_command();
288 assert!(display.starts_with("node scripts/model-policy.mjs set"));
289 assert!(display.contains("--scope all"));
290 }
291
292 #[test]
293 fn assignments_for_agent_filters() {
294 let doc = ModelAssignments {
295 manifest_version: 1,
296 generated_by: "test".to_string(),
297 policy_sha256: "0".repeat(64),
298 capabilities: HashMap::new(),
299 assignments: vec![
300 ModelAssignment {
301 agent_id: "a".to_string(),
302 harness: "codex".to_string(),
303 model: Some("gpt-5.4".to_string()),
304 model_provider: None,
305 model_fallback_from: None,
306 model_warning: None,
307 reasoning_effort: Some("high".to_string()),
308 model_source: "all".to_string(),
309 reasoning_source: "all".to_string(),
310 },
311 ModelAssignment {
312 agent_id: "b".to_string(),
313 harness: "codex".to_string(),
314 model: None,
315 model_provider: None,
316 model_fallback_from: None,
317 model_warning: None,
318 reasoning_effort: None,
319 model_source: "default".to_string(),
320 reasoning_source: "default".to_string(),
321 },
322 ],
323 };
324 let for_a = doc.for_agent("a");
325 assert_eq!(for_a.len(), 1);
326 assert_eq!(for_a[0].model.as_deref(), Some("gpt-5.4"));
327 assert!(doc.for_agent("missing").is_empty());
328 }
329
330 #[test]
331 fn parses_generated_assignments_shape() {
332 let json = r#"{
333 "manifest_version": 1,
334 "generated_by": "scripts/model-policy.mjs",
335 "policy_sha256": "abc",
336 "capabilities": {
337 "codex": { "model": true, "reasoning_effort": true }
338 },
339 "assignments": [
340 {
341 "agent_id": "x",
342 "harness": "codex",
343 "model": null,
344 "model_provider": null,
345 "reasoning_effort": "high",
346 "model_source": "default",
347 "reasoning_source": "all"
348 }
349 ]
350 }"#;
351 let parsed: ModelAssignments = serde_json::from_str(json).expect("parse");
352 assert_eq!(parsed.assignments.len(), 1);
353 assert!(parsed.assignments[0].model.is_none());
354 assert!(parsed.assignments[0].model_fallback_from.is_none());
358 assert!(parsed.assignments[0].model_warning.is_none());
359 }
360
361 #[test]
362 fn parses_assignment_with_lifecycle_warning() {
363 let json = r#"{
364 "manifest_version": 1,
365 "generated_by": "scripts/model-policy.mjs",
366 "policy_sha256": "abc",
367 "capabilities": {
368 "codex": { "model": true, "reasoning_effort": true }
369 },
370 "assignments": [
371 {
372 "agent_id": "x",
373 "harness": "codex",
374 "model": "gpt-5.5",
375 "model_provider": null,
376 "model_fallback_from": "gpt-5-2025-08-07",
377 "model_warning": "model \"gpt-5-2025-08-07\" was retired by the provider — projecting documented successor \"gpt-5.5\"; migrate the policy rule",
378 "reasoning_effort": "high",
379 "model_source": "agent:x",
380 "reasoning_source": "agent:x"
381 }
382 ]
383 }"#;
384 let parsed: ModelAssignments = serde_json::from_str(json).expect("parse");
385 let a = &parsed.assignments[0];
386 assert_eq!(a.model.as_deref(), Some("gpt-5.5"));
387 assert_eq!(a.model_fallback_from.as_deref(), Some("gpt-5-2025-08-07"));
388 assert!(a.model_warning.as_deref().unwrap().contains("was retired"));
389 }
390}