use crate::Result;
use chrono::Utc;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use sha2::{Digest, Sha256};
use std::path::{Path, PathBuf};
use tokio::fs;
pub struct ArtifactFreezer {
base_dir: PathBuf,
}
impl ArtifactFreezer {
pub fn new() -> Self {
Self {
base_dir: PathBuf::from(".mockforge/frozen"),
}
}
pub fn with_base_dir<P: AsRef<Path>>(base_dir: P) -> Self {
Self {
base_dir: base_dir.as_ref().to_path_buf(),
}
}
pub fn base_dir(&self) -> &Path {
&self.base_dir
}
pub async fn freeze(&self, request: &FreezeRequest) -> Result<FrozenArtifact> {
let path = if let Some(custom_path) = &request.path {
PathBuf::from(custom_path)
} else {
let timestamp = Utc::now().format("%Y%m%d_%H%M%S");
let extension = if request.format == "yaml" || request.format == "yml" {
"yaml"
} else {
"json"
};
self.base_dir
.join(format!("{}_{}.{}", request.artifact_type, timestamp, extension))
};
if let Some(parent) = path.parent() {
fs::create_dir_all(parent).await.map_err(|e| {
crate::Error::io_with_context("create frozen artifacts directory", e.to_string())
})?;
}
let output_hash = if request.metadata.is_some() {
let content_str = serde_json::to_string(&request.content)?;
let mut hasher = Sha256::new();
hasher.update(content_str.as_bytes());
Some(format!("{:x}", hasher.finalize()))
} else {
None
};
let mut frozen_content = request.content.clone();
if let Some(obj) = frozen_content.as_object_mut() {
let mut metadata_json = serde_json::json!({
"frozen_at": Utc::now().to_rfc3339(),
"artifact_type": request.artifact_type,
"source": "ai_generated",
"format": request.format,
});
if let Some(ref metadata) = request.metadata {
if let Some(ref provider) = metadata.llm_provider {
metadata_json["llm_provider"] = Value::String(provider.clone());
}
if let Some(ref model) = metadata.llm_model {
metadata_json["llm_model"] = Value::String(model.clone());
}
if let Some(ref version) = metadata.llm_version {
metadata_json["llm_version"] = Value::String(version.clone());
}
if let Some(ref prompt_hash) = metadata.prompt_hash {
metadata_json["prompt_hash"] = Value::String(prompt_hash.clone());
}
if let Some(ref output_hash) = output_hash {
metadata_json["output_hash"] = Value::String(output_hash.clone());
}
if let Some(ref prompt) = metadata.original_prompt {
metadata_json["original_prompt"] = Value::String(prompt.clone());
}
}
obj.insert("_frozen_metadata".to_string(), metadata_json);
}
let content_str = if request.format == "yaml" || request.format == "yml" {
serde_yaml::to_string(&frozen_content).map_err(|e| {
crate::Error::internal(format!("Failed to serialize to YAML: {}", e))
})?
} else {
serde_json::to_string_pretty(&frozen_content).map_err(|e| {
crate::Error::internal(format!("Failed to serialize to JSON: {}", e))
})?
};
fs::write(&path, content_str)
.await
.map_err(|e| crate::Error::io_with_context("write frozen artifact", e.to_string()))?;
Ok(FrozenArtifact {
artifact_type: request.artifact_type.clone(),
content: frozen_content,
format: request.format.clone(),
path: path.to_string_lossy().to_string(),
metadata: request.metadata.clone(),
output_hash,
})
}
pub async fn auto_freeze_if_enabled(
&self,
request: &FreezeRequest,
deterministic_config: &crate::ai_studio::config::DeterministicModeConfig,
) -> Result<Option<FrozenArtifact>> {
if deterministic_config.enabled && deterministic_config.is_auto_freeze_enabled() {
Ok(Some(self.freeze(request).await?))
} else {
Ok(None)
}
}
pub async fn verify_frozen_artifact(&self, artifact: &FrozenArtifact) -> Result<bool> {
let content_str = serde_json::to_string(&artifact.content)?;
let mut hasher = Sha256::new();
hasher.update(content_str.as_bytes());
let current_hash = format!("{:x}", hasher.finalize());
if let Some(ref stored_hash) = artifact.output_hash {
Ok(current_hash == *stored_hash)
} else {
Ok(true)
}
}
pub async fn freeze_batch(&self, requests: &[FreezeRequest]) -> Result<Vec<FrozenArtifact>> {
let mut results = Vec::new();
for request in requests {
results.push(self.freeze(request).await?);
}
Ok(results)
}
pub async fn load_frozen(
&self,
artifact_type: &str,
identifier: Option<&str>,
) -> Result<Option<FrozenArtifact>> {
let _search_pattern = if let Some(id) = identifier {
format!("{}_*_{}", artifact_type, id)
} else {
format!("{}_*", artifact_type)
};
let mut entries = fs::read_dir(&self.base_dir).await.map_err(|e| {
crate::Error::io_with_context("read frozen artifacts directory", e.to_string())
})?;
let mut latest_match: Option<FrozenArtifact> = None;
let mut latest_time = chrono::DateTime::<Utc>::MIN_UTC;
while let Some(entry) = entries
.next_entry()
.await
.map_err(|e| crate::Error::io_with_context("read directory entry", e.to_string()))?
{
let path = entry.path();
if path.is_file() {
let file_name = path.file_name().and_then(|n| n.to_str()).unwrap_or("");
let matches = if let Some(id) = identifier {
file_name.contains(artifact_type) && file_name.contains(id)
} else {
file_name.starts_with(&format!("{}_", artifact_type))
};
if matches {
let content = fs::read_to_string(&path).await.map_err(|e| {
crate::Error::io_with_context("read frozen artifact", e.to_string())
})?;
let content_value: Value = if path.extension().and_then(|e| e.to_str())
== Some("yaml")
|| path.extension().and_then(|e| e.to_str()) == Some("yml")
{
serde_yaml::from_str(&content).map_err(|e| {
crate::Error::internal(format!("Failed to parse YAML: {}", e))
})?
} else {
serde_json::from_str(&content).map_err(|e| {
crate::Error::internal(format!("Failed to parse JSON: {}", e))
})?
};
let frozen_time = content_value
.get("_frozen_metadata")
.and_then(|m| m.get("frozen_at"))
.and_then(|t| t.as_str())
.and_then(|s| chrono::DateTime::parse_from_rfc3339(s).ok())
.map(|dt| dt.with_timezone(&Utc))
.unwrap_or_else(|| {
Utc::now()
});
if frozen_time > latest_time {
latest_time = frozen_time;
latest_match = Some(FrozenArtifact {
artifact_type: artifact_type.to_string(),
content: content_value,
format: if path.extension().and_then(|e| e.to_str()) == Some("yaml")
|| path.extension().and_then(|e| e.to_str()) == Some("yml")
{
"yaml".to_string()
} else {
"json".to_string()
},
path: path.to_string_lossy().to_string(),
metadata: None,
output_hash: None,
});
}
}
}
}
Ok(latest_match)
}
}
impl Default for ArtifactFreezer {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FreezeRequest {
pub artifact_type: String,
pub content: Value,
pub format: String,
pub path: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub metadata: Option<FreezeMetadata>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FreezeMetadata {
#[serde(skip_serializing_if = "Option::is_none")]
pub llm_provider: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub llm_model: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub llm_version: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub prompt_hash: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub output_hash: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub original_prompt: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FrozenArtifact {
pub artifact_type: String,
pub content: Value,
pub format: String,
pub path: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub metadata: Option<FreezeMetadata>,
#[serde(skip_serializing_if = "Option::is_none")]
pub output_hash: Option<String>,
}