use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::Path;
const PRIVATE_LOCK_FILENAME: &str = "agpm.private.lock";
const PRIVATE_LOCK_VERSION: u32 = 1;
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct PrivateLockedResource {
pub name: String,
#[serde(default, skip_serializing_if = "HashMap::is_empty")]
pub applied_patches: HashMap<String, toml::Value>,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct PrivateLockFile {
pub version: u32,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub agents: Vec<PrivateLockedResource>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub snippets: Vec<PrivateLockedResource>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub commands: Vec<PrivateLockedResource>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub scripts: Vec<PrivateLockedResource>,
#[serde(default, skip_serializing_if = "Vec::is_empty", rename = "mcp-servers")]
pub mcp_servers: Vec<PrivateLockedResource>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub hooks: Vec<PrivateLockedResource>,
}
impl Default for PrivateLockFile {
fn default() -> Self {
Self::new()
}
}
impl PrivateLockFile {
pub fn new() -> Self {
Self {
version: PRIVATE_LOCK_VERSION,
agents: Vec::new(),
snippets: Vec::new(),
commands: Vec::new(),
scripts: Vec::new(),
mcp_servers: Vec::new(),
hooks: Vec::new(),
}
}
pub fn load(project_dir: &Path) -> Result<Option<Self>> {
let path = project_dir.join(PRIVATE_LOCK_FILENAME);
if !path.exists() {
return Ok(None);
}
let content = std::fs::read_to_string(&path)
.with_context(|| format!("Failed to read {}", path.display()))?;
let lock: Self = toml::from_str(&content)
.with_context(|| format!("Failed to parse {}", path.display()))?;
if lock.version > PRIVATE_LOCK_VERSION {
anyhow::bail!(
"Private lockfile version {} is newer than supported version {}. \
Please upgrade AGPM.",
lock.version,
PRIVATE_LOCK_VERSION
);
}
Ok(Some(lock))
}
pub fn save(&self, project_dir: &Path) -> Result<()> {
let path = project_dir.join(PRIVATE_LOCK_FILENAME);
if self.is_empty() {
if path.exists() {
std::fs::remove_file(&path)
.with_context(|| format!("Failed to remove {}", path.display()))?;
}
return Ok(());
}
let content = serialize_private_lockfile_with_inline_patches(self)?;
std::fs::write(&path, content)
.with_context(|| format!("Failed to write {}", path.display()))?;
Ok(())
}
pub fn is_empty(&self) -> bool {
self.agents.is_empty()
&& self.snippets.is_empty()
&& self.commands.is_empty()
&& self.scripts.is_empty()
&& self.mcp_servers.is_empty()
&& self.hooks.is_empty()
}
pub fn total_patches(&self) -> usize {
self.agents.len()
+ self.snippets.len()
+ self.commands.len()
+ self.scripts.len()
+ self.mcp_servers.len()
+ self.hooks.len()
}
pub fn add_private_patches(
&mut self,
resource_type: &str,
name: &str,
patches: HashMap<String, toml::Value>,
) {
if patches.is_empty() {
return;
}
let vec = match resource_type {
"agents" => &mut self.agents,
"snippets" => &mut self.snippets,
"commands" => &mut self.commands,
"scripts" => &mut self.scripts,
"mcp-servers" => &mut self.mcp_servers,
"hooks" => &mut self.hooks,
_ => return,
};
vec.retain(|r| r.name != name);
vec.push(PrivateLockedResource {
name: name.to_string(),
applied_patches: patches,
});
}
pub fn get_patches(
&self,
resource_type: &str,
name: &str,
) -> Option<&HashMap<String, toml::Value>> {
let vec = match resource_type {
"agents" => &self.agents,
"snippets" => &self.snippets,
"commands" => &self.commands,
"scripts" => &self.scripts,
"mcp-servers" => &self.mcp_servers,
"hooks" => &self.hooks,
_ => return None,
};
vec.iter().find(|r| r.name == name).map(|r| &r.applied_patches)
}
}
fn serialize_private_lockfile_with_inline_patches(lockfile: &PrivateLockFile) -> Result<String> {
use toml_edit::{DocumentMut, Item};
let toml_str =
toml::to_string_pretty(lockfile).context("Failed to serialize private lockfile to TOML")?;
let mut doc: DocumentMut = toml_str.parse().context("Failed to parse TOML document")?;
let resource_types = ["agents", "snippets", "commands", "scripts", "hooks", "mcp-servers"];
for resource_type in &resource_types {
if let Some(Item::ArrayOfTables(array)) = doc.get_mut(resource_type) {
for table in array.iter_mut() {
if let Some(Item::Table(patches_table)) = table.get_mut("applied_patches") {
let mut inline = toml_edit::InlineTable::new();
for (key, val) in patches_table.iter() {
if let Some(v) = val.as_value() {
inline.insert(key, v.clone());
}
}
table.insert("applied_patches", toml_edit::value(inline));
}
}
}
}
Ok(doc.to_string())
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::TempDir;
#[test]
fn test_new_lockfile_is_empty() {
let lock = PrivateLockFile::new();
assert!(lock.is_empty());
assert_eq!(lock.total_patches(), 0);
}
#[test]
fn test_add_private_patches() {
let mut lock = PrivateLockFile::new();
let patches = HashMap::from([
("model".to_string(), toml::Value::String("haiku".into())),
("temp".to_string(), toml::Value::String("0.9".into())),
]);
lock.add_private_patches("agents", "my-agent", patches);
assert!(!lock.is_empty());
assert_eq!(lock.total_patches(), 1);
assert!(lock.agents.iter().any(|r| r.name == "my-agent"));
}
#[test]
fn test_empty_patches_not_added() {
let mut lock = PrivateLockFile::new();
lock.add_private_patches("agents", "my-agent", HashMap::new());
assert!(lock.is_empty());
}
#[test]
fn test_save_and_load() {
let temp_dir = TempDir::new().unwrap();
let mut lock = PrivateLockFile::new();
let patches = HashMap::from([("model".to_string(), toml::Value::String("haiku".into()))]);
lock.add_private_patches("agents", "test", patches);
lock.save(temp_dir.path()).unwrap();
let loaded = PrivateLockFile::load(temp_dir.path()).unwrap();
assert!(loaded.is_some());
assert_eq!(loaded.unwrap(), lock);
}
#[test]
fn test_empty_lockfile_deletes_file() {
let temp_dir = TempDir::new().unwrap();
let lock_path = temp_dir.path().join(PRIVATE_LOCK_FILENAME);
std::fs::write(&lock_path, "test").unwrap();
assert!(lock_path.exists());
let lock = PrivateLockFile::new();
lock.save(temp_dir.path()).unwrap();
assert!(!lock_path.exists());
}
#[test]
fn test_load_nonexistent_returns_none() {
let temp_dir = TempDir::new().unwrap();
let result = PrivateLockFile::load(temp_dir.path()).unwrap();
assert!(result.is_none());
}
#[test]
fn test_get_patches() {
let mut lock = PrivateLockFile::new();
let patches = HashMap::from([("model".to_string(), toml::Value::String("haiku".into()))]);
lock.add_private_patches("agents", "test", patches.clone());
let retrieved = lock.get_patches("agents", "test");
assert!(retrieved.is_some());
assert_eq!(retrieved.unwrap(), &patches);
let missing = lock.get_patches("agents", "nonexistent");
assert!(missing.is_none());
}
}