cookie_agent_config/
agent_document.rs

1use cookie_agent_identity::AgentId;
2use cookie_agent_models::Sha256Digest;
3pub use cookie_agent_protocol::AgentDocumentSource;
4use serde::Serialize;
5use sha2::{Digest as _, Sha256};
6use std::path::Path;
7
8use crate::{AgentFrontmatter, ConfigError};
9
10const MAX_FRONTMATTER_BYTES: usize = 128 * 1024;
11const MAX_BODY_BYTES: usize = 128 * 1024;
12const MAX_LIST: usize = 256;
13
14#[derive(Clone, Debug)]
15pub struct AgentDocument {
16    pub id: AgentId,
17    pub frontmatter: AgentFrontmatter,
18    pub body: String,
19    pub source: AgentDocumentSource,
20    pub document_fingerprint: Sha256Digest,
21    pub prompt_fingerprint: Sha256Digest,
22}
23
24pub(crate) fn parse_agent(
25    id: AgentId,
26    bytes: &[u8],
27    source: AgentDocumentSource,
28    path: &Path,
29) -> Result<AgentDocument, ConfigError> {
30    let text = std::str::from_utf8(bytes)
31        .map_err(|_| invalid_document(path, "content is not UTF-8"))?
32        .replace("\r\n", "\n")
33        .replace('\r', "\n");
34    if text.contains("${env:") {
35        return Err(invalid_document(
36            path,
37            "environment interpolation is not allowed",
38        ));
39    }
40    let rest = text
41        .strip_prefix("---\n")
42        .ok_or_else(|| invalid_document(path, "frontmatter must start with `---`"))?;
43    let closing = rest
44        .find("\n---\n")
45        .ok_or_else(|| invalid_document(path, "frontmatter must end with `---`"))?;
46    let yaml = &rest[..closing];
47    let raw_body = &rest[closing + "\n---\n".len()..];
48    if yaml.len() > MAX_FRONTMATTER_BYTES || raw_body.len() > MAX_BODY_BYTES || forbidden_yaml(yaml)
49    {
50        return Err(invalid_document(
51            path,
52            "frontmatter is too large or uses unsupported YAML features",
53        ));
54    }
55    let yaml_value: serde_yaml::Value =
56        serde_yaml::from_str(yaml).map_err(|error| yaml_error(path, &error))?;
57    validate_yaml_limits(&yaml_value, 0)
58        .map_err(|_| invalid_document(path, "YAML resource limit exceeded"))?;
59    reject_schema_field(&yaml_value, yaml, path)?;
60    reject_renamed_models_field(&yaml_value, &id)?;
61    reject_removed_tools_field(&yaml_value, &id)?;
62    validate_permission_expressions(&yaml_value, &id)?;
63    let frontmatter: AgentFrontmatter =
64        serde_yaml::from_str(yaml).map_err(|error| yaml_error(path, &error))?;
65    let body = format!("{}\n", raw_body.trim_end_matches('\n'));
66    if !body.chars().any(|character| !character.is_whitespace()) {
67        return Err(ConfigError::EmptyPrompt(id));
68    }
69    let document_fingerprint = hash(
70        "cookie-agent/agent-document/v1",
71        &(id.as_str(), yaml, &body),
72    );
73    let prompt_fingerprint = hash("cookie-agent/system-prompt/v1", &body);
74    Ok(AgentDocument {
75        id,
76        frontmatter,
77        body,
78        source,
79        document_fingerprint,
80        prompt_fingerprint,
81    })
82}
83
84fn reject_schema_field(
85    value: &serde_yaml::Value,
86    yaml: &str,
87    path: &Path,
88) -> Result<(), ConfigError> {
89    let schema = serde_yaml::Value::String("schema".to_owned());
90    if value
91        .as_mapping()
92        .is_some_and(|frontmatter| frontmatter.contains_key(&schema))
93    {
94        let line = yaml
95            .lines()
96            .position(|line| line.trim_start().starts_with("schema:"))
97            .map_or(2, |index| index + 2);
98        return Err(ConfigError::AgentSchemaRemoved {
99            path: path.to_owned(),
100            line,
101        });
102    }
103    Ok(())
104}
105
106fn invalid_document(path: &Path, message: impl Into<String>) -> ConfigError {
107    ConfigError::AgentDocument {
108        path: path.to_owned(),
109        message: message.into(),
110    }
111}
112
113fn yaml_error(path: &Path, error: &serde_yaml::Error) -> ConfigError {
114    let location = error
115        .location()
116        .map(|location| {
117            format!(
118                " at line {}, column {}",
119                location.line() + 1,
120                location.column()
121            )
122        })
123        .unwrap_or_default();
124    let rendered = error.to_string();
125    let message = extract_serde_field(&rendered, "unknown field")
126        .map(|field| format!("unknown field `{field}`{location}; remove it"))
127        .or_else(|| {
128            extract_serde_field(&rendered, "missing field")
129                .map(|field| format!("missing required field `{field}`{location}"))
130        })
131        .unwrap_or_else(|| format!("malformed YAML frontmatter{location}"));
132    invalid_document(path, message)
133}
134
135fn extract_serde_field(message: &str, marker: &str) -> Option<String> {
136    let rest = message.split_once(marker)?.1.trim_start();
137    let quote = rest.chars().next()?;
138    if !matches!(quote, '`' | '\'' | '"') {
139        return None;
140    }
141    rest[quote.len_utf8()..]
142        .split_once(quote)
143        .map(|(field, _)| field.to_owned())
144}
145
146fn reject_removed_tools_field(
147    value: &serde_yaml::Value,
148    agent: &AgentId,
149) -> Result<(), ConfigError> {
150    let tools = serde_yaml::Value::String("tools".to_owned());
151    if value
152        .as_mapping()
153        .is_some_and(|frontmatter| frontmatter.contains_key(&tools))
154    {
155        return Err(ConfigError::AgentToolsRemoved(agent.clone()));
156    }
157    Ok(())
158}
159
160fn reject_renamed_models_field(
161    value: &serde_yaml::Value,
162    agent: &AgentId,
163) -> Result<(), ConfigError> {
164    let legacy = serde_yaml::Value::String("model_fallback".to_owned());
165    if value
166        .as_mapping()
167        .is_some_and(|frontmatter| frontmatter.contains_key(&legacy))
168    {
169        return Err(ConfigError::AgentModelsRenamed(agent.clone()));
170    }
171    Ok(())
172}
173
174fn validate_permission_expressions(
175    value: &serde_yaml::Value,
176    agent: &AgentId,
177) -> Result<(), ConfigError> {
178    let Some(frontmatter) = value.as_mapping() else {
179        return Ok(());
180    };
181    let permissions_key = serde_yaml::Value::String("permissions".to_owned());
182    let Some(permissions) = frontmatter
183        .get(&permissions_key)
184        .and_then(serde_yaml::Value::as_mapping)
185    else {
186        return Ok(());
187    };
188    for (action, permission) in permissions {
189        let Some(action) = action.as_str() else {
190            continue;
191        };
192        let Some(resources) = permission.as_mapping() else {
193            continue;
194        };
195        for resource in resources.keys().filter_map(serde_yaml::Value::as_str) {
196            let remainder = resource.replace("${workspace_dir}", "");
197            let malformed_expression = remainder.contains("${");
198            let unsupported_action =
199                resource.contains("${workspace_dir}") && !matches!(action, "read" | "write");
200            if malformed_expression || unsupported_action {
201                return Err(ConfigError::AgentPermissionExpression(agent.clone()));
202            }
203        }
204    }
205    Ok(())
206}
207
208fn forbidden_yaml(yaml: &str) -> bool {
209    yaml.lines().any(|line| {
210        if line.trim_start().starts_with("<<:") {
211            return true;
212        }
213        let mut single = false;
214        let mut double = false;
215        let mut escaped = false;
216        for (index, character) in line.char_indices() {
217            if escaped {
218                escaped = false;
219                continue;
220            }
221            if double && character == '\\' {
222                escaped = true;
223                continue;
224            }
225            match character {
226                '\'' if !double => single = !single,
227                '"' if !single => double = !double,
228                '#' if !single && !double => break,
229                '&' | '!' | '*' if !single && !double => {
230                    let boundary = index == 0
231                        || line[..index]
232                            .chars()
233                            .next_back()
234                            .is_some_and(char::is_whitespace);
235                    if boundary {
236                        return true;
237                    }
238                }
239                _ => {}
240            }
241        }
242        false
243    })
244}
245
246fn validate_yaml_limits(value: &serde_yaml::Value, depth: usize) -> Result<(), ConfigError> {
247    if depth > 16 {
248        return Err(ConfigError::AgentYamlLimit);
249    }
250    match value {
251        serde_yaml::Value::Sequence(values) => {
252            if values.len() > MAX_LIST {
253                return Err(ConfigError::AgentYamlLimit);
254            }
255            for value in values {
256                validate_yaml_limits(value, depth + 1)?;
257            }
258        }
259        serde_yaml::Value::Mapping(values) => {
260            if values.len() > MAX_LIST {
261                return Err(ConfigError::AgentYamlLimit);
262            }
263            for (key, value) in values {
264                validate_yaml_limits(key, depth + 1)?;
265                validate_yaml_limits(value, depth + 1)?;
266            }
267        }
268        serde_yaml::Value::Tagged(_) => return Err(ConfigError::AgentYamlLimit),
269        serde_yaml::Value::String(value) if value.len() > 128 * 1024 => {
270            return Err(ConfigError::AgentYamlLimit);
271        }
272        _ => {}
273    }
274    Ok(())
275}
276
277fn hash(domain: &str, value: &impl Serialize) -> Sha256Digest {
278    let mut hasher = Sha256::new();
279    hasher.update(domain.as_bytes());
280    hasher.update([0]);
281    hasher.update(serde_json::to_vec(value).expect("strings always serialize to JSON"));
282    Sha256Digest::new(format!("{:x}", hasher.finalize()))
283        .expect("SHA-256 output is always a valid digest")
284}
cookie_agent_config/agent_document.rs

cookie_agent_config/
agent_document.rs
