1use std::collections::{BTreeMap, BTreeSet};
2
3use cookie_agent_identity::{AgentId, ConfiguredVariantRef, ModelKey, WildcardPattern};
4pub use cookie_agent_protocol::{AgentMode, PermissionAction, PermissionEffect, PermissionRule};
5use indexmap::IndexMap;
6use serde::{Deserialize, Serialize};
7
8use crate::{AgentDocument, ConfigError};
9
10const MAX_LIST: usize = 256;
11pub const BUILT_IN_DEFAULT_AGENT_ID: &str = "default";
12pub const BUILT_IN_APPROVAL_AGENT_ID: &str = "approval";
13pub const BUILT_IN_COMPACTION_AGENT_ID: &str = "compaction";
14pub const BUILT_IN_TITLE_AGENT_ID: &str = "title";
15pub const PARENT_MODEL_EXPRESSION: &str = "${parent_model}";
16
17#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
18#[serde(untagged)]
19pub enum PermissionValue {
20 Effect(PermissionEffect),
21 Resources(IndexMap<WildcardPattern, PermissionEffect>),
22}
23
24impl<'de> Deserialize<'de> for PermissionValue {
25 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
26 where
27 D: serde::Deserializer<'de>,
28 {
29 struct Visitor;
30 impl<'de> serde::de::Visitor<'de> for Visitor {
31 type Value = PermissionValue;
32
33 fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
34 formatter.write_str("a permission effect or ordered resource map")
35 }
36
37 fn visit_str<E>(self, value: &str) -> Result<Self::Value, E>
38 where
39 E: serde::de::Error,
40 {
41 let effect = match value {
42 "allow" => PermissionEffect::Allow,
43 "ask" => PermissionEffect::Ask,
44 "deny" => PermissionEffect::Deny,
45 _ => return Err(E::custom("invalid permission effect")),
46 };
47 Ok(PermissionValue::Effect(effect))
48 }
49
50 fn visit_map<A>(self, mut map: A) -> Result<Self::Value, A::Error>
51 where
52 A: serde::de::MapAccess<'de>,
53 {
54 let mut resources = IndexMap::new();
55 while let Some((resource, effect)) =
56 map.next_entry::<WildcardPattern, PermissionEffect>()?
57 {
58 if resources.insert(resource, effect).is_some() {
59 return Err(serde::de::Error::custom("duplicate permission resource"));
60 }
61 }
62 Ok(PermissionValue::Resources(resources))
63 }
64 }
65 deserializer.deserialize_any(Visitor)
66 }
67}
68
69fn deserialize_permissions<'de, D>(
70 deserializer: D,
71) -> Result<IndexMap<PermissionAction, PermissionValue>, D::Error>
72where
73 D: serde::Deserializer<'de>,
74{
75 struct Visitor;
76 impl<'de> serde::de::Visitor<'de> for Visitor {
77 type Value = IndexMap<PermissionAction, PermissionValue>;
78
79 fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
80 formatter.write_str("an ordered action permission map")
81 }
82
83 fn visit_map<A>(self, mut map: A) -> Result<Self::Value, A::Error>
84 where
85 A: serde::de::MapAccess<'de>,
86 {
87 let mut permissions = IndexMap::new();
88 while let Some((action, value)) =
89 map.next_entry::<PermissionAction, PermissionValue>()?
90 {
91 if permissions.insert(action, value).is_some() {
92 return Err(serde::de::Error::custom("duplicate permission action"));
93 }
94 }
95 Ok(permissions)
96 }
97 }
98 deserializer.deserialize_map(Visitor)
99}
100
101impl PermissionValue {
102 pub fn rules(&self, action: PermissionAction) -> Vec<PermissionRule> {
103 match self {
104 Self::Effect(effect) => vec![PermissionRule {
105 action,
106 resource: WildcardPattern::new("*").expect("static wildcard"),
107 effect: *effect,
108 }],
109 Self::Resources(resources) => resources
110 .iter()
111 .map(|(resource, effect)| PermissionRule {
112 action,
113 resource: resource.clone(),
114 effect: *effect,
115 })
116 .collect(),
117 }
118 }
119}
120
121#[derive(Clone, Debug, Eq, Ord, PartialEq, PartialOrd, Serialize)]
122pub enum AgentModelRef {
123 Model(ModelKey),
124 ParentModel,
125}
126
127impl<'de> Deserialize<'de> for AgentModelRef {
128 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
129 where
130 D: serde::Deserializer<'de>,
131 {
132 let value = String::deserialize(deserializer)?;
133 if value == PARENT_MODEL_EXPRESSION {
134 Ok(Self::ParentModel)
135 } else {
136 value
137 .parse::<ModelKey>()
138 .map(Self::Model)
139 .map_err(serde::de::Error::custom)
140 }
141 }
142}
143
144#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
145#[serde(deny_unknown_fields)]
146pub struct AgentModelFallback {
147 pub model: AgentModelRef,
148 pub variant: Option<ConfiguredVariantRef>,
149 pub cache: Option<crate::CacheConfig>,
150}
151
152#[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
153#[serde(default, deny_unknown_fields)]
154pub struct AgentLimits {
155 pub timeout_ms: u64,
156 pub max_output_tokens: u64,
157}
158
159#[derive(Clone, Debug, Serialize)]
160#[serde(deny_unknown_fields)]
161pub struct AgentFrontmatter {
162 pub description: String,
163 pub mode: AgentMode,
164 pub enabled: bool,
165 pub models: Vec<AgentModelFallback>,
166 #[serde(default)]
167 pub limits: AgentLimits,
168 #[serde(default, deserialize_with = "deserialize_permissions")]
169 pub permissions: IndexMap<PermissionAction, PermissionValue>,
170 #[serde(default)]
173 pub agent_md: Option<bool>,
174}
175
176impl<'de> Deserialize<'de> for AgentFrontmatter {
177 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
178 where
179 D: serde::Deserializer<'de>,
180 {
181 #[derive(Default, Deserialize)]
182 #[serde(default, deny_unknown_fields)]
183 struct RawLimits {
184 timeout_ms: Option<u64>,
185 max_output_tokens: Option<u64>,
186 }
187
188 #[derive(Deserialize)]
189 #[serde(deny_unknown_fields)]
190 struct RawFrontmatter {
191 description: String,
192 mode: AgentMode,
193 enabled: bool,
194 models: Vec<AgentModelFallback>,
195 #[serde(default)]
196 limits: RawLimits,
197 #[serde(default, deserialize_with = "deserialize_permissions")]
198 permissions: IndexMap<PermissionAction, PermissionValue>,
199 #[serde(default)]
200 agent_md: Option<bool>,
201 }
202
203 let raw = RawFrontmatter::deserialize(deserializer)?;
204 let internal = raw.mode == AgentMode::Internal;
205 Ok(Self {
206 description: raw.description,
207 mode: raw.mode,
208 enabled: raw.enabled,
209 models: raw.models,
210 limits: AgentLimits {
211 timeout_ms: raw
212 .limits
213 .timeout_ms
214 .unwrap_or(if internal { 30_000 } else { 0 }),
215 max_output_tokens: raw.limits.max_output_tokens.unwrap_or(0),
219 },
220 permissions: raw.permissions,
221 agent_md: raw.agent_md,
222 })
223 }
224}
225
226#[derive(Clone, Copy, Debug)]
231pub struct AgentMaterializationInput<'a> {
232 pub document: &'a AgentDocument,
233 pub root_eligible: bool,
234}
235
236#[derive(Clone, Debug)]
238pub struct AgentRegistry {
239 agents: BTreeMap<AgentId, AgentDocument>,
240}
241
242impl AgentRegistry {
243 pub fn validate(agents: BTreeMap<AgentId, AgentDocument>) -> Result<Self, ConfigError> {
244 Self::validate_ref(&agents)?;
245 Ok(Self { agents })
246 }
247
248 pub(crate) fn validate_ref(
249 agents: &BTreeMap<AgentId, AgentDocument>,
250 ) -> Result<(), ConfigError> {
251 for document in agents.values() {
252 validate_agent_document(document, agents)?;
253 }
254 Ok(())
255 }
256
257 pub(crate) fn from_validated(agents: BTreeMap<AgentId, AgentDocument>) -> Self {
258 Self { agents }
259 }
260
261 #[must_use]
262 pub fn get(&self, id: &AgentId) -> Option<&AgentDocument> {
263 self.agents.get(id)
264 }
265
266 pub fn agents(&self) -> impl ExactSizeIterator<Item = (&AgentId, &AgentDocument)> {
267 self.agents.iter()
268 }
269
270 #[must_use]
271 pub fn documents(&self) -> &BTreeMap<AgentId, AgentDocument> {
272 &self.agents
273 }
274
275 #[must_use]
276 pub fn into_documents(self) -> BTreeMap<AgentId, AgentDocument> {
277 self.agents
278 }
279
280 pub fn materialization_inputs(
281 &self,
282 ) -> impl ExactSizeIterator<Item = AgentMaterializationInput<'_>> {
283 self.agents
284 .values()
285 .map(|document| AgentMaterializationInput {
286 document,
287 root_eligible: authored_root_eligible(document),
288 })
289 }
290}
291
292fn authored_root_eligible(document: &AgentDocument) -> bool {
293 document.frontmatter.enabled
294 && matches!(
295 document.frontmatter.mode,
296 AgentMode::Primary | AgentMode::All
297 )
298 && !document.frontmatter.models.is_empty()
299}
300
301fn validate_agent_document(
302 document: &AgentDocument,
303 all: &BTreeMap<AgentId, AgentDocument>,
304) -> Result<(), ConfigError> {
305 if document.id.as_str() == BUILT_IN_DEFAULT_AGENT_ID {
306 return Err(ConfigError::ReservedAgentId(document.id.clone()));
307 }
308 let frontmatter = &document.frontmatter;
309 if matches!(
310 document.id.as_str(),
311 BUILT_IN_APPROVAL_AGENT_ID | BUILT_IN_COMPACTION_AGENT_ID | BUILT_IN_TITLE_AGENT_ID
312 ) && frontmatter.mode != AgentMode::Internal
313 {
314 return Err(ConfigError::AgentField {
315 agent: document.id.clone(),
316 field: "mode",
317 });
318 }
319 if frontmatter.description.is_empty()
320 || frontmatter.description.len() > 512
321 || frontmatter.description.chars().any(char::is_control)
322 {
323 return Err(ConfigError::AgentField {
324 agent: document.id.clone(),
325 field: "description",
326 });
327 }
328 if frontmatter.agent_md.is_some()
331 && matches!(frontmatter.mode, AgentMode::Subagent | AgentMode::Internal)
332 {
333 return Err(ConfigError::AgentField {
334 agent: document.id.clone(),
335 field: "agent_md",
336 });
337 }
338 if matches!(frontmatter.mode, AgentMode::Primary) && frontmatter.models.is_empty() {
339 return Err(ConfigError::PrimaryFallback(document.id.clone()));
340 }
341 if frontmatter.models.len() > MAX_LIST {
342 return Err(ConfigError::AgentLimit(document.id.clone()));
343 }
344 let mut fallback_models = BTreeSet::new();
345 for fallback in &frontmatter.models {
346 if matches!(fallback.model, AgentModelRef::ParentModel)
347 && (frontmatter.mode != AgentMode::Internal || fallback.variant.is_some())
348 {
349 return Err(ConfigError::AgentField {
350 agent: document.id.clone(),
351 field: "models",
352 });
353 }
354 if !fallback_models.insert(fallback.model.clone()) {
355 return Err(ConfigError::AgentField {
356 agent: document.id.clone(),
357 field: "models",
358 });
359 }
360 }
361 if frontmatter.mode != AgentMode::Internal && frontmatter.limits.timeout_ms != 0 {
362 return Err(ConfigError::AgentTimeoutInternalOnly(document.id.clone()));
363 }
364 let permission_rules = frontmatter
365 .permissions
366 .values()
367 .map(|value| match value {
368 PermissionValue::Effect(_) => 1,
369 PermissionValue::Resources(resources) => resources.len(),
370 })
371 .sum::<usize>();
372 if permission_rules > MAX_LIST {
373 return Err(ConfigError::AgentLimit(document.id.clone()));
374 }
375 if let Some(delegate) = frontmatter.permissions.get(&PermissionAction::Delegate) {
376 let rules = delegate.rules(PermissionAction::Delegate);
377 let mut expanded = BTreeSet::new();
381 for rule in rules {
382 if rule.effect == PermissionEffect::Deny {
383 continue;
384 }
385 let eligible = |target_document: &AgentDocument| {
386 target_document.frontmatter.enabled
387 && matches!(
388 target_document.frontmatter.mode,
389 AgentMode::Subagent | AgentMode::All
390 )
391 };
392 let pattern = rule.resource.as_str();
393 if pattern.contains('*') || pattern.contains('?') {
394 let matches = all
398 .iter()
399 .filter(|(id, target_document)| {
400 crate::simple_wildcard_match(pattern, id.as_str())
401 && eligible(target_document)
402 })
403 .map(|(id, _)| id.clone())
404 .collect::<Vec<_>>();
405 if matches.is_empty() {
406 return Err(ConfigError::UnmatchedDelegationPattern {
407 agent: document.id.clone(),
408 pattern: pattern.to_owned(),
409 });
410 }
411 expanded.extend(matches);
412 continue;
413 }
414 let target =
415 AgentId::new(pattern).map_err(|_| ConfigError::Delegation(document.id.clone()))?;
416 let target_document =
417 all.get(&target)
418 .ok_or_else(|| ConfigError::UnknownDelegationTarget {
419 agent: document.id.clone(),
420 target: target.clone(),
421 })?;
422 if !eligible(target_document) {
423 return Err(ConfigError::IneligibleDelegationTarget {
424 agent: document.id.clone(),
425 target: target.clone(),
426 });
427 }
428 expanded.insert(target);
429 }
430 if expanded.len() > MAX_LIST {
431 return Err(ConfigError::AgentLimit(document.id.clone()));
432 }
433 }
434 for (action, value) in &frontmatter.permissions {
435 if !matches!(action, PermissionAction::Read | PermissionAction::Write)
436 && value.rules(*action).iter().any(|rule| {
437 rule.resource
438 .as_str()
439 .contains(WildcardPattern::WORKSPACE_DIR_EXPRESSION)
440 })
441 {
442 return Err(ConfigError::AgentPermissionExpression(document.id.clone()));
443 }
444 }
445 Ok(())
446}