Skip to main content

mars_schema/
model_policy.rs

1//! Model-policy definition and validation shared by profiles and consumer config.
2
3use crate::agent::{ApprovalMode, EffortLevel, SandboxMode};
4use crate::network::NetworkPolicy;
5use serde::ser::SerializeMap;
6use serde::{Deserialize, Serialize};
7
8#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)]
9#[serde(rename_all = "kebab-case")]
10pub enum ModelPolicyMatchType {
11    Model,
12    Alias,
13    ModelGlob,
14}
15
16#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
17#[serde(deny_unknown_fields)]
18pub struct ModelPolicyOverride {
19    #[serde(default, skip_serializing_if = "Option::is_none")]
20    pub effort: Option<String>,
21    #[serde(default, skip_serializing_if = "Option::is_none")]
22    pub sandbox: Option<String>,
23    #[serde(default, skip_serializing_if = "Option::is_none")]
24    pub approval: Option<String>,
25    #[serde(default, skip_serializing_if = "Option::is_none")]
26    pub network: Option<String>,
27}
28
29impl ModelPolicyOverride {
30    pub fn is_empty(&self) -> bool {
31        self.effort.is_none()
32            && self.sandbox.is_none()
33            && self.approval.is_none()
34            && self.network.is_none()
35    }
36
37    pub fn get(&self, key: &str) -> Option<&str> {
38        match key {
39            "effort" => self.effort.as_deref(),
40            "sandbox" => self.sandbox.as_deref(),
41            "approval" => self.approval.as_deref(),
42            "network" => self.network.as_deref(),
43            _ => None,
44        }
45    }
46}
47
48/// Shared model-policy rule type used by profile frontmatter, agent overlays,
49/// and settings-level model policies.
50#[derive(Debug, Clone, PartialEq)]
51pub struct ModelPolicyRule {
52    pub match_type: ModelPolicyMatchType,
53    pub match_value: String,
54    pub no_fallback: bool,
55    pub overrides: ModelPolicyOverride,
56    /// A compact, authored guide explaining when this model is a good choice.
57    pub when: Option<String>,
58}
59
60#[derive(Debug, Clone, PartialEq, Eq)]
61pub enum ModelPolicyRuleParseError {
62    RuleMustBeMapping {
63        found: String,
64    },
65    MatchMissing,
66    MatchMustBeMapping {
67        found: String,
68    },
69    MatchMustContainExactlyOne {
70        found: String,
71    },
72    MatchKeyMustBeString {
73        found: String,
74    },
75    UnknownMatchKey {
76        key: String,
77    },
78    MatchValueMustBeString {
79        key: String,
80        found: String,
81    },
82    MatchValueEmpty {
83        key: String,
84    },
85    OverrideMustBeMapping {
86        found: String,
87    },
88    UnknownOverrideKey {
89        key: String,
90    },
91    /// A host set is layered by settings, profile, overlay and CLI, never by a model rule.
92    NetworkHostsNotAllowed {
93        key: String,
94    },
95    OverrideValueInvalid {
96        key: String,
97        found: String,
98        allowed: &'static str,
99    },
100    NoFallbackMustBeBoolean {
101        found: String,
102    },
103    WhenInvalid {
104        found: String,
105    },
106}
107
108impl ModelPolicyRuleParseError {
109    fn deserialize_message(&self) -> String {
110        match self {
111            Self::MatchMustContainExactlyOne { .. }
112            | Self::MatchMissing
113            | Self::MatchMustBeMapping { .. } => {
114                "model policy `match` must contain exactly one of model, alias, model-glob"
115                    .to_string()
116            }
117            Self::MatchKeyMustBeString { .. } => {
118                "model policy `match` key must be a string".to_string()
119            }
120            Self::MatchValueMustBeString { .. } => {
121                "model policy `match` value must be a string".to_string()
122            }
123            Self::MatchValueEmpty { .. } => {
124                "model policy `match` value must be a non-empty string".to_string()
125            }
126            Self::UnknownMatchKey { key } => format!(
127                "unknown model policy match key `{key}`; expected model, alias, or model-glob"
128            ),
129            Self::OverrideMustBeMapping { .. } => {
130                "model policy `override` must be a mapping".to_string()
131            }
132            Self::UnknownOverrideKey { key } => format!(
133                "unknown model policy override `{key}`; expected effort, sandbox, approval, or network"
134            ),
135            Self::NetworkHostsNotAllowed { key } => format!(
136                "model policy override `{key}` is not supported; set it on settings, the profile, an agent overlay, or the CLI"
137            ),
138            Self::OverrideValueInvalid {
139                key,
140                found,
141                allowed,
142            } => format!(
143                "model policy override `{key}` has invalid value `{found}`; allowed: {allowed}"
144            ),
145            Self::NoFallbackMustBeBoolean { .. } => {
146                "model policy `no-fallback` must be a boolean".to_string()
147            }
148            Self::WhenInvalid { .. } => {
149                "model policy `when` must be a non-empty one-line string".to_string()
150            }
151            Self::RuleMustBeMapping { .. } => "model policy rule must be a mapping".to_string(),
152        }
153    }
154}
155
156pub(crate) fn parse_model_policy_rule_value(
157    value: &serde_yaml::Value,
158) -> Result<ModelPolicyRule, ModelPolicyRuleParseError> {
159    let rule = value
160        .as_mapping()
161        .ok_or_else(|| ModelPolicyRuleParseError::RuleMustBeMapping {
162            found: crate::ordered::plain(value),
163        })?;
164
165    let match_value = rule.get(serde_yaml::Value::String("match".to_string()));
166    let match_mapping = match match_value {
167        Some(value) => {
168            value
169                .as_mapping()
170                .ok_or_else(|| ModelPolicyRuleParseError::MatchMustBeMapping {
171                    found: crate::ordered::plain(value),
172                })?
173        }
174        None => return Err(ModelPolicyRuleParseError::MatchMissing),
175    };
176
177    let mut entries = match_mapping.iter();
178    let Some((match_key, match_value)) = entries.next() else {
179        return Err(ModelPolicyRuleParseError::MatchMustContainExactlyOne {
180            found: crate::ordered::plain(&serde_yaml::Value::Mapping(match_mapping.clone())),
181        });
182    };
183    if entries.next().is_some() {
184        return Err(ModelPolicyRuleParseError::MatchMustContainExactlyOne {
185            found: crate::ordered::plain(&serde_yaml::Value::Mapping(match_mapping.clone())),
186        });
187    }
188
189    let key =
190        match_key
191            .as_str()
192            .ok_or_else(|| ModelPolicyRuleParseError::MatchKeyMustBeString {
193                found: crate::ordered::plain(match_key),
194            })?;
195    let value =
196        match_value
197            .as_str()
198            .ok_or_else(|| ModelPolicyRuleParseError::MatchValueMustBeString {
199                key: key.to_string(),
200                found: crate::ordered::plain(match_value),
201            })?;
202    let match_value = value.to_string();
203    if match_value.is_empty() {
204        return Err(ModelPolicyRuleParseError::MatchValueEmpty {
205            key: key.to_string(),
206        });
207    }
208
209    let match_type = match key {
210        "model" => ModelPolicyMatchType::Model,
211        "alias" => ModelPolicyMatchType::Alias,
212        "model-glob" => ModelPolicyMatchType::ModelGlob,
213        _ => {
214            return Err(ModelPolicyRuleParseError::UnknownMatchKey {
215                key: key.to_string(),
216            });
217        }
218    };
219
220    let overrides = parse_overrides(rule.get(serde_yaml::Value::String("override".to_string())))?;
221
222    let no_fallback = match rule.get(serde_yaml::Value::String("no-fallback".to_string())) {
223        None | Some(serde_yaml::Value::Null) => false,
224        Some(serde_yaml::Value::Bool(value)) => *value,
225        Some(value) => {
226            return Err(ModelPolicyRuleParseError::NoFallbackMustBeBoolean {
227                found: crate::ordered::plain(value),
228            });
229        }
230    };
231
232    let when = match rule.get(serde_yaml::Value::String("when".to_string())) {
233        None | Some(serde_yaml::Value::Null) => None,
234        Some(serde_yaml::Value::String(value))
235            if !value.is_empty() && !value.contains(['\n', '\r']) =>
236        {
237            Some(value.clone())
238        }
239        Some(value) => {
240            return Err(ModelPolicyRuleParseError::WhenInvalid {
241                found: crate::ordered::plain(value),
242            });
243        }
244    };
245
246    Ok(ModelPolicyRule {
247        match_type,
248        match_value,
249        no_fallback,
250        overrides,
251        when,
252    })
253}
254
255fn parse_overrides(
256    value: Option<&serde_yaml::Value>,
257) -> Result<ModelPolicyOverride, ModelPolicyRuleParseError> {
258    let Some(value) = value.filter(|value| !value.is_null()) else {
259        return Ok(ModelPolicyOverride::default());
260    };
261    let mapping =
262        value
263            .as_mapping()
264            .ok_or_else(|| ModelPolicyRuleParseError::OverrideMustBeMapping {
265                found: crate::ordered::plain(value),
266            })?;
267    let mut overrides = ModelPolicyOverride::default();
268    for (key, value) in mapping {
269        let key = key
270            .as_str()
271            .ok_or_else(|| ModelPolicyRuleParseError::UnknownOverrideKey {
272                key: crate::ordered::plain(key),
273            })?;
274        if matches!(key, "network-allow" | "network-deny") {
275            return Err(ModelPolicyRuleParseError::NetworkHostsNotAllowed {
276                key: key.to_string(),
277            });
278        }
279        let raw =
280            value
281                .as_str()
282                .ok_or_else(|| ModelPolicyRuleParseError::OverrideValueInvalid {
283                    key: key.to_string(),
284                    found: crate::ordered::plain(value),
285                    allowed: allowed_override_values(key),
286                })?;
287        let invalid = || ModelPolicyRuleParseError::OverrideValueInvalid {
288            key: key.to_string(),
289            found: raw.to_string(),
290            allowed: allowed_override_values(key),
291        };
292        match key {
293            "effort" => {
294                overrides.effort =
295                    Some(EffortLevel::parse(raw).ok_or_else(invalid)?.as_str().into());
296            }
297            "sandbox" => {
298                overrides.sandbox =
299                    Some(SandboxMode::parse(raw).ok_or_else(invalid)?.as_str().into());
300            }
301            "approval" => {
302                overrides.approval = Some(
303                    ApprovalMode::parse(raw)
304                        .ok_or_else(invalid)?
305                        .as_str()
306                        .into(),
307                );
308            }
309            "network" => {
310                overrides.network = Some(
311                    NetworkPolicy::parse(raw)
312                        .ok_or_else(invalid)?
313                        .as_str()
314                        .into(),
315                );
316            }
317            _ => {
318                return Err(ModelPolicyRuleParseError::UnknownOverrideKey {
319                    key: key.to_string(),
320                });
321            }
322        }
323    }
324    Ok(overrides)
325}
326
327fn allowed_override_values(key: &str) -> &'static str {
328    match key {
329        "effort" => EffortLevel::ALLOWED,
330        "sandbox" => SandboxMode::ALLOWED,
331        "approval" => ApprovalMode::ALLOWED,
332        "network" => NetworkPolicy::ALLOWED,
333        _ => "effort, sandbox, approval, network",
334    }
335}
336
337impl<'de> Deserialize<'de> for ModelPolicyRule {
338    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
339    where
340        D: serde::Deserializer<'de>,
341    {
342        let value = serde_yaml::Value::deserialize(deserializer)?;
343        parse_model_policy_rule_value(&value)
344            .map_err(|err| serde::de::Error::custom(err.deserialize_message()))
345    }
346}
347
348impl Serialize for ModelPolicyRule {
349    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
350    where
351        S: serde::Serializer,
352    {
353        let mut map = serializer.serialize_map(None)?;
354        let match_key = match self.match_type {
355            ModelPolicyMatchType::Model => "model",
356            ModelPolicyMatchType::Alias => "alias",
357            ModelPolicyMatchType::ModelGlob => "model-glob",
358        };
359
360        let mut match_clause = serde_yaml::Mapping::new();
361        match_clause.insert(
362            serde_yaml::Value::String(match_key.to_string()),
363            serde_yaml::Value::String(self.match_value.clone()),
364        );
365        map.serialize_entry("match", &match_clause)?;
366        if !self.overrides.is_empty() {
367            map.serialize_entry("override", &self.overrides)?;
368        }
369        if self.no_fallback {
370            map.serialize_entry("no-fallback", &self.no_fallback)?;
371        }
372        if let Some(when) = &self.when {
373            map.serialize_entry("when", when)?;
374        }
375        map.end()
376    }
377}