Skip to main content

agent_runtime_definition_builder/
lib.rs

1//! Ergonomic builders for Runtime Agent packages.
2
3use std::sync::Arc;
4
5use runtime_extension_api::{
6    AgentPackage, CompletionMode, ContextHook, ContextSource, Guard, Observer, PackageManifest,
7    Tool,
8};
9use runtime_types::PackageId;
10use sha2::{Digest, Sha256};
11use thiserror::Error;
12
13#[derive(Debug, Error)]
14pub enum PackageBuildError {
15    #[error("invalid package id: {0}")]
16    InvalidId(String),
17    #[error("package version must not be empty")]
18    EmptyVersion,
19    #[error("runtime type must not be empty")]
20    EmptyRuntimeType,
21    #[error("system prompt must not be empty")]
22    EmptySystemPrompt,
23    #[error("duplicate tool: {0}")]
24    DuplicateTool(String),
25    #[error("required completion tool is not registered: {0}")]
26    MissingCompletionTool(String),
27}
28
29pub struct AgentPackageBuilder {
30    id: String,
31    version: String,
32    runtime_type: String,
33    system_prompt: String,
34    tools: Vec<Arc<dyn Tool>>,
35    context_sources: Vec<Arc<dyn ContextSource>>,
36    context_hooks: Vec<Arc<dyn ContextHook>>,
37    guards: Vec<Arc<dyn Guard>>,
38    observers: Vec<Arc<dyn Observer>>,
39    completion_mode: CompletionMode,
40}
41
42impl AgentPackageBuilder {
43    pub fn new(id: impl Into<String>, version: impl Into<String>) -> Self {
44        Self {
45            id: id.into(),
46            version: version.into(),
47            runtime_type: "generic".into(),
48            system_prompt: String::new(),
49            tools: Vec::new(),
50            context_sources: Vec::new(),
51            context_hooks: Vec::new(),
52            guards: Vec::new(),
53            observers: Vec::new(),
54            completion_mode: CompletionMode::ModelOrTool,
55        }
56    }
57
58    pub fn runtime_type(mut self, value: impl Into<String>) -> Self {
59        self.runtime_type = value.into();
60        self
61    }
62    pub fn system_prompt(mut self, value: impl Into<String>) -> Self {
63        self.system_prompt = value.into();
64        self
65    }
66    pub fn tool(mut self, value: Arc<dyn Tool>) -> Self {
67        self.tools.push(value);
68        self
69    }
70    pub fn context_source(mut self, value: Arc<dyn ContextSource>) -> Self {
71        self.context_sources.push(value);
72        self
73    }
74    pub fn context_hook(mut self, value: Arc<dyn ContextHook>) -> Self {
75        self.context_hooks.push(value);
76        self
77    }
78    pub fn guard(mut self, value: Arc<dyn Guard>) -> Self {
79        self.guards.push(value);
80        self
81    }
82    pub fn observer(mut self, value: Arc<dyn Observer>) -> Self {
83        self.observers.push(value);
84        self
85    }
86    pub fn completion_mode(mut self, value: CompletionMode) -> Self {
87        self.completion_mode = value;
88        self
89    }
90
91    pub fn build(self) -> Result<Arc<AgentPackage>, PackageBuildError> {
92        let id = PackageId::new(self.id)
93            .map_err(|error| PackageBuildError::InvalidId(error.to_string()))?;
94        if self.version.trim().is_empty() {
95            return Err(PackageBuildError::EmptyVersion);
96        }
97        if self.runtime_type.trim().is_empty() {
98            return Err(PackageBuildError::EmptyRuntimeType);
99        }
100        if self.system_prompt.trim().is_empty() {
101            return Err(PackageBuildError::EmptySystemPrompt);
102        }
103        let mut names = std::collections::BTreeSet::new();
104        for tool in &self.tools {
105            let name = tool.descriptor().name;
106            if !names.insert(name.clone()) {
107                return Err(PackageBuildError::DuplicateTool(name));
108            }
109        }
110        if let CompletionMode::RequiredTool(name) = &self.completion_mode
111            && !names.contains(name)
112        {
113            return Err(PackageBuildError::MissingCompletionTool(name.clone()));
114        }
115        let digest = package_digest(
116            id.as_str(),
117            &self.version,
118            &self.runtime_type,
119            &self.system_prompt,
120            names.iter().map(String::as_str),
121        );
122        Ok(Arc::new(AgentPackage {
123            manifest: PackageManifest {
124                id,
125                version: self.version,
126                runtime_type: self.runtime_type,
127                schema_version: 1,
128                digest,
129            },
130            system_prompt: self.system_prompt,
131            tools: self.tools,
132            context_sources: self.context_sources,
133            context_hooks: self.context_hooks,
134            guards: self.guards,
135            observers: self.observers,
136            completion_mode: self.completion_mode,
137        }))
138    }
139}
140
141fn package_digest<'a>(
142    id: &str,
143    version: &str,
144    runtime_type: &str,
145    prompt: &str,
146    tools: impl Iterator<Item = &'a str>,
147) -> String {
148    let mut digest = Sha256::new();
149    for value in [id, version, runtime_type, prompt] {
150        digest.update(value.len().to_be_bytes());
151        digest.update(value.as_bytes());
152    }
153    for tool in tools {
154        digest.update(tool.len().to_be_bytes());
155        digest.update(tool.as_bytes());
156    }
157    format!("sha256:{:x}", digest.finalize())
158}
159
160#[cfg(test)]
161mod tests {
162    use super::*;
163
164    #[test]
165    fn package_digest_is_reproducible() {
166        let one = package_digest("a", "1", "generic", "prompt", ["x", "y"].into_iter());
167        let two = package_digest("a", "1", "generic", "prompt", ["x", "y"].into_iter());
168        assert_eq!(one, two);
169    }
170}