Skip to main content

cookie_agent_config/
agents.rs

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    /// Per-agent AGENTS.md loading. `None` defers to `[agent_md] enabled`;
171    /// a value overrides it for root runs of this agent.
172    #[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                // Internal agents default to no document cap: at run time a zero
216                // document cap inherits the owner run's cap, mirroring how they
217                // inherit `${parent_model}`.
218                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/// Validated authored-agent input for downstream runtime materialization.
227///
228/// `root_eligible` covers only authored policy. P7 computes `runnable_as_root`
229/// by additionally requiring at least one available materialized fallback.
230#[derive(Clone, Copy, Debug)]
231pub struct AgentMaterializationInput<'a> {
232    pub document: &'a AgentDocument,
233    pub root_eligible: bool,
234}
235
236/// Atomically validated authored agent documents, independent of model execution.
237#[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    // Only root runs load AGENTS.md, so the override is meaningless on agents
329    // that can never run as a root.
330    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        // Expanded targets land in AgentDescriptor and FrozenDelegationPolicy,
378        // both limited to 256 IDs; validate the deduplicated union exactly as
379        // the runtime expansion produces it.
380        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                // Wildcard delegation uses the same runtime matching as every
395                // other permission; statically the pattern must match at
396                // least one known eligible subagent so it cannot be dead.
397                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}