1use std::collections::BTreeSet;
2use std::path::{Path, PathBuf};
3
4use agentkit_core::{Item, ItemKind, MetadataMap, Part, TextPart};
5use async_trait::async_trait;
6use futures_lite::StreamExt;
7use serde_json::Value;
8use thiserror::Error;
9
10const DEFAULT_AGENTS_FILE: &str = "AGENTS.md";
11const DEFAULT_SKILL_FILE: &str = "SKILL.md";
12
13#[derive(Clone, Copy, Debug, PartialEq, Eq)]
14pub enum AgentsMdMode {
15 Nearest,
16 All,
17}
18
19#[async_trait]
20pub trait ContextSource: Send + Sync {
21 async fn load(&self) -> Result<Vec<Item>, ContextError>;
22}
23
24#[derive(Default)]
25pub struct ContextLoader {
26 sources: Vec<Box<dyn ContextSource>>,
27}
28
29impl ContextLoader {
30 pub fn new() -> Self {
31 Self::default()
32 }
33
34 pub fn with_source(mut self, source: impl ContextSource + 'static) -> Self {
35 self.sources.push(Box::new(source));
36 self
37 }
38
39 pub async fn load(&self) -> Result<Vec<Item>, ContextError> {
40 let mut items = Vec::new();
41
42 for source in &self.sources {
43 items.extend(source.load().await?);
44 }
45
46 Ok(items)
47 }
48}
49
50#[derive(Clone, Debug)]
51pub struct AgentsMd {
52 start_dir: PathBuf,
53 mode: AgentsMdMode,
54 file_name: String,
55 explicit_paths: Vec<PathBuf>,
56 search_dirs: Vec<PathBuf>,
57}
58
59impl AgentsMd {
60 pub fn discover(start_dir: impl Into<PathBuf>) -> Self {
61 Self {
62 start_dir: start_dir.into(),
63 mode: AgentsMdMode::Nearest,
64 file_name: DEFAULT_AGENTS_FILE.into(),
65 explicit_paths: Vec::new(),
66 search_dirs: Vec::new(),
67 }
68 }
69
70 pub fn discover_all(start_dir: impl Into<PathBuf>) -> Self {
71 Self::discover(start_dir).with_mode(AgentsMdMode::All)
72 }
73
74 pub fn with_mode(mut self, mode: AgentsMdMode) -> Self {
75 self.mode = mode;
76 self
77 }
78
79 pub fn with_file_name(mut self, file_name: impl Into<String>) -> Self {
80 self.file_name = file_name.into();
81 self
82 }
83
84 pub fn with_path(mut self, path: impl Into<PathBuf>) -> Self {
85 self.explicit_paths.push(path.into());
86 self
87 }
88
89 pub fn with_search_dir(mut self, dir: impl Into<PathBuf>) -> Self {
90 self.search_dirs.push(dir.into());
91 self
92 }
93
94 pub async fn resolve(&self) -> Result<Option<PathBuf>, ContextError> {
95 Ok(self.resolve_all().await?.into_iter().next())
96 }
97
98 pub async fn resolve_all(&self) -> Result<Vec<PathBuf>, ContextError> {
99 let mut paths = Vec::new();
100
101 for path in &self.explicit_paths {
102 if path_exists(path).await? {
103 paths.push(path.clone());
104 }
105 }
106
107 for dir in &self.search_dirs {
108 let candidate = dir.join(&self.file_name);
109 if path_exists(&candidate).await? {
110 paths.push(candidate);
111 }
112 }
113
114 paths.extend(
115 find_in_ancestors_with_mode(
116 &self.start_dir,
117 &self.file_name,
118 self.mode == AgentsMdMode::All,
119 )
120 .await?,
121 );
122
123 let mut seen = BTreeSet::new();
124 paths.retain(|path| seen.insert(path.clone()));
125 if self.mode == AgentsMdMode::Nearest {
126 Ok(paths.into_iter().rev().take(1).collect())
127 } else {
128 Ok(paths)
129 }
130 }
131}
132
133#[async_trait]
134impl ContextSource for AgentsMd {
135 async fn load(&self) -> Result<Vec<Item>, ContextError> {
136 let paths = self.resolve_all().await?;
137 let mut items = Vec::with_capacity(paths.len());
138
139 for path in paths {
140 let body = async_fs::read_to_string(&path).await.map_err(|error| {
141 ContextError::ReadFailed {
142 path: path.clone(),
143 error,
144 }
145 })?;
146
147 items.push(context_item(
148 format!(
149 "[Loaded AGENTS]\nPath: {}\n\n{}",
150 path.display(),
151 body.trim_end()
152 ),
153 metadata_for("agents_md", &path, None),
154 ));
155 }
156
157 Ok(items)
158 }
159}
160
161#[derive(Clone, Debug)]
162pub struct SkillsDirectory {
163 roots: Vec<PathBuf>,
164 skill_file_name: String,
165}
166
167impl SkillsDirectory {
168 pub fn from_dir(root: impl Into<PathBuf>) -> Self {
169 Self {
170 roots: vec![root.into()],
171 skill_file_name: DEFAULT_SKILL_FILE.into(),
172 }
173 }
174
175 pub fn with_dir(mut self, root: impl Into<PathBuf>) -> Self {
176 self.roots.push(root.into());
177 self
178 }
179
180 pub fn with_skill_file_name(mut self, skill_file_name: impl Into<String>) -> Self {
181 self.skill_file_name = skill_file_name.into();
182 self
183 }
184}
185
186#[async_trait]
187impl ContextSource for SkillsDirectory {
188 async fn load(&self) -> Result<Vec<Item>, ContextError> {
189 let mut skill_paths = Vec::new();
190 for root in &self.roots {
191 if !path_exists(root).await? {
192 continue;
193 }
194 skill_paths.extend(collect_skill_files(root, &self.skill_file_name).await?);
195 }
196 skill_paths.sort();
197 skill_paths.dedup();
198
199 let mut items = Vec::with_capacity(skill_paths.len());
200
201 for path in skill_paths {
202 let body = async_fs::read_to_string(&path).await.map_err(|error| {
203 ContextError::ReadFailed {
204 path: path.clone(),
205 error,
206 }
207 })?;
208 let skill_name = path
209 .parent()
210 .and_then(Path::file_name)
211 .map(|value| value.to_string_lossy().into_owned());
212
213 items.push(context_item(
214 format!(
215 "[Loaded Skill]\nName: {}\nPath: {}\n\n{}",
216 skill_name.clone().unwrap_or_else(|| "unknown".into()),
217 path.display(),
218 body.trim_end()
219 ),
220 metadata_for("skill", &path, skill_name),
221 ));
222 }
223
224 Ok(items)
225 }
226}
227
228fn context_item(text: String, metadata: MetadataMap) -> Item {
229 Item {
230 id: None,
231 kind: ItemKind::Context,
232 parts: vec![Part::Text(TextPart {
233 text,
234 metadata: MetadataMap::new(),
235 })],
236 metadata,
237 }
238}
239
240fn metadata_for(source_kind: &str, path: &Path, name: Option<String>) -> MetadataMap {
241 let mut metadata = MetadataMap::new();
242 metadata.insert(
243 "agentkit.context.source".into(),
244 Value::String(source_kind.into()),
245 );
246 metadata.insert(
247 "agentkit.context.path".into(),
248 Value::String(path.display().to_string()),
249 );
250 if let Some(name) = name {
251 metadata.insert("agentkit.context.name".into(), Value::String(name));
252 }
253 metadata
254}
255
256async fn path_exists(path: &Path) -> Result<bool, ContextError> {
257 match async_fs::metadata(path).await {
258 Ok(_) => Ok(true),
259 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(false),
260 Err(error) => Err(ContextError::InspectFailed {
261 path: path.to_path_buf(),
262 error,
263 }),
264 }
265}
266
267async fn find_in_ancestors_with_mode(
268 start_dir: &Path,
269 file_name: &str,
270 include_all: bool,
271) -> Result<Vec<PathBuf>, ContextError> {
272 let mut current = start_dir.to_path_buf();
273 let mut matches = Vec::new();
274
275 loop {
276 let candidate = current.join(file_name);
277 if path_exists(&candidate).await? {
278 matches.push(candidate);
279 if !include_all {
280 break;
281 }
282 }
283 let Some(parent) = current.parent() else {
284 break;
285 };
286 current = parent.to_path_buf();
287 }
288
289 matches.reverse();
290 Ok(matches)
291}
292
293async fn collect_skill_files(
294 root: &Path,
295 skill_file_name: &str,
296) -> Result<Vec<PathBuf>, ContextError> {
297 let mut pending = vec![root.to_path_buf()];
298 let mut skill_paths = Vec::new();
299
300 while let Some(dir_path) = pending.pop() {
301 let mut read_dir =
302 async_fs::read_dir(&dir_path)
303 .await
304 .map_err(|error| ContextError::InspectFailed {
305 path: dir_path.clone(),
306 error,
307 })?;
308
309 while let Some(entry) = read_dir.next().await {
310 let entry = entry.map_err(|error| ContextError::InspectFailed {
311 path: dir_path.clone(),
312 error,
313 })?;
314 let path = entry.path();
315 let file_type =
316 entry
317 .file_type()
318 .await
319 .map_err(|error| ContextError::InspectFailed {
320 path: path.clone(),
321 error,
322 })?;
323
324 if file_type.is_dir() {
325 pending.push(path);
326 continue;
327 }
328
329 if file_type.is_file() && path.file_name().is_some_and(|name| name == skill_file_name) {
330 skill_paths.push(path);
331 }
332 }
333 }
334
335 skill_paths.sort();
336 Ok(skill_paths)
337}
338
339#[derive(Debug, Error)]
340pub enum ContextError {
341 #[error("failed to inspect {path}: {error}")]
342 InspectFailed {
343 path: PathBuf,
344 #[source]
345 error: std::io::Error,
346 },
347 #[error("failed to read {path}: {error}")]
348 ReadFailed {
349 path: PathBuf,
350 #[source]
351 error: std::io::Error,
352 },
353}
354
355#[cfg(test)]
356mod tests {
357 use std::time::{SystemTime, UNIX_EPOCH};
358
359 use super::*;
360
361 #[tokio::test]
362 async fn discovers_agents_file_in_ancestors() {
363 let root = temp_path("agentkit-context-agents");
364 let nested = root.join("nested/project");
365 async_fs::create_dir_all(&nested).await.unwrap();
366 let agents_path = root.join("AGENTS.md");
367 async_fs::write(&agents_path, "project = lantern")
368 .await
369 .unwrap();
370
371 let items = AgentsMd::discover(&nested).load().await.unwrap();
372 assert_eq!(items.len(), 1);
373 assert_eq!(items[0].kind, ItemKind::Context);
374 assert_eq!(
375 items[0].metadata.get("agentkit.context.source"),
376 Some(&Value::String("agents_md".into()))
377 );
378
379 async_fs::remove_dir_all(&root).await.unwrap();
380 }
381
382 #[tokio::test]
383 async fn discovers_all_agents_files_when_requested() {
384 let root = temp_path("agentkit-context-agents-all");
385 let nested = root.join("nested/project");
386 async_fs::create_dir_all(&nested).await.unwrap();
387 async_fs::write(root.join("AGENTS.md"), "project = lantern")
388 .await
389 .unwrap();
390 async_fs::write(root.join("nested/AGENTS.md"), "team = orbit")
391 .await
392 .unwrap();
393
394 let items = AgentsMd::discover_all(&nested).load().await.unwrap();
395 assert_eq!(items.len(), 2);
396
397 async_fs::remove_dir_all(&root).await.unwrap();
398 }
399
400 #[tokio::test]
401 async fn loads_agents_from_explicit_search_paths() {
402 let root = temp_path("agentkit-context-agents-explicit");
403 let nested = root.join("nested/project");
404 let shared = root.join("shared");
405 async_fs::create_dir_all(&nested).await.unwrap();
406 async_fs::create_dir_all(&shared).await.unwrap();
407 async_fs::write(shared.join("AGENTS.md"), "policy = explicit")
408 .await
409 .unwrap();
410
411 let items = AgentsMd::discover(&nested)
412 .with_search_dir(&shared)
413 .load()
414 .await
415 .unwrap();
416 assert_eq!(items.len(), 1);
417 assert!(
418 items[0]
419 .metadata
420 .get("agentkit.context.path")
421 .and_then(Value::as_str)
422 .is_some_and(|path| path.ends_with("/shared/AGENTS.md"))
423 );
424
425 async_fs::remove_dir_all(&root).await.unwrap();
426 }
427
428 #[tokio::test]
429 async fn loads_skills_recursively() {
430 let root = temp_path("agentkit-context-skills");
431 let skill_dir = root.join("skills/release-notes");
432 async_fs::create_dir_all(&skill_dir).await.unwrap();
433 async_fs::write(skill_dir.join("SKILL.md"), "# Release Notes")
434 .await
435 .unwrap();
436
437 let items = SkillsDirectory::from_dir(root.join("skills"))
438 .load()
439 .await
440 .unwrap();
441 assert_eq!(items.len(), 1);
442 assert_eq!(
443 items[0].metadata.get("agentkit.context.name"),
444 Some(&Value::String("release-notes".into()))
445 );
446
447 async_fs::remove_dir_all(&root).await.unwrap();
448 }
449
450 #[tokio::test]
451 async fn loads_skills_from_multiple_roots() {
452 let root = temp_path("agentkit-context-skills-multi");
453 let root_a = root.join("skills-a/release-notes");
454 let root_b = root.join("skills-b/deploy");
455 async_fs::create_dir_all(&root_a).await.unwrap();
456 async_fs::create_dir_all(&root_b).await.unwrap();
457 async_fs::write(root_a.join("SKILL.md"), "# Release Notes")
458 .await
459 .unwrap();
460 async_fs::write(root_b.join("SKILL.md"), "# Deploy")
461 .await
462 .unwrap();
463
464 let items = SkillsDirectory::from_dir(root.join("skills-a"))
465 .with_dir(root.join("skills-b"))
466 .load()
467 .await
468 .unwrap();
469 assert_eq!(items.len(), 2);
470
471 async_fs::remove_dir_all(&root).await.unwrap();
472 }
473
474 fn temp_path(prefix: &str) -> PathBuf {
475 let suffix = SystemTime::now()
476 .duration_since(UNIX_EPOCH)
477 .unwrap()
478 .as_nanos();
479 std::env::temp_dir().join(format!("{prefix}-{suffix}"))
480 }
481}