Skip to main content

spool/
enhancement_trace.rs

1use crate::lifecycle_store::lifecycle_root_from_config;
2use serde::{Deserialize, Serialize};
3use std::fs;
4use std::path::{Path, PathBuf};
5use std::time::{SystemTime, UNIX_EPOCH};
6use ts_rs::TS;
7
8const TRACE_FILE_NAME: &str = "prompt-optimize.latest.json";
9const TRACE_SCHEMA_VERSION: &str = "prompt-optimize-trace.v1";
10
11#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, TS)]
12#[ts(export, export_to = "../frontend/src/lib/types/generated/")]
13pub struct PromptOptimizeTrace {
14    pub schema_version: String,
15    pub recorded_at: String,
16    pub source: String,
17    pub cwd: String,
18    pub task: String,
19    pub target: String,
20    pub profile: String,
21    #[serde(default, skip_serializing_if = "Option::is_none")]
22    pub provider: Option<String>,
23    #[serde(default, skip_serializing_if = "Option::is_none")]
24    pub session_id: Option<String>,
25    #[serde(default, skip_serializing_if = "Option::is_none")]
26    pub matched_project_id: Option<String>,
27    pub note_count: usize,
28    pub used_vault_root: String,
29}
30
31impl PromptOptimizeTrace {
32    #[allow(clippy::too_many_arguments)]
33    pub fn new(
34        cwd: impl Into<String>,
35        task: impl Into<String>,
36        target: impl Into<String>,
37        profile: impl Into<String>,
38        provider: Option<String>,
39        session_id: Option<String>,
40        matched_project_id: Option<String>,
41        note_count: usize,
42        used_vault_root: impl Into<String>,
43    ) -> Self {
44        Self {
45            schema_version: TRACE_SCHEMA_VERSION.to_string(),
46            recorded_at: timestamp_string(),
47            source: "mcp.prompt_optimize".to_string(),
48            cwd: cwd.into(),
49            task: task.into(),
50            target: target.into(),
51            profile: profile.into(),
52            provider,
53            session_id,
54            matched_project_id,
55            note_count,
56            used_vault_root: used_vault_root.into(),
57        }
58    }
59}
60
61pub fn write_latest_prompt_optimize_trace(
62    config_path: &Path,
63    trace: &PromptOptimizeTrace,
64) -> anyhow::Result<()> {
65    let path = trace_file_path(config_path);
66    if let Some(parent) = path.parent() {
67        fs::create_dir_all(parent)?;
68    }
69    let temp_path = path.with_extension("tmp");
70    fs::write(&temp_path, serde_json::to_vec_pretty(trace)?)?;
71    fs::rename(temp_path, path)?;
72    Ok(())
73}
74
75pub fn read_latest_prompt_optimize_trace(
76    config_path: &Path,
77) -> anyhow::Result<Option<PromptOptimizeTrace>> {
78    let path = trace_file_path(config_path);
79    if !path.exists() {
80        return Ok(None);
81    }
82    let trace = serde_json::from_slice::<PromptOptimizeTrace>(&fs::read(path)?)?;
83    if trace.schema_version != TRACE_SCHEMA_VERSION {
84        return Ok(None);
85    }
86    Ok(Some(trace))
87}
88
89fn trace_file_path(config_path: &Path) -> PathBuf {
90    let config_dir = config_path.parent().unwrap_or_else(|| Path::new("."));
91    lifecycle_root_from_config(config_dir).join(TRACE_FILE_NAME)
92}
93
94fn timestamp_string() -> String {
95    let seconds = SystemTime::now()
96        .duration_since(UNIX_EPOCH)
97        .unwrap_or_default()
98        .as_secs();
99    format!("unix:{seconds}")
100}
101
102#[cfg(test)]
103mod tests {
104    use super::{
105        PromptOptimizeTrace, read_latest_prompt_optimize_trace, write_latest_prompt_optimize_trace,
106    };
107    use std::fs;
108    use tempfile::tempdir;
109
110    #[test]
111    fn prompt_optimize_trace_should_round_trip_from_config_path() {
112        let temp = tempdir().unwrap();
113        let config_path = temp.path().join("spool.toml");
114        fs::write(&config_path, "[vault]\nroot = \"/tmp\"\n").unwrap();
115
116        let trace = PromptOptimizeTrace::new(
117            "/tmp/repo",
118            "continue task",
119            "codex",
120            "project",
121            Some("codex".to_string()),
122            Some("codex:session-1".to_string()),
123            Some("spool".to_string()),
124            3,
125            "/tmp/vault",
126        );
127        write_latest_prompt_optimize_trace(&config_path, &trace).unwrap();
128
129        let loaded = read_latest_prompt_optimize_trace(&config_path)
130            .unwrap()
131            .unwrap();
132        assert_eq!(loaded.source, "mcp.prompt_optimize");
133        assert_eq!(loaded.cwd, "/tmp/repo");
134        assert_eq!(loaded.provider.as_deref(), Some("codex"));
135        assert_eq!(loaded.session_id.as_deref(), Some("codex:session-1"));
136        assert_eq!(loaded.note_count, 3);
137    }
138}