use super::installer::{ContextPackage, markdown_files};
use super::skill::frontmatter_description;
use bake::{Error, Result};
use socketry_markdown::{
ParseOptions,
mdast::{Heading, Link, Node, Paragraph, Text},
to_mdast,
};
use std::collections::HashMap;
use std::fs;
use std::path::{Path, PathBuf};
#[derive(Clone, Debug)]
struct ContextDocument {
path: PathBuf,
title: String,
description: Option<String>,
}
struct SourceHeading {
index: usize,
level: u8,
title: String,
start: usize,
body_start: usize,
}
#[derive(Clone, Debug)]
pub struct AgentIndex {
root: PathBuf,
context_path: PathBuf,
context_link_path: PathBuf,
package_descriptions: HashMap<String, String>,
}
impl AgentIndex {
pub fn new(root: impl Into<PathBuf>) -> Self {
let root = root.into();
Self {
context_path: root.join(".agents/context"),
root,
context_link_path: PathBuf::from(".agents/context"),
package_descriptions: HashMap::new(),
}
}
pub fn with_packages(mut self, packages: &[ContextPackage]) -> Self {
self.package_descriptions = packages
.iter()
.filter_map(|package| {
package
.description
.as_ref()
.map(|description| (package.selector().to_owned(), description.clone()))
})
.collect();
self
}
pub fn context_path(&self) -> &Path {
&self.context_path
}
pub fn generate_context_section(&self) -> Result<String> {
let mut sections = vec![
"This section links to documentation from installed packages. It is automatically generated and can be refreshed with `cargo bake agent:context:install`.".to_owned(),
String::new(),
"**Before working on a package, read the relevant context files below. They contain package-specific guidance and workflows.**".to_owned(),
String::new(),
"If these files are missing or dependencies have changed, run `cargo bake agent:context:install` to install them.".to_owned(),
String::new(),
];
let packages = self.collect_context_packages()?;
if packages.is_empty() {
sections.push(
"No context files found. Run `cargo bake agent:context:install` to install context from dependencies.".to_owned(),
);
return Ok(sections.join("\n"));
}
for (package_name, files) in packages {
sections.push(format!("### {package_name}"));
sections.push(String::new());
sections.push(
self.package_descriptions
.get(&package_name)
.cloned()
.unwrap_or_else(|| format!("Context files for {package_name}")),
);
sections.push(String::new());
for document in files {
append_document(
&mut sections,
&self.context_link_path.join(&package_name),
&document.path,
&document.title,
document.description.as_deref(),
);
}
}
while sections.last().is_some_and(String::is_empty) {
sections.pop();
}
Ok(sections.join("\n"))
}
pub fn update_agents_md(&self, path: impl AsRef<Path>) -> Result<()> {
let path = path.as_ref();
let path = if path.is_absolute() {
path.to_path_buf()
} else {
self.root.join(path)
};
let context = self.generate_context_section()?;
let existing = match fs::read_to_string(&path) {
Ok(contents) => Some(contents),
Err(error) if error.kind() == std::io::ErrorKind::NotFound => None,
Err(error) => {
return Err(Error::new(format!(
"cannot read {}: {error}",
path.display()
)));
}
};
let updated = match existing {
Some(contents) => update_existing(&contents, &context),
None => format!("# Agent\n\n## Context\n\n{context}\n"),
};
if let Some(parent) = path.parent() {
fs::create_dir_all(parent).map_err(|error| {
Error::new(format!("cannot create {}: {error}", parent.display()))
})?;
}
fs::write(&path, updated)
.map_err(|error| Error::new(format!("cannot write {}: {error}", path.display())))
}
fn collect_context_packages(&self) -> Result<Vec<(String, Vec<ContextDocument>)>> {
let entries = match fs::read_dir(&self.context_path) {
Ok(entries) => entries,
Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()),
Err(error) => {
return Err(Error::new(format!(
"cannot read {}: {error}",
self.context_path.display()
)));
}
};
let mut packages = Vec::new();
for entry in entries {
let entry = entry?;
let file_type = entry.file_type()?;
if !file_type.is_dir() {
continue;
}
let package_path = entry.path();
let mut files = Vec::new();
for path in markdown_files(&package_path)? {
let (title, description) = extract_content(&path)?;
let relative_path = path.strip_prefix(&package_path).map_err(|error| {
Error::new(format!("cannot make context path relative: {error}"))
})?;
files.push(ContextDocument {
path: relative_path.to_path_buf(),
title,
description,
});
}
files.sort_by_key(|document| canonical_order(&document.path));
if !files.is_empty() {
packages.push((entry.file_name().to_string_lossy().into_owned(), files));
}
}
packages.sort_by(|left, right| left.0.cmp(&right.0));
Ok(packages)
}
}
fn update_existing(contents: &str, context: &str) -> String {
let Ok(root) = to_mdast(contents, &ParseOptions::default()) else {
return contents.to_owned();
};
let Some(children) = root.children() else {
return contents.to_owned();
};
let headings: Vec<_> = children
.iter()
.enumerate()
.filter_map(|(index, node)| {
let Node::Heading(heading) = node else {
return None;
};
let position = node.position()?;
Some(SourceHeading {
index,
level: heading.depth,
title: heading_text(node),
start: position.start.offset,
body_start: after_heading_line(contents, position.end.offset),
})
})
.collect();
let Some(agent_heading) = headings
.iter()
.find(|heading| heading.level == 1 && heading.title.eq_ignore_ascii_case("agent"))
else {
return format!("# Agent\n\n## Context\n\n{context}\n\n{contents}");
};
let agent_end_index = headings
.iter()
.find(|heading| heading.index > agent_heading.index && heading.level <= 1)
.map_or(children.len(), |heading| heading.index);
let context_heading = headings.iter().find(|heading| {
heading.index > agent_heading.index
&& heading.index < agent_end_index
&& heading.level == 2
&& heading.title.eq_ignore_ascii_case("context")
});
let newline = if contents.contains("\r\n") {
"\r\n"
} else {
"\n"
};
let context = context.replace("\r\n", "\n").replace('\n', newline);
let mut updated = String::with_capacity(contents.len() + context.len() + 32);
if let Some(context_heading) = context_heading {
let end = headings
.iter()
.find(|heading| heading.index > context_heading.index && heading.level <= 2)
.map_or(contents.len(), |heading| heading.start);
updated.push_str(&contents[..context_heading.start]);
updated.push_str("## Context");
updated.push_str(newline);
updated.push_str(newline);
updated.push_str(&context);
if end < contents.len() {
updated.push_str(newline);
updated.push_str(newline);
} else if contents.ends_with('\n') || !context.is_empty() {
updated.push_str(newline);
}
updated.push_str(&contents[end..]);
} else {
updated.push_str(&contents[..agent_heading.body_start]);
updated.push_str(newline);
updated.push_str("## Context");
updated.push_str(newline);
updated.push_str(newline);
updated.push_str(&context);
if agent_heading.body_start < contents.len() {
updated.push_str(newline);
updated.push_str(newline);
} else if contents.ends_with('\n') || !context.is_empty() {
updated.push_str(newline);
}
updated.push_str(&contents[agent_heading.body_start..]);
}
updated
}
fn after_heading_line(contents: &str, offset: usize) -> usize {
let Some(rest) = contents.get(offset..) else {
return offset;
};
if rest.starts_with("\r\n") {
offset + 2
} else if rest.starts_with('\r') || rest.starts_with('\n') {
offset + 1
} else {
offset
}
}
fn heading_text(node: &Node) -> String {
let text = node.text_content();
text.strip_suffix("\r\n")
.or_else(|| text.strip_suffix('\n'))
.or_else(|| text.strip_suffix('\r'))
.unwrap_or(text.as_str())
.to_owned()
}
fn append_document(
sections: &mut Vec<String>,
link_root: &Path,
relative_path: &Path,
title: &str,
description: Option<&str>,
) {
let heading = Node::Heading(Heading {
children: vec![Node::Link(Link {
children: vec![Node::Text(Text {
value: title.to_owned(),
position: None,
})],
position: None,
url: markdown_link(&link_root.join(relative_path)),
title: None,
})],
position: None,
depth: 4,
});
sections.push(heading.to_markdown().trim_end().to_owned());
sections.push(String::new());
if let Some(description) = description.filter(|description| !description.is_empty()) {
let paragraph = Node::Paragraph(Paragraph {
children: vec![Node::Text(Text {
value: description.to_owned(),
position: None,
})],
position: None,
});
sections.push(paragraph.to_markdown().trim_end().to_owned());
sections.push(String::new());
}
}
fn markdown_link(path: &Path) -> String {
path_to_string(path)
.replace('%', "%25")
.replace(' ', "%20")
.replace('#', "%23")
.replace('?', "%3F")
.replace('(', "%28")
.replace(')', "%29")
}
fn canonical_order(path: &Path) -> (usize, String, String) {
const CANONICAL: &[&str] = &[
"getting-started",
"overview",
"usage",
"configuration",
"migration",
"troubleshooting",
"debugging",
];
let name = path
.file_stem()
.and_then(|name| name.to_str())
.unwrap_or_default()
.to_ascii_lowercase();
let order = CANONICAL
.iter()
.position(|canonical| *canonical == name)
.unwrap_or(CANONICAL.len());
(order, name, path_to_string(path))
}
fn extract_content(path: &Path) -> Result<(String, Option<String>)> {
let content = fs::read_to_string(path)
.map_err(|error| Error::new(format!("cannot read {}: {error}", path.display())))?;
let mut options = ParseOptions::default();
options.constructs.frontmatter = true;
let root = to_mdast(&content, &options)
.map_err(|error| Error::new(format!("could not parse {}: {error}", path.display())))?;
let children = root
.children()
.ok_or_else(|| Error::new(format!("{} is not a Markdown document", path.display())))?;
let title = children
.iter()
.find_map(|node| match node {
Node::Heading(_) => {
let title = heading_text(node);
(!title.trim().is_empty()).then_some(title)
}
_ => None,
})
.unwrap_or_else(|| {
path.file_stem()
.and_then(|stem| stem.to_str())
.unwrap_or("Documentation")
.replace('-', " ")
});
let first_paragraph = children.iter().find_map(|node| match node {
Node::Paragraph(_) => Some(node.text_content()),
_ => None,
});
let description = frontmatter_description(&root, path)?
.or_else(|| first_paragraph.as_deref().and_then(first_sentence));
Ok((title, description))
}
fn first_sentence(paragraph: &str) -> Option<String> {
let paragraph = paragraph.trim();
if paragraph.is_empty() {
return None;
}
for (index, character) in paragraph.char_indices() {
if matches!(character, '.' | '!' | '?')
&& paragraph[index + character.len_utf8()..]
.chars()
.next()
.is_none_or(char::is_whitespace)
{
return Some(paragraph[..index + character.len_utf8()].to_owned());
}
}
Some(paragraph.to_owned())
}
fn path_to_string(path: &Path) -> String {
path.components()
.map(|component| component.as_os_str().to_string_lossy())
.collect::<Vec<_>>()
.join("/")
}