Skip to main content

agent_runtime_definition_loader/
lib.rs

1//! Filesystem loader and immutable catalog for declarative Agent definitions.
2
3use std::{
4    collections::BTreeMap,
5    fs,
6    path::{Path, PathBuf},
7    sync::Arc,
8};
9
10use async_trait::async_trait;
11use runtime_control_tools::CompleteTool;
12use runtime_extension_api::{AgentDefinition, CompletionMode};
13use runtime_extension_sdk::AgentDefinitionBuilder;
14use runtime_kernel_api::{AgentDefinitionResolver, KernelFailure};
15use runtime_ports::{OperationContext, ResolvedExecutionContext};
16use serde::Deserialize;
17use sha2::{Digest, Sha256};
18use thiserror::Error;
19
20const MAX_RESOURCE_BYTES: u64 = 1024 * 1024;
21
22#[derive(Debug, Error)]
23pub enum DefinitionLoadError {
24    #[error("definition I/O failed: {0}")]
25    Io(#[from] std::io::Error),
26    #[error("definition manifest is invalid: {0}")]
27    Manifest(#[from] toml::de::Error),
28    #[error("definition contract is invalid: {0}")]
29    Invalid(String),
30}
31
32#[derive(Deserialize)]
33#[serde(deny_unknown_fields)]
34struct Manifest {
35    schema_version: String,
36    id: String,
37    version: String,
38    runtime_api_version: u32,
39    runtime_type: String,
40    system: String,
41    tools: Tools,
42    #[serde(default)]
43    eval: Option<Eval>,
44}
45
46#[derive(Deserialize)]
47#[serde(deny_unknown_fields)]
48struct Tools {
49    #[serde(default)]
50    required: Vec<String>,
51    #[serde(default)]
52    enabled: Vec<String>,
53}
54
55#[derive(Deserialize)]
56#[serde(deny_unknown_fields)]
57struct Eval {
58    suite: String,
59}
60
61pub fn load_definition(root: &Path) -> Result<Arc<AgentDefinition>, DefinitionLoadError> {
62    Ok(load_definition_report(root)?.definition)
63}
64
65pub struct DefinitionLoadReport {
66    pub definition: Arc<AgentDefinition>,
67}
68
69pub fn load_definition_report(root: &Path) -> Result<DefinitionLoadReport, DefinitionLoadError> {
70    let root = root.canonicalize()?;
71    if !root.is_dir() {
72        return Err(DefinitionLoadError::Invalid(
73            "definition root is not a directory".into(),
74        ));
75    }
76    let manifest_path = safe_file(&root, Path::new("agent.toml"))?;
77    let manifest_bytes = bounded_read(&manifest_path)?;
78    let manifest: Manifest = toml::from_slice(&manifest_bytes)?;
79    if manifest.schema_version != "v1" || manifest.runtime_api_version != 3 {
80        return Err(DefinitionLoadError::Invalid(
81            "definition requires schema v1 and Runtime API version 3".into(),
82        ));
83    }
84    if let Some(eval) = &manifest.eval {
85        let _ = safe_file(&root, Path::new(&eval.suite))?;
86    }
87    let system_path = safe_file(&root, Path::new(&manifest.system))?;
88    let system_bytes = bounded_read(&system_path)?;
89    let system = String::from_utf8(system_bytes.clone())
90        .map_err(|_| DefinitionLoadError::Invalid("system prompt must be UTF-8".into()))?;
91    let mut tool_names = manifest.tools.enabled.clone();
92    tool_names.extend(manifest.tools.required.clone());
93    tool_names.sort();
94    tool_names.dedup();
95    if tool_names != ["runtime.complete"] {
96        return Err(DefinitionLoadError::Invalid(
97            "declarative v1 definitions may enable only the built-in runtime.complete tool; add typed tools through an extension catalog".into(),
98        ));
99    }
100    let completion = if manifest
101        .tools
102        .required
103        .iter()
104        .any(|name| name == "runtime.complete")
105    {
106        CompletionMode::RequiredTool("runtime.complete".into())
107    } else {
108        CompletionMode::ModelOrTool
109    };
110    let definition = AgentDefinitionBuilder::new(&manifest.id, &manifest.version)
111        .runtime_type(&manifest.runtime_type)
112        .system_prompt(system)
113        .tool(Arc::new(CompleteTool))
114        .completion_mode(completion)
115        .build()
116        .map_err(|error| DefinitionLoadError::Invalid(error.to_string()))?;
117    let mut definition = Arc::into_inner(definition).ok_or_else(|| {
118        DefinitionLoadError::Invalid("definition builder returned a shared definition".into())
119    })?;
120    let mut digest = Sha256::new();
121    for bytes in [&manifest_bytes, &system_bytes] {
122        digest.update(bytes.len().to_be_bytes());
123        digest.update(bytes);
124    }
125    definition.manifest.digest = format!("sha256:{:x}", digest.finalize());
126    Ok(DefinitionLoadReport {
127        definition: Arc::new(definition),
128    })
129}
130
131pub struct DefinitionCatalogResolver {
132    definitions: BTreeMap<(String, String), Arc<AgentDefinition>>,
133}
134
135impl DefinitionCatalogResolver {
136    pub fn new(
137        definitions: impl IntoIterator<Item = Arc<AgentDefinition>>,
138    ) -> Result<Self, DefinitionLoadError> {
139        let mut catalog = BTreeMap::new();
140        for definition in definitions {
141            let key = (
142                definition.manifest.id.to_string(),
143                definition.manifest.version.clone(),
144            );
145            if catalog.insert(key.clone(), definition).is_some() {
146                return Err(DefinitionLoadError::Invalid(format!(
147                    "duplicate definition {}@{}",
148                    key.0, key.1
149                )));
150            }
151        }
152        if catalog.is_empty() {
153            return Err(DefinitionLoadError::Invalid(
154                "definition catalog is empty".into(),
155            ));
156        }
157        Ok(Self {
158            definitions: catalog,
159        })
160    }
161}
162
163#[async_trait]
164impl AgentDefinitionResolver for DefinitionCatalogResolver {
165    async fn resolve(
166        &self,
167        _: &OperationContext,
168        resolved: &ResolvedExecutionContext,
169    ) -> Result<Arc<AgentDefinition>, KernelFailure> {
170        let definition = self
171            .definitions
172            .get(&(
173                resolved.definition_id.clone(),
174                resolved.definition_version.clone(),
175            ))
176            .cloned()
177            .ok_or_else(|| {
178                KernelFailure::new("DEFINITION_NOT_FOUND", "Agent definition is not loaded")
179            })?;
180        if definition.manifest.runtime_type != resolved.runtime_type
181            || resolved
182                .definition_digest
183                .as_ref()
184                .is_some_and(|digest| digest != &definition.manifest.digest)
185        {
186            return Err(KernelFailure::new(
187                "DEFINITION_BINDING_MISMATCH",
188                "Agent definition does not match the immutable Runtime binding",
189            ));
190        }
191        Ok(definition)
192    }
193}
194
195fn safe_file(root: &Path, relative: &Path) -> Result<PathBuf, DefinitionLoadError> {
196    if relative.as_os_str().is_empty()
197        || relative.is_absolute()
198        || relative
199            .components()
200            .any(|part| !matches!(part, std::path::Component::Normal(_)))
201    {
202        return Err(DefinitionLoadError::Invalid(format!(
203            "unsafe resource path {}",
204            relative.display()
205        )));
206    }
207    let path = root.join(relative).canonicalize()?;
208    if !path.starts_with(root) || !path.is_file() {
209        return Err(DefinitionLoadError::Invalid(format!(
210            "resource escapes definition root: {}",
211            relative.display()
212        )));
213    }
214    Ok(path)
215}
216
217fn bounded_read(path: &Path) -> Result<Vec<u8>, DefinitionLoadError> {
218    if fs::metadata(path)?.len() > MAX_RESOURCE_BYTES {
219        return Err(DefinitionLoadError::Invalid(format!(
220            "resource exceeds {MAX_RESOURCE_BYTES} bytes: {}",
221            path.display()
222        )));
223    }
224    Ok(fs::read(path)?)
225}
226
227#[cfg(test)]
228mod tests {
229    use super::*;
230
231    #[test]
232    fn loads_a_minimal_non_coding_definition_and_rejects_escape() {
233        let temp = tempfile::tempdir().unwrap();
234        fs::write(temp.path().join("system.md"), "Finish the task.").unwrap();
235        fs::write(
236            temp.path().join("agent.toml"),
237            r#"schema_version="v1"
238id="com.example.fixture"
239version="0.1.0"
240runtime_api_version=3
241runtime_type="generic"
242system="system.md"
243[tools]
244required=["runtime.complete"]
245enabled=["runtime.complete"]
246"#,
247        )
248        .unwrap();
249        let definition = load_definition(temp.path()).unwrap();
250        assert_eq!(definition.manifest.runtime_type, "generic");
251        assert!(definition.tool("runtime.complete").is_some());
252        fs::write(
253            temp.path().join("agent.toml"),
254            fs::read_to_string(temp.path().join("agent.toml"))
255                .unwrap()
256                .replace("system.md", "../outside.md"),
257        )
258        .unwrap();
259        assert!(load_definition(temp.path()).is_err());
260    }
261}